diff --git a/CLAUDE.md b/CLAUDE.md index ccab334..c3cc181 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -91,4 +91,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`). +**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. **Full design contract, with every config option documented: `docs/v0.3.0-design.md` — read it before touching `giant/config.py` or `giant/model/network.py`.** + **Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet. diff --git a/docs/v0.3.0-design.md b/docs/v0.3.0-design.md new file mode 100644 index 0000000..5b4408d --- /dev/null +++ b/docs/v0.3.0-design.md @@ -0,0 +1,932 @@ +# 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](#1-why) +2. [Decisions register](#2-decisions-register) +3. [Config format reference](#3-config-format-reference) +4. [Migration: v0.2 -> v0.3](#4-migration-v02---v03) +5. [network.py refactor](#5-networkpy-refactor) +6. [Stage-2 autoregressive design](#6-stage-2-autoregressive-design) +7. [Training loop](#7-training-loop) +8. [Data and setup-cache changes](#8-data-and-setup-cache-changes) +9. [Config machinery changes](#9-config-machinery-changes) +10. [Callers that need updating](#10-callers-that-need-updating) +11. [Open questions](#11-open-questions) +12. [Implementation order](#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. | +| 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. + +```toml +[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. + +```toml +[stage1_model] +active = true +generator = "flow" +hidden_dim = 256 +n_res_blocks = 6 +dropout = 0.1 +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 out** — `giant 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 `ResBlock`s 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.1` | Dropout inside each `ResBlock`. | +| `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. + +```toml +[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. + +```toml +[stage2_model] +active = true +decoder = "autoregressive" +generator = "wgan" +hidden_dim = 256 +n_res_blocks = 6 +dropout = 0.1 +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.1` | Dropout inside each `ResBlock`. | +| `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 (`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, documented in §3 as a later possibility, not a decision. `"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 + +```python +((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: + +```python +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. + +```toml +[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 = false +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 | `false` | Opt-in W&B logging. | +| `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** + +```python +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** + +```python +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: + +```python +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 free** — `pipeline.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 0–2 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 `train.py` one trainer object +per active stage: + +```python +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 (`train.py:262`) generalizes cleanly. +- `metrics.csv` and W&B metric names gain a stage prefix. +- 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 N−1 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_dims`** — **deleted.** 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; re-measure and note + the new calibration point. + +--- + +## 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/train.py` | Per-stage trainers (§7). | +| `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. Open questions + +1. **"Other" bucket at rollout.** A predicted "other" class has no concrete PDG, + so no mass/charge for that secondary's own downstream conditioning. + `other_policy` is the config surface; the mechanism is a judgement call + (`"sample"` from the empirical within-bucket distribution seems least biased). + **Not decided in the meeting.** +2. **Charge conservation.** Out of scope for v0.3.0 (decision 6), but worth + recording *why* it interacts with decision 5: a categorical head makes a hard + mechanism tractable — mask classes whose charge cannot fit the remaining charge + budget before the softmax, analogous to the energy simplex. This is **only** + possible with a categorical type; the continuous `(mass, charge)` target admits + no such mask. That is a second, independent argument for the §5 switch beyond + the species-collapse one. +3. **Does `"stop_token"` deserve building in v0.3.0?** The meeting kept the `n_sec` + head and recorded the implicit stop as a later possibility. The config key + exists; the implementation can be deferred. +4. **`estimate_batch_size` calibration** for AR stage 2 (§9). +5. **Nearest-neighbour decode has no reject option.** Under + `particle_type.target = "embedding"`, an output far from every table row still + snaps to its nearest neighbour — there is no equivalent of `"onehot"`'s "other" + bin, and no confidence signal. Worth logging the L1 distance distribution at + rollout: a heavy tail means the decoder is emitting vectors off the embedding + manifold, which would be the direct analogue of the species-collapse symptom + this redesign is chasing. +6. **Should `flow` and `ddpm` share one sub-table?** Both need `time_dim`; only + ddpm needs `n_steps`. Currently duplicated as `[stage*_model.flow]` and + `[stage*_model.ddpm]`. A single `[stage*_model.diffusion]` would avoid the + duplication at the cost of a name that fits flow matching poorly. +5. **Differentiability sanity check.** The meeting's §5 notes that the original + argument for the continuous type target rested on avoiding a non-differentiable + categorical *sampling* step — but per-token training loss is differentiable + either way under teacher forcing, and full shower-rollout backprop is already + structurally non-differentiable once secondaries spawn branches. The note flags + this as "worth confirming explicitly with Jan rather than assuming it". + +--- + +## 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. **`train.py`** — 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 map** — `loader.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. +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`. +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. + +Steps 1–3 are pure refactor with a bit-identical acceptance criterion. Steps 4–7 +are the actual physics change.