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>
build_models/build_critics parsed model_config into frozen dataclasses
(ConditioningConfig, Stage2ModelConfig, ...) but then threw the parsed
sub-objects away and passed the original raw dicts (conditioning["particle"],
s2_spec.particle_type.to_dict()) down into ConditionEncoder/StageModel/etc,
which re-read them with their own hardcoded .get(key, default) fallbacks —
each an independent copy of a fact the dataclass already stated once. Worst
instance: giant/training/trainers.py:236 converted an already-parsed
ParticleTypeConfig back into a dict for no reason.
Threads ConditioningAxisConfig (particle_cfg/material_cfg) and
ParticleTypeConfig (particle_type_cfg) as the actual dataclass instances
through every signature that used to type them dict: ConditionEncoder,
StageModel/CriticModel, resolve_type_n_classes/stage2_type_dim/
stage2_trunk_sec_dim, giant/model/builders.py, giant/sample.py,
giant/training/stage2_inputs.py, giant/training/trainers.py (StageSpec/
StageTrainer), giant/pipeline.py, giant/rollout.py, giant/validate.py — so ty
now catches a misspelled field instead of it silently falling back. No
config-schema change: config.toml/checkpoint model_config keep the same
nested-dict shape; only what happens after the existing X.from_dict(...)
parse changes.
User-confirmed scope decision: both axes (particle_cfg/material_cfg and
particle_type_cfg), not just the more heavily-duplicated particle_type_cfg
axis, and not stopping at the two most literal parse-then-discard round
trips — matching the issue's own proposal.
Preserved-default decision: StageModel's particle_type_cfg=None sentinel
(hit only by direct/test construction — build_models always passes an
explicit particle_type) still resolves to ParticleTypeConfig(target=
"physical"), not ParticleTypeConfig()'s own target="onehot" config-file
default — switching it would have silently grown an unused, gradient-less
type_head on every test that constructs Stage2OneShot/Stage2Autoregressive
without particle_type_cfg=, breaking their "every param has a grad" checks.
New tests in tests/test_network.py: ConditionEncoder/StageModel store the
exact ConditioningAxisConfig/ParticleTypeConfig instance passed in (identity,
not just equality) — no internal dict round-trip — and build_models's output
carries real dataclass instances end to end, not the plain dicts it produced
before this fix.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
generator ∈ {"flow", "ddpm", "wgan"} was tested as a bare string in ~45
sites across models.py, sample.py, builders.py, trainers.py, and
stage2_inputs.py, each independently re-deriving one of five consequences
of the choice (needs a time embedding? what does the trunk take as input?
is the type slice folded into the trunk output? which sampler? which
loss?). giant/model/objectives.py adds an Objective ABC + OBJECTIVE_REGISTRY
+ build_objective factory, mirroring routers.py's Router pattern, and every
bare-string site now goes through it (needs_time, is_adversarial,
folds_type_slice, trunk_in_dim, build_schedule, stage1_loss/stage2_loss).
Per discussion: FlowDDPMStageTrainer and WGANStageTrainer stay separate
classes rather than merging into one StageTrainer as the issue's sketch
proposed — their training loops are genuinely different shapes (single loss
vs. dual G/D step with gradient penalty/n_critic/ST-Gumbel), and trainers.py
is the least-covered-by-fast-tests part of the codebase, so a full merge
was judged out of proportion to this issue's risk budget.
FlowDDPMStageTrainer's own loss dispatch (flow vs ddpm, one-shot vs AR) does
move onto the objective, so a future non-adversarial objective (rectified
flow, consistency distillation) is still a one-file, zero-trainer-edits
addition.
No config-schema change — stage{1,2}_model.generator stays the persisted
string, just looked up in the registry instead of string-compared. An
unrecognized generator value now fails fast with a clear ValueError instead
of silently falling through some bare-string checks and not others (same
behavior build_router/build_trunk already have for their own type keys).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The design doc and its followups doc are no longer needed as a live
reference now that the v0.3.0 redesign is implemented — comments and
docstrings across the codebase cited it extensively (file path, "design
doc §X.Y", "decision N", or bare "§X.Y" section numbers) as design
rationale. Removed docs/ and edited every citing comment/docstring to
drop the now-dangling reference while keeping the substantive
explanation next to it. CLAUDE.md's v0.3.0 roadmap bullet loses its
trailing pointer to the deleted file.
Verified: no remaining "docs/v0.3.0", "design doc", "decision N", or
"§N.N" references (repo-wide grep); ruff and ty clean; full test suite
on the heaviest-touched modules (network, sample, rollout, migration,
config, train) passes.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Every metric name used to exist in four places: the dict keys each
StageTrainer returned, the hardcoded _metrics_fields() column list, the
~110-line metrics_row assembly in train(), and the tqdm/summary
formatting. The two had to be kept in exact correspondence by hand or
csv.DictWriter would raise.
Each metric is now declared once, as a MetricSpec on the trainer that
computes it. MetricsCollector derives the CSV header and W&B payload from
those declarations and owns all accumulation, so train() no longer carries
a running sum, and every isinstance(tr, WGANStageTrainer) branch is gone —
replaced by four trainer hooks (batch_loss, summary, val_objective,
supports_val_loss).
giant/train.py (1875 lines) becomes giant/training/:
trainers.py StageSpec + shared StageTrainer base + the two subclasses
metrics.py MetricSpec, MetricsCollector
stage2_inputs.py the pure AR/teacher-forcing tensor helpers, moved verbatim
loop.py train() (225 lines, was ~514) + graceful shutdown
checkpoint.py build/load, lifted out of train()'s closures
The trainers shared ~15 identical constructor arguments and copy-pasted
their cosine-warmup lambda, EMA setup, state_dict/load_state_dict,
resume_lr and train_mode/eval_mode. StageSpec resolves one stage's config
once (constructors go from 24 and 22 keyword arguments to (spec, model,
device)), the base class holds the rest, and build_stage_trainers drops
from ~100 lines to 15.
Metric columns are renamed to a uniform stage/split/metric scheme
(stage1/train/loss, stage2/train/d_loss, stage1/lr, stage1/router/entropy,
val/loss, ...). Old metrics.csv files and W&B history are not comparable.
The checkpoint format is unchanged.
BEHAVIOR CHANGE — WGAN best-checkpoint selection. The old code meant to
score a WGAN stage on its marginal KL, but the guard
`{n: kl for n in wgan_names if n not in val_loss_per_stage}` could never
fire: val_loss_per_stage was pre-seeded with 0.0 for every stage, so a
WGAN stage contributed a flat 0.0 and the KL was written to metrics.csv
without ever influencing best.pt. val_objective now returns it as
intended. On the test harness's default flow+wgan config val_loss went
from 2.182 (stage 1 only) to 15.137 (stage 1 + KL 12.954), and which epoch
won changed. Runs before this commit picked their best checkpoint on the
non-adversarial stages alone. Written up in docs/v0.3.0-followups.md.
Verified: 699 tests pass; ruff, ruff format and ty clean. Baseline-vs-
refactor metrics.csv compared across five configs (flow+wgan, AR+onehot,
routed, both-flow, AR-flow) — every comparable value bit-identical except
val/loss where the fix applies. Resume appends without a duplicate header
and reproduces a HEAD worktree's per-epoch losses and LRs exactly across
the resume boundary. A refactored last.pt loads through
cli.py:_load_model_weights in both raw and ema modes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
AttentionHistory (giant/model/network.py) adds causal self-attention over
the emitted-secondary prefix as the alternative to MarkovHistory, with a
parallel forward() for training and an init_cache()/step() KV-cache path
for sample.py's per-slot AR inference loop, wired into
Stage2Autoregressive via history="attention".
giant/train.py adds _stage2_tf_prob and _assemble_stage2_ar_inputs_scheduled,
mixing ground-truth history with a detached sample_secondaries_ar self-sample
per slot so teacher_forcing="scheduled"/"never" close the train/inference gap
teacher_forcing="always" always avoided; wired into both stage-2 AR trainers.
config.py's validate_config no longer rejects these two previously
unimplemented schema values.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- giant/sample.py: fix every sampler's call convention against
Stage1Model/Stage2OneShot's actual forward signatures (was still
calling model(x, t, cond_cont, cond_cat) positionally); add
sample_secondaries_ar (free-running AR loop, unsnapped history feature)
and sample_stage1/sample_stage2/resolve_n_sec dispatch helpers that read
each stage's generator_kind/decoder off the model instance itself.
- giant/particles.py: decode_topn_class (argmax + other_policy) and
decode_embedding_nearest (L1-snap + distance) turn a secondary's
"onehot"/"embedding" type prediction into a concrete PDG.
- giant/rollout.py: decode_secondary_identity routes all three
particle_type.target values to real mass/charge; per-stage generator
dispatch (drops the single shared `mode` string, adds ddpm support);
L1DistCollector accumulates the §11.3 embedding-distance diagnostic.
- giant/cli.py: drop the onehot/embedding-target rejection gate (narrowed
to the still-unimplemented conditioning.particle/material.type=onehot
axis); fix the dead model_cfg.get("mode") bug in predict/rollout.
- giant/analysis/: new type_embedding_l1_distance PlotSpec, wired through
the rollout YAML sidecar (no live-model call needed, unlike
router_gating -- the histogram is already pre-aggregated at rollout
time).
- Un-xfail every test that was blocked on this step (test_rollout.py,
test_flow.py, test_wgan.py, test_phase2.py, test_router.py,
test_validate.py); add test_sample.py, test_type_embedding_distance.py.
Known follow-up: giant/validate.py still unpacks the training val-batch
as a stale 6-tuple and doesn't use the new per-stage dispatch, so
marginal validation during training degrades gracefully with a warning
rather than working -- not in this step's scope.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds --mode wgan alongside flow/ddpm: both stages get a WGAN-GP
generator/critic pair (giant.model.wgan) instead of flow matching, so
inference is a single forward pass per stage rather than a 10-step ODE
integration — the fast-eval architecture noted in the roadmap.
predict/rollout auto-detect the mode from the checkpoint's model_config.
Best-checkpoint selection for wgan uses marginal-KL against the EMA
generators every epoch, since a critic loss isn't a monotone quality
signal. --router is not supported together with --mode wgan.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds model.conditioning = "physical" | "embedding": physical mode routes
particle mass/charge and material Z_eff/A_eff/density/X0/lambda_int through
small MLPs to replace the learned PDG/material embedding tables, so the
surrogate generalizes to PDG codes/materials outside the training vocab
instead of memorizing it. "embedding" stays available as the comparison
baseline (old checkpoints without the key default to it).
Stage 2 now regresses a secondary's mass/charge directly against a fixed
physics-derived target instead of a learned/snapped embedding, and uses no
snapping at inference — the model's raw predicted (mass, charge) is the
secondary's physical identity, including for its own further rollout steps.
A separate reporting-only nearest-known-PDG lookup (never fed back into the
model) populates output pdg columns / the embedding-mode rollout fallback.
giant/materials.py's table is populated with Geant4's own built-in NIST
constants (Z_eff, A_eff, density, X0, lambda_int), extracted directly from
the Geant4 11.4.1 build vendored in minicalosim via G4NistManager rather
than hand-typed literature values. G4_LYSO is left unfilled: confirmed (both
by runtime lookup and by searching minicalosim's history) that it's never
actually a constructed Geant4 material there, only documentation/UI color-map
text.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
cli.py and scripts/train.py duplicated ~140 lines of training setup and had
drifted (scripts/train.py forgot to save model_config, breaking predict on
those checkpoints). Extract shared logic into giant/constants.py (X_DIM,
target names), giant/config.py (device/git/TOML/seeding helpers, run
metadata), and giant/pipeline.py (the actual training-job orchestration),
so both entry points become thin CLI wrappers around the same code path.
Also adds --seed/--resume support (checkpoints now carry optimizer/scheduler
state, epoch, and best_val_loss), a richer [meta] section in the saved
config.toml (git hash, seed, versions, timestamp, invocation, dataset
stats), and a metrics.csv (train/val loss, lr, epoch time) written every
epoch and append-safe across resumes.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
step_length already encodes |post_pos - pre_pos| by definition, so a raw
post_pos target would duplicate that magnitude and could drift inconsistent
with step_length during sampling. Instead add travel_dir, a unit vector
(local frame) giving only the direction of pre_pos->post_pos; post_pos is
reconstructed at inference as pre_pos + step_length * travel_dir, keeping
the two self-consistent. Target grows from 6D to 9D.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
larsandClaude Opus 5