Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way:
a one-shot n_sec_head classifier over conditioning alone, run before any
secondary token existed, with the AR loop then always executing k_max slots
and discarding the tail. This adds a real per-slot EOS mechanism instead:
- Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts
P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head
(n_sec.mode = "stop_token" builds no n_sec_head at all).
- sample_secondaries_ar accepts n_sec_pred=None to drive generation off the
stop head instead of a pre-resolved count: each row stops the first slot
its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the
default, threshold at 0 — or "sample", a Bernoulli draw), and the whole
batch loop breaks once every row has stopped, so cost scales with the
realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly
(the scheduled-sampling self-sample path) is unchanged.
- resolve_n_sec returns None for a stop-token decoder instead of raising;
rollout.py/cli.py/validate.py now derive the realized count from
sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling,
rather than resolving it up front — a no-op reordering under every other
n_sec.mode, where sec_valid was already built from n_sec_pred.
- Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds
the per-slot target/mask (one slot wider than the existing token-content
sec_mask, since the stop slot itself needs supervision) and
StageTrainer._stop_loss trains it with masked BCE, gated on stop_head
exactly like _n_sec_loss gates on n_sec_head. Wired into both the
flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks
stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop
head replaces n_sec_head under this mode, so no new weight key.
- validate_config now accepts stop_token (requires decoder="autoregressive"
and n_sec.owner="stage2") instead of always rejecting it.
Decisions made during planning: stop_sampling defaults to "greedy" for
deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's
HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding
new config keys, since the two heads never coexist.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
lars
c1c4957e2f