Merge pull request 'Add configs/baseline.toml as the kept reference model' (#58) from add/baseline-config into master

Reviewed-on: #58

configs/baseline.toml

@@ -0,0 +1,129 @@
+# GIANT reference baseline (v0.3 schema).
+#
+# The fixed comparison point every future architecture variant is measured
+# against. Chosen so that each experimental axis the roadmap cares about
+# (routed trunk, WGAN generators, attention history, shared conditioning,
+# embedding/onehot conditioning) is a *single* edit away from this file.
+#
+# Rationale for the choices below, from the runs already on record
+# (analysis_runs/ + the `giant` W&B project):
+#
+#   * flow, not wgan, for both stages. Ranking the five existing rollouts by
+#     mean Jensen-Shannon divergence against the Geant4 reference, the plain
+#     non-routed flow model wins (0.172) over the routed flow runs
+#     (0.197/0.200) and both WGAN runs (0.218/0.234) — and it beats them by
+#     ~7x on per-event total deposited energy and by 3-10x on every
+#     per-PDG marginal. WGAN stays a variant, not the reference.
+#
+#   * no router. The routed runs are not better, and soft-mixing 10 small
+#     experts costs ~10x per-pass throughput at train time (29k samples/s vs
+#     the WGAN runs' 52-116k), which is what made those runs take ~110 h for
+#     30 epochs.
+#
+#   * hidden_dim 512 / 6 blocks per stage. The best-scoring rollout so far
+#     was hidden_dim 1024, but at 4x the trunk FLOPs of 512. 512/6 sits in
+#     the same weight class as the variants it will be compared against and
+#     leaves headroom to train it properly rather than cheaply.
+#
+#   * dropout 0.0. Training set is ~5e8 steps against <1e7 parameters;
+#     capacity overfitting is not the binding constraint, and every recent
+#     run used 0.0.
+#
+# Known weak spots this baseline is expected to *exhibit* (they are the
+# reason for the comparisons, not a reason to retune this file): every model
+# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
+# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
+# head accuracy sits at 0.863-0.867 regardless of size or objective.
+
+[meta]
+# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
+# particle_type.target = "physical" and the v0.2 default sizes, while still
+# passing validate_config.
+config_version = 3
+
+[conditioning]
+# Physical-property MLPs rather than learned vocab embeddings: computable for
+# any PDG code / material, which is what the held-out-species and
+# held-out-material generalization comparisons need.
+out_dim      = 128
+share_stages = false
+
+# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
+# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
+# identical to the architecture that produced the results cited above.
+[conditioning.particle]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[conditioning.material]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[stage1_model]
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+
+[stage2_model]
+# The v0.3 pivot: autoregressive in descending-energy order with a
+# categorical species target, which is the agreed response to the 2026-08-03
+# secondary-species failure. Flow (not the schema default wgan) so the
+# baseline varies only the decoder relative to the best v0.2 result.
+#
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
+# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
+# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
+# k_max * steps = 150 stage-2 calls per physics step. That makes this block
+# the dominant cost on both sides:
+#   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
+#   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
+# Accepted deliberately: one-shot is the configuration whose secondary
+# species distribution failed, and that failure is what v0.3 exists to fix.
+decoder      = "autoregressive"
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+k_max        = 15
+
+[stage2_model.autoregressive]
+history         = "markov"
+teacher_forcing = "always"
+
+[stage2_model.particle_type]
+target = "onehot"
+# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
+# "other" bucket keeps essentially all real secondary species out of "other"
+# without making the head expensive.
+n_classes    = 32
+other_policy = "sample"
+
+[train]
+epochs        = 50
+# Sized for ONE NVIDIA L40S on deepthought2 (46068 MiB; the box has two, and
+# CLAUDE.md's shared-machine rule allows a single GPU). From a measured
+# linear fit of this exact config's training step on the local RTX 4070:
+#   peak reserved MiB = 0.9736 * batch_size + 115
+# so 36864 reserves ~36.0 GiB, i.e. 78% of the card, leaving ~10 GiB of
+# headroom for fragmentation and the CUDA context. Throughput is already
+# flat above bs~4096 on the 4070, so this is chosen for occupancy on the
+# larger card, not for step efficiency — and it sits next to the 43008/32768
+# of the runs lr = 3e-4 was proven at.
+batch_size    = 36864
+lr            = 3e-4
+warmup_epochs = 3
+weight_decay  = 0.01
+ema_decay     = 0.9999
+val_fraction  = 0.1
+num_workers   = 4
+seed          = 0
+# The marginal/KL pass is expensive (~5000 s on top of an epoch), so keep it
+# to every 10th epoch; the cheap per-epoch val loss still runs every epoch.
+validate_every = 10
+validate_steps = 10
+wandb          = true
+wandb_project  = "giant"