sample_secondaries_ar ran all k_max=15 slots for every row regardless of
each row's own predicted secondary count, even though the baseline
checkpoint's rollout measured only 0.382 secondaries/step — so ~97% of
stage-2 model calls generated tokens sec_valid then masked away.
Compact the loop to the still-active row set at each slot: drop a row the
moment its n_sec_pred is exhausted (or, under n_sec.mode="stop_token", the
moment its own stop logit fires), so slot k's model calls cost O(active
rows) instead of O(B). Exact — rows are independent given their own
history — verified by comparing the compacted path against a new
full_length=True escape hatch that reproduces the original uncompacted
behavior bit-for-bit under deterministic noise.
full_length=True is required by
_assemble_stage2_ar_inputs_scheduled's scheduled-sampling self-sample,
whose training contract needs a real prediction at every slot up to
k_max regardless of a row's own count, so training behavior is
unchanged.
AttentionHistory's KV cache and MarkovHistory's O(1) state are kept
aligned to the shrinking active set via a new
HistoryEncoder.select_cache / Stage2Autoregressive.select_history_cache.
Also fixes a latent bug the refactor surfaced: derived_n_sec (stop-token
mode) could be overwritten by a later spurious re-fire of the stop logit
on a row that had already stopped; now tracked via an explicit `finished`
mask so only the first stop slot is recorded, matching the documented
contract.
No architecture or checkpoint-format change — every existing v0.3.0
Stage2Autoregressive checkpoint (flow/wgan, markov/attention,
head/stop_token) picks up the speedup automatically on its next
`giant rollout`/`giant predict`, no retraining needed.
Measured (CPU, hidden_dim=512/6 blocks, k_max=15, batch 512, mean
n_sec≈0.38 matching the baseline checkpoint's own rollout): 17.6-22.9x
fewer wall-clock seconds for the AR loop alone (attention/markov history
respectively). Directional only — baseline.toml's GPU inference-cost
comment is updated accordingly, flagged stale pending a real rollout
re-measurement via eval_cost_per_step.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HPt7bVLZYFJe5cG6V7ahqC
Branches off configs/baseline.toml with both stages on WGAN-GP and
stage-2 n_sec.mode = stop_token, 30 epochs — combines two unbenchmarked
roadmap axes (post-v0.3.0 WGAN, and the AR stop-token multiplicity mode)
into one variant that stays a single edit away from baseline for
attribution.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KzPghrmFcAJYrWUvHApY9N
Replaces the extrapolated pre-v0.3 weak-spot claims in baseline.toml's
header (which had the wrong sign on step-count error) with measured
numbers from the first full rollout validation of this exact config,
and adds a matching Roadmap entry in CLAUDE.md.
A fixed comparison point for future architecture variants, so each
experimental axis (routed trunk, WGAN generators, attention history,
shared conditioning) is a single edit away from one known config.
flow/flow autoregressive, hidden_dim 512 / 6 blocks per stage, physical
conditioning, no router, 7.70M params. Chosen by ranking the five runs in
analysis_runs/ by mean Jensen-Shannon divergence against the Geant4
reference: unrouted flow wins (0.172) over routed flow (0.197/0.200) and
both WGAN runs (0.218/0.234), with the lead concentrated in per-event
total deposited energy and the per-PDG marginals.
batch_size 36864 is sized for one L40S on deepthought2 from a measured
linear fit of this config's training step (reserved MiB = 0.9736 * bs +
115), giving ~36 GiB, 78% of the card.
The comments record two measured facts that are easy to get wrong:
WGAN is slower to *train* than flow (n_critic plus the gradient-penalty
double-backward), its advantage being inference-only; and
sample_secondaries_ar loops over all k_max slots unconditionally rather
than short-circuiting on n_sec, which is what makes the autoregressive
decoder the dominant cost on both axes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Scaled-up variant of the already-trained wgan_h128_b4_physical.toml
benchmark config, same training recipe (mode/epochs/lr/warmup),
increased model capacity.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Extends the 10-expert EnergyRouter + physical-conditioning benchmark
config with the new opt-in gumbel combine weights, isolating the
learn_centers/learn_temperature axis: none, centers only, and
centers+temperature.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
lars