lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Add v0.3.0 design doc: Stage-2 autoregressive redesign
CI / Format (ruff format) (push) Successful in 27s
Details
CI / Lint (ruff check) (push) Successful in 29s
Details
CI / Sync project version with tag (push) Has been skipped
Details
CI / Type check (ty) (push) Successful in 28s
Details
CI / Tests (push) Successful in 56s
Details

Browse Source 
Design contract for the v0.3.0 config break and network.py refactor,
following the 2026-08-04 meeting with Jan. Not implemented yet.

The 2026-08-03 WGAN rollout benchmark failed specifically at the
secondary-species level (zero photons, ~4M hallucinated -14 muon
antineutrinos). The response pivots Stage 2 to autoregressive generation
in descending-energy order with teacher forcing, and reverts the particle
type to a categorical representation.

That needs a config break: [conditioning] / [stage1_model] /
[stage2_model] / [train] replace the single global train.mode and
[model] block, so per-stage generators (stage1 flow + stage2 wgan),
stage-2-only training, and one-shot-vs-autoregressive comparison all
become expressible. The particle and material conditioning axes are
configured independently and mix freely, each with physical / embedding /
onehot modes; the stage-2 type target mirrors the same three names, with
conditioning.particle.emb_dim sizing both so the two share one class map.

network.py collapses from ten permutation classes (stage x objective x
routed) into composable parts — encoder x trunk x objective — which also
makes routed WGAN work for the first time; it was only ever rejected
because no routed WGAN generator class existed.

The doc specifies every config option, the v0.2 migration (shim for both
configs and checkpoints, gated on a bit-identical output diff), the
refactor, and the implementation order.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-08-04 15:25:32 +02:00
parent 2fe887b49a
commit f390884f67
2 changed files with 934 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+2
View File
@@ -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.
docs/v0.3.0-design.md
+932
View File
@@ -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.