n_sec mode='head': allow sampled instead of argmax secondary counts #86

New Issue

2026-08-28T10:48:17+02:00

lars commented

2026-08-28 10:48:17 +02:00

Problem

For stage2_model.n_sec.mode = "head", the secondary count is resolved deterministically as an argmax over the classifier logits (giant/sample.py:510-511):

logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
return logits.argmax(dim=-1)

The head (Stage2Autoregressive.predict_n_sec, giant/model/models.py:652) is a linear head over the base conditioning embedding emitting (B, K_MAX+1) logits, trained with plain F.cross_entropy (giant/training/trainers.py:539).

Taking the argmax collapses the multiplicity distribution onto its conditional mode: at fixed pre-step conditioning the count is fully deterministic, so all spread in n_sec across a rollout comes from variation in the conditioning, never from the head. Real secondary multiplicity is genuinely stochastic at fixed pre-step state, so this should systematically under-disperse n_sec — and being a mode rather than a mean, it biases low wherever the true conditional distribution is right-skewed, which for multiplicity it typically is.

mode = "stop_token" already has the corresponding knob: n_sec.stop_sampling ("greedy" | "sample", giant/sample.py:348) chooses between thresholding and a Bernoulli draw. mode = "head" has no equivalent.

Proposal

Add a categorical sampling branch to resolve_n_sec, i.e.

probs = logits.softmax(dim=-1)
return torch.multinomial(probs, 1).squeeze(-1)

Preferred config shape: generalize the existing key into one n_sec.sampling = "greedy" | "sample" covering both modes (greedy = argmax for head, threshold at 0 for stop_token; sample = categorical draw for head, Bernoulli for stop_token), rather than adding a second parallel key. stop_sampling would need to be kept as a deprecated alias, or migrated, since it appears in existing checkpoints' model_config — see giant/config.py:431,440,448,1515 and giant/_migration.py.

This is weight-neutral and shape-neutral: it changes only sampling-time behavior, so it also belongs on the inference-safe override allowlist (see the companion issue on inference-time config overrides).

Acceptance

n_sec.sampling (or equivalent) selects greedy vs sampled n_sec under mode = "head"
stop_token behavior unchanged by default; existing stop_sampling values in old checkpoints still honored
Config validation rejects bad values, matching the existing stop_sampling check
Tests parametrized over greedy/sample, in the style of tests/test_sample.py:270-292
Default stays greedy so existing runs are unaffected

Follow-up

Worth checking against a rollout: if a head-mode rollout shows a too-narrow secondary-count distribution vs Geant4 in the analysis pipeline, this is a prime suspect.

## Problem For `stage2_model.n_sec.mode = "head"`, the secondary count is resolved deterministically as an argmax over the classifier logits (`giant/sample.py:510-511`): ```python logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out) return logits.argmax(dim=-1) ``` The head (`Stage2Autoregressive.predict_n_sec`, `giant/model/models.py:652`) is a linear head over the base conditioning embedding emitting `(B, K_MAX+1)` logits, trained with plain `F.cross_entropy` (`giant/training/trainers.py:539`). Taking the argmax collapses the multiplicity distribution onto its conditional mode: at fixed pre-step conditioning the count is fully deterministic, so all spread in `n_sec` across a rollout comes from variation in the conditioning, never from the head. Real secondary multiplicity is genuinely stochastic at fixed pre-step state, so this should systematically under-disperse `n_sec` — and being a mode rather than a mean, it biases low wherever the true conditional distribution is right-skewed, which for multiplicity it typically is. `mode = "stop_token"` already has the corresponding knob: `n_sec.stop_sampling` (`"greedy"` | `"sample"`, `giant/sample.py:348`) chooses between thresholding and a Bernoulli draw. `mode = "head"` has no equivalent. ## Proposal Add a categorical sampling branch to `resolve_n_sec`, i.e. ```python probs = logits.softmax(dim=-1) return torch.multinomial(probs, 1).squeeze(-1) ``` Preferred config shape: generalize the existing key into one `n_sec.sampling = "greedy" | "sample"` covering **both** modes (greedy = argmax for `head`, threshold at 0 for `stop_token`; sample = categorical draw for `head`, Bernoulli for `stop_token`), rather than adding a second parallel key. `stop_sampling` would need to be kept as a deprecated alias, or migrated, since it appears in existing checkpoints' `model_config` — see `giant/config.py:431,440,448,1515` and `giant/_migration.py`. This is weight-neutral and shape-neutral: it changes only sampling-time behavior, so it also belongs on the inference-safe override allowlist (see the companion issue on inference-time config overrides). ## Acceptance - [x] `n_sec.sampling` (or equivalent) selects greedy vs sampled `n_sec` under `mode = "head"` - [x] `stop_token` behavior unchanged by default; existing `stop_sampling` values in old checkpoints still honored - [x] Config validation rejects bad values, matching the existing `stop_sampling` check - [x] Tests parametrized over greedy/sample, in the style of `tests/test_sample.py:270-292` - [x] Default stays `greedy` so existing runs are unaffected ## Follow-up Worth checking against a rollout: if a `head`-mode rollout shows a too-narrow secondary-count distribution vs Geant4 in the analysis pipeline, this is a prime suspect.

lars referenced this issue from a commit

2026-08-28 11:04:54 +02:00

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

lars referenced this issue

2026-08-28 11:08:05 +02:00

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

lars referenced this issue from a commit

2026-08-28 11:18:08 +02:00

Merge pull request 'Add sampled n_sec under n_sec.mode = 'head' (gitea #86)' (#88) from fix/issue-86 into master

lars closed this issue

2026-08-28 11:18:47 +02:00

Sign in to join this conversation.