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>
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>
lars