diff --git a/configs/baseline.toml b/configs/baseline.toml new file mode 100644 index 0000000..3bd095e --- /dev/null +++ b/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"