giant/docs/v0.3.0-design.md

# 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. [Settled scope and open questions](#11-settled-scope-and-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.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 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.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.

```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.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 (`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

```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 = 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**

```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. **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/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. 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.

### 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 Still open

- **Differentiability sanity check with Jan.** 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". Decision 5 rests on this reasoning, so it is worth raising before step 4
  rather than after.

---

## 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`. 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 1–3 are pure refactor with a bit-identical acceptance criterion. Steps 4–7
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.