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

---

1. Why

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

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

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

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

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

---

2. Decisions register

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

2.1 Consequence of decision 5

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

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

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

2.2 Rejected alternatives worth remembering

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

2.3 Correction to an earlier claim

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

---

3. Config format reference

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

3.1 [conditioning]

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

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

[conditioning]
out_dim      = 128
share_stages = false

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

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

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

[conditioning.particle] and [conditioning.material]

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

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

The three modes:

mode	particle input	material input	width	learned parameters
"physical"	log(mass), charge (PARTICLE_PHYS_DIM = 2)	Z_eff, A_eff, log(density), log(X0), log(λ_int) (MATERIAL_PHYS_DIM = 5)	emb_dim	one n_layers-deep MLP with emb_dim neurons
"embedding"	dense vocab index	dense vocab index	emb_dim	nn.Embedding(vocab, emb_dim)
"onehot"	top emb_dim - 1 PDG codes by training-set count, plus one "other" bin	top emb_dim - 1 materials by count, plus "other"	emb_dim	none — a fixed vector

• "physical" reads columns already present in cond_cont[:, COND_DIM_BASE:] (see giant.data.transforms.build_features). It is computable for any PDG code or material, which is what allows generalization beyond the training menu.
• "embedding" memorizes the training menu — the generalization-comparison baseline, and the only mode that supports stage2_model.particle_type.target = "embedding" (§3.3).
• "onehot" is a fixed, unlearned representation. It is not a reparameterization of "embedding": the difference is the vocabulary cap. With 237 PDG codes and emb_dim = 16, "embedding" gives 237 distinct learned vectors while "onehot" gives 16 classes. Needs the same data-derived top-N map as the stage-2 type target (§8).

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

3.2 [stage1_model]

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

[stage1_model]
active       = true
generator    = "flow"
hidden_dim   = 256
n_res_blocks = 6
dropout      = 0.0
lambda       = 1.0

Key	Type	Default	Meaning
active	bool	true	false skips building and training stage 1 entirely. The resulting checkpoint holds only stage 2 and cannot be rolled 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 ResBlocks in the trunk, and in every expert under a routed trunk. Was model.n_blocks; renamed for clarity since v0.3.0 also has attention layers in stage 2.
dropout	float	0.0	Dropout inside each ResBlock. Default changed in v0.3.0 (was 0.1).
lambda	float	1.0	Weight of this stage's loss in the total. Meaningful when both stages are active and non-adversarial; a WGAN stage's adversarial loss drives its own optimizer, so lambda scales only its non-adversarial auxiliary terms.

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

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

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

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

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

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

[stage1_model.router]

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

[stage1_model.router]
enabled = false
type    = "energy"
n_experts = 4
temperature = 0.5
learn_centers = true
learn_width = false
learn_temperature = false
width_min_ratio = 0.1
width_max_ratio = 10.0
lambda_balance = 0.0
lambda_entropy = 0.0
lambda_proc    = 0.0
gumbel = false
gumbel_tau_start = 1.0
gumbel_tau_end   = 0.1
emb_dim = 8
hidden_dim = 64

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

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

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

3.3 [stage2_model]

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

[stage2_model]
active         = true
decoder        = "autoregressive"
generator      = "wgan"
hidden_dim     = 256
n_res_blocks   = 6
dropout        = 0.0
lambda         = 1.0
k_max          = 15
context_dim    = 64
stage1_context = "truth"

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

[stage2_model.n_sec]

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

[stage2_model.particle_type]

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

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

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

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

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

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

target = "embedding" in detail

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

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

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

Inverse map. The natural exact-match form

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

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

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

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

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

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

[stage2_model.flow] / [stage2_model.ddpm]

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

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

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

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

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

[stage2_model.router]

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

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

3.4 [train]

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

[train]
epochs = 100
batch_size = 4096
lr = 3e-4
weight_decay = 0.01
ema_decay = 0.9999
warmup_epochs = 5
val_fraction = 0.1
max_val_batches = 200
num_workers = 4
seed = 0
validate_every = 10
validate_steps = 10
wandb = true
wandb_project = "giant"
wandb_run_name = ""
wandb_log_every = 50

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

3.5 Removed from v0.2

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

---

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

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

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

4.1 The awkward one: n_sec head ownership

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

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

4.2 Non-inheriting expert dims

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

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

4.3 Migration test

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

---

5. network.py refactor

5.1 What it looks like today

Ten classes that are permutations of three independent choices:

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

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

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

5.2 Proposed decomposition — one axis per config block

(a) [conditioning] -> encoders

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

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

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

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

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

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

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

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

The generator choice controls exactly two things:

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

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

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

5.3 Resulting class list

# building blocks
SinusoidalEmbedding, ResBlock, ConditionEncoder, ContextAdapter

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

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

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

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

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

Ten classes become four stage classes plus reusable parts, and routed WGAN comes for 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:

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.
• stage2_model.generator = "ddpm". The value is accepted by the schema (§3.3 lists "flow" | "ddpm" | "wgan" with no caveat) but FlowDDPMStageTrainer.__init__ (giant/train.py) raises NotImplementedError for stage 2 — only "flow" and "wgan" have a stage-2 secondary-decoder loss implemented. stage1_model.generator = "ddpm" is unaffected; this restriction is stage-2-only. Landing stage-2 ddpm later is a trainer addition, not a config break.

11.3 Tracked as implementation work

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

11.4 Accepted with a validation obligation: differentiability

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

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

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

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

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

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

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

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

---

12. Implementation order

1. config.py — new DEFAULT_CONFIG, recursive merge/write, migrate_config, tests. Nothing else can land first.
2. network.py — the §5 decomposition, with Stage2OneShot reproducing v0.2 exactly. Gate on the §4.3 migration test: load a v0.2 checkpoint through the shim and diff outputs against v0.2 code.
3. 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. Add the gradient-magnitude instrumentation of §11.4 test 1 here, so the differentiability assumption accrues evidence during the architecture comparison instead of needing its own run later.
6. sample.py / rollout.py — AR generation and class -> PDG decode, so an AR model can actually be rolled out and put through giant analyze. Includes the L1-distance diagnostic of §11.3.
7. history = "attention", scheduled sampling — then run the meeting's §7 comparison (one-shot vs autoregressive, standalone) and only chain the winner behind Stage 1.
8. estimate_batch_size recalibration — last, once the architectures are settled. Measure on real hardware with the example configs and record the new calibration points in config.py (§9, §11.3).

Steps 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.