Add sampled n_sec under n_sec.mode = 'head' (gitea #86)

Taking argmax over the n_sec classifier logits collapses secondary
multiplicity onto its conditional mode at fixed pre-step conditioning,
under-dispersing n_sec in rollouts and biasing low wherever the true
conditional count distribution is right-skewed (typical for
multiplicity).

Generalizes stage2_model.n_sec.stop_sampling (previously stop_token-only)
into stage2_model.n_sec.sampling, covering both "head" (greedy: argmax;
sample: categorical draw via torch.multinomial) and "stop_token" (unchanged:
greedy threshold / Bernoulli draw) modes. stop_sampling is kept as a
deprecated alias in NSecConfig.from_dict and migrate_config, since it
appears in existing checkpoints' model_config. Default stays "greedy" so
existing runs/checkpoints are unaffected.

giant/config.py

@@ -424,11 +424,15 @@ class NSecConfig:
     # n_sec head was trained against Stage 1's own ConditionEncoder output and so has
     # to stay attached there, not just be labeled as such).
     owner: str = "stage2"
-    # mode="stop_token" only: how sample_secondaries_ar turns a slot's stop logit into a
-    # stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample":
-    # a Bernoulli draw at sigmoid(logit) (a real sample from the learned length
-    # distribution, at the cost of an extra RNG draw per slot).
-    stop_sampling: str = "greedy"
+    # How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an
+    # actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits
+    # (deterministic — the conditional mode, not a sample); "sample" is a categorical draw
+    # from softmax(logits) (a real sample from the learned count distribution). mode=
+    # "stop_token": "greedy" is sigmoid(stop_logit) >= 0.5 per slot (deterministic);
+    # "sample" is a Bernoulli draw at sigmoid(stop_logit) per slot. Renamed from
+    # "stop_sampling" (gitea #86), which is still accepted as a deprecated alias since it
+    # appears in existing checkpoints' model_config.
+    sampling: str = "greedy"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "NSecConfig":
@@ -437,7 +441,7 @@ class NSecConfig:
             mode=d.get("mode", "head"),
             lambda_weight=d.get("lambda", 0.1),
             owner=d.get("owner", "stage2"),
-            stop_sampling=d.get("stop_sampling", "greedy"),
+            sampling=d.get("sampling", d.get("stop_sampling", "greedy")),
         )
 
     def to_dict(self) -> dict:
@@ -445,7 +449,7 @@ class NSecConfig:
             "mode": self.mode,
             "lambda": self.lambda_weight,
             "owner": self.owner,
-            "stop_sampling": self.stop_sampling,
+            "sampling": self.sampling,
         }
 
 
@@ -1075,6 +1079,27 @@ def _set_path(d: dict, dotted: str, value) -> None:
     cur[parts[-1]] = value
 
 
+def _pop_path(d: dict, dotted: str) -> None:
+    """Remove a dotted path from a nested dict, if present. No-op if any
+    component along the path is missing."""
+    parts = dotted.split(".")
+    cur = d
+    for part in parts[:-1]:
+        if not isinstance(cur, dict) or part not in cur:
+            return
+        cur = cur[part]
+    if isinstance(cur, dict):
+        cur.pop(parts[-1], None)
+
+
+# Config keys renamed within v0.3 itself (not part of the v0.2->v0.3 migration
+# above) — normalized by migrate_config so a config.toml still using an older
+# v0.3 key name keeps passing validate_config_keys.
+_RENAMED_KEYS = {
+    "stage2_model.n_sec.stop_sampling": "stage2_model.n_sec.sampling",  # gitea #86
+}
+
+
 def _deep_merge(base: dict, override: dict) -> dict:
     """Recursively merge `override` onto a copy of `base`.
 
@@ -1271,11 +1296,21 @@ def migrate_config(cfg: dict) -> dict:
     (which additionally carries n_sec_head ownership and needs
     `network.build_models`'s cooperation) is a separate migration surface,
     deferred to the network.py refactor.
-    """
-    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
-        return copy.deepcopy(cfg)
 
+    Independently of the v0.2/v0.3 branch below, `_RENAMED_KEYS` normalizes
+    keys renamed within v0.3 itself (e.g. `stop_sampling` -> `sampling`,
+    gitea #86) so a config.toml written against an older v0.3 key name still
+    passes `validate_config_keys`.
+    """
     cfg = copy.deepcopy(cfg)
+    for old_path, new_path in _RENAMED_KEYS.items():
+        if _get_path(cfg, old_path) is not None and _get_path(cfg, new_path) is None:
+            _set_path(cfg, new_path, _get_path(cfg, old_path))
+        _pop_path(cfg, old_path)
+
+    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
+        return cfg
+
     old_train = cfg.pop("train", {})
     old_model = cfg.pop("model", {})
     old_router = dict(old_model.pop("router", {}))
@@ -1512,9 +1547,9 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
                 "conditioning to hang an EOS decision off"
             )
 
-    stop_sampling = _get_path(cfg, "stage2_model.n_sec.stop_sampling")
-    if stop_sampling not in ("greedy", "sample"):
-        raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
+    n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
+    if n_sec_sampling not in ("greedy", "sample"):
+        raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")
 
     precision = _get_path(cfg, "train.precision")
     if precision not in ("fp32", "bf16"):

giant/model/builders.py

@@ -138,7 +138,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
                 build_stop_head=stop_token,
-                stop_sampling=s2_spec.n_sec.stop_sampling,
+                n_sec_sampling=s2_spec.n_sec.sampling,
                 stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
             )
         else:
@@ -168,6 +168,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 cond_enc=shared_cond_enc,
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
+                n_sec_sampling=s2_spec.n_sec.sampling,
             )
 
     return result

giant/model/models.py

@@ -370,6 +370,7 @@ class Stage2OneShot(StageModel):
         cond_enc: ConditionEncoder | None = None,
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
+        n_sec_sampling: str = "greedy",
     ) -> None:
         super().__init__(
             pdg_vocab,
@@ -383,6 +384,7 @@ class Stage2OneShot(StageModel):
             particle_type_cfg=particle_type_cfg,
             cond_enc=cond_enc,
         )
+        self.n_sec_sampling = n_sec_sampling
         self._build_context_fusion(x_dim, context_dim, cond_out_dim)
         target = self.particle_type_cfg.target
         type_head_out_dim = None if target == "physical" else k_max * self.type_dim
@@ -498,7 +500,7 @@ class Stage2Autoregressive(StageModel):
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
         build_stop_head: bool = False,
-        stop_sampling: str = "greedy",
+        n_sec_sampling: str = "greedy",
         stop_head_cfg: dict | None = None,
     ) -> None:
         super().__init__(
@@ -514,7 +516,7 @@ class Stage2Autoregressive(StageModel):
             cond_enc=cond_enc,
         )
         self.history_kind = history
-        self.stop_sampling = stop_sampling
+        self.n_sec_sampling = n_sec_sampling
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
             nn.Linear(cond_out_dim + context_dim, cond_out_dim),

giant/model/summary.py

@@ -140,7 +140,7 @@ def _fingerprint(modules: dict[str, nn.Module]) -> list:
     exist, every parameter's/buffer's shape+dtype (never values — those are
     randomly initialized and irrelevant to *structure*), and every plain
     scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
-    .stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
+    .n_sec_sampling`, `EnergyRouter.temperature`) — this is what makes a
     non-parametric key's effect on construction observable."""
     sig = []
     for stage_name, module in modules.items():

giant/sample.py

@@ -261,7 +261,7 @@ def sample_secondaries_ar(
     slot's own stop logit (`predict_stop`, evaluated on the same prefix
     conditioning as the token itself — see `predict_type`'s docstring for
     why this needs no extra state) decides whether generation should have
-    already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
+    already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
     0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
     `n_sec_pred` is the first slot index where this fires; once every row in
     the batch has fired, the loop breaks before spending a model call on the
@@ -345,7 +345,7 @@ def sample_secondaries_ar(
                 slot_idx,
                 hist=hist,
             ).squeeze(1)
-            if sec_decoder.stop_sampling == "sample":
+            if sec_decoder.n_sec_sampling == "sample":
                 stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
             else:
                 stop_now = stop_logit >= 0.0
@@ -496,7 +496,12 @@ def resolve_n_sec(
 
     Raises if neither stage owns any n_sec mechanism at all — the only way
     that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
-    rollout-/predict-capable checkpoint."""
+    rollout-/predict-capable checkpoint.
+
+    `n_sec.mode = "head"` resolves the classifier logits per
+    `sec_decoder.n_sec_sampling`: "greedy" (default) takes the conditional
+    mode via argmax; "sample" draws a real sample from the learned count
+    distribution via `torch.multinomial` on the softmax — see gitea #86."""
     if n_sec_pred is not None:
         return n_sec_pred
     if getattr(sec_decoder, "stop_head", None) is not None:
@@ -508,4 +513,6 @@ def resolve_n_sec(
             "'truth' is standalone-evaluation-only"
         )
     logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
+    if sec_decoder.n_sec_sampling == "sample":
+        return torch.multinomial(logits.softmax(dim=-1), 1).squeeze(-1)
     return logits.argmax(dim=-1)

tests/test_config.py

@@ -171,17 +171,40 @@ def test_n_sec_config_owner_defaults_to_stage2():
 def test_n_sec_config_owner_round_trips():
     n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
     assert n_sec.owner == "stage1"
-    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "sampling": "greedy"}
 
 
-def test_n_sec_config_stop_sampling_defaults_to_greedy():
-    assert gconfig.NSecConfig().stop_sampling == "greedy"
+def test_n_sec_config_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().sampling == "greedy"
 
 
-def test_n_sec_config_stop_sampling_round_trips():
+def test_n_sec_config_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "sampling": "sample"})
+    assert n_sec.sampling == "sample"
+    assert n_sec.to_dict()["sampling"] == "sample"
+
+
+def test_n_sec_config_stop_sampling_alias_still_honored():
+    """gitea #86: stop_sampling was renamed to sampling; old checkpoints'
+    model_config still carries the old key and must keep working."""
     n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
-    assert n_sec.stop_sampling == "sample"
-    assert n_sec.to_dict()["stop_sampling"] == "sample"
+    assert n_sec.sampling == "sample"
+    assert "stop_sampling" not in n_sec.to_dict()
+
+
+def test_n_sec_config_sampling_key_wins_over_stop_sampling_alias():
+    n_sec = gconfig.NSecConfig.from_dict({"sampling": "sample", "stop_sampling": "greedy"})
+    assert n_sec.sampling == "sample"
+
+
+def test_migrate_config_renames_stop_sampling_key():
+    cfg = {
+        "meta": {"config_version": gconfig.CONFIG_VERSION},
+        "stage2_model": {"n_sec": {"stop_sampling": "sample"}},
+    }
+    migrated = gconfig.migrate_config(cfg)
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.sampling") == "sample"
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.stop_sampling") is None
 
 
 # ---------------------------------------------------------------------------
@@ -851,13 +874,13 @@ def test_validate_config_stop_token_rejected_for_stage1_owner():
         assert "stop_token" in str(e) and "owner" in str(e)
 
 
-def test_validate_config_bad_stop_sampling_rejected():
-    cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
+def test_validate_config_bad_n_sec_sampling_rejected():
+    cfg = _cfg_with(**{"stage2_model.n_sec.sampling": "bogus"})
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "stop_sampling" in str(e)
+        assert "sampling" in str(e)
 
 
 def test_validate_config_default_precision_is_fp32():

tests/test_sample.py

@@ -14,6 +14,7 @@ from giant.model.network import (
     stage2_trunk_sec_dim,
 )
 from giant.sample import (
+    resolve_n_sec,
     sample_flow,
     sample_secondaries,
     sample_secondaries_ar,
@@ -72,6 +73,7 @@ def _stage2_ar(
     mat: int = 2,
     k_max: int = 5,
     history: str = "markov",
+    n_sec_sampling: str = "greedy",
 ) -> Stage2Autoregressive:
     particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     return Stage2Autoregressive(
@@ -89,6 +91,7 @@ def _stage2_ar(
         history=history,
         attn_n_heads=2,
         attn_n_layers=1,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -99,7 +102,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
 def _stage2_ar_stop_token(
     target: str,
     generator: str,
-    stop_sampling: str = "greedy",
+    n_sec_sampling: str = "greedy",
     emb_dim: int = 6,
     pdg: int = 3,
     mat: int = 2,
@@ -120,7 +123,7 @@ def _stage2_ar_stop_token(
         particle_type_cfg=ParticleTypeConfig(target=target),
         build_n_sec_head=False,
         build_stop_head=True,
-        stop_sampling=stop_sampling,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -267,14 +270,14 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
 # ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_sampling):
     """A stop_head pinned to a large positive logit fires at slot 0 for
     every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
     a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
     generating any token."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, 50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -283,13 +286,13 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(sto
     assert not sec_valid.any()
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_sampling):
     """A stop_head pinned to a large negative logit never fires under either
     policy, so every row is capped at k_max (the safety cap, not a modeling
     ceiling)."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, -50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -336,3 +339,60 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
     stage1_out = torch.randn(3, X_DIM)
     with pytest.raises(AssertionError):
         sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+
+
+# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
+
+
+def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
+    """Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
+    returns `bias` (broadcast over the batch) as logits regardless of
+    conditioning — mirrors `_force_stop_head_logit`."""
+    assert decoder.n_sec_head is not None
+    last_linear = decoder.n_sec_head[-1]
+    with torch.no_grad():
+        last_linear.weight.zero_()
+        last_linear.bias.copy_(bias)
+
+
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
+    """A logit vector overwhelmingly favoring one class gives the same
+    answer under both policies — greedy because it's the argmax, sample
+    because softmax puts ~all mass on it."""
+    B, k_max = 8, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
+    bias = torch.full((k_max + 1,), -50.0)
+    bias[2] = 50.0
+    _force_n_sec_head_bias(decoder, bias)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))
+
+
+def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
+    B, k_max = 32, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() == 1
+
+
+def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
+    """Under a flat logit vector, a categorical draw across a large batch
+    should hit more than one class — the whole point of gitea #86: greedy
+    always collapses to one, sample should not."""
+    torch.manual_seed(0)
+    B, k_max = 256, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() > 1