Files
giant/docs/v0.3.0-design.md
T
lars 8019a80563
CI / Format (ruff format) (push) Successful in 27s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 39s
CI / Type check (ty) (push) Successful in 43s
CI / Format (ruff format) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (push) Successful in 2m35s
CI / Tests (pull_request) Successful in 2m36s
Refactor train.py into giant/training/ around a metrics collector
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>
2026-08-07 17:03:20 +02:00

58 KiB
Raw Blame History

GIANT v0.3.0 — Stage-2 autoregressive redesign

Status: design agreed, not implemented. Branch v0.3.0-stage2-autoregressive. Date: 2026-08-04. Source: ~/knowledge-base/meetings/2026-08-04-jan-stage2-autoregressive-architecture.md (meeting with Jan), plus the design decisions taken in the session that produced this document.

This document is the implementation contract for v0.3.0. It specifies the new config format option by option, the network.py refactor, and the order in which to build it. Read it before touching giant/config.py or giant/model/network.py.


Table of contents

  1. Why
  2. Decisions register
  3. Config format reference
  4. Migration: v0.2 -> v0.3
  5. network.py refactor
  6. Stage-2 autoregressive design
  7. Training loop
  8. Data and setup-cache changes
  9. Config machinery changes
  10. Callers that need updating
  11. Settled scope and open questions
  12. Implementation order

1. Why

The 2026-08-03 WGAN rollout benchmark (~/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md) was good at the primary-step level and failed at the secondary-species level: zero photon secondaries generated, ~4M hallucinated -14 (muon antineutrino) secondaries — a species essentially absent from Geant4.

Stage 1 is not implicated; the meeting scoped everything to Stage 2. Two changes were agreed together:

  • Autoregressive generation over secondaries in decreasing energy order, replacing the one-shot masked K_MAX=15 prediction, trained with teacher forcing.
  • Categorical particle type with a data-derived "other" bucket, reversing the 2026-07-17 move to a continuous (log-mass, charge) target. Working hypothesis: the continuous target is part of what let the generator collapse onto degenerate species.

The meeting also set a methodology: compare Stage-2 architectures standalone (trained directly on secondary columns, no Stage-1 forward pass) before chaining the winner behind Stage 1. Iterating on the compounding-rollout-error problem is far cheaper that way than paying for a full two-stage run per candidate.

That methodology is what forces the config refactor: v0.2 has a single global train.mode and a single [model] block, with no way to express "train stage 2 only", "stage 1 flow + stage 2 WGAN", or "one-shot vs autoregressive stage 2".


2. Decisions register

# Decision Rationale
1 n_sec head moves from stage 1 to stage 2 Stage 1 becomes the pure 9D primary step. A stage-2-only run is then self-contained (it can predict its own multiplicity), and §3's "implicit stop token" alternative gets a natural home in the same config block.
2 Full mixed per-stage objectives stage1 = flow + stage2 = wgan must actually run — Stage 1 is good as flow, Stage 2 is what is being iterated on. The train loop becomes one trainer object per stage, each owning its optimizers and update cadence.
3 Migration shim for configs and checkpoints Nothing on /ceph goes dead. v0.2 config.toml files and v0.2 model_config dicts are translated on load.
4 Nested objective sub-tables [stage1_model.wgan], [stage2_model.ddpm] rather than flat wgan_n_critic keys — self-documenting about which keys the active generator ignores, and validation can warn on a populated sub-table that is never read.
5 Particle type is adversarial: straight-through Gumbel into the critic The critic sees a relaxed one-hot alongside energy/direction, so the joint (species, kinematics) distribution is learned rather than factorized. Makes the collapse hypothesis directly testable instead of assumed. Accepts a biased gradient through the type path — assumed negligible, with a validation obligation in §11.4.
6 No charge conservation in v0.3.0 Explicitly "not yet worked out" in the meeting. No [stage2_model.conservation] block at all. Energy conservation stays exact and implicit in the stick-breaking encoding.
7 Explicit stage-prefixed CLI flags --stage2-hidden-dim etc., no generic --set path=value. Discoverable via --help and tab-completable; the cost is a flag list kept in sync with DEFAULT_CONFIG by hand.
8 AR history is a config axis, markov default Markov (previous token + remaining budget + slot index) is the baseline that makes the meeting's §6 "is attention useful" question answerable by ablation rather than by comparing differently-shaped models. Both sit behind one history_encoder(prefix) -> vector interface.

2.1 Consequence of decision 5

Straight-through Gumbel needs a critic to receive the relaxed one-hot. So the type mechanism is implied by the generator, and needs no config key of its own:

stage2_model.generator Type mechanism
"wgan" The type slice goes into the critic's input alongside energy/direction. Adversarial; learns the joint. Under particle_type.target = "onehot" it is relaxed through ST-Gumbel first; "physical" and "embedding" are already continuous and feed the critic directly.
"flow" / "ddpm" No critic exists -> the type slice trains against its own target, weighted by particle_type.lambda (same pattern as n_sec today): cross-entropy for "onehot", regression for "physical" / "embedding". Non-adversarial.

There is therefore no particle_type.adversarial key — the mechanism follows from generator × particle_type.target.

2.2 Rejected alternatives worth remembering

  • tie_to_stage1 as a tri-state (none/gate/full). Sharing experts between stages is impossible — the stage-1 trunk's input is the 9D target vector, stage 2's is a token vector of a different width. Sharing the gate is the only meaningful tying, so the key is a bool.
  • Generic --set path.to.key=value CLI overrides. Considered and rejected in favour of explicit flags (decision 7).
  • Per-expert trunk sizing (expert_hidden_dim / expert_n_blocks). Removed in v0.3.0: experts always use the stage's own hidden_dim / n_res_blocks. This was already the effective behaviour — v0.2's 0 sentinel meant "inherit" and nothing ever set it otherwise. Consequence worth knowing: a routed model is n_experts × the parameters of the monolith at equal per-row eval cost (top-1 dispatch runs one full-size trunk), so routing buys specialization, not a per-call speedup.

2.3 Correction to an earlier claim

An earlier draft of this design asserted that a one-hot conditioning mode is "mathematically identical to an nn.Embedding lookup" and should be dropped. That is wrong for the mode specified in §3.1. One-hot here is a fixed, unlearned vector of width emb_dim covering the top emb_dim - 1 species by training-set count plus an "other" bin. The difference from embedding is the vocabulary cap, not the parameterization: with 237 PDG codes and emb_dim = 16, embedding gives 237 distinct learned vectors while onehot gives 16 classes. That is a real capacity difference and a real statement about how rare species are treated, so all three modes are kept.


3. Config format reference

Six top-level blocks: [conditioning], [stage1_model], [stage2_model], [train], plus per-stage sub-tables and the existing [meta] (written by save_config, never hand-authored).

3.1 [conditioning]

One block shared by both stages. Each stage still builds its own encoder instance (separate weights) unless share_stages = true.

The particle and material axes are configured independently and may mix freely — e.g. material physical with particle embedding is a valid and intended combination.

[conditioning]
out_dim      = 128
share_stages = false

[conditioning.particle]
type     = "physical"
emb_dim  = 16
n_layers = 1

[conditioning.material]
type     = "physical"
emb_dim  = 16
n_layers = 1
Key Type Default Meaning
out_dim int 128 Width of the fused conditioning vector produced by the encoder's fusion MLP, consumed by every downstream trunk. New in v0.3.0 — v0.2 hardcoded this as cond_out_dim = 128 in every constructor signature, unreachable from config.
share_stages bool false false: stage 1 and stage 2 each construct their own ConditionEncoder with identical config but independent weights (v0.2 behaviour). true: one instance, shared by reference. Shared weights halve the conditioning parameter count and force a common representation; independent weights let each stage specialize its view of the pre-step state.

[conditioning.particle] and [conditioning.material]

Identical key sets, applied to the two identity axes independently.

Key Type Default Meaning
type "physical" | "embedding" | "onehot" "physical" How this axis's identity becomes an emb_dim-wide vector. See the table below.
emb_dim int 16 Width of this axis's vector. Under "onehot" it also sets the class count — see below.
n_layers int 1 Depth of the sub-MLP under "physical". 1 is the single-layer net with emb_dim neurons. Ignored under "embedding" / "onehot".

The three modes:

mode particle input material input width learned parameters
"physical" log(mass), charge (PARTICLE_PHYS_DIM = 2) Z_eff, A_eff, log(density), log(X0), log(λ_int) (MATERIAL_PHYS_DIM = 5) emb_dim one n_layers-deep MLP with emb_dim neurons
"embedding" dense vocab index dense vocab index emb_dim nn.Embedding(vocab, emb_dim)
"onehot" top emb_dim - 1 PDG codes by training-set count, plus one "other" bin top emb_dim - 1 materials by count, plus "other" emb_dim none — a fixed vector
  • "physical" reads columns already present in cond_cont[:, COND_DIM_BASE:] (see giant.data.transforms.build_features). It is computable for any PDG code or material, which is what allows generalization beyond the training menu.
  • "embedding" memorizes the training menu — the generalization-comparison baseline, and the only mode that supports stage2_model.particle_type.target = "embedding" (§3.3).
  • "onehot" is a fixed, unlearned representation. It is not a reparameterization of "embedding": the difference is the vocabulary cap. With 237 PDG codes and emb_dim = 16, "embedding" gives 237 distinct learned vectors while "onehot" gives 16 classes. Needs the same data-derived top-N map as the stage-2 type target (§8).

Interaction: particle.type = "physical" is incompatible with router types "pdg" and "process", which build their own dataset-scoped nn.Embedding(pdg_vocab, ...) regardless of the trunk's conditioning mode. Pairing them silently reintroduces a training-menu-scoped lookup at the routing layer, defeating the point of physical conditioning. Rejected loudly at build time — _check_router_conditioning_compat in network.py, which carries over but must now read conditioning.particle.type rather than a single global mode.

3.2 [stage1_model]

Stage 1 predicts the 9D primary post-step vector (giant/constants.py:LOCAL_TARGET_NAMES). As of decision 1 it carries no n_sec head.

[stage1_model]
active       = true
generator    = "flow"
hidden_dim   = 256
n_res_blocks = 6
dropout      = 0.0
lambda       = 1.0
Key Type Default Meaning
active bool true false skips building and training stage 1 entirely. The resulting checkpoint holds only stage 2 and cannot be rolled outgiant rollout must refuse it with a clear error. Used for the meeting's Stage-2-only architecture comparison.
generator "flow" | "ddpm" | "wgan" "flow" The generative objective. "flow": conditional flow matching (Lipman et al. 2022), ~10 ODE steps at inference. "ddpm": cosine-schedule diffusion baseline. "wgan": WGAN-GP, single forward pass at inference. Replaces the global train.mode. Objective-specific knobs live in the matching sub-table below.
hidden_dim int 256 Trunk width — also the width of every expert under a routed trunk.
n_res_blocks int 6 Number of ResBlocks in the trunk, and in every expert under a routed trunk. Was model.n_blocks; renamed for clarity since v0.3.0 also has attention layers in stage 2.
dropout float 0.0 Dropout inside each ResBlock. Default changed in v0.3.0 (was 0.1).
lambda float 1.0 Weight of this stage's loss in the total. Meaningful when both stages are active and non-adversarial; a WGAN stage's adversarial loss drives its own optimizer, so lambda scales only its non-adversarial auxiliary terms.

[stage1_model.flow] — read only when generator = "flow"

Key Type Default Meaning
time_dim int 64 Width of the SinusoidalEmbedding for the flow time variable, concatenated into the trunk's conditioning. New in v0.3.0 — v0.2 hardcoded 64.

[stage1_model.ddpm] — read only when generator = "ddpm"

Key Type Default Meaning
time_dim int 64 As above, for the diffusion time variable.
n_steps int 1000 Cosine-schedule diffusion steps. New in v0.3.0 — v0.2 hardcoded this in CosineSchedule.

[stage1_model.wgan] — read only when generator = "wgan"

Key Type Default Meaning
noise_dim int 64 Width of the generator's input noise vector. There is no time variable, hence no time_dim.
n_critic int 5 Critic updates per generator update (Gulrajani et al. 2017).
gp_weight float 10.0 Gradient-penalty coefficient.
critic_lr float 0.0 Critic learning rate. 0.0 means "inherit train.lr" — not None, since the TOML writer has no null literal to round-trip.
critic_hidden_dim int 0 Critic trunk width. 0 = inherit stage1_model.hidden_dim. New in v0.3.0 — v0.2 always sized the critic from the generator.
critic_n_res_blocks int 0 Critic trunk depth. 0 = inherit stage1_model.n_res_blocks.

[stage1_model.router]

Content carries over from v0.2's [model.router] unchanged. Reproduced here in full because the block is now per-stage and its semantics are easy to lose.

[stage1_model.router]
enabled = false
type    = "energy"
n_experts = 4
temperature = 0.5
learn_centers = true
learn_width = false
learn_temperature = false
width_min_ratio = 0.1
width_max_ratio = 10.0
lambda_balance = 0.0
lambda_entropy = 0.0
lambda_proc    = 0.0
gumbel = false
gumbel_tau_start = 1.0
gumbel_tau_end   = 0.1
emb_dim = 8
hidden_dim = 64
Key Type Default Meaning
enabled bool false Replace the monolithic trunk with a mixture of per-expert trunks: soft-mixed over all experts at train time, top-1 dispatched at eval time (each row runs exactly one expert). Every expert is hidden_dim × n_res_blocks — v0.3.0 removes the per-expert sizing keys, so a routed model costs n_experts × the monolith's parameters at equal per-row eval cost. Routing buys specialization, not a per-call speedup.
type str "energy" Router impl from ROUTER_REGISTRY. "energy": soft turn-on gate over normalized pre-step log-energy. "pdg": gate over a learned PDG embedding. "process": own classifier over pre-step conditioning predicting the step-ending physics process. "composed": joint outer-product gating over multiple axes via axis{i}_{field} keys.
n_experts int 4 Number of experts. For type = "process" this doubles as the number of process classes. Ignored for type = "composed" (each axis has its own).
temperature float 0.5 Softmax denominator for the distance-based gate. As tau -> 0 the gate hardens to nearest-center (Voronoi) selection, which is exactly what eval-time top1 uses. Energy/pdg routers only.
learn_centers bool true Whether gate centers are nn.Parameter or a fixed buffer.
learn_width bool false Give each expert its own learnable width, sigmoid-bounded to [width_min_ratio, width_max_ratio] * temperature. Mutually exclusive with learn_temperature.
learn_temperature bool false Make the single shared temperature learnable, same bounding. Mutually exclusive with learn_width.
width_min_ratio float 0.1 Lower bound multiplier for the above. Must bracket 1.0 with width_max_ratio so enabling either mode is a no-op at init.
width_max_ratio float 10.0 Upper bound multiplier. Deliberately bounded rather than softplus/exp: an unbounded width lets one expert's logit -d²/width -> 0 almost everywhere, so it wins nearly every row regardless of distance — the same "experts overlap instead of partitioning" failure the router design exists to avoid.
lambda_balance float 0.0 Importance-CV² load-balancing auxiliary loss weight (Shazeer et al. 2017). The 2026-07-22 benchmark failure ran with 0.0; do not repeat that.
lambda_entropy float 0.0 Entropy-regularization weight penalizing uniform/collapsed gating. Secondary guard against all experts' widths co-inflating together, which lambda_balance cannot see (usage shares stay even throughout that failure). Use with caution: indiscriminate entropy penalties also suppress legitimate ambiguity near a decision boundary.
lambda_proc float 0.0 Supervised process-classification CE weight. "process" router only; 0.0 still trains a working router (the gate gets gradient through the downstream loss) but only > 0 grounds it in the true process label.
gumbel bool false Straight-through Gumbel-softmax train-time combine weights: the forward pass samples a hard one-hot combination (matching eval-time top-1 dispatch exactly) while the backward pass still flows smooth gradient to every expert. Targets the train/eval mismatch.
gumbel_tau_start float 1.0 Gumbel temperature at step 0, annealed linearly over training.
gumbel_tau_end float 0.1 Gumbel temperature at the final step.
emb_dim int 8 The router's own pdg (and material) embedding width, separate from the trunk's ConditionEncoder. "pdg" / "process" routers only.
hidden_dim int 64 The "process" router's internal classifier hidden width.

Composed routers use flat axis{i}_{field} keys instead of type/n_experts — e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type = "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. Indices must be contiguous from 0. Flat keys keep the block a table of scalars, which the merge machinery relies on.

centers_init is not authored by hand: giant/pipeline.py populates it for type = "energy" from real data quantiles collected during the existing normalizer-fitting pass, then writes it into the checkpoint's model_config.

3.3 [stage2_model]

Stage 2 predicts n_sec and the per-secondary energy/direction/type.

[stage2_model]
active         = true
decoder        = "autoregressive"
generator      = "wgan"
hidden_dim     = 256
n_res_blocks   = 6
dropout        = 0.0
lambda         = 1.0
k_max          = 15
context_dim    = 64
stage1_context = "truth"
Key Type Default Meaning
active bool true false trains stage 1 alone. The checkpoint then has no secondary decoder; giant rollout must refuse it, giant predict still works.
decoder "one_shot" | "autoregressive" "autoregressive" "one_shot": predict all k_max slots simultaneously with padded slots masked from the loss — v0.2 behaviour, kept as the baseline arm of the meeting's §7 comparison. "autoregressive": emit one secondary at a time in descending-energy order.
generator "flow" | "ddpm" | "wgan" "wgan" As stage 1. Under "autoregressive" this is the objective for each token: a WGAN token costs one forward pass, a flow token costs ~10 ODE steps. See the cost note in §6.4.
hidden_dim int 256 Trunk width.
n_res_blocks int 6 Trunk depth.
dropout float 0.0 Dropout inside each ResBlock. Default changed in v0.3.0 (was 0.1).
lambda float 1.0 Weight of stage 2's loss in the total. Was train.lambda_s2.
k_max int 15 Maximum secondary slots. Under "one_shot" this is the fixed output width; under "autoregressive" it is a safety cap on the generation loop. Was the global constant K_MAX in giant/constants.py (max observed n_sec is 14 in the PbWO4 dataset, so 15 covers it with one spare).
context_dim int 64 Width of the projected stage-1 outcome fed into stage 2's conditioning. Was the hardcoded stage1_proj_dim = 64.
stage1_context "truth" | "sampled" "truth" What stage 2 conditions on during training. "truth": the ground-truth stage-1 target vector, detached — v0.2 behaviour (v0.2 train.py:254 passes x1_s1.detach()), i.e. stage-level teacher forcing. "sampled": stage 1's own sampled output, closing the train/inference gap at the cost of a sampling pass per batch and a moving target early in training.

[stage2_model.n_sec]

Key Type Default Meaning
mode "head" | "stop_token" | "truth" "head" "head": a classifier over {0..k_max} on the condition encoding alone (no diffusion noise), so it is callable independently at inference — v0.2 behaviour, and what the meeting's §3 explicitly decided to keep. "stop_token": an EOS-style implicit stop — accepted by the schema but not implemented in v0.3.0, raising a clear "not implemented" error if set (§11.2). The key exists now so landing the mechanism later is not a config break. "truth": take n_sec from ground truth — only valid for standalone stage-2 evaluation, never for rollout.
lambda float 0.1 Cross-entropy weight for the head. Was train.lambda_nsec.

[stage2_model.particle_type]

The three targets mirror the three conditioning modes of §3.1, and use the same names.

Key Type Default Meaning
target "onehot" | "physical" | "embedding" "onehot" What the token's type slice is. See the table below.
lambda float 1.0 Loss weight. Under generator = "flow"/"ddpm" this weights the cross-entropy ("onehot") or regression ("physical"/"embedding") term; under "wgan" the type is adversarial (§2.1) and this weights only any auxiliary term.
other_policy "sample" | "modal" | "drop" "sample" How a predicted "other" class becomes a concrete PDG code at rollout, needed because a secondary's mass/charge feed its own downstream conditioning. "sample": draw from the empirical within-bucket distribution recorded at map-build time. "modal": always the most common member. "drop": discard the secondary. Read only under target = "onehot". Not decided in the meeting — see §11.

There is no n_classes key. The class count under target = "onehot" is conditioning.particle.emb_dim — the same number that sizes the particle axis everywhere else. One knob sets the model's particle-type resolution, and the stage-2 onehot classes are by construction the same classes the conditioning onehot uses, so an emitted secondary's type is directly consumable as the conditioning of its own next step with no re-mapping.

Note the coupling this creates: under conditioning.particle.type = "physical", emb_dim primarily means "sub-MLP output width", yet it still sets the stage-2 class count. Intentional, but worth knowing when tuning either.

target token type slice width training target inverse map at rollout
"physical" regressed (log mass, charge) 2 the true PDG's physics values (giant.particles.particle_mass_charge) none needed — mass/charge are used directly; giant.particles.nearest_known_pdg gives a reporting-only label. v0.2 behaviour.
"onehot" class logits conditioning.particle.emb_dim true class index argmax -> class -> PDG (via other_policy for the "other" bin)
"embedding" an emb_dim-wide vector conditioning.particle.emb_dim emb.weight[class].detach() L1-nearest row of emb.weight — see below

So the type slice is conditioning.particle.emb_dim wide for both "onehot" and "embedding", and 2 only for "physical".

target = "embedding" in detail

Stage 2 emits a vector that should equal the conditioning's own particle embedding for the secondary's species — the same nn.Embedding table [conditioning.particle] builds, not a second one.

Requires conditioning.particle.type = "embedding". There is no table to match against under "physical" or "onehot"; reject at config-validation time with an explicit error.

Why detached: the regression target is emb.weight[class].detach(), so the embedding table receives gradient only through the conditioning path, never through the stage-2 output loss. Without the detach the target moves as the decoder chases it — the exact failure mode that motivated abandoning the learned type target in the first place (see decisions/physical-property-conditioning). The detach is what makes this option viable again.

Inverse map. The natural exact-match form

((out - emb.weight).abs().sum(1) < 1e-6).nonzero()

is correct as a round-trip assertion in tests (encode a known PDG, decode, recover the same PDG) but must not be used at inference: a generative model's continuous output essentially never lands within 1e-6 of a table row, so it returns an empty tensor almost always. Inference needs the nearest row:

pdg_idx = (out.unsqueeze(-2) - emb.weight).abs().sum(-1).argmin(-1)   # L1 nearest

Note this decode is unbounded in vocabulary, unlike "onehot" — every PDG code in the training vocab is reachable, and there is no "other" bucket, hence no other_policy. The trade-off is that nearest-neighbour decode has no notion of confidence: an output far from every row still snaps to something.

[stage2_model.autoregressive] — read only when decoder = "autoregressive"

Key Type Default Meaning
order "energy_desc" "energy_desc" Canonical generation order. Descending energy is the ordering already flagged as natural in the Phase-2 note's open questions, and the one the existing stick-breaking encoding assumes. Single-valued for now; the key exists so an alternative ordering is not a config break.
history "markov" | "attention" "markov" How token i+1 sees tokens ≤ i. "markov": previous token plus running scalars (remaining energy budget, slot index) — a fixed-width summary. "attention": causal self-attention over all emitted tokens. See §6.2 for the trade-off.
teacher_forcing "always" | "scheduled" | "never" "always" "always": condition on the ground-truth previous secondary throughout training (the meeting's confirmed plan). "scheduled": scheduled sampling — interpolate toward conditioning on the model's own prediction. "never": free-running from the start.
tf_p_start float 1.0 Under "scheduled", P(use ground truth) at epoch 0.
tf_p_end float 1.0 Under "scheduled", P(use ground truth) at the final epoch. Linear interpolation between the two.
attn_n_heads int 4 Attention heads. Read only under history = "attention".
attn_n_layers int 2 Causal self-attention layers. Read only under history = "attention".

[stage2_model.flow] / [stage2_model.ddpm]

Same keys as their [stage1_model.*] counterparts (time_dim; plus n_steps for ddpm).

[stage2_model.wgan] — read only when generator = "wgan"

Same keys as [stage1_model.wgan] (noise_dim, n_critic, gp_weight, critic_lr, critic_hidden_dim, critic_n_res_blocks), plus:

Key Type Default Meaning
gumbel_tau_start float 1.0 Straight-through Gumbel temperature for the particle-type one-hot at step 0, annealed linearly. Read only under particle_type.target = "onehot" (the other two targets are continuous and need no relaxation). Distinct from router.gumbel_tau_start, which anneals expert-combination weights — two unrelated Gumbel relaxations that must not share a key.
gumbel_tau_end float 0.1 Same, at the final step.

Under "autoregressive" a fresh noise_dim draw is made per token.

[stage2_model.router]

Key Type Default Meaning
tie_to_stage1 bool false true: stage 2 shares stage 1's Router module instance, so expert i in stage 1 and expert i in stage 2 gate on identical conditions by construction. Every other key in this block is then ignored. false: an independent router — note this is v0.2's actual behaviour, which built two separate routers from one config, so stage-1 expert i and stage-2 expert i had no semantic relationship despite identical hyperparameters.

All other keys are as [stage1_model.router]. Invalid when stage1_model.active = false and tie_to_stage1 = true — reject at config-validation time.

3.4 [train]

Optimizer, schedule, data split and logging only. Everything model-shaped moved into the stage blocks.

[train]
epochs = 100
batch_size = 4096
lr = 3e-4
weight_decay = 0.01
ema_decay = 0.9999
warmup_epochs = 5
val_fraction = 0.1
max_val_batches = 200
num_workers = 4
seed = 0
validate_every = 10
validate_steps = 10
wandb = true
wandb_project = "giant"
wandb_run_name = ""
wandb_log_every = 50
Key Type Default Meaning
epochs int 100 Training epochs.
batch_size int 4096 Steps per batch. auto on the CLI estimates from free VRAM.
lr float 3e-4 AdamW learning rate for every stage's generator.
weight_decay float 0.01 AdamW weight decay.
ema_decay float 0.9999 EMA of model weights used for sampling; 0 disables. Maintained per stage.
warmup_epochs int 5 Linear LR warmup.
val_fraction float 0.1 Fraction of events (not steps) held out — the split is by event_id to avoid leaking correlated steps from the same shower.
max_val_batches int 200 Cap on the per-epoch val-loss pass; 0 = full val set. Distinct from validate_every's marginal/KL pass.
num_workers int 4 DataLoader workers. Warns above ~1/4 of the machine's CPUs — portal machines are shared.
seed int 0 Seeds Python/numpy/torch and the event split.
validate_every int 10 Epochs between full marginal/KL validation passes.
validate_steps int 10 Sampler steps used during those passes.
wandb bool true W&B per-epoch metric logging. Default changed in v0.3.0 (was opt-in false): the v0.3.0 work is a sequence of architecture comparisons, and a run that was not logged is not comparable. Set false for throwaway/debug runs.
wandb_project str "giant" W&B project.
wandb_run_name str "" "" means "use the checkpoint out_dir name" — not None, since the TOML writer has no null literal.
wandb_log_every int 50 Optimizer steps between batch-granularity metric logs. A single epoch can be tens of thousands of steps; per-epoch metrics always log in full.

3.5 Removed from v0.2

Key Fate
train.mode Split into stage1_model.generator / stage2_model.generator.
train.lambda_nsec -> stage2_model.n_sec.lambda.
train.lambda_s2 -> stage2_model.lambda.
train.n_critic, gp_weight, critic_lr -> stage{1,2}_model.wgan.*.
[model] (whole block) Split across [conditioning] and the two stage blocks.

4. Migration: v0.2 -> v0.3

[meta] config_version = 3 tags the new format; absent means v0.2. migrate_config(cfg) -> cfg applies the table below and is called on both config.toml load and checkpoint model_config load, so nothing on /ceph goes dead (decision 3).

v0.2 v0.3
train.mode stage1_model.generator and stage2_model.generator (same value)
train.lambda_nsec stage2_model.n_sec.lambda
train.lambda_s2 stage2_model.lambda
train.n_critic / gp_weight / critic_lr stage1_model.wgan.* and stage2_model.wgan.*
model.hidden_dim / n_blocks / dropout stage{1,2}_model.hidden_dim / n_res_blocks / dropout
model.emb_dim conditioning.particle.emb_dim and conditioning.material.emb_dim (v0.2 had one shared value)
model.conditioning conditioning.particle.type and conditioning.material.type (v0.2 had one shared mode)
model.noise_dim stage{1,2}_model.wgan.noise_dim
model.router.* stage1_model.router.*, copied verbatim to stage2_model.router.* with tie_to_stage1 = false (preserves v0.2's two-independent-routers behaviour)
model.expert_hidden_dim / expert_n_blocks dropped. v0.2's 0 sentinel meant "inherit from the monolith", which is now unconditional. A v0.2 config with a non-zero value must fail loudly rather than silently resize the experts — see §4.3.
model.k_max stage2_model.k_max
conditioning.out_dim = 128 (v0.2's hardcoded value)
conditioning.{particle,material}.n_layers = 2 (v0.2's hardcoded depth; note the v0.3 default is 1)
stage{1,2}_model.flow.time_dim = 64 / .ddpm.time_dim = 64 (v0.2's hardcoded value)
stage2_model.context_dim = 64 (v0.2's hardcoded stage1_proj_dim)
stage2_model.decoder = "one_shot" (v0.2 had no other option)
stage2_model.particle_type.target = "physical" (v0.2 behaviour)
stage{1,2}_model.active = true

4.1 The awkward one: n_sec head ownership

Decision 1 moves the n_sec head to stage 2, but v0.2 checkpoints carry n_sec_head weights inside the stage-1 module (DenoisingMLP.n_sec_head, WGANGenerator.n_sec_head, RoutedDenoisingMLP.n_sec_head). The shim must keep those loading where they are.

Handling: migrate_config sets an internal stage2_model.n_sec.legacy_owner = "stage1" that build_models honours by attaching the head to the stage-1 module. Never written by new runs, never CLI-settable, never documented as a user-facing option.

4.2 Non-inheriting expert dims

v0.3.0 drops expert_hidden_dim / expert_n_blocks (§2.2). Migration must distinguish two cases:

  • value is 0 (the "inherit" sentinel, and what every real run used) — drop the key silently, behaviour is unchanged.
  • value is non-zero — fail loudly. Silently resizing those experts to hidden_dim would change the architecture, so the checkpoint's weights would no longer match. Such a checkpoint can only be loaded by v0.2.

4.3 Migration test

The acceptance criterion for the whole shim: load a v0.2 checkpoint through migrate_config + the new build_models, and diff its outputs against v0.2 code on the same input batch. Bit-identical, or the refactor has changed something it should not have. Pick one flow checkpoint and one WGAN checkpoint from /ceph.


5. network.py refactor

5.1 What it looks like today

Ten classes that are permutations of three independent choices:

flow/ddpm wgan generator wgan critic
stage 1 DenoisingMLP WGANGenerator Critic
stage 1, routed RoutedDenoisingMLP
stage 2 SecondaryDecoder WGANSecondaryGenerator SecondaryCritic
stage 2, routed RoutedSecondaryDecoder

The empty cells are the entire reason giant/pipeline.py:275 hard-rejects --mode wgan --router: no routed WGAN generator class was ever written. There is no deeper reason — the routed trunk is orthogonal to the objective.

Every one of those classes repeats the same body: build a condition encoder, optionally a time embedding, project input, run blocks, project output.

5.2 Proposed decomposition — one axis per config block

(a) [conditioning] -> encoders

ConditionEncoder(type, emb_dim, n_layers, out_dim)   # behaviour unchanged, now configurable
ContextAdapter(in_dim, context_dim)                  # stage-1 outcome -> context vector

SecondaryConditionEncoder disappears as a class. It was ConditionEncoder + a stage1_proj linear + a fuse layer; those compose at the stage-model level instead.

(b) [stage*_model] + .router -> trunks, behind one interface

class Trunk(nn.Module):
    def forward(self, x, cond, cond_cont=None, cond_cat=None) -> Tensor: ...

MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)

build_trunk(stage_cfg, in_dim, out_dim, cond_dim) -> Trunk

ExpertTrunk and _route_forward carry over unchanged. One required change: MonolithicTrunk and ExpertTrunk need a separate out_dim — today ExpertTrunk hardcodes out_proj = nn.Linear(hidden_dim, in_dim), i.e. out_dim == in_dim. That stops working the moment stage 2's per-token output is 4 + type_dim wide while its input is a noise vector of width noise_dim.

(c) [stage*_model].generator -> a thin wrapper, not a class family

The generator choice controls exactly two things:

  • whether SinusoidalEmbedding(t) is concatenated into cond (flow/ddpm) or not (wgan)
  • whether the trunk's x is the diffused/interpolated x_t (flow/ddpm) or a noise draw z (wgan)

That is small enough for one wrapper, collapsing six of today's ten classes:

class GenerativeTrunk(nn.Module):
    """cond-encode -> (optional time-embed) -> trunk. Backs flow, ddpm and wgan."""
    def forward(self, x, cond_cont, cond_cat, t=None, context=None) -> Tensor: ...

5.3 Resulting class list

# building blocks
SinusoidalEmbedding, ResBlock, ConditionEncoder, ContextAdapter

# trunks
Trunk (ABC), MonolithicTrunk, RoutedTrunk, ExpertTrunk

# routers — carried over unchanged
Router, EnergyRouter, PdgRouter, ProcessRouter, ComposedRouter
ROUTER_REGISTRY, register_router, build_router, build_composed_router

# history encoders — new, stage-2 AR only
HistoryEncoder (ABC), MarkovHistory, AttentionHistory

# stage models
Stage1Model            # 9D primary step
Stage2OneShot          # k_max slots at once (v0.2 behaviour)
Stage2Autoregressive   # one token at a time
CriticModel            # stage-1 or stage-2 critic, generator-agnostic

# factories
build_models(cfg)  -> {"stage1": ... | None, "stage2": ... | None}
build_critics(cfg) -> {"stage1": ... | None, "stage2": ... | None}

Ten classes become four stage classes plus reusable parts, and routed WGAN comes for freepipeline.py's rejection can be deleted.

5.4 Factory signature change

build_models returns a dict, not a tuple: active = false on either stage means that key is None. Every caller unpacking stage1, sec_decoder = build_models(...) must be updated (pipeline.py:390, cli.py:871, cli.py:1254).


6. Stage-2 autoregressive design

6.1 Token layout

Secondaries are emitted one at a time in descending energy order. Per-token output width is 4 + type_dim:

slice meaning
[0] stick-breaking logit — fraction of the remaining energy budget
[1:4] local-frame direction, normalized to a unit vector
[4:4+type_dim] particle type — see below

type_dim follows particle_type.target (§3.3): 2 for "physical", and conditioning.particle.emb_dim for both "onehot" and "embedding". Only "onehot" involves a relaxation — ST-Gumbel under wgan, cross-entropy under flow/ddpm; the other two are continuous and feed the critic (or the regression loss) directly.

Per-token conditioning is:

base condition encoding          (ConditionEncoder, [conditioning].out_dim)
+ stage-1 context                (ContextAdapter, stage2_model.context_dim)
+ history_encoder(prefix)        (§6.2)
+ running scalars                (remaining energy budget, slot index)

6.2 History encoders

One interface, history_encoder(prefix) -> fixed-width vector:

  • MarkovHistory — the previous token's (energy_fraction, direction, type_embedding). Fixed-width, one small MLP.
  • AttentionHistory — causal self-attention over all emitted tokens, taking the last position. attn_n_heads × attn_n_layers.

The trade-off (this is decision 8's reasoning, recorded so it does not have to be re-derived):

  1. Markov is less impoverished than it sounds. The remaining-energy budget is an exact sufficient statistic for the conservation constraint — stick-breaking needs nothing else from history. Slot index likewise. What markov genuinely cannot see is set composition: "I have already emitted two photons and an electron." That matters for correlated production — pair production emits exactly e⁺e⁻, a brems cascade correlates species across the set. Note that the meeting's charge-conservation idea was exactly such a hand-engineered summary statistic, and it is out of scope for v0.3.0 (decision 6), so markov does not get that crutch.
  2. Sequences are short and front-loaded. n_sec is 02 for most steps, max 14. For K ≤ 2, attention is markov. They diverge only in the high-multiplicity tail — physically interesting but data-poor, so the attention path trains on few examples where it actually matters.
  3. Cost is not where intuition puts it. Under teacher forcing both train in a single parallel pass over all K tokens (every token's input is ground truth, so nothing is sequential). At inference both need K sequential forwards; attention additionally needs a KV cache to avoid re-encoding the prefix. Attention's marginal FLOPs over ≤15 tokens are rounding error.
  4. Exposure bias cuts against attention. Attention conditions on the entire generated prefix, so one off-manifold early token poisons every later token through the context. Markov's fixed summary sees only one bad token, and the budget scalars stay exact regardless. Given the explicitly-flagged train/ inference gap and the known compounding-rollout-error problem, the more expressive history is also the more fragile one.

Hence: markov is the default and the baseline; attention is a flag.

6.3 Energy budget under AR

No re-derivation needed, despite the meeting's action item suggesting otherwise. The existing stick-breaking is already sequential in spirit — each slot takes a fraction of what remains — so it carries over to per-token generation directly by feeding "remaining budget" as a per-token conditioning scalar. Conservation stays exact by construction: the valid slots' energies sum to e_sec, which the Stage-1 simplex already guarantees sums correctly with edep and post_E.

6.4 Cost warning

AR costs K sequential forward passes per step where one-shot costs one. Against the ~10× native-Geant4 eval budget that motivated the whole fast-eval track, this is the number to watch — not the history-encoder choice. With generator = "flow" it is worse still: ~10 ODE steps per token, so ~150 forwards per step in the worst case. generator = "wgan" (one pass per token) is the only configuration that plausibly meets the budget; flow AR is for quality comparison.


7. Training loop

Decision 2 (full mixed per-stage objectives) makes giant/training/ one trainer object per active stage:

class StageTrainer:                 # owns optimizers, EMA, update cadence
    def step(self, batch, global_step) -> dict[str, float]: ...

FlowTrainer, DDPMTrainer, WGANTrainer(critic, n_critic, gp_weight)
  • Non-adversarial stages contribute lambda * loss to one backward pass.
  • A WGAN stage runs its own critic inner loop on the same batch, with a generator update every n_critic-th batch — today's _wgan_train_step cadence.
  • A mixed run (flow + wgan) steps stage 1 every batch while stage 2 does 5 critic updates then a generator update. Independent optimizers, independent EMA.
  • Router auxiliary losses (lambda_balance / lambda_proc / lambda_entropy) become per-stage, summed over whichever stages are routed. Today's hasattr(model, "router") check (v0.2 train.py:262) generalizes cleanly.
  • metrics.csv and W&B metric names gain a stage prefix. Implemented as <stage>/<split>/<metric> (stage1/train/loss, stage2/train/d_loss, stage1/lr, stage1/router/entropy) plus an unprefixed run-level tail (val/loss, val/marginal_kl, grad_norm, gpu_mem_mb, samples_per_sec, is_best, epoch_time_s). Each name is declared once, as a MetricSpec on the StageTrainer that computes it; the CSV header and W&B payload are derived from those declarations (giant/training/metrics.py).
  • With stage1_model.active = false, stage 2 still needs its stage-1 context: it comes from the ground-truth target already in the batch (x1_s1), which is exactly what v0.2 does anyway. Stage-2-only training is therefore a cheap ablation, not new plumbing — drop the stage-1 loss, skip building stage 1.

8. Data and setup-cache changes

Three config options need a "top N1 by training-set count plus other" map, but they resolve to at most two distinct maps per run — one per axis — because the class count always comes from that axis's emb_dim:

consumer axis N
conditioning.particle.type = "onehot" PDG conditioning.particle.emb_dim
stage2_model.particle_type.target = "onehot" PDG conditioning.particle.emb_dim
conditioning.material.type = "onehot" material conditioning.material.emb_dim

The two PDG consumers therefore share one map — which is the point of dropping n_classes: a secondary's emitted type is directly consumable as the conditioning of its own next step, with no re-mapping between two class systems. Build one shared helper, structurally identical to today's proc_map:

  • build_topn_map_from_files(files, column, n_classes=...) in giant/data/loader.py, next to build_process_map_from_files
  • a setup_cache section keyed by (axis, N), same shape as cache.proc_maps (which is keyed by n_experts). N is still part of the key so the sidecar stays reusable across runs with different emb_dim, even though a single run only ever needs one N per axis.
  • persisted into the checkpoint beside pdg_map / mat_map
  • inverted at rollout to recover a concrete PDG -> mass/charge per secondary

Record the empirical within-bucket distribution at map-build time as well — other_policy = "sample" needs it.

particle_type.target = "embedding" needs no map: it reaches the full training vocab through the conditioning's embedding table (§3.3).

giant/constants.py: K_MAX and SEC_DIM stop being authoritative constants (they become stage2_model.k_max and a derived quantity). Keep them as defaults only, and audit the ~12 modules importing K_MAX for places that assume it is global truth.


9. Config machinery changes

All in giant/config.py:

  • merge_cli_overrides — replace the hand-written one-level router merge with a generic recursive deep-merge. The new layout is three levels deep (stage1_model.router.axis0_type).
  • save_config — recursive TOML writer; today it handles exactly one nesting level (see its nested_sections list).
  • default_out_dir_name_OUT_DIR_NAME_CANDIDATES entries become dotted paths ("stage2_model.decoder") instead of (section, field) pairs. Add candidates for the new discriminating fields: decoder, per-stage generator, particle_type.target, autoregressive.history.
  • resolve_expert_dimsdeleted. Experts always take the stage's hidden_dim / n_res_blocks (§2.2). Its two callers (pipeline.py:351 and the CLI's batch-size auto-estimate) read the stage keys directly, as does pipeline.py's "experts are NxM, different from model.hidden_dim" warning, which becomes unreachable and should go.
  • migrate_config(cfg) -> cfg — §4's table, applied on both config.toml load and checkpoint model_config load. New [meta] config_version = 3.
  • Conditioning enum — gains a third member onehot, and now feeds two keys (conditioning.particle.type, conditioning.material.type) rather than one. Shared with scripts/dwarf.py, so both CLIs stay in sync.
  • Cross-block validation — a new validate_config(cfg) pass, since v0.3.0 has constraints no single block can check: particle_type.target = "embedding" requires conditioning.particle.type = "embedding"; router types "pdg"/"process" require conditioning.particle.type != "physical"; stage2_model.router.tie_to_stage1 requires stage1_model.active; n_sec.mode = "truth" is invalid for a rollout-capable checkpoint.
  • estimate_batch_size — its calibration constants assume the v0.2 architecture ("post-Phase-2, including the Stage-2 secondary decoder and n_sec head"). AR stage 2 changes the activation-memory profile. Recalibrate last (§12 step 8), measured on real hardware with the example configs.

10. Callers that need updating

File Why
giant/pipeline.py Builds model_config; now per-stage. Delete the wgan+router rejection at :275. Router centers_init seeding becomes per-stage.
giant/training/ Per-stage trainers (§7), metric collection, checkpointing, the epoch loop.
giant/cli.py Stage-prefixed flags for train and new-run; build_models now returns a dict (:871, :1254); :1339 writes model_config into the rollout sidecar.
giant/sample.py Sampler picked per stage from stage*_model.generator; new AR sampling loop with KV cache under history = "attention".
giant/rollout.py AR secondary generation; categorical class -> PDG decode; other_policy handling.
giant/validate.py Stage-2 marginals gain a type-class marginal.
giant/analysis/render.py _router_summary(model_config) at :44 reads model_config["router"].
giant/analysis/router_gating.py Same, at :76.
giant/constants.py K_MAX / SEC_DIM demoted to defaults (§8).
configs/*.toml All eight shipped configs are v0.2-format; regenerate or rely on the shim.
condor-gpu-train-rollout branch Submits giant train flags; needs rebasing onto the new flag surface.

11. Settled scope and open questions

11.1 Settled — build as specified

  • other_policy switch. The three-way "sample" / "modal" / "drop" design in §3.3 is approved as the config surface for turning a predicted "other" class into a concrete PDG at rollout. "sample" (draw from the empirical within-bucket distribution recorded at map-build time) stays the default as the least-biased option.
  • [stage*_model.flow] and [stage*_model.ddpm] stay separate tables. The duplicated time_dim is accepted; a merged [stage*_model.diffusion] was considered and rejected because the name fits flow matching poorly.

11.2 Deferred — not implemented in v0.3.0

  • n_sec.mode = "stop_token". The value is accepted by the schema but raises a clear "not implemented in v0.3.0" error if set. The key existing now means landing the implicit-stop mechanism later is not a config break. "head" is what v0.3.0 builds, per the meeting's §3 decision.
  • Charge conservation. No key at all — [stage2_model.conservation] is absent from v0.3.0 (decision 6), unlike stop_token above. Deliberately left undesigned: the mechanism should be worked out on its own terms when it is taken up, not pre-shaped by choices made for this refactor. Adding the block later is a config addition, not a break.
  • stage2_model.generator = "ddpm". The value is accepted by the schema (§3.3 lists "flow" | "ddpm" | "wgan" with no caveat) but FlowDDPMStageTrainer.__init__ (giant/training/trainers.py) raises NotImplementedError for stage 2 — only "flow" and "wgan" have a stage-2 secondary-decoder loss implemented. stage1_model.generator = "ddpm" is unaffected; this restriction is stage-2-only. Landing stage-2 ddpm later is a trainer addition, not a config break.

11.3 Tracked as implementation work

  • Log the L1 distance distribution at rollout under particle_type.target = "embedding". Nearest-neighbour decode has no reject option — an output far from every table row still snaps to its nearest neighbour, with no "other" bin and no confidence signal. A heavy tail in that distribution means the decoder is emitting vectors off the embedding manifold, which is the direct analogue of the species-collapse symptom this redesign exists to fix. Belongs with the giant/rollout.py decode work (step 6) and should surface as a giant analyze diagnostic plot alongside router_gating.
  • estimate_batch_size recalibration for AR stage 2 is the last step of the implementation flow (§12 step 8), measured on real hardware using the example configs — not guessed from the existing calibration constants.

11.4 Accepted with a validation obligation: differentiability

Position taken: differentiability through the categorical type path is broken, and that is accepted. The expected contribution of the broken path to the total gradient is small enough to ignore. This is an assumption, not a result — it has to be demonstrated later.

Recording it precisely, since "broken" covers three distinct things:

where status
Per-token training loss under teacher forcing Fine. Softmax cross-entropy needs no sampling; the loss is differentiable in the logits.
ST-Gumbel into the critic (generator = "wgan", §2.1) Biased, not absent. The forward pass is a hard one-hot; the backward pass pretends it was the soft sample. Gradient flows, but it is not the gradient of what was actually computed.
Full shower-rollout backprop Structurally broken regardless. Already non-differentiable once secondaries spawn branches, independent of the type representation — so the categorical switch costs nothing that was not already lost.

The claim being accepted is about the middle row: the straight-through estimator's bias, propagated into the shared trunk, is expected to be negligible against the gradient from the continuous paths (stick-breaking energy, direction, and — when stage 1 is active — the 9D primary target). The third row is the reason this is tolerable at all: end-to-end differentiability was never available.

Validation obligation. Do not treat this as settled until one of the following has actually been run. Cheapest first:

  1. Gradient-magnitude accounting. Instrument a training run to log the norm of the trunk gradient contributed through the type slice against the norm from the continuous slices. "Negligible" should mean a stable, small ratio — not merely small at initialization. This is the direct measurement of the claim and costs almost nothing to add.
  2. Detached-type ablation. Train with the type path detached from the shared trunk entirely (type head still learns; no type gradient reaches the trunk) against the ST-Gumbel default. Comparable species marginals and kinematics mean the coupling was weak, which is the same conclusion by a different route.
  3. Estimator swap, only if 12 are inconclusive: compare ST-Gumbel against an unbiased-but-high-variance estimator (e.g. REINFORCE with a baseline) on a short run. Agreement in the learned marginals means the bias did not matter.

Option 1 should be added when the AR trunk lands (step 5), so the evidence accrues during the architecture comparison rather than needing a dedicated run afterwards.

Why it still matters that this is written down: decision 5 (adversarial type via ST-Gumbel) rests on this assumption. If the ratio in test 1 turns out not to be small, the fallback is not a redesign — it is particle_type.target = "physical" or a non-adversarial CE head, both of which already exist as config options.


12. Implementation order

  1. config.py — new DEFAULT_CONFIG, recursive merge/write, migrate_config, tests. Nothing else can land first.
  2. network.py — the §5 decomposition, with Stage2OneShot reproducing v0.2 exactly. Gate on the §4.3 migration test: load a v0.2 checkpoint through the shim and diff outputs against v0.2 code.
  3. giant/training/ — per-stage trainers; active = false paths. At this point Stage-2-only training works and the meeting's step 2 (one-shot WGAN baseline, trained standalone) is runnable.
  4. Type maploader.py + setup_cache.py + checkpoint persistence + particle_type.target = "onehot" in Stage2OneShot. This is the meeting's action item 1, and it is testable against the one-shot baseline before any AR work.
  5. Stage2Autoregressive with history = "markov", teacher_forcing = "always". The meeting's step 3. Add the gradient-magnitude instrumentation of §11.4 test 1 here, so the differentiability assumption accrues evidence during the architecture comparison instead of needing its own run later.
  6. sample.py / rollout.py — AR generation and class -> PDG decode, so an AR model can actually be rolled out and put through giant analyze. Includes the L1-distance diagnostic of §11.3.
  7. history = "attention", scheduled sampling — then run the meeting's §7 comparison (one-shot vs autoregressive, standalone) and only chain the winner behind Stage 1.
  8. estimate_batch_size recalibration — last, once the architectures are settled. Measure on real hardware with the example configs and record the new calibration points in config.py (§9, §11.3).

Steps 13 are pure refactor with a bit-identical acceptance criterion. Steps 47 are the actual physics change. Step 8 is measurement, deliberately last: the activation-memory profile is not knowable until the AR trunk and its history encoder are final.