Merge pull request 'Fix/issue 40' (#64) from fix/issue-40 into master

Reviewed-on: #64

configs/baseline.toml

@@ -0,0 +1,129 @@
+# GIANT reference baseline (v0.3 schema).
+#
+# The fixed comparison point every future architecture variant is measured
+# against. Chosen so that each experimental axis the roadmap cares about
+# (routed trunk, WGAN generators, attention history, shared conditioning,
+# embedding/onehot conditioning) is a *single* edit away from this file.
+#
+# Rationale for the choices below, from the runs already on record
+# (analysis_runs/ + the `giant` W&B project):
+#
+#   * flow, not wgan, for both stages. Ranking the five existing rollouts by
+#     mean Jensen-Shannon divergence against the Geant4 reference, the plain
+#     non-routed flow model wins (0.172) over the routed flow runs
+#     (0.197/0.200) and both WGAN runs (0.218/0.234) — and it beats them by
+#     ~7x on per-event total deposited energy and by 3-10x on every
+#     per-PDG marginal. WGAN stays a variant, not the reference.
+#
+#   * no router. The routed runs are not better, and soft-mixing 10 small
+#     experts costs ~10x per-pass throughput at train time (29k samples/s vs
+#     the WGAN runs' 52-116k), which is what made those runs take ~110 h for
+#     30 epochs.
+#
+#   * hidden_dim 512 / 6 blocks per stage. The best-scoring rollout so far
+#     was hidden_dim 1024, but at 4x the trunk FLOPs of 512. 512/6 sits in
+#     the same weight class as the variants it will be compared against and
+#     leaves headroom to train it properly rather than cheaply.
+#
+#   * dropout 0.0. Training set is ~5e8 steps against <1e7 parameters;
+#     capacity overfitting is not the binding constraint, and every recent
+#     run used 0.0.
+#
+# Known weak spots this baseline is expected to *exhibit* (they are the
+# reason for the comparisons, not a reason to retune this file): every model
+# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
+# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
+# head accuracy sits at 0.863-0.867 regardless of size or objective.
+
+[meta]
+# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
+# particle_type.target = "physical" and the v0.2 default sizes, while still
+# passing validate_config.
+config_version = 3
+
+[conditioning]
+# Physical-property MLPs rather than learned vocab embeddings: computable for
+# any PDG code / material, which is what the held-out-species and
+# held-out-material generalization comparisons need.
+out_dim      = 128
+share_stages = false
+
+# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
+# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
+# identical to the architecture that produced the results cited above.
+[conditioning.particle]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[conditioning.material]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[stage1_model]
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+
+[stage2_model]
+# The v0.3 pivot: autoregressive in descending-energy order with a
+# categorical species target, which is the agreed response to the 2026-08-03
+# secondary-species failure. Flow (not the schema default wgan) so the
+# baseline varies only the decoder relative to the best v0.2 result.
+#
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
+# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
+# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
+# k_max * steps = 150 stage-2 calls per physics step. That makes this block
+# the dominant cost on both sides:
+#   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
+#   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
+# Accepted deliberately: one-shot is the configuration whose secondary
+# species distribution failed, and that failure is what v0.3 exists to fix.
+decoder      = "autoregressive"
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+k_max        = 15
+
+[stage2_model.autoregressive]
+history         = "markov"
+teacher_forcing = "always"
+
+[stage2_model.particle_type]
+target = "onehot"
+# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
+# "other" bucket keeps essentially all real secondary species out of "other"
+# without making the head expensive.
+n_classes    = 32
+other_policy = "sample"
+
+[train]
+epochs        = 50
+# Sized for ONE NVIDIA L40S on deepthought2 (46068 MiB; the box has two, and
+# CLAUDE.md's shared-machine rule allows a single GPU). From a measured
+# linear fit of this exact config's training step on the local RTX 4070:
+#   peak reserved MiB = 0.9736 * batch_size + 115
+# so 36864 reserves ~36.0 GiB, i.e. 78% of the card, leaving ~10 GiB of
+# headroom for fragmentation and the CUDA context. Throughput is already
+# flat above bs~4096 on the 4070, so this is chosen for occupancy on the
+# larger card, not for step efficiency — and it sits next to the 43008/32768
+# of the runs lr = 3e-4 was proven at.
+batch_size    = 36864
+lr            = 3e-4
+warmup_epochs = 3
+weight_decay  = 0.01
+ema_decay     = 0.9999
+val_fraction  = 0.1
+num_workers   = 4
+seed          = 0
+# The marginal/KL pass is expensive (~5000 s on top of an epoch), so keep it
+# to every 10th epoch; the cheap per-epoch val loss still runs every epoch.
+validate_every = 10
+validate_steps = 10
+wandb          = true
+wandb_project  = "giant"

giant/analysis/reduce.py

@@ -29,8 +29,17 @@ from giant.constants import TERM_ESCAPED
 
 
 def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
-    """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
-    return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
+    """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins.
+
+    Out-of-range values clamp into the edge bins, and the clamp deliberately
+    happens in f64 *before* the integer cast: a rollout is free to emit a wildly
+    out-of-range outlier (a step_length of 1e10 mm, say) or an inf, whose
+    unclamped bin index overflows i32 and makes the cast fail outright. NaN has
+    no edge to clamp to, so it becomes null and is dropped by the callers below
+    — the same thing ``np.histogram`` does with it.
+    """
+    idx = ((value - lo) / (hi - lo) * nbins).floor().clip(0, nbins - 1)
+    return pl.when(idx.is_nan()).then(None).otherwise(idx).cast(pl.Int32)
 
 
 def hist1d(
@@ -50,6 +59,7 @@ def hist1d(
     group = pl.lit(0, dtype=pl.Int64) if group is None else group
     res = (
         lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
+        .drop_nulls("_b")
         .group_by("_g", "_b")
         .agg(pl.len().alias("_n"))
         .collect(engine="streaming")
@@ -188,6 +198,7 @@ def profile_partial(
             _bin_expr(coord, lo, hi, nbins).alias("_b"),
             weight.alias("_w"),
         )
+        .drop_nulls("_b")
         .group_by("event_id", "_b")
         .agg(pl.col("_w").sum().alias("_ws"))
         .collect(engine="streaming")

giant/cli.py

@@ -1043,8 +1043,12 @@ def predict(
             # A fresh v0.3.0 Stage1Model owns no n_sec_head —
             # sample_stage1 returns n_sec_pred=None then, so ask stage 2.
             n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred)
-            sec_cont, sec_type, _sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
-            n_sec_pred_np = n_sec_pred.cpu().numpy()
+            sec_cont, sec_type, sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+            # A stop-token decoder resolves n_sec_pred=None above — read the
+            # real count back off sec_valid_pred instead (a no-op round trip
+            # under every other n_sec.mode, where sec_valid_pred was built
+            # FROM n_sec_pred in the first place).
+            n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy()
 
         pred = stage1_norm.cpu().numpy()  # normalised
 

giant/config.py

@@ -408,18 +408,27 @@ class Stage2RouterConfig(RouterConfig):
 class NSecConfig:
     # "head": a classifier over {0..k_max} on the condition encoding alone
     # (no diffusion noise), callable independently at inference.
-    # "stop_token": an EOS-style implicit stop — accepted by the schema but
-    # not implemented in v0.3.0 (see validate_config).
+    # "stop_token": an EOS-style per-slot stop head on the autoregressive
+    # secondary decoder (Stage2Autoregressive only — see validate_config),
+    # evaluated against the generated prefix instead of conditioning alone.
+    # Replaces n_sec_head entirely: the two are mutually exclusive, so this
+    # mode builds no n_sec_head and stage2_model.n_sec.lambda instead weights
+    # the stop head's BCE term.
     # "truth": take n_sec from ground truth — standalone stage-2 evaluation
     # only, never for rollout.
     mode: str = "head"
-    lambda_weight: float = 0.1  # dict key "lambda" — cross-entropy weight for the head
+    lambda_weight: float = 0.1  # dict key "lambda" — cross-entropy/BCE weight for the head
     # Which stage's module physically owns the n_sec_head weights: "stage2" (default,
     # fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1"
     # (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose
     # 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"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "NSecConfig":
@@ -428,10 +437,16 @@ 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"),
         )
 
     def to_dict(self) -> dict:
-        return {"mode": self.mode, "lambda": self.lambda_weight, "owner": self.owner}
+        return {
+            "mode": self.mode,
+            "lambda": self.lambda_weight,
+            "owner": self.owner,
+            "stop_sampling": self.stop_sampling,
+        }
 
 
 @dataclass(frozen=True)
@@ -1368,11 +1383,23 @@ def validate_config(cfg: dict) -> None:
         )
 
     if _get_path(cfg, "stage2_model.n_sec.mode") == "stop_token":
-        raise ValueError(
-            "stage2_model.n_sec.mode = 'stop_token' is accepted by the schema "
-            "but not implemented in v0.3.0 — use 'head' (default) or 'truth' "
-            "(standalone stage-2 evaluation only, never for rollout)"
-        )
+        if _get_path(cfg, "stage2_model.decoder") != "autoregressive":
+            raise ValueError(
+                "stage2_model.n_sec.mode = 'stop_token' requires "
+                "stage2_model.decoder = 'autoregressive' — there is no "
+                "per-token loop to stop under 'one_shot'"
+            )
+        if _get_path(cfg, "stage2_model.n_sec.owner") != "stage2":
+            raise ValueError(
+                "stage2_model.n_sec.mode = 'stop_token' requires "
+                "stage2_model.n_sec.owner = 'stage2' — a migrated v0.2 "
+                "checkpoint's stage-1 n_sec_head has no per-token "
+                "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'")
 
     if _get_path(cfg, "stage2_model.stage1_context") == "sampled":
         raise ValueError(

giant/data/transforms.py

@@ -14,6 +14,14 @@ _EPS = 1e-8
 # the conservation it slightly softens is physically negligible (~0.001%).
 _SIMPLEX_FLOOR = 1e-5
 
+# Upper clip for a raw predicted log_mass before inv_log_transform: exp(y)
+# must stay well inside float32 range (~3.4e38, i.e. y < ~88.7) or it
+# overflows to inf, which — like the negative-mass case below — blows up the
+# next log_transform call once that mass is fed back in as conditioning.
+# 80.0 leaves comfortable headroom while still being far beyond any physical
+# particle mass a converged model would ever predict.
+_LOG_MASS_MAX = 80.0
+
 
 def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
     x = np.asarray(x, dtype=np.float32)
@@ -685,9 +693,18 @@ def decode_secondaries(
     log_mass = sec_cont[:, :, 4]  # (N, K)
     charge = sec_cont[:, :, 5]  # (N, K)
 
-    # mass is non-negative by construction (inv_log_transform of a real
-    # number is always > 0); clip to 0 for padded/invalid slots rather than
-    # leaving a spurious small positive floor from the log inverse.
+    # log_mass is a raw model prediction, not itself the output of
+    # log_transform, so it can land far outside the range that round-trips
+    # cleanly through inv_log_transform: too negative and exp(log_mass)
+    # undershoots _EPS, making inv_log_transform go slightly negative; too
+    # positive and exp(log_mass) overflows float32 to inf. Either one then
+    # blows up the next log_transform call on this track's mass once it's
+    # fed back in as conditioning for a further rollout step
+    # (giant/rollout.py -> build_cond_features -> _physical_cond_columns).
+    # Clip to a range whose inverse is guaranteed finite and >= 0 before
+    # that can happen; clip to 0 separately for padded/invalid slots rather
+    # than leaving a spurious small positive floor.
+    log_mass = np.clip(log_mass, np.log(_EPS), _LOG_MASS_MAX)
     sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32)
     sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32)
 

giant/model/builders.py

@@ -106,6 +106,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
         # above.
         time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64
         n_sec_owner = s2_spec.n_sec.owner
+        stop_token = s2_spec.n_sec.mode == "stop_token"
         k_max = s2_spec.k_max
         particle_type_cfg = s2_spec.particle_type
 
@@ -128,7 +129,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 router=stage2_router,
                 trunk_type=s2_spec.trunk.type,
                 block_conditioning=s2_spec.trunk.block_conditioning,
-                build_n_sec_head=n_sec_owner != "stage1",
+                build_n_sec_head=n_sec_owner != "stage1" and not stop_token,
                 particle_type_cfg=particle_type_cfg,
                 history=ar_cfg.history,
                 attn_n_heads=ar_cfg.attn_n_heads,
@@ -136,6 +137,9 @@ 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(),
+                build_stop_head=stop_token,
+                stop_sampling=s2_spec.n_sec.stop_sampling,
+                stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
             )
         else:
             sec_dim = stage2_trunk_sec_dim(

giant/model/models.py

@@ -132,13 +132,15 @@ class StageModel(nn.Module):
         n_sec_head_cfg: dict | None,
         type_head_out_dim: int | None,
         type_head_cfg: dict | None,
+        build_stop_head: bool = False,
+        stop_head_cfg: dict | None = None,
     ) -> None:
         """Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`,
-        `self.type_head`. Called by a subclass's `__init__` after it has set
-        up its own conditioning-assembly modules — `merged_cond_dim` below
-        must match the width that assembly (`_cond_embed`/`_base_cond`/
-        `_token_cond`, or plain `cond_enc` for `Stage1Model`) actually
-        produces.
+        `self.type_head`, `self.stop_head`. Called by a subclass's `__init__`
+        after it has set up its own conditioning-assembly modules —
+        `merged_cond_dim` below must match the width that assembly
+        (`_cond_embed`/`_base_cond`/`_token_cond`, or plain `cond_enc` for
+        `Stage1Model`) actually produces.
 
         `n_sec_head` is built iff `n_sec_head_k_max is not None` (output
         width `n_sec_head_k_max + 1`) — `Stage1Model` passes this only for a
@@ -148,7 +150,11 @@ class StageModel(nn.Module):
         classes — passes `None` exactly when `particle_type_cfg.target ==
         "physical"`) *and* the objective doesn't fold the type slice into its
         own trunk output (checked here, since `objective` is already needed
-        for the trunk itself).
+        for the trunk itself). `stop_head` is built iff `build_stop_head` —
+        only `Stage2Autoregressive` ever passes `True` (`n_sec.mode ==
+        "stop_token"`, mutually exclusive with `n_sec_head`), a single
+        `cond_out_dim -> 1` logit per call, same `HeadConfig` shape rules as
+        the other two heads.
         """
         objective = build_objective(self.generator_kind)
         has_time = objective.needs_time
@@ -176,6 +182,11 @@ class StageModel(nn.Module):
             head_cfg = HeadConfig.from_dict(type_head_cfg)
             hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
             self.type_head = build_mlp_head(cond_out_dim, type_head_out_dim, hidden, head_cfg.depth)
+        self.stop_head = None
+        if build_stop_head:
+            head_cfg = HeadConfig.from_dict(stop_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth)
 
     def _require_n_sec_head(self) -> None:
         if self.n_sec_head is None:
@@ -195,6 +206,14 @@ class StageModel(nn.Module):
                 "directly instead)"
             )
 
+    def _require_stop_head(self) -> None:
+        if self.stop_head is None:
+            raise RuntimeError(
+                f"this {type(self).__name__} has no stop_head — only a "
+                "Stage2Autoregressive built with stage2_model.n_sec.mode = "
+                "'stop_token' owns one"
+            )
+
 
 class Stage1Model(StageModel):
     """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
@@ -431,7 +450,13 @@ class Stage2Autoregressive(StageModel):
     `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend
     on token position; `_token_cond` additionally fuses in the history
     encoding and two running scalars (remaining energy-budget fraction,
-    normalized slot index), and feeds `forward`/`predict_type`/the trunk.
+    normalized slot index), and feeds `forward`/`predict_type`/`predict_stop`/
+    the trunk.
+
+    `n_sec.mode = "stop_token"` (`build_stop_head=True`) replaces
+    `predict_n_sec`'s one-shot classifier with `predict_stop`'s per-token EOS
+    logit instead — the two heads are mutually exclusive (`build_n_sec_head`
+    is `False` whenever this is `True`, see `giant.model.builders`).
     """
 
     def __init__(
@@ -461,6 +486,9 @@ class Stage2Autoregressive(StageModel):
         cond_enc: ConditionEncoder | None = None,
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
+        build_stop_head: bool = False,
+        stop_sampling: str = "greedy",
+        stop_head_cfg: dict | None = None,
     ) -> None:
         super().__init__(
             pdg_vocab,
@@ -475,6 +503,7 @@ class Stage2Autoregressive(StageModel):
             cond_enc=cond_enc,
         )
         self.history_kind = history
+        self.stop_sampling = stop_sampling
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
             nn.Linear(cond_out_dim + context_dim, cond_out_dim),
@@ -516,6 +545,8 @@ class Stage2Autoregressive(StageModel):
             n_sec_head_cfg=n_sec_head_cfg,
             type_head_out_dim=type_head_out_dim,
             type_head_cfg=type_head_cfg,
+            build_stop_head=build_stop_head,
+            stop_head_cfg=stop_head_cfg,
         )
 
     def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
@@ -638,6 +669,38 @@ class Stage2Autoregressive(StageModel):
         B, K, _ = c_emb.shape
         return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
 
+    def predict_stop(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """`(B, K)` raw stop logits — `n_sec.mode = "stop_token"` only.
+        Evaluated on slot `k`'s own conditioning (which carries slot `k-1`'s
+        history, same as `predict_type`), so this is `P(n_sec == k |
+        prefix)`: a high logit at slot `k` means "stop before generating a
+        token here" — the caller (`giant.sample.sample_secondaries_ar`)
+        checks it before spending a model call on that slot's token."""
+        self._require_stop_head()
+        assert self.stop_head is not None
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+            hist=hist,
+        )
+        B, K, _ = c_emb.shape
+        return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K)
+
 
 class CriticModel(nn.Module):
     """Generator-agnostic WGAN-GP critic body: a scalar realism score, for

giant/rollout.py

@@ -681,7 +681,6 @@ def _step_chunk(
     post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local)
 
     n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)
-    n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
 
     # --- Secondaries ---
     # No snapping for "physical"/history-facing state elsewhere in the
@@ -691,7 +690,13 @@ def _step_chunk(
     # decode_secondary_identity's docstring for how each
     # particle_type.target differs on whether PDG resolution is a real
     # identity decision or just a reporting label.
-    sec_cont, sec_type, _valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+    sec_cont, sec_type, sec_valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+    # A stop-token decoder resolves n_sec_pred=None above — the real count
+    # only exists once sample_stage2 has actually generated (or stopped
+    # generating) tokens, so read it back off sec_valid here. Under every
+    # other n_sec.mode sec_valid was built FROM n_sec_pred, so this is a
+    # no-op round trip in those cases.
+    n_sec_np = sec_valid.sum(dim=-1).cpu().numpy().astype(np.int64)
     sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
         sec_decoder,
         sec_cont,

giant/sample.py

@@ -229,14 +229,14 @@ def sample_secondaries_ar(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     stage1_out: torch.Tensor,
-    n_sec_pred: torch.Tensor,
+    n_sec_pred: torch.Tensor | None,
     steps: int = 10,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """`Stage2Autoregressive` inference loop: one token at a time, in
-    descending-energy slot order, `k_max` sequential calls. Unlike training
-    (teacher forcing — a single parallel pass over ground-truth tokens, see
-    `giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no
-    ground truth at inference: each token's conditioning is built
+    descending-energy slot order, up to `k_max` sequential calls. Unlike
+    training (teacher forcing — a single parallel pass over ground-truth
+    tokens, see `giant.training.stage2_inputs._assemble_stage2_ar_inputs`),
+    there is no ground truth at inference: each token's conditioning is built
     free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the
     train/inference gap that is the cost of markov history's
     expressiveness.
@@ -244,7 +244,30 @@ def sample_secondaries_ar(
     A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
     the "K sequential forwards" cost applies per-token here, not
     once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
-    physics step.
+    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
+    every row in the batch has stopped).
+
+    `n_sec_pred`, if given, fixes each row's secondary count up front (as
+    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
+    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
+    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
+    running self-sample regardless of the decoder's own stop head — the
+    scheduled-sampling training contract does not truncate). This always
+    runs the full `k_max`-iteration loop, masking by the given count at the
+    end exactly as before.
+
+    `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
+    (`n_sec.mode = "stop_token"`): before generating each slot's token, that
+    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
+    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
+    next slot's token — the average-case cost win the docstring above
+    describes. A row that never fires within `k_max` is capped there
+    (`K_MAX` stays a safety cap, not a modeling ceiling).
 
     Under `history="attention"` the history encoding is computed once per
     slot via `Stage2Autoregressive.history_step` (a KV-cache append)
@@ -294,6 +317,16 @@ def sample_secondaries_ar(
     remaining = torch.ones(B, device=device)
     history_cache = sec_decoder.init_history_cache()
 
+    use_stop_token = n_sec_pred is None
+    if use_stop_token:
+        assert getattr(sec_decoder, "stop_head", None) is not None, (
+            "sample_secondaries_ar called with n_sec_pred=None on a decoder "
+            "with no stop_head — only valid under stage2_model.n_sec.mode = "
+            "'stop_token'"
+        )
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+        derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
+
     for k in range(k_max):
         has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
         history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
@@ -301,6 +334,26 @@ def sample_secondaries_ar(
         slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
         hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
+        if use_stop_token:
+            stop_logit = sec_decoder.predict_stop(
+                cond_cont,
+                cond_cat,
+                stage1_out,
+                history_feat,
+                has_prev,
+                remaining_frac,
+                slot_idx,
+                hist=hist,
+            ).squeeze(1)
+            if sec_decoder.stop_sampling == "sample":
+                stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
+            else:
+                stop_now = stop_logit >= 0.0
+            derived_n_sec[stop_now & ~finished] = k
+            finished = finished | stop_now
+            if finished.all():
+                break
+
         if objective.is_adversarial:
             z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
             token = sec_decoder(
@@ -362,7 +415,8 @@ def sample_secondaries_ar(
         prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
         remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
 
-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
+    resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
     return sec_cont, sec_type, sec_valid
 
 
@@ -397,7 +451,7 @@ def sample_stage2(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     stage1_out: torch.Tensor,
-    n_sec_pred: torch.Tensor,
+    n_sec_pred: torch.Tensor | None,
     steps: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Dispatches on `decoder` (one-shot vs autoregressive — the class
@@ -406,9 +460,15 @@ def sample_stage2(
     built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*`
     is the only stage-2 training path that exists for the non-adversarial
     case, so there's nothing to dispatch to here.
+
+    `n_sec_pred=None` (from `resolve_n_sec` on a stop-token decoder) is only
+    meaningful for the autoregressive path — see `sample_secondaries_ar`'s
+    docstring; the one-shot samplers have no per-token stop mechanism to
+    derive a count from, so `n_sec_pred` must already be resolved for them.
     """
     if isinstance(sec_decoder, Stage2Autoregressive):
         return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
+    assert n_sec_pred is not None, "one-shot stage-2 decoders need a resolved n_sec_pred"
     if build_objective(sec_decoder.generator_kind).is_adversarial:
         return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
     return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
@@ -421,20 +481,31 @@ def resolve_n_sec(
     cond_cat: torch.Tensor,
     stage1_out: torch.Tensor,
     n_sec_pred: torch.Tensor | None,
-) -> torch.Tensor:
+) -> torch.Tensor | None:
     """`n_sec_pred` is already populated when `stage1_model` owns a legacy
     `n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s
-    docstring); otherwise ask stage 2, which owns it by default. Raises if
-    neither stage owns a head at all — the only way that happens is
-    `stage2_model.n_sec.mode` other than `"head"` (`"truth"`/`"stop_token"`),
-    neither of which is a valid rollout-/predict-capable checkpoint."""
+    docstring); otherwise ask stage 2, which owns it by default.
+
+    Returns `None` when `sec_decoder` owns a `stop_head` (`n_sec.mode =
+    "stop_token"`) instead of an `n_sec_head` — there is nothing to resolve
+    up front in that case, since the count only exists once
+    `sample_secondaries_ar` has actually generated (or stopped generating)
+    tokens; the caller passes this `None` straight through to `sample_stage2`
+    and reads the real count back off its returned `sec_valid`
+    (`sec_valid.sum(-1)`) afterwards.
+
+    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."""
     if n_sec_pred is not None:
         return n_sec_pred
+    if getattr(sec_decoder, "stop_head", None) is not None:
+        return None
     if getattr(sec_decoder, "n_sec_head", None) is None:
         raise RuntimeError(
-            "checkpoint has no n_sec_head on either stage — needs "
-            "stage2_model.n_sec.mode = 'head' (the default); 'truth' is "
-            "standalone-evaluation-only and 'stop_token' isn't implemented"
+            "checkpoint has no n_sec_head/stop_head on either stage — needs "
+            "stage2_model.n_sec.mode = 'head' (the default) or 'stop_token'; "
+            "'truth' is standalone-evaluation-only"
         )
     logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
     return logits.argmax(dim=-1)

giant/tools/dwarf.py

@@ -442,43 +442,65 @@ def warm_cache(
         Path,
         typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
     ],
+    config: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--config",
+            "-c",
+            help="TOML config file to warm for — same file the `giant train` run(s) will use. "
+            "Mutually exclusive with the flags below (put val-fraction/seed/conditioning/router "
+            "settings in the file itself, so warming and training can't disagree on them)",
+        ),
+    ] = None,
     val_fraction: Annotated[
-        float,
+        Optional[float],
         typer.Option(
             "--val-fraction",
             "-f",
-            help="Must match the `giant train` run(s) to warm for",
+            help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
         ),
-    ] = 0.1,
+    ] = None,
     seed: Annotated[
-        int,
-        typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
-    ] = 0,
+        Optional[int],
+        typer.Option(
+            "--seed",
+            "-s",
+            help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
+        ),
+    ] = None,
     particle_conditioning: Annotated[
-        Conditioning,
+        Optional[Conditioning],
         typer.Option(
             "--particle-conditioning",
-            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for",
+            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for. "
+            "Not allowed together with --config",
         ),
-    ] = Conditioning.physical,
+    ] = None,
     material_conditioning: Annotated[
-        Conditioning,
+        Optional[Conditioning],
         typer.Option(
             "--material-conditioning",
             help="Must match the `giant train` run(s)' conditioning.material.type "
             "to warm for — independent of --particle-conditioning "
-            "(the two axes may differ)",
+            "(the two axes may differ). Not allowed together with --config",
         ),
-    ] = Conditioning.physical,
+    ] = None,
     router: Annotated[
-        bool,
+        Optional[bool],
         typer.Option(
             "--router/--no-router",
-            help="Warm the process vocabulary too (only takes effect with --router-type process)",
+            help="Warm the process vocabulary too (only takes effect with --router-type process). "
+            "Not allowed together with --config",
         ),
-    ] = False,
-    router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
-    n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
+    ] = None,
+    router_type: Annotated[
+        Optional[str],
+        typer.Option("--router-type", help="Router implementation name. Not allowed together with --config"),
+    ] = None,
+    n_experts: Annotated[
+        Optional[int],
+        typer.Option("--n-experts", help="Number of routed experts. Not allowed together with --config"),
+    ] = None,
     rebuild: Annotated[
         bool,
         typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
@@ -487,17 +509,38 @@ def warm_cache(
     """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
 
     Warms the vocab maps, event-id split index, and the normalizer entry for
-    the given --val-fraction/--seed/--particle-conditioning/
-    --material-conditioning, so a later `giant train` run (or a `dwarf
-    hparam-scan` sweep, which shares one such entry across every run) skips
-    straight to training. See giant/data/setup_cache.py.
+    either --config, or the given --val-fraction/--seed/
+    --particle-conditioning/--material-conditioning/--router* flags, so a
+    later `giant train` run (or a `dwarf hparam-scan` sweep, which shares one
+    such entry across every run) skips straight to training. See
+    giant/data/setup_cache.py.
     """
+    flag_overrides = {
+        "--val-fraction": val_fraction,
+        "--seed": seed,
+        "--particle-conditioning": particle_conditioning,
+        "--material-conditioning": material_conditioning,
+        "--router/--no-router": router,
+        "--router-type": router_type,
+        "--n-experts": n_experts,
+    }
+    if config is not None:
+        given = [name for name, value in flag_overrides.items() if value is not None]
+        if given:
+            typer.echo(
+                f"error: --config cannot be combined with {', '.join(given)} "
+                "— put these settings in the config file instead",
+                err=True,
+            )
+            raise typer.Exit(1)
+
     run_warm_setup_cache(
         data=str(data),
+        config_path=config,
         val_fraction=val_fraction,
         seed=seed,
-        particle_conditioning=particle_conditioning.value,
-        material_conditioning=material_conditioning.value,
+        particle_conditioning=particle_conditioning.value if particle_conditioning is not None else None,
+        material_conditioning=material_conditioning.value if material_conditioning is not None else None,
         router_enabled=router,
         router_type=router_type,
         n_experts=n_experts,

giant/tools/warm_setup_cache.py

@@ -10,63 +10,86 @@ for the sidecar itself.
 from pathlib import Path
 
 from giant import config as gconfig
-from giant.constants import K_MAX
 from giant.pipeline import run_setup_stage
 
 
 def run_warm_setup_cache(
     data: str,
-    val_fraction: float = 0.1,
-    seed: int = 0,
-    particle_conditioning: str = "physical",
-    material_conditioning: str = "physical",
-    router_enabled: bool = False,
-    router_type: str = "energy",
-    n_experts: int = 4,
+    config_path: Path | None = None,
+    val_fraction: float | None = None,
+    seed: int | None = None,
+    particle_conditioning: str | None = None,
+    material_conditioning: str | None = None,
+    router_enabled: bool | None = None,
+    router_type: str | None = None,
+    n_experts: int | None = None,
     rebuild: bool = False,
     echo=print,
 ) -> None:
     """Populate (or refresh) the setup cache sidecar for `data`.
 
-    `val_fraction`/`seed`/`particle_conditioning`/`material_conditioning`
-    select the normalizer cache entry
-    (`giant.data.setup_cache.normalizer_key`) — pass the same values a later
-    `giant train` invocation will use so it hits this warmed entry. The two
-    conditioning axes are independent and may differ.
-    `router_enabled`/`router_type`/`n_experts` only matter for
-    `router_type == "process"` (warms that `n_experts`'s process map); the
-    energy-router quantile summary is always collected regardless, so a
-    later `--router-type energy` run never needs to rescan just to seed
-    centers.
+    Two mutually exclusive ways to select what to warm for (enforced by the
+    caller, `giant.tools.dwarf.warm_cache` — this function just trusts
+    whichever combination it's given):
+
+    - `config_path`: the same TOML `giant train --config` takes. Every value
+      `run_setup_stage` needs (`train.val_fraction`/`seed`,
+      `conditioning.particle`/`material.type`, both stages' `router`,
+      `stage2_model.particle_type.n_classes`, ...) is read from the one
+      resulting merged `cfg`, so a later `giant train --config <same file>`
+      run resolves to exactly the same cache keys — see gitea #59.
+    - The individual flags below: `val_fraction`/`seed`/
+      `particle_conditioning`/`material_conditioning` select the normalizer
+      cache entry (`giant.data.setup_cache.normalizer_key`) — pass the same
+      values a later `giant train` invocation will use so it hits this
+      warmed entry. The two conditioning axes are independent and may
+      differ. `router_enabled`/`router_type`/`n_experts` only matter for
+      `router_type == "process"` (warms that `n_experts`'s process map); the
+      energy-router quantile summary is always collected regardless, so a
+      later `--router-type energy` run never needs to rescan just to seed
+      centers.
+
+    Any flag left `None` is omitted from the merge, so it falls back to
+    `DEFAULT_CONFIG`'s own value (or the config file's, if `config_path` is
+    given) instead of silently overriding it — see gitea #59.
     """
-    router_cfg = {
-        "enabled": router_enabled,
-        "type": router_type,
-        "n_experts": n_experts,
-    }
-    # Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so
-    # run_setup_stage always sees every key it might read (e.g.
-    # conditioning.particle.emb_dim, stage2_model.particle_type.target) at
-    # its real default, not silently missing/None — see issues.md Issue 1.
+    overrides: dict = {}
+
+    conditioning_overrides: dict = {}
+    if particle_conditioning is not None:
+        conditioning_overrides["particle"] = {"type": particle_conditioning}
+    if material_conditioning is not None:
+        conditioning_overrides["material"] = {"type": material_conditioning}
+    if conditioning_overrides:
+        overrides["conditioning"] = conditioning_overrides
+
     # This CLI only ever configures one router (matching today's single
-    # --router-type flag), so it's placed on stage1_model; stage2_model's
-    # stays disabled.
-    cfg = gconfig.merge_cli_overrides(
-        gconfig.DEFAULT_CONFIG,
-        None,
-        {
-            "conditioning": {
-                "particle": {"type": particle_conditioning},
-                "material": {"type": material_conditioning},
-            },
-            "stage1_model": {"router": router_cfg},
-            "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
-        },
-    )
+    # --router-type flag), so it's placed on stage1_model; stage2_model's is
+    # left to DEFAULT_CONFIG/the config file rather than forced disabled.
+    router_overrides: dict = {}
+    if router_enabled is not None:
+        router_overrides["enabled"] = router_enabled
+    if router_type is not None:
+        router_overrides["type"] = router_type
+    if n_experts is not None:
+        router_overrides["n_experts"] = n_experts
+    if router_overrides:
+        overrides["stage1_model"] = {"router": router_overrides}
+
+    train_overrides: dict = {}
+    if val_fraction is not None:
+        train_overrides["val_fraction"] = val_fraction
+    if seed is not None:
+        train_overrides["seed"] = seed
+    if train_overrides:
+        overrides["train"] = train_overrides
+
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, overrides)
+    gconfig.validate_config(cfg)
     run_setup_stage(
         Path(data),
-        val_fraction=val_fraction,
-        seed=seed,
+        val_fraction=cfg["train"]["val_fraction"],
+        seed=cfg["train"]["seed"],
         cfg=cfg,
         cache_setup=True,
         rebuild_setup_cache=rebuild,

giant/training/stage2_inputs.py

@@ -140,6 +140,27 @@ def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
     return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
 
 
+def _stop_target_and_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+    """`(target, mask)`, both `(B, K_MAX)`, for `n_sec.mode = "stop_token"`'s
+    per-slot EOS head (`Stage2Autoregressive.predict_stop`).
+
+    `predict_stop` is evaluated on slot `k`'s own (pre-token) conditioning —
+    "should generation have already stopped by here" — so `target[k] = 1`
+    exactly at `k == n_sec` (the first invalid slot: `sample_secondaries_ar`
+    checks this before spending a model call generating that slot's token),
+    `0` elsewhere. `mask` is `k <= n_sec` — one slot *wider* than
+    `StageTrainer._sec_mask`'s `k < n_sec` token-content mask, since the stop
+    slot itself (`k == n_sec`) must be supervised even though there is no
+    real secondary there. A row with `n_sec == k_max` has no in-range stop
+    slot at all: `mask` covers the full `k_max` range (every generated token
+    is real) and `target` is all-zero — `sample_secondaries_ar` correctly
+    never breaks early for it, running into the `k_max` safety cap instead."""
+    idx = torch.arange(k_max, device=device).unsqueeze(0)
+    target = (idx == n_sec.unsqueeze(1)).float()
+    mask = idx <= n_sec.unsqueeze(1)
+    return target, mask
+
+
 def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]:
     """`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR
     conditioning tensors that don't depend on *which* history representation

giant/training/trainers.py

@@ -34,6 +34,7 @@ from giant.training.stage2_inputs import (
     _gumbel_tau,
     _relax_onehot_type_slice,
     _stage2_tf_prob,
+    _stop_target_and_mask,
 )
 
 
@@ -89,6 +90,7 @@ class StageSpec:
     # loss weights
     lambda_weight: float = 1.0
     n_sec_lambda: float = 0.1
+    n_sec_mode: str = "head"
 
     # particle-type target (stage 2 only)
     particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
@@ -144,6 +146,7 @@ class StageSpec:
             decoder=s2_spec.decoder if is_stage2 else "one_shot",
             lambda_weight=stage_spec.lambda_weight,
             n_sec_lambda=s2_spec.n_sec.lambda_weight,
+            n_sec_mode=s2_spec.n_sec.mode,
             particle_type=s2_spec.particle_type,
             particle_type_n_classes=resolve_type_n_classes(
                 s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"]
@@ -374,10 +377,9 @@ class StageTrainer:
         stage1-vs-stage2 `predict_n_sec` signature split, shared by the
         non-adversarial and WGAN trainers.
 
-        Gated on `n_sec_head is None`, not on `n_sec.mode`: a future
-        `mode="stop_token"` model (currently rejected in
-        `validate_config`) carries no head and would train its EOS signal in
-        the generator/AR loss path instead, so this correctly stays zero.
+        Gated on `n_sec_head is None`, not on `n_sec.mode`: a `mode =
+        "stop_token"` model carries no head at all (see `_stop_loss` for its
+        EOS signal instead), so this correctly stays zero for it.
         """
         if self.model.n_sec_head is None:
             zero = torch.zeros((), device=device)
@@ -391,6 +393,47 @@ class StageTrainer:
         nsec_acc = (logits.argmax(dim=-1) == n_sec).float().mean()
         return l_nsec, nsec_acc
 
+    def _stop_loss(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_ctx: torch.Tensor,
+        n_sec: torch.Tensor,
+        device: torch.device,
+        ar_inputs: dict[str, torch.Tensor] | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """`(l_stop, stop_acc)` for `n_sec.mode = "stop_token"`'s per-slot EOS
+        head (`Stage2Autoregressive.predict_stop`) — zeros when this stage
+        owns no `stop_head` (every other `n_sec.mode`), the same gating
+        convention `_n_sec_loss` uses for `n_sec_head`. The two heads are
+        mutually exclusive (`giant.model.builders`), so exactly one of
+        `_n_sec_loss`/`_stop_loss` is ever non-zero for a given stage.
+
+        Masked BCE against `_stop_target_and_mask`'s per-slot target — one
+        slot wider than `sec_mask` (the stop slot itself, `k == n_sec`, needs
+        supervision even though it holds no real secondary)."""
+        stop_head = getattr(self.model, "stop_head", None)
+        if stop_head is None:
+            zero = torch.zeros((), device=device)
+            return zero, zero
+        assert ar_inputs is not None
+        logits = self.model.predict_stop(
+            cond_cont,
+            cond_cat,
+            stage1_ctx,
+            ar_inputs["history_feat"],
+            ar_inputs["has_prev"],
+            ar_inputs["remaining_frac"],
+            ar_inputs["slot_idx"],
+        )
+        target, mask = _stop_target_and_mask(n_sec, logits.size(1), device)
+        mask_f = mask.float()
+        denom = mask_f.sum().clamp(min=1)
+        bce = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
+        l_stop = (bce * mask_f).sum() / denom
+        stop_acc = (((logits >= 0).float() == target).float() * mask_f).sum() / denom
+        return l_stop, stop_acc
+
     @staticmethod
     def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
         """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
@@ -477,10 +520,12 @@ class FlowDDPMStageTrainer(StageTrainer):
                 "loss",
                 "loss_gen",
                 "loss_nsec",
+                "loss_stop",
                 "loss_balance",
                 "loss_proc",
                 "loss_entropy",
                 "nsec_acc",
+                "stop_acc",
                 "loss_type",
                 "type_acc",
                 "grad_norm",
@@ -493,6 +538,8 @@ class FlowDDPMStageTrainer(StageTrainer):
                 "loss_gen",
                 "loss_nsec",
                 "nsec_acc",
+                "loss_stop",
+                "stop_acc",
                 "loss_type",
                 "type_acc",
             )
@@ -586,6 +633,7 @@ class FlowDDPMStageTrainer(StageTrainer):
 
         l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs)
         l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+        l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)
 
         l_type, type_acc = self._type_loss(
             cond_cont,
@@ -606,7 +654,11 @@ class FlowDDPMStageTrainer(StageTrainer):
             if self.spec.lambda_entropy > 0:
                 l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
 
-        total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type
+        total = (
+            self.spec.lambda_weight * l_gen
+            + self.spec.n_sec_lambda * (l_nsec + l_stop)
+            + self.particle_type_lambda * l_type
+        )
         if self.spec.lambda_balance > 0:
             total = total + self.spec.lambda_balance * l_balance
         if self.spec.lambda_proc > 0:
@@ -618,12 +670,14 @@ class FlowDDPMStageTrainer(StageTrainer):
             "loss": total,
             "loss_gen": l_gen,
             "loss_nsec": l_nsec,
+            "loss_stop": l_stop,
             "loss_type": l_type,
             "type_acc": type_acc,
             "loss_balance": l_balance,
             "loss_proc": l_proc,
             "loss_entropy": l_entropy,
             "nsec_acc": nsec_acc,
+            "stop_acc": stop_acc,
         }
 
     def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
@@ -723,6 +777,8 @@ class WGANStageTrainer(StageTrainer):
             "gp_loss",
             "loss_nsec",
             "nsec_acc",
+            "loss_stop",
+            "stop_acc",
             "grad_norm_d",
             "grad_norm_g",
         ]
@@ -735,11 +791,15 @@ class WGANStageTrainer(StageTrainer):
         self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")]
 
     def _stage2_real_and_fake(self, batch_tensors: _Stage2RealFakeBatch, stage1_ctx, global_step, device):
-        """Build `(real, fake_raw, mask, critic_fn)` for stage 2, covering
-        both decoders and all three particle-type targets. `fake_raw` still
-        needs the caller's straight-through relaxation under
+        """Build `(real, fake_raw, mask, critic_fn, ar_inputs)` for stage 2,
+        covering both decoders and all three particle-type targets. `fake_raw`
+        still needs the caller's straight-through relaxation under
         `particle_type.target = "onehot"`, and neither tensor is masked-and-
-        multiplied on the fake side yet."""
+        multiplied on the fake side yet. `ar_inputs` is `None` under
+        `decoder = "one_shot"`; under `"autoregressive"` it's the same dict
+        `_ar_inputs` built to condition `self.model` above — returned so the
+        caller's `_stop_loss` reuses it instead of paying for a second
+        (possibly self-sampling) `_ar_inputs` call."""
         cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
         B = cond_cont.size(0)
         type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes)
@@ -752,9 +812,10 @@ class WGANStageTrainer(StageTrainer):
         def critic_fn(x):
             return self.critic(x, cond_cont, cond_cat, stage1_ctx)
 
+        ar_inputs = None
         if self.decoder == "autoregressive":
             epoch = global_step // self.spec.steps_per_epoch
-            ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
+            ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
             real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False).reshape(B, -1) * mask
             z = torch.randn(B, k_max, self.model.noise_dim, device=device)
             fake_raw = self.model(
@@ -762,17 +823,17 @@ class WGANStageTrainer(StageTrainer):
                 cond_cont,
                 cond_cat,
                 stage1_ctx,
-                ar["history_feat"],
-                ar["has_prev"],
-                ar["remaining_frac"],
-                ar["slot_idx"],
+                ar_inputs["history_feat"],
+                ar_inputs["has_prev"],
+                ar_inputs["remaining_frac"],
+                ar_inputs["slot_idx"],
             ).reshape(B, -1)
         else:
             real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True) * mask
             z = torch.randn(B, self.model.noise_dim, device=device)
             fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx)
 
-        return real, fake_raw, mask, critic_fn
+        return real, fake_raw, mask, critic_fn, ar_inputs
 
     def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         (
@@ -788,6 +849,7 @@ class WGANStageTrainer(StageTrainer):
         stage1_ctx = x1_s1.detach()
         grad_probe: dict[str, float] = {}
 
+        ar_inputs = None
         if not self.is_stage2:
             real = x1_s1
 
@@ -798,7 +860,7 @@ class WGANStageTrainer(StageTrainer):
             fake = self.model(z, cond_cont, cond_cat)
             mask = None
         else:
-            real, fake_raw, mask, critic_fn = self._stage2_real_and_fake(
+            real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake(
                 _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
                 stage1_ctx,
                 global_step,
@@ -840,18 +902,19 @@ class WGANStageTrainer(StageTrainer):
         # --- generator (+ n_sec) step ---
         did_g_step = global_step % self.n_critic == 0
         l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+        l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)
 
-        # On a non-generator-step batch with no n_sec_head on this stage
-        # (n_sec now defaults to stage 2), there's nothing for
+        # On a non-generator-step batch with no n_sec_head/stop_head on this
+        # stage (n_sec now defaults to stage 2), there's nothing for
         # the generator optimizer to do this batch — g_loss would otherwise
         # be a graph-less zero tensor, which .backward() rejects outright.
-        skip_g_step = not did_g_step and self.model.n_sec_head is None
+        skip_g_step = not did_g_step and self.model.n_sec_head is None and self.model.stop_head is None
         if did_g_step:
             g_loss_adv = generator_loss(critic_fn, fake)
-            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
+            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * (l_nsec + l_stop)
         else:
             g_loss_adv = torch.zeros((), device=device)
-            g_loss = self.spec.n_sec_lambda * l_nsec
+            g_loss = self.spec.n_sec_lambda * (l_nsec + l_stop)
         if skip_g_step:
             grad_norm_g = 0.0
         else:
@@ -869,6 +932,8 @@ class WGANStageTrainer(StageTrainer):
             "gp_loss": gp.item(),
             "loss_nsec": l_nsec.item(),
             "nsec_acc": nsec_acc.item(),
+            "loss_stop": l_stop.item(),
+            "stop_acc": stop_acc.item(),
             "did_g_step": did_g_step,
             "grad_norm": grad_norm_d + grad_norm_g,
             "grad_norm_d": grad_norm_d,

giant/validate.py

@@ -129,10 +129,7 @@ def validate_marginals(
             continue
 
         n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
-        n_sec_pred_np = n_sec_pred.cpu().numpy()
         n_sec_np = n_sec.numpy()
-        all_n_sec_real.append(n_sec_np)
-        all_n_sec_pred.append(n_sec_pred_np)
 
         real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None]  # (B, k_max)
         real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
@@ -140,6 +137,13 @@ def validate_marginals(
         sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
             sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
         )
+        # A stop-token decoder resolves n_sec_pred=None above — read the real
+        # count back off sec_valid_pred instead (a no-op round trip under
+        # every other n_sec.mode, where sec_valid_pred was built FROM
+        # n_sec_pred in the first place).
+        n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy()
+        all_n_sec_real.append(n_sec_np)
+        all_n_sec_pred.append(n_sec_pred_np)
         gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
         gen_valid = sec_valid_pred.cpu().numpy()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.1"
+version = "0.3.2"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"

tests/test_analysis_reduce.py

@@ -102,6 +102,31 @@ def test_hist1d_overall_and_grouped():
     assert hg[11].sum() == 4
 
 
+def test_hist1d_clamps_extreme_values_and_drops_nan():
+    # A rollout can emit a wildly out-of-range step_length (or an inf/NaN); the
+    # fixed-edge binning must clamp rather than overflow the i32 bin cast.
+    lf = pl.DataFrame({"x": [5.0, 1.0725e10, float("inf"), -float("inf"), float("nan"), None]}).lazy()
+    edges = np.linspace(0.0, 50.0, 6)  # width 10
+    h = R.hist1d(lf, pl.col("x"), edges)
+    # 5 -> bin 0; 1e10 and +inf -> top bin; -inf -> bin 0; NaN/null dropped
+    assert h[0].tolist() == [2, 0, 0, 0, 2]
+
+
+def test_profile_partial_clamps_extreme_values_and_drops_nan():
+    lf = pl.DataFrame(
+        {
+            "event_id": [1, 1, 1, 1],
+            "z": [5.0, 1.0725e10, float("nan"), 45.0],
+            "w": [1.0, 2.0, 4.0, 8.0],
+        }
+    ).lazy()
+    edges = np.linspace(0.0, 50.0, 6)
+    ev, mat = R.profile_partial(lf, pl.col("z"), edges, pl.col("w"))
+    assert ev.tolist() == [1]
+    # 1e10 clamps into the top bin alongside 45; the NaN row's weight is dropped
+    assert mat[0].tolist() == [1.0, 0.0, 0.0, 0.0, 10.0]
+
+
 def test_physical_steps_drops_synthetic_rollout_rows_only():
     lf = _rollout_frame()
     phys = physical_steps(lf, Side.rollout).collect()

tests/test_config.py

@@ -150,7 +150,17 @@ 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"}
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+
+
+def test_n_sec_config_stop_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().stop_sampling == "greedy"
+
+
+def test_n_sec_config_stop_sampling_round_trips():
+    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"
 
 
 # ---------------------------------------------------------------------------
@@ -725,13 +735,49 @@ def test_validate_config_tie_to_stage1_requires_stage1_active():
         assert "tie_to_stage1" in str(e)
 
 
-def test_validate_config_stop_token_not_implemented():
+def test_validate_config_stop_token_accepted_under_autoregressive():
+    """DEFAULT_CONFIG's stage2_model.decoder is already "autoregressive"
+    (see test_stage2_model_config_defaults_match_documented_v030_intent), so
+    mode="stop_token" alone must not raise."""
     cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_stop_token_rejected_under_one_shot():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "stop_token",
+            "stage2_model.decoder": "one_shot",
+        }
+    )
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "stop_token" in str(e)
+        assert "stop_token" in str(e) and "autoregressive" in str(e)
+
+
+def test_validate_config_stop_token_rejected_for_stage1_owner():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "stop_token",
+            "stage2_model.n_sec.owner": "stage1",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        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"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stop_sampling" in str(e)
 
 
 def test_validate_config_stage1_context_sampled_not_implemented():

tests/test_dwarf.py

@@ -124,6 +124,11 @@ def test_warm_cache_router_process_warms_proc_map(tmp_path):
         [
             "warm-cache",
             str(data),
+            # router.type="process" is incompatible with the default
+            # conditioning.particle.type="physical" (validate_config, now
+            # enforced by warm-cache too — see gitea #59).
+            "--particle-conditioning",
+            "embedding",
             "--router",
             "--router-type",
             "process",
@@ -167,3 +172,63 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
     assert loaded is not None
     assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
     assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
+
+
+def test_warm_cache_config_warms_particle_type_n_classes(tmp_path):
+    """gitea #59: a config setting stage2_model.particle_type.n_classes away
+    from its 0 (= inherit conditioning.particle.emb_dim) default must warm
+    the pdg top-N map under that n_classes, not the emb_dim default, so a
+    later `giant train --config <same file>` run hits it instead of quietly
+    re-scanning every parquet file."""
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    config_path = tmp_path / "config.toml"
+    config_path.write_text("[meta]\nconfig_version = 3\n\n[stage2_model.particle_type]\nn_classes = 32\n")
+
+    runner.invoke(app, ["warm-cache", str(data), "--config", str(config_path)])
+    result = runner.invoke(app, ["warm-cache", str(data), "--config", str(config_path)])
+
+    assert result.exit_code == 0, result.output
+    assert "pdg top-N map: cache hit" in result.output
+    assert "32 classes" in result.output
+
+
+def test_warm_cache_config_rejects_val_fraction_flag(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    config_path = tmp_path / "config.toml"
+    config_path.write_text("[meta]\nconfig_version = 3\n")
+
+    result = runner.invoke(
+        app,
+        ["warm-cache", str(data), "--config", str(config_path), "--val-fraction", "0.2"],
+    )
+
+    assert result.exit_code != 0
+    assert "--config" in result.output
+    assert "--val-fraction" in result.output
+
+
+def test_warm_cache_config_rejects_router_flags(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    config_path = tmp_path / "config.toml"
+    config_path.write_text("[meta]\nconfig_version = 3\n")
+
+    result = runner.invoke(
+        app,
+        [
+            "warm-cache",
+            str(data),
+            "--config",
+            str(config_path),
+            "--router",
+            "--router-type",
+            "process",
+            "--n-experts",
+            "3",
+        ],
+    )
+
+    assert result.exit_code != 0
+    assert "--config" in result.output
+    assert "--router/--no-router" in result.output
+    assert "--router-type" in result.output
+    assert "--n-experts" in result.output

tests/test_phase2.py

@@ -627,3 +627,53 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
     _, _, sec_mass, sec_charge, _ = decode_secondaries(sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm)
     assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
     assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
+
+
+def test_decode_secondaries_extreme_negative_log_mass_stays_nonnegative():
+    from giant.data.transforms import decode_secondaries, log_transform
+
+    N = 1
+    e_sec = np.array([5.0], dtype=np.float32)
+    n_sec = np.array([1])
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
+    sec_cont[0, 0, 0] = 10.0  # stick logit -> ~all of e_sec
+    sec_cont[0, 0, 1:4] = [0, 0, 1]
+    sec_cont[0, 0, 4] = -50.0  # raw model prediction: extremely negative log_mass
+    sec_cont[0, 0, 5] = 1.0
+
+    _, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+
+    # A raw model prediction isn't itself the output of log_transform, so
+    # naively applying inv_log_transform can undershoot zero (see
+    # decode_secondaries) — which then crashes the next log_transform call
+    # once this mass is fed back in as conditioning during rollout. The
+    # float32 residual from clipping can land a hair below zero, but must
+    # stay well above -eps so log_transform(mass) stays finite.
+    assert sec_mass[0, 0] > -1e-8
+    log_transform(sec_mass[0, 0])
+
+
+def test_decode_secondaries_extreme_positive_log_mass_stays_finite():
+    from giant.data.transforms import decode_secondaries, log_transform
+
+    N = 1
+    e_sec = np.array([5.0], dtype=np.float32)
+    n_sec = np.array([1])
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
+    sec_cont[0, 0, 0] = 10.0  # stick logit -> ~all of e_sec
+    sec_cont[0, 0, 1:4] = [0, 0, 1]
+    sec_cont[0, 0, 4] = 200.0  # raw model prediction: extremely positive log_mass
+    sec_cont[0, 0, 5] = 1.0
+
+    _, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+
+    # Mirror image of the extreme-negative case above: exp(log_mass)
+    # overflows float32 to inf for an unclipped raw prediction this large,
+    # which then crashes the next log_transform call the same way a
+    # negative mass would.
+    assert np.isfinite(sec_mass[0, 0])
+    log_transform(sec_mass[0, 0])

tests/test_rollout.py

@@ -366,6 +366,7 @@ def _models_v3(
     k_max=6,
     emb_dim=4,
     stage2_has_n_sec_head=True,
+    stop_token=False,
 ):
     particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
     material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
@@ -417,7 +418,8 @@ def _models_v3(
             noise_dim=8,
             k_max=k_max,
             particle_type_cfg=particle_type_cfg,
-            build_n_sec_head=stage2_has_n_sec_head,
+            build_n_sec_head=stage2_has_n_sec_head and not stop_token,
+            build_stop_head=stop_token,
         )
     return s1.eval(), s2.eval()
 
@@ -481,6 +483,23 @@ def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props):
         assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
 
 
+@pytest.mark.parametrize("generator2", ["flow", "wgan"])
+def test_rollout_stop_token_end_to_end(fake_material_props, generator2):
+    """gitea #40: n_sec.mode='stop_token' (no n_sec_head on either stage —
+    resolve_n_sec must return None and let sample_stage2's AR loop derive
+    the count from its own stop head) must still run to completion, produce
+    secondaries, and conserve energy exactly like the 'head' mode."""
+    s1, s2 = _models_v3(decoder="autoregressive", generator2=generator2, stop_token=True)
+    rec = _run_v3(s1, s2)
+    assert len(rec["event_id"]) > 0
+    seeds = _seeds()
+    for i, ev in enumerate(seeds["event_id"]):
+        m = rec["event_id"] == ev
+        dep = rec["edep"][m].sum()
+        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
+        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
+
+
 def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
     """Neither stage owning n_sec_head only happens for a
     stage2_model.n_sec.mode other than "head" — not a valid rollout-capable

tests/test_sample.py

@@ -96,6 +96,45 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
     return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
 
 
+def _stage2_ar_stop_token(
+    target: str,
+    generator: str,
+    stop_sampling: str = "greedy",
+    emb_dim: int = 6,
+    pdg: int = 3,
+    mat: int = 2,
+    k_max: int = 5,
+) -> Stage2Autoregressive:
+    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
+    return Stage2Autoregressive(
+        pdg_vocab=pdg,
+        mat_vocab=mat,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        generator=generator,
+        time_dim=16,
+        noise_dim=8,
+        k_max=k_max,
+        particle_type_cfg=ParticleTypeConfig(target=target),
+        build_n_sec_head=False,
+        build_stop_head=True,
+        stop_sampling=stop_sampling,
+    ).eval()
+
+
+def _force_stop_head_logit(decoder: Stage2Autoregressive, logit: float) -> None:
+    """Zeroes stop_head's weights and pins its bias, so predict_stop returns
+    `logit` for every row/slot regardless of conditioning — makes the AR
+    loop's stop decision deterministic for testing."""
+    assert decoder.stop_head is not None
+    last_linear = decoder.stop_head[-1]
+    with torch.no_grad():
+        last_linear.weight.zero_()
+        last_linear.bias.fill_(logit)
+
+
 # ── Stage-1 n_sec ownership ──────────────────────────────────────────────────
 
 
@@ -223,3 +262,77 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
     sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
     assert sec_valid.tolist() == [[False], [True], [True]]
+
+
+# ── 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):
+    """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)
+    _force_stop_head_logit(decoder, 50.0)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    assert sec_valid.shape == (B, k_max)
+    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):
+    """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)
+    _force_stop_head_logit(decoder, -50.0)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    assert sec_valid.all()
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_stop_token_valid_mask_is_always_a_prefix(generator):
+    """Without forcing the stop head, per-row stop timing varies — but
+    sec_valid must always be a contiguous prefix (slot k valid implies every
+    slot < k is also valid), matching the "head"/"truth" contract."""
+    B, k_max = 6, 5
+    decoder = _stage2_ar_stop_token("physical", generator, k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    n = sec_valid.sum(dim=-1)
+    expected = torch.arange(k_max).unsqueeze(0) < n.unsqueeze(1)
+    assert torch.equal(sec_valid, expected)
+
+
+def test_sample_secondaries_ar_stop_token_explicit_n_sec_pred_ignores_stop_head():
+    """The scheduled-sampling training contract: passing n_sec_pred
+    explicitly (as _assemble_stage2_ar_inputs_scheduled's self-sample call
+    does, with ground-truth n_sec) must run the full k_max loop and mask by
+    the given count, even though the decoder owns a stop_head that would
+    otherwise stop early."""
+    B, k_max = 3, 5
+    decoder = _stage2_ar_stop_token("physical", "flow", k_max=k_max)
+    _force_stop_head_logit(decoder, 50.0)  # would stop immediately if consulted
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 2, k_max])
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
+    for i, n in enumerate(n_sec_pred.tolist()):
+        assert sec_valid[i, :n].all()
+        assert not sec_valid[i, n:].any()
+
+
+def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
+    decoder = _stage2_ar("physical", "flow", k_max=5)  # head mode: no stop_head
+    cond_cont, cond_cat = _cond(3)
+    stage1_out = torch.randn(3, X_DIM)
+    with pytest.raises(AssertionError):
+        sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)

tests/test_train.py

@@ -20,7 +20,7 @@ from giant.constants import (
     X_DIM,
 )
 from giant.data.dataset import StepBatch
-from giant.model.network import build_critics, build_models
+from giant.model.network import Stage2Autoregressive, build_critics, build_models
 from giant.training import (
     FlowDDPMStageTrainer,
     StageSpec,
@@ -40,6 +40,7 @@ from giant.training.stage2_inputs import (
     _shift_prev,
     _stage2_tf_prob,
     _stick_fraction,
+    _stop_target_and_mask,
     _type_repr,
 )
 
@@ -120,6 +121,23 @@ def test_ar_has_prev_false_only_at_slot_zero():
     assert has_prev.tolist() == [[False, True, True, True, True]]
 
 
+def test_stop_target_and_mask_hand_computed():
+    # k_max=5; n_sec=0 (no real secondaries, stop slot is 0), n_sec=2
+    # (stop slot is 2), n_sec=5 (== k_max: no in-range stop slot at all).
+    n_sec = torch.tensor([0, 2, 5])
+    target, mask = _stop_target_and_mask(n_sec, 5, torch.device("cpu"))
+    assert target.tolist() == [
+        [1, 0, 0, 0, 0],
+        [0, 0, 1, 0, 0],
+        [0, 0, 0, 0, 0],
+    ]
+    assert mask.tolist() == [
+        [True, False, False, False, False],
+        [True, True, True, False, False],
+        [True, True, True, True, True],
+    ]
+
+
 # --- _stage2_tf_prob (v0.3.0 step 7) -----------
 
 
@@ -740,6 +758,71 @@ def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
         assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows)
 
 
+# --- n_sec.mode = "stop_token" (gitea #40) ----------------------------------
+
+
+def _stop_token_cfg():
+    cfg = _base_cfg()
+    cfg["stage2_model"]["decoder"] = "autoregressive"
+    cfg["stage2_model"]["n_sec"] = {"mode": "stop_token", "lambda": 0.1}
+    return cfg
+
+
+@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
+def test_build_stage_trainers_stop_token_step_runs(stage2_generator):
+    """A stop_token AR stage-2 trainer.step() must run and emit a finite
+    loss_stop for both non-adversarial (flow) and WGAN generators — the two
+    trainer subclasses wire the stop head's BCE term in independently."""
+    cfg = _stop_token_cfg()
+    cfg["stage2_model"]["generator"] = stage2_generator
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
+    trainer = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    stats = trainer.step(batch, torch.device("cpu"), global_step=1)
+    assert math.isfinite(stats["loss_stop"])
+    assert math.isfinite(stats["stop_acc"])
+
+
+def test_stop_token_model_has_stop_head_not_n_sec_head():
+    cfg = _stop_token_cfg()
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    stage2 = models["stage2"]
+    assert isinstance(stage2, Stage2Autoregressive)
+    assert stage2.n_sec_head is None
+    assert stage2.stop_head is not None
+
+
+def test_head_mode_model_has_n_sec_head_not_stop_head():
+    """Sanity check on the other side of the gate — the default 'head' mode
+    must be unaffected by the stop_head plumbing."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["decoder"] = "autoregressive"
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    stage2 = models["stage2"]
+    assert isinstance(stage2, Stage2Autoregressive)
+    assert stage2.n_sec_head is not None
+    assert stage2.stop_head is None
+
+
+def test_train_end_to_end_stop_token():
+    """Full train() run with n_sec.mode='stop_token' must complete and write
+    a checkpoint + metrics.csv with finite losses throughout."""
+    cfg = _stop_token_cfg()
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        with open(out_dir / "metrics.csv", newline="") as f:
+            rows = list(csv.DictReader(f))
+        assert len(rows) == cfg["train"]["epochs"]
+        assert all(math.isfinite(float(r["stage2/train/loss_stop"])) for r in rows)
+
+
 def test_wgan_physical_omits_grad_norm_slice_columns():
     cfg = _base_cfg()  # _base_cfg's stage2_model.particle_type.target is "physical"
     with tempfile.TemporaryDirectory() as tmp:

uv.lock

@@ -633,7 +633,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.1"
+version = "0.3.2"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },