EnergyRouter's gate sharpness was a single fixed temperature shared by
every expert, with no way for an expert to independently learn how much
of the energy axis it covers. Adds two mutually exclusive, default-off
modes: learn_width (per-expert learnable width) and learn_temperature
(single learnable shared scalar), both bounded via a sigmoid
interpolation warm-started to reproduce today's fixed-temperature gate
exactly at init, to compare against each other without risking the
unbounded-width collapse failure mode. Also promotes gate_stats's
entropy into a generic, optional Router.entropy_loss (lambda_entropy) as
a secondary guard against all experts' widths co-inflating together.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds Router.gate_stats (per-router gate entropy + per-expert utilization),
logged both per-batch (entropy only, train loop) and per-epoch (full
stats, over the whole val set) — the router-collapse failure mode from
the roadmap's rollout postmortem is now visible during training instead
of only after a full rollout+analysis run. Also splits WGAN critic/
generator grad norms instead of summing them, logs critic LR, n_sec head
accuracy, GPU peak memory + samples/sec, model parameter counts (in
wandb.config), and an is_best flag — all wired into both metrics.csv and
W&B.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
expert_hidden_dim/expert_n_blocks were hardcoded to 128/3 in
DEFAULT_CONFIG, independent of model.hidden_dim/n_blocks, so a routed
run always got fixed 128/3-wide experts no matter what --hidden-dim/
--n-blocks was passed. They now default to 0 ("unset"), which
resolve_expert_dims() resolves by inheriting the model dims; an
explicit override still works and now warns when it diverges from
model.hidden_dim/n_blocks, since the checkpoint dir name won't
reflect it.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The 2026-07-22 rollout benchmark's router_gating diagnostic showed the
10-expert EnergyRouter's default linspace(-2, 2, n_experts) init assumes a
roughly uniform z-normalized energy distribution, leaving experts heavily
overlapping instead of partitioning the range. Add an optional
centers_init kwarg (backward compatible, defaults to the old linspace) and
have giant train estimate it from a reservoir sample of the real energy
column, collected during the existing normalizer-fitting pass.
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>
Route on several independent axes at once (e.g. energy x pdg), each with
its own expert count and hyperparameters. The joint gate is the outer
product of per-axis softmax gates, so it stays a partition of unity and
top1/balance_loss factor per-axis. Config uses flat axis{i}_{field} keys
in model.router (TOML/CLI friendly), also settable via repeatable
--router-axis flags.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Routes on the pre-step PDG code, which — unlike ProcessRouter's process
label — is already known at gate time (a conditioning input), so no
supervision is needed and classify_loss falls back to the zero default.
Generalizes EnergyRouter's soft-turn-on-then-Voronoi trick from a 1-D
distance to a small learned PDG embedding space: its own embedding table
maps each PDG code to a point, and n_experts learnable centers partition
that space.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Brings in the rollout-validation fixes developed alongside Phase 2
(exact e_sec budget rescaling in decode_secondaries, filtering
synthetic termination rows out of load_rollout_vs_truth, Tier 4 truth
overlay, --energy-gev support in dwarf make-root) and reconciles them
with this branch's mixture-of-experts routing work: build_features/
build_models/dataset plumbing keep the ProcessRouter's proc_map/
proc_idx threading, and create_root_files.py's job_seed folds in both
the per-job seed derivation and the new energy_gev component.
Routes on the physics process (Compton, phot, brems, ...) that ends a
step, supervised by a small classifier since process is a post-step
outcome unobservable at gate time. Threads a process label end-to-end
through the data pipeline (loader, build_features, dataset batches,
training loss/checkpointing) alongside the existing EnergyRouter.
Both stages can now route through a pluggable Router (EnergyRouter as the
first implementation, a soft turn-on gate over pre-step log-energy) into
several small ExpertTrunks instead of one monolithic trunk. Trains as a
differentiable soft mixture and dispatches to a single expert per row at
eval time, which is the source of the per-call speedup this prototype is
after (issue #5's ~10x native-Geant4 budget). Disabled by default, so
existing configs/checkpoints are unaffected; build_models() centralizes
routed-vs-monolith construction across train/predict/rollout.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Replaces the independent log_delta_e/log_edep targets with 2 additive-log-ratio
coordinates over the deposit/secondary/post-energy simplex (fractions of pre_E
summing to 1), so edep + e_sec + post_E == pre_E holds by construction after
decoding (softmax) rather than being learned approximately. Requires e_sec
(secondary energy) as a new conditioning input and a steps_to_parquet.py pass
to derive it from child track first-step energies.
Wire a dropout hyperparameter (default 0.1) through the config, model,
training pipeline, and CLI. Persisted in saved model_config so checkpoints
reconstruct the architecture correctly.
Co-Authored-By: Claude Opus 4.8 <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>
lars