Fix/issue 40 #64

2026-08-17T10:47:02+02:00

lars commented

2026-08-17 10:47:02 +02:00

No description provided.

lars added 2 commits 2026-08-17 10:47:03 +02:00

Implement n_sec.mode = "stop_token" for the AR secondary decoder (gitea #40 ) c1c4957e2f

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>

Bump patch version to 0.3.2

CI / Lint (ruff check) (push) Successful in 27s

Details

CI / Format (ruff format) (push) Successful in 28s

Details

CI / Sync project version with tag (push) Has been skipped

Details

CI / Type check (ty) (push) Successful in 26s

Details

CI / Lint (ruff check) (pull_request) Successful in 29s

Details

CI / Format (ruff format) (pull_request) Successful in 28s

Details

CI / Type check (ty) (pull_request) Successful in 38s

Details

CI / Sync project version with tag (pull_request) Has been skipped

Details

CI / Tests (push) Successful in 3m51s

Details

CI / Tests (pull_request) Successful in 2m24s

Details

b42fa95d1a

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

lars merged commit 59eccbb5cb into master

2026-08-17 10:55:33 +02:00

lars deleted branch fix/issue-40

2026-08-17 10:55:33 +02:00

Sign in to join this conversation.

1 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: lars/giant#64