From 878e9ddca3eb41e0dac12a66fa5410a4245b83a7 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 10 Aug 2026 11:19:02 +0200 Subject: [PATCH] Delete docs/v0.3.0-design.md and strip all references to it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The design doc and its followups doc are no longer needed as a live reference now that the v0.3.0 redesign is implemented — comments and docstrings across the codebase cited it extensively (file path, "design doc §X.Y", "decision N", or bare "§X.Y" section numbers) as design rationale. Removed docs/ and edited every citing comment/docstring to drop the now-dangling reference while keeping the substantive explanation next to it. CLAUDE.md's v0.3.0 roadmap bullet loses its trailing pointer to the deleted file. Verified: no remaining "docs/v0.3.0", "design doc", "decision N", or "§N.N" references (repo-wide grep); ruff and ty clean; full test suite on the heaviest-touched modules (network, sample, rollout, migration, config, train) passes. Co-Authored-By: Claude Sonnet 5 --- CLAUDE.md | 2 +- docs/v0.3.0-design.md | 1002 --------------------- docs/v0.3.0-followups.md | 53 -- giant/analysis/catalog.py | 2 +- giant/analysis/condor.py | 2 +- giant/analysis/render.py | 4 +- giant/analysis/router_gating.py | 4 +- giant/analysis/type_embedding_distance.py | 2 +- giant/cli.py | 16 +- giant/config.py | 14 +- giant/data/dataset.py | 3 +- giant/data/loader.py | 11 +- giant/data/setup_cache.py | 6 +- giant/data/transforms.py | 35 +- giant/model/network.py | 97 +- giant/model/schedule.py | 6 +- giant/particles.py | 20 +- giant/pipeline.py | 20 +- giant/rollout.py | 35 +- giant/sample.py | 52 +- giant/training/loop.py | 7 +- giant/training/stage2_inputs.py | 28 +- giant/training/trainers.py | 36 +- giant/validate.py | 2 +- scripts/check_migration_v02_v03.py | 10 +- scripts/dwarf.py | 2 +- scripts/warm_setup_cache.py | 4 +- tests/legacy/network_v02_snapshot.py | 8 +- tests/test_cli_train_overrides.py | 5 +- tests/test_config.py | 8 +- tests/test_loader.py | 3 +- tests/test_migration_v02_v03.py | 28 +- tests/test_network.py | 18 +- tests/test_pipeline.py | 10 +- tests/test_rollout.py | 20 +- tests/test_router.py | 8 +- tests/test_sample.py | 6 +- tests/test_train.py | 12 +- tests/test_type_embedding_distance.py | 2 +- tests/test_validate.py | 4 +- uv.lock | 2 +- 41 files changed, 263 insertions(+), 1346 deletions(-) delete mode 100644 docs/v0.3.0-design.md delete mode 100644 docs/v0.3.0-followups.md diff --git a/CLAUDE.md b/CLAUDE.md index c3cc181..7c89dab 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -91,6 +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`.** +**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. **Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet. diff --git a/docs/v0.3.0-design.md b/docs/v0.3.0-design.md deleted file mode 100644 index 81f9698..0000000 --- a/docs/v0.3.0-design.md +++ /dev/null @@ -1,1002 +0,0 @@ -# GIANT v0.3.0 — Stage-2 autoregressive redesign - -**Status:** design agreed, not implemented. Branch `v0.3.0-stage2-autoregressive`. -**Date:** 2026-08-04. -**Source:** `~/knowledge-base/meetings/2026-08-04-jan-stage2-autoregressive-architecture.md` -(meeting with Jan), plus the design decisions taken in the session that produced -this document. - -This document is the implementation contract for v0.3.0. It specifies the new -config format option by option, the `network.py` refactor, and the order in which -to build it. Read it before touching `giant/config.py` or `giant/model/network.py`. - ---- - -## Table of contents - -1. [Why](#1-why) -2. [Decisions register](#2-decisions-register) -3. [Config format reference](#3-config-format-reference) -4. [Migration: v0.2 -> v0.3](#4-migration-v02---v03) -5. [network.py refactor](#5-networkpy-refactor) -6. [Stage-2 autoregressive design](#6-stage-2-autoregressive-design) -7. [Training loop](#7-training-loop) -8. [Data and setup-cache changes](#8-data-and-setup-cache-changes) -9. [Config machinery changes](#9-config-machinery-changes) -10. [Callers that need updating](#10-callers-that-need-updating) -11. [Settled scope and open questions](#11-settled-scope-and-open-questions) -12. [Implementation order](#12-implementation-order) - ---- - -## 1. Why - -The 2026-08-03 WGAN rollout benchmark -(`~/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md`) was -good at the primary-step level and **failed at the secondary-species level**: -zero photon secondaries generated, ~4M hallucinated `-14` (muon antineutrino) -secondaries — a species essentially absent from Geant4. - -Stage 1 is not implicated; the meeting scoped everything to Stage 2. Two changes -were agreed together: - -- **Autoregressive generation** over secondaries in decreasing energy order, - replacing the one-shot masked `K_MAX=15` prediction, trained with teacher - forcing. -- **Categorical particle type** with a data-derived "other" bucket, reversing the - 2026-07-17 move to a continuous `(log-mass, charge)` target. Working hypothesis: - the continuous target is part of what let the generator collapse onto degenerate - species. - -The meeting also set a **methodology**: compare Stage-2 architectures *standalone* -(trained directly on secondary columns, no Stage-1 forward pass) before chaining -the winner behind Stage 1. Iterating on the compounding-rollout-error problem is -far cheaper that way than paying for a full two-stage run per candidate. - -That methodology is what forces the config refactor: v0.2 has a single global -`train.mode` and a single `[model]` block, with no way to express "train stage 2 -only", "stage 1 flow + stage 2 WGAN", or "one-shot vs autoregressive stage 2". - ---- - -## 2. Decisions register - -| # | Decision | Rationale | -|---|----------|-----------| -| 1 | **`n_sec` head moves from stage 1 to stage 2** | Stage 1 becomes the pure 9D primary step. A stage-2-only run is then self-contained (it can predict its own multiplicity), and §3's "implicit stop token" alternative gets a natural home in the same config block. | -| 2 | **Full mixed per-stage objectives** | `stage1 = flow` + `stage2 = wgan` must actually run — Stage 1 is good as flow, Stage 2 is what is being iterated on. The train loop becomes one trainer object per stage, each owning its optimizers and update cadence. | -| 3 | **Migration shim for configs *and* checkpoints** | Nothing on `/ceph` goes dead. v0.2 `config.toml` files and v0.2 `model_config` dicts are translated on load. | -| 4 | **Nested objective sub-tables** | `[stage1_model.wgan]`, `[stage2_model.ddpm]` rather than flat `wgan_n_critic` keys — self-documenting about which keys the active generator ignores, and validation can warn on a populated sub-table that is never read. | -| 5 | **Particle type is adversarial: straight-through Gumbel into the critic** | The critic sees a relaxed one-hot alongside energy/direction, so the joint (species, kinematics) distribution is learned rather than factorized. Makes the collapse hypothesis directly testable instead of assumed. 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. - -```toml -[conditioning] -out_dim = 128 -share_stages = false - -[conditioning.particle] -type = "physical" -emb_dim = 16 -n_layers = 1 - -[conditioning.material] -type = "physical" -emb_dim = 16 -n_layers = 1 -``` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `out_dim` | int | `128` | Width of the fused conditioning vector produced by the encoder's fusion MLP, consumed by every downstream trunk. **New in v0.3.0** — v0.2 hardcoded this as `cond_out_dim = 128` in every constructor signature, unreachable from config. | -| `share_stages` | bool | `false` | `false`: stage 1 and stage 2 each construct their own `ConditionEncoder` with identical config but independent weights (v0.2 behaviour). `true`: one instance, shared by reference. Shared weights halve the conditioning parameter count and force a common representation; independent weights let each stage specialize its view of the pre-step state. | - -#### `[conditioning.particle]` and `[conditioning.material]` - -Identical key sets, applied to the two identity axes independently. - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `type` | `"physical"` \| `"embedding"` \| `"onehot"` | `"physical"` | How this axis's identity becomes an `emb_dim`-wide vector. See the table below. | -| `emb_dim` | int | `16` | Width of this axis's vector. Under `"onehot"` it **also sets the class count** — see below. | -| `n_layers` | int | `1` | Depth of the sub-MLP under `"physical"`. `1` is the single-layer net with `emb_dim` neurons. Ignored under `"embedding"` / `"onehot"`. | - -**The three modes:** - -| mode | particle input | material input | width | learned parameters | -|------|----------------|----------------|-------|--------------------| -| `"physical"` | `log(mass)`, `charge` (`PARTICLE_PHYS_DIM = 2`) | `Z_eff`, `A_eff`, `log(density)`, `log(X0)`, `log(λ_int)` (`MATERIAL_PHYS_DIM = 5`) | `emb_dim` | one `n_layers`-deep MLP with `emb_dim` neurons | -| `"embedding"` | dense vocab index | dense vocab index | `emb_dim` | `nn.Embedding(vocab, emb_dim)` | -| `"onehot"` | top `emb_dim - 1` PDG codes by training-set count, plus one "other" bin | top `emb_dim - 1` materials by count, plus "other" | `emb_dim` | **none** — a fixed vector | - -- `"physical"` reads columns already present in `cond_cont[:, COND_DIM_BASE:]` - (see `giant.data.transforms.build_features`). It is computable for **any** PDG - code or material, which is what allows generalization beyond the training menu. -- `"embedding"` memorizes the training menu — the generalization-comparison - baseline, and the only mode that supports - `stage2_model.particle_type.target = "embedding"` (§3.3). -- `"onehot"` is a *fixed, unlearned* representation. It is **not** a - reparameterization of `"embedding"`: the difference is the vocabulary cap. With - 237 PDG codes and `emb_dim = 16`, `"embedding"` gives 237 distinct learned - vectors while `"onehot"` gives 16 classes. Needs the same data-derived top-N - map as the stage-2 type target (§8). - -**Interaction:** `particle.type = "physical"` is incompatible with router types -`"pdg"` and `"process"`, which build their own dataset-scoped -`nn.Embedding(pdg_vocab, ...)` regardless of the trunk's conditioning mode. -Pairing them silently reintroduces a training-menu-scoped lookup at the routing -layer, defeating the point of physical conditioning. Rejected loudly at build time -— `_check_router_conditioning_compat` in `network.py`, which carries over but must -now read `conditioning.particle.type` rather than a single global mode. - -### 3.2 `[stage1_model]` - -Stage 1 predicts the 9D primary post-step vector -(`giant/constants.py:LOCAL_TARGET_NAMES`). As of decision 1 it carries **no** -`n_sec` head. - -```toml -[stage1_model] -active = true -generator = "flow" -hidden_dim = 256 -n_res_blocks = 6 -dropout = 0.0 -lambda = 1.0 -``` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `active` | bool | `true` | `false` skips building and training stage 1 entirely. The resulting checkpoint holds only stage 2 and **cannot be rolled out** — `giant rollout` must refuse it with a clear error. Used for the meeting's Stage-2-only architecture comparison. | -| `generator` | `"flow"` \| `"ddpm"` \| `"wgan"` | `"flow"` | The generative objective. `"flow"`: conditional flow matching (Lipman et al. 2022), ~10 ODE steps at inference. `"ddpm"`: cosine-schedule diffusion baseline. `"wgan"`: WGAN-GP, single forward pass at inference. Replaces the global `train.mode`. Objective-specific knobs live in the matching sub-table below. | -| `hidden_dim` | int | `256` | Trunk width — also the width of **every expert** under a routed trunk. | -| `n_res_blocks` | int | `6` | Number of `ResBlock`s in the trunk, and in every expert under a routed trunk. Was `model.n_blocks`; renamed for clarity since v0.3.0 also has attention layers in stage 2. | -| `dropout` | float | `0.0` | Dropout inside each `ResBlock`. **Default changed in v0.3.0** (was `0.1`). | -| `lambda` | float | `1.0` | Weight of this stage's loss in the total. Meaningful when both stages are active and non-adversarial; a WGAN stage's adversarial loss drives its own optimizer, so `lambda` scales only its non-adversarial auxiliary terms. | - -#### `[stage1_model.flow]` — read only when `generator = "flow"` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `time_dim` | int | `64` | Width of the `SinusoidalEmbedding` for the flow time variable, concatenated into the trunk's conditioning. **New in v0.3.0** — v0.2 hardcoded 64. | - -#### `[stage1_model.ddpm]` — read only when `generator = "ddpm"` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `time_dim` | int | `64` | As above, for the diffusion time variable. | -| `n_steps` | int | `1000` | Cosine-schedule diffusion steps. **New in v0.3.0** — v0.2 hardcoded this in `CosineSchedule`. | - -#### `[stage1_model.wgan]` — read only when `generator = "wgan"` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `noise_dim` | int | `64` | Width of the generator's input noise vector. There is no time variable, hence no `time_dim`. | -| `n_critic` | int | `5` | Critic updates per generator update (Gulrajani et al. 2017). | -| `gp_weight` | float | `10.0` | Gradient-penalty coefficient. | -| `critic_lr` | float | `0.0` | Critic learning rate. `0.0` means "inherit `train.lr`" — not `None`, since the TOML writer has no null literal to round-trip. | -| `critic_hidden_dim` | int | `0` | Critic trunk width. `0` = inherit `stage1_model.hidden_dim`. **New in v0.3.0** — v0.2 always sized the critic from the generator. | -| `critic_n_res_blocks` | int | `0` | Critic trunk depth. `0` = inherit `stage1_model.n_res_blocks`. | - -#### `[stage1_model.router]` - -Content carries over from v0.2's `[model.router]` unchanged. Reproduced here in -full because the block is now per-stage and its semantics are easy to lose. - -```toml -[stage1_model.router] -enabled = false -type = "energy" -n_experts = 4 -temperature = 0.5 -learn_centers = true -learn_width = false -learn_temperature = false -width_min_ratio = 0.1 -width_max_ratio = 10.0 -lambda_balance = 0.0 -lambda_entropy = 0.0 -lambda_proc = 0.0 -gumbel = false -gumbel_tau_start = 1.0 -gumbel_tau_end = 0.1 -emb_dim = 8 -hidden_dim = 64 -``` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `enabled` | bool | `false` | Replace the monolithic trunk with a mixture of per-expert trunks: soft-mixed over all experts at train time, **top-1 dispatched at eval time** (each row runs exactly one expert). Every expert is `hidden_dim` × `n_res_blocks` — v0.3.0 removes the per-expert sizing keys, so a routed model costs `n_experts` × the monolith's parameters at equal per-row eval cost. Routing buys specialization, not a per-call speedup. | -| `type` | str | `"energy"` | Router impl from `ROUTER_REGISTRY`. `"energy"`: soft turn-on gate over normalized pre-step log-energy. `"pdg"`: gate over a learned PDG embedding. `"process"`: own classifier over pre-step conditioning predicting the step-ending physics process. `"composed"`: joint outer-product gating over multiple axes via `axis{i}_{field}` keys. | -| `n_experts` | int | `4` | Number of experts. For `type = "process"` this doubles as the number of process classes. Ignored for `type = "composed"` (each axis has its own). | -| `temperature` | float | `0.5` | Softmax denominator for the distance-based gate. As `tau -> 0` the gate hardens to nearest-center (Voronoi) selection, which is exactly what eval-time `top1` uses. Energy/pdg routers only. | -| `learn_centers` | bool | `true` | Whether gate centers are `nn.Parameter` or a fixed buffer. | -| `learn_width` | bool | `false` | Give each expert its own learnable width, sigmoid-bounded to `[width_min_ratio, width_max_ratio] * temperature`. Mutually exclusive with `learn_temperature`. | -| `learn_temperature` | bool | `false` | Make the single shared `temperature` learnable, same bounding. Mutually exclusive with `learn_width`. | -| `width_min_ratio` | float | `0.1` | Lower bound multiplier for the above. Must bracket 1.0 with `width_max_ratio` so enabling either mode is a no-op at init. | -| `width_max_ratio` | float | `10.0` | Upper bound multiplier. **Deliberately bounded rather than `softplus`/`exp`:** an unbounded width lets one expert's logit `-d²/width -> 0` almost everywhere, so it wins nearly every row regardless of distance — the same "experts overlap instead of partitioning" failure the router design exists to avoid. | -| `lambda_balance` | float | `0.0` | Importance-CV² load-balancing auxiliary loss weight (Shazeer et al. 2017). **The 2026-07-22 benchmark failure ran with `0.0`; do not repeat that.** | -| `lambda_entropy` | float | `0.0` | Entropy-regularization weight penalizing uniform/collapsed gating. Secondary guard against all experts' widths co-inflating together, which `lambda_balance` cannot see (usage shares stay even throughout that failure). Use with caution: indiscriminate entropy penalties also suppress legitimate ambiguity near a decision boundary. | -| `lambda_proc` | float | `0.0` | Supervised process-classification CE weight. `"process"` router only; `0.0` still trains a working router (the gate gets gradient through the downstream loss) but only `> 0` grounds it in the true `process` label. | -| `gumbel` | bool | `false` | Straight-through Gumbel-softmax train-time combine weights: the forward pass samples a hard one-hot combination (matching eval-time top-1 dispatch exactly) while the backward pass still flows smooth gradient to every expert. Targets the train/eval mismatch. | -| `gumbel_tau_start` | float | `1.0` | Gumbel temperature at step 0, annealed linearly over training. | -| `gumbel_tau_end` | float | `0.1` | Gumbel temperature at the final step. | -| `emb_dim` | int | `8` | The router's **own** pdg (and material) embedding width, separate from the trunk's `ConditionEncoder`. `"pdg"` / `"process"` routers only. | -| `hidden_dim` | int | `64` | The `"process"` router's internal classifier hidden width. | - -**Composed routers** use flat `axis{i}_{field}` keys instead of `type`/`n_experts` -— e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`, -`axis1_n_experts = 3`, `axis1_emb_dim = 8`. Indices must be contiguous from 0. -Flat keys keep the block a table of scalars, which the merge machinery relies on. - -**`centers_init`** is not authored by hand: `giant/pipeline.py` populates it for -`type = "energy"` from real data quantiles collected during the existing -normalizer-fitting pass, then writes it into the checkpoint's `model_config`. - -### 3.3 `[stage2_model]` - -Stage 2 predicts `n_sec` and the per-secondary energy/direction/type. - -```toml -[stage2_model] -active = true -decoder = "autoregressive" -generator = "wgan" -hidden_dim = 256 -n_res_blocks = 6 -dropout = 0.0 -lambda = 1.0 -k_max = 15 -context_dim = 64 -stage1_context = "truth" -``` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `active` | bool | `true` | `false` trains stage 1 alone. The checkpoint then has no secondary decoder; `giant rollout` must refuse it, `giant predict` still works. | -| `decoder` | `"one_shot"` \| `"autoregressive"` | `"autoregressive"` | `"one_shot"`: predict all `k_max` slots simultaneously with padded slots masked from the loss — v0.2 behaviour, kept as the baseline arm of the meeting's §7 comparison. `"autoregressive"`: emit one secondary at a time in descending-energy order. | -| `generator` | `"flow"` \| `"ddpm"` \| `"wgan"` | `"wgan"` | As stage 1. Under `"autoregressive"` this is the objective for **each token**: a WGAN token costs one forward pass, a flow token costs ~10 ODE steps. See the cost note in §6.4. | -| `hidden_dim` | int | `256` | Trunk width. | -| `n_res_blocks` | int | `6` | Trunk depth. | -| `dropout` | float | `0.0` | Dropout inside each `ResBlock`. **Default changed in v0.3.0** (was `0.1`). | -| `lambda` | float | `1.0` | Weight of stage 2's loss in the total. Was `train.lambda_s2`. | -| `k_max` | int | `15` | Maximum secondary slots. Under `"one_shot"` this is the fixed output width; under `"autoregressive"` it is a safety cap on the generation loop. Was the global constant `K_MAX` in `giant/constants.py` (max observed `n_sec` is 14 in the PbWO4 dataset, so 15 covers it with one spare). | -| `context_dim` | int | `64` | Width of the projected stage-1 outcome fed into stage 2's conditioning. Was the hardcoded `stage1_proj_dim = 64`. | -| `stage1_context` | `"truth"` \| `"sampled"` | `"truth"` | What stage 2 conditions on during training. `"truth"`: the ground-truth stage-1 target vector, detached — v0.2 behaviour (v0.2 `train.py:254` passes `x1_s1.detach()`), i.e. stage-level teacher forcing. `"sampled"`: stage 1's own sampled output, closing the train/inference gap at the cost of a sampling pass per batch and a moving target early in training. | - -#### `[stage2_model.n_sec]` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `mode` | `"head"` \| `"stop_token"` \| `"truth"` | `"head"` | `"head"`: a classifier over `{0..k_max}` on the condition encoding alone (no diffusion noise), so it is callable independently at inference — v0.2 behaviour, and what the meeting's §3 explicitly decided to keep. `"stop_token"`: an EOS-style implicit stop — **accepted by the schema but not implemented in v0.3.0**, raising a clear "not implemented" error if set (§11.2). The key exists now so landing the mechanism later is not a config break. `"truth"`: take `n_sec` from ground truth — only valid for standalone stage-2 evaluation, never for rollout. | -| `lambda` | float | `0.1` | Cross-entropy weight for the head. Was `train.lambda_nsec`. | - -#### `[stage2_model.particle_type]` - -**The three targets mirror the three conditioning modes of §3.1**, and use the -same names. - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `target` | `"onehot"` \| `"physical"` \| `"embedding"` | `"onehot"` | What the token's type slice *is*. See the table below. | -| `lambda` | float | `1.0` | Loss weight. Under `generator = "flow"`/`"ddpm"` this weights the cross-entropy (`"onehot"`) or regression (`"physical"`/`"embedding"`) term; under `"wgan"` the type is adversarial (§2.1) and this weights only any auxiliary term. | -| `other_policy` | `"sample"` \| `"modal"` \| `"drop"` | `"sample"` | How a predicted "other" class becomes a concrete PDG code at rollout, needed because a secondary's mass/charge feed its own downstream conditioning. `"sample"`: draw from the empirical within-bucket distribution recorded at map-build time. `"modal"`: always the most common member. `"drop"`: discard the secondary. Read only under `target = "onehot"`. **Not decided in the meeting** — see §11. | - -**There is no `n_classes` key.** The class count under `target = "onehot"` is -`conditioning.particle.emb_dim` — the same number that sizes the particle axis -everywhere else. One knob sets the model's particle-type resolution, and the -stage-2 onehot classes are by construction the same classes the conditioning -onehot uses, so an emitted secondary's type is directly consumable as the -conditioning of its own next step with no re-mapping. - -Note the coupling this creates: under `conditioning.particle.type = "physical"`, -`emb_dim` primarily means "sub-MLP output width", yet it still sets the stage-2 -class count. Intentional, but worth knowing when tuning either. - -| target | token type slice | width | training target | inverse map at rollout | -|--------|------------------|-------|-----------------|------------------------| -| `"physical"` | regressed `(log mass, charge)` | `2` | the true PDG's physics values (`giant.particles.particle_mass_charge`) | none needed — mass/charge are used directly; `giant.particles.nearest_known_pdg` gives a reporting-only label. v0.2 behaviour. | -| `"onehot"` | class logits | `conditioning.particle.emb_dim` | true class index | `argmax` -> class -> PDG (via `other_policy` for the "other" bin) | -| `"embedding"` | an `emb_dim`-wide vector | `conditioning.particle.emb_dim` | `emb.weight[class].detach()` | L1-nearest row of `emb.weight` — see below | - -So the type slice is `conditioning.particle.emb_dim` wide for **both** `"onehot"` -and `"embedding"`, and `2` only for `"physical"`. - -#### `target = "embedding"` in detail - -Stage 2 emits a vector that should equal **the conditioning's own particle -embedding** for the secondary's species — the same `nn.Embedding` table -`[conditioning.particle]` builds, not a second one. - -**Requires `conditioning.particle.type = "embedding"`.** There is no table to -match against under `"physical"` or `"onehot"`; reject at config-validation time -with an explicit error. - -**Why detached:** the regression target is `emb.weight[class].detach()`, so the -embedding table receives gradient **only through the conditioning path**, never -through the stage-2 output loss. Without the detach the target moves as the -decoder chases it — the exact failure mode that motivated abandoning the learned -type target in the first place (see `decisions/physical-property-conditioning`). -The detach is what makes this option viable again. - -**Inverse map.** The natural exact-match form - -```python -((out - emb.weight).abs().sum(1) < 1e-6).nonzero() -``` - -is correct as a **round-trip assertion in tests** (encode a known PDG, decode, -recover the same PDG) but **must not be used at inference**: a generative model's -continuous output essentially never lands within `1e-6` of a table row, so it -returns an empty tensor almost always. Inference needs the nearest row: - -```python -pdg_idx = (out.unsqueeze(-2) - emb.weight).abs().sum(-1).argmin(-1) # L1 nearest -``` - -Note this decode is **unbounded in vocabulary**, unlike `"onehot"` — every PDG -code in the training vocab is reachable, and there is no "other" bucket, hence no -`other_policy`. The trade-off is that nearest-neighbour decode has no notion of -confidence: an output far from every row still snaps to something. - -#### `[stage2_model.autoregressive]` — read only when `decoder = "autoregressive"` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `order` | `"energy_desc"` | `"energy_desc"` | Canonical generation order. Descending energy is the ordering already flagged as natural in the Phase-2 note's open questions, and the one the existing stick-breaking encoding assumes. Single-valued for now; the key exists so an alternative ordering is not a config break. | -| `history` | `"markov"` \| `"attention"` | `"markov"` | How token *i+1* sees tokens ≤ *i*. `"markov"`: previous token plus running scalars (remaining energy budget, slot index) — a fixed-width summary. `"attention"`: causal self-attention over all emitted tokens. See §6.2 for the trade-off. | -| `teacher_forcing` | `"always"` \| `"scheduled"` \| `"never"` | `"always"` | `"always"`: condition on the ground-truth previous secondary throughout training (the meeting's confirmed plan). `"scheduled"`: scheduled sampling — interpolate toward conditioning on the model's own prediction. `"never"`: free-running from the start. | -| `tf_p_start` | float | `1.0` | Under `"scheduled"`, P(use ground truth) at epoch 0. | -| `tf_p_end` | float | `1.0` | Under `"scheduled"`, P(use ground truth) at the final epoch. Linear interpolation between the two. | -| `attn_n_heads` | int | `4` | Attention heads. Read only under `history = "attention"`. | -| `attn_n_layers` | int | `2` | Causal self-attention layers. Read only under `history = "attention"`. | - -#### `[stage2_model.flow]` / `[stage2_model.ddpm]` - -Same keys as their `[stage1_model.*]` counterparts (`time_dim`; plus `n_steps` -for ddpm). - -#### `[stage2_model.wgan]` — read only when `generator = "wgan"` - -Same keys as `[stage1_model.wgan]` (`noise_dim`, `n_critic`, `gp_weight`, -`critic_lr`, `critic_hidden_dim`, `critic_n_res_blocks`), plus: - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `gumbel_tau_start` | float | `1.0` | Straight-through Gumbel temperature for the **particle-type one-hot** at step 0, annealed linearly. Read only under `particle_type.target = "onehot"` (the other two targets are continuous and need no relaxation). Distinct from `router.gumbel_tau_start`, which anneals expert-combination weights — two unrelated Gumbel relaxations that must not share a key. | -| `gumbel_tau_end` | float | `0.1` | Same, at the final step. | - -Under `"autoregressive"` a fresh `noise_dim` draw is made **per token**. - -#### `[stage2_model.router]` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `tie_to_stage1` | bool | `false` | `true`: stage 2 shares **stage 1's `Router` module instance**, so expert *i* in stage 1 and expert *i* in stage 2 gate on identical conditions by construction. Every other key in this block is then ignored. `false`: an independent router — note this is v0.2's actual behaviour, which built *two separate routers from one config*, so stage-1 expert *i* and stage-2 expert *i* had no semantic relationship despite identical hyperparameters. | - -All other keys are as `[stage1_model.router]`. Invalid when `stage1_model.active -= false` and `tie_to_stage1 = true` — reject at config-validation time. - -### 3.4 `[train]` - -Optimizer, schedule, data split and logging only. Everything model-shaped moved -into the stage blocks. - -```toml -[train] -epochs = 100 -batch_size = 4096 -lr = 3e-4 -weight_decay = 0.01 -ema_decay = 0.9999 -warmup_epochs = 5 -val_fraction = 0.1 -max_val_batches = 200 -num_workers = 4 -seed = 0 -validate_every = 10 -validate_steps = 10 -wandb = true -wandb_project = "giant" -wandb_run_name = "" -wandb_log_every = 50 -``` - -| Key | Type | Default | Meaning | -|-----|------|---------|---------| -| `epochs` | int | `100` | Training epochs. | -| `batch_size` | int | `4096` | Steps per batch. `auto` on the CLI estimates from free VRAM. | -| `lr` | float | `3e-4` | AdamW learning rate for every stage's generator. | -| `weight_decay` | float | `0.01` | AdamW weight decay. | -| `ema_decay` | float | `0.9999` | EMA of model weights used for sampling; `0` disables. Maintained per stage. | -| `warmup_epochs` | int | `5` | Linear LR warmup. | -| `val_fraction` | float | `0.1` | Fraction of **events** (not steps) held out — the split is by `event_id` to avoid leaking correlated steps from the same shower. | -| `max_val_batches` | int | `200` | Cap on the per-epoch val-loss pass; `0` = full val set. Distinct from `validate_every`'s marginal/KL pass. | -| `num_workers` | int | `4` | DataLoader workers. Warns above ~1/4 of the machine's CPUs — portal machines are shared. | -| `seed` | int | `0` | Seeds Python/numpy/torch and the event split. | -| `validate_every` | int | `10` | Epochs between full marginal/KL validation passes. | -| `validate_steps` | int | `10` | Sampler steps used during those passes. | -| `wandb` | bool | `true` | W&B per-epoch metric logging. **Default changed in v0.3.0** (was opt-in `false`): the v0.3.0 work is a sequence of architecture comparisons, and a run that was not logged is not comparable. Set `false` for throwaway/debug runs. | -| `wandb_project` | str | `"giant"` | W&B project. | -| `wandb_run_name` | str | `""` | `""` means "use the checkpoint out_dir name" — not `None`, since the TOML writer has no null literal. | -| `wandb_log_every` | int | `50` | Optimizer steps between batch-granularity metric logs. A single epoch can be tens of thousands of steps; per-epoch metrics always log in full. | - -### 3.5 Removed from v0.2 - -| Key | Fate | -|-----|------| -| `train.mode` | Split into `stage1_model.generator` / `stage2_model.generator`. | -| `train.lambda_nsec` | -> `stage2_model.n_sec.lambda`. | -| `train.lambda_s2` | -> `stage2_model.lambda`. | -| `train.n_critic`, `gp_weight`, `critic_lr` | -> `stage{1,2}_model.wgan.*`. | -| `[model]` (whole block) | Split across `[conditioning]` and the two stage blocks. | - ---- - -## 4. Migration: v0.2 -> v0.3 - -`[meta] config_version = 3` tags the new format; **absent means v0.2**. -`migrate_config(cfg) -> cfg` applies the table below and is called on both -`config.toml` load and checkpoint `model_config` load, so nothing on `/ceph` goes -dead (decision 3). - -| v0.2 | v0.3 | -|------|------| -| `train.mode` | `stage1_model.generator` **and** `stage2_model.generator` (same value) | -| `train.lambda_nsec` | `stage2_model.n_sec.lambda` | -| `train.lambda_s2` | `stage2_model.lambda` | -| `train.n_critic` / `gp_weight` / `critic_lr` | `stage1_model.wgan.*` **and** `stage2_model.wgan.*` | -| `model.hidden_dim` / `n_blocks` / `dropout` | `stage{1,2}_model.hidden_dim` / `n_res_blocks` / `dropout` | -| `model.emb_dim` | `conditioning.particle.emb_dim` **and** `conditioning.material.emb_dim` (v0.2 had one shared value) | -| `model.conditioning` | `conditioning.particle.type` **and** `conditioning.material.type` (v0.2 had one shared mode) | -| `model.noise_dim` | `stage{1,2}_model.wgan.noise_dim` | -| `model.router.*` | `stage1_model.router.*`, copied verbatim to `stage2_model.router.*` with `tie_to_stage1 = false` (preserves v0.2's two-independent-routers behaviour) | -| `model.expert_hidden_dim` / `expert_n_blocks` | **dropped.** v0.2's `0` sentinel meant "inherit from the monolith", which is now unconditional. A v0.2 config with a *non-zero* value must fail loudly rather than silently resize the experts — see §4.3. | -| `model.k_max` | `stage2_model.k_max` | -| — | `conditioning.out_dim = 128` (v0.2's hardcoded value) | -| — | `conditioning.{particle,material}.n_layers = 2` (v0.2's hardcoded depth; note the v0.3 *default* is 1) | -| — | `stage{1,2}_model.flow.time_dim = 64` / `.ddpm.time_dim = 64` (v0.2's hardcoded value) | -| — | `stage2_model.context_dim = 64` (v0.2's hardcoded `stage1_proj_dim`) | -| — | `stage2_model.decoder = "one_shot"` (v0.2 had no other option) | -| — | `stage2_model.particle_type.target = "physical"` (v0.2 behaviour) | -| — | `stage{1,2}_model.active = true` | - -### 4.1 The awkward one: `n_sec` head ownership - -Decision 1 moves the `n_sec` head to stage 2, but **v0.2 checkpoints carry -`n_sec_head` weights inside the stage-1 module** (`DenoisingMLP.n_sec_head`, -`WGANGenerator.n_sec_head`, `RoutedDenoisingMLP.n_sec_head`). The shim must keep -those loading where they are. - -Handling: `migrate_config` sets an internal -`stage2_model.n_sec.legacy_owner = "stage1"` that `build_models` honours by -attaching the head to the stage-1 module. Never written by new runs, never -CLI-settable, never documented as a user-facing option. - -### 4.2 Non-inheriting expert dims - -v0.3.0 drops `expert_hidden_dim` / `expert_n_blocks` (§2.2). Migration must -distinguish two cases: - -- value is `0` (the "inherit" sentinel, and what every real run used) — drop the - key silently, behaviour is unchanged. -- value is non-zero — **fail loudly.** Silently resizing those experts to - `hidden_dim` would change the architecture, so the checkpoint's weights would no - longer match. Such a checkpoint can only be loaded by v0.2. - -### 4.3 Migration test - -The acceptance criterion for the whole shim: **load a v0.2 checkpoint through -`migrate_config` + the new `build_models`, and diff its outputs against v0.2 code -on the same input batch.** Bit-identical, or the refactor has changed something it -should not have. Pick one flow checkpoint and one WGAN checkpoint from `/ceph`. - ---- - -## 5. network.py refactor - -### 5.1 What it looks like today - -Ten classes that are permutations of three independent choices: - -| | flow/ddpm | wgan generator | wgan critic | -|---|---|---|---| -| **stage 1** | `DenoisingMLP` | `WGANGenerator` | `Critic` | -| **stage 1, routed** | `RoutedDenoisingMLP` | — | — | -| **stage 2** | `SecondaryDecoder` | `WGANSecondaryGenerator` | `SecondaryCritic` | -| **stage 2, routed** | `RoutedSecondaryDecoder` | — | — | - -The empty cells are the entire reason `giant/pipeline.py:275` hard-rejects -`--mode wgan --router`: no routed WGAN generator class was ever written. There is -no deeper reason — the routed trunk is orthogonal to the objective. - -Every one of those classes repeats the same body: build a condition encoder, -optionally a time embedding, project input, run blocks, project output. - -### 5.2 Proposed decomposition — one axis per config block - -**(a) `[conditioning]` -> encoders** - -```python -ConditionEncoder(type, emb_dim, n_layers, out_dim) # behaviour unchanged, now configurable -ContextAdapter(in_dim, context_dim) # stage-1 outcome -> context vector -``` - -`SecondaryConditionEncoder` **disappears as a class**. It was -`ConditionEncoder` + a `stage1_proj` linear + a fuse layer; those compose at the -stage-model level instead. - -**(b) `[stage*_model]` + `.router` -> trunks, behind one interface** - -```python -class Trunk(nn.Module): - def forward(self, x, cond, cond_cont=None, cond_cat=None) -> Tensor: ... - -MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) -RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) - -build_trunk(stage_cfg, in_dim, out_dim, cond_dim) -> Trunk -``` - -`ExpertTrunk` and `_route_forward` carry over unchanged. **One required change:** -`MonolithicTrunk` and `ExpertTrunk` need a separate `out_dim` — today -`ExpertTrunk` hardcodes `out_proj = nn.Linear(hidden_dim, in_dim)`, i.e. -`out_dim == in_dim`. That stops working the moment stage 2's per-token output is -`4 + type_dim` wide while its input is a noise vector of width `noise_dim`. - -**(c) `[stage*_model].generator` -> a thin wrapper, not a class family** - -The generator choice controls exactly two things: - -- whether `SinusoidalEmbedding(t)` is concatenated into `cond` (flow/ddpm) or not (wgan) -- whether the trunk's `x` is the diffused/interpolated `x_t` (flow/ddpm) or a noise draw `z` (wgan) - -That is small enough for one wrapper, collapsing six of today's ten classes: - -```python -class GenerativeTrunk(nn.Module): - """cond-encode -> (optional time-embed) -> trunk. Backs flow, ddpm and wgan.""" - def forward(self, x, cond_cont, cond_cat, t=None, context=None) -> Tensor: ... -``` - -### 5.3 Resulting class list - -``` -# building blocks -SinusoidalEmbedding, ResBlock, ConditionEncoder, ContextAdapter - -# trunks -Trunk (ABC), MonolithicTrunk, RoutedTrunk, ExpertTrunk - -# routers — carried over unchanged -Router, EnergyRouter, PdgRouter, ProcessRouter, ComposedRouter -ROUTER_REGISTRY, register_router, build_router, build_composed_router - -# history encoders — new, stage-2 AR only -HistoryEncoder (ABC), MarkovHistory, AttentionHistory - -# stage models -Stage1Model # 9D primary step -Stage2OneShot # k_max slots at once (v0.2 behaviour) -Stage2Autoregressive # one token at a time -CriticModel # stage-1 or stage-2 critic, generator-agnostic - -# factories -build_models(cfg) -> {"stage1": ... | None, "stage2": ... | None} -build_critics(cfg) -> {"stage1": ... | None, "stage2": ... | None} -``` - -Ten classes become four stage classes plus reusable parts, and **routed WGAN comes -for free** — `pipeline.py`'s rejection can be deleted. - -### 5.4 Factory signature change - -`build_models` returns a **dict, not a tuple**: `active = false` on either stage -means that key is `None`. Every caller unpacking -`stage1, sec_decoder = build_models(...)` must be updated -(`pipeline.py:390`, `cli.py:871`, `cli.py:1254`). - ---- - -## 6. Stage-2 autoregressive design - -### 6.1 Token layout - -Secondaries are emitted one at a time in descending energy order. Per-token output -width is `4 + type_dim`: - -| slice | meaning | -|-------|---------| -| `[0]` | stick-breaking logit — fraction of the **remaining** energy budget | -| `[1:4]` | local-frame direction, normalized to a unit vector | -| `[4:4+type_dim]` | particle type — see below | - -`type_dim` follows `particle_type.target` (§3.3): `2` for `"physical"`, and -`conditioning.particle.emb_dim` for both `"onehot"` and `"embedding"`. -Only `"onehot"` involves a relaxation — ST-Gumbel under `wgan`, cross-entropy -under `flow`/`ddpm`; the other two are continuous and feed the critic (or the -regression loss) directly. - -Per-token conditioning is: - -``` -base condition encoding (ConditionEncoder, [conditioning].out_dim) -+ stage-1 context (ContextAdapter, stage2_model.context_dim) -+ history_encoder(prefix) (§6.2) -+ running scalars (remaining energy budget, slot index) -``` - -### 6.2 History encoders - -One interface, `history_encoder(prefix) -> fixed-width vector`: - -- **`MarkovHistory`** — the previous token's `(energy_fraction, direction, - type_embedding)`. Fixed-width, one small MLP. -- **`AttentionHistory`** — causal self-attention over all emitted tokens, taking - the last position. `attn_n_heads` × `attn_n_layers`. - -**The trade-off** (this is decision 8's reasoning, recorded so it does not have to -be re-derived): - -1. **Markov is less impoverished than it sounds.** The remaining-energy budget is - an *exact sufficient statistic* for the conservation constraint — stick-breaking - needs nothing else from history. Slot index likewise. What markov genuinely - cannot see is set *composition*: "I have already emitted two photons and an - electron." That matters for correlated production — pair production emits - exactly e⁺e⁻, a brems cascade correlates species across the set. Note that the - meeting's charge-conservation idea was exactly such a hand-engineered summary - statistic, and it is out of scope for v0.3.0 (decision 6), so markov does not - get that crutch. -2. **Sequences are short and front-loaded.** `n_sec` is 0–2 for most steps, max 14. - For K ≤ 2, attention *is* markov. They diverge only in the high-multiplicity - tail — physically interesting but data-poor, so the attention path trains on - few examples where it actually matters. -3. **Cost is not where intuition puts it.** Under teacher forcing both train in a - single parallel pass over all K tokens (every token's input is ground truth, so - nothing is sequential). At inference both need K sequential forwards; attention - additionally needs a KV cache to avoid re-encoding the prefix. Attention's - marginal FLOPs over ≤15 tokens are rounding error. -4. **Exposure bias cuts against attention.** Attention conditions on the entire - generated prefix, so one off-manifold early token poisons every later token - through the context. Markov's fixed summary sees only one bad token, and the - budget scalars stay exact regardless. Given the explicitly-flagged train/ - inference gap and the known compounding-rollout-error problem, the more - expressive history is also the more fragile one. - -Hence: markov is the default and the baseline; attention is a flag. - -### 6.3 Energy budget under AR - -**No re-derivation needed**, despite the meeting's action item suggesting -otherwise. The existing stick-breaking is already sequential in spirit — each slot -takes a fraction of what remains — so it carries over to per-token generation -directly by feeding "remaining budget" as a per-token conditioning scalar. -Conservation stays exact by construction: the valid slots' energies sum to `e_sec`, -which the Stage-1 simplex already guarantees sums correctly with `edep` and -`post_E`. - -### 6.4 Cost warning - -AR costs **K sequential forward passes per step** where one-shot costs one. -Against the ~10× native-Geant4 eval budget that motivated the whole fast-eval -track, this is the number to watch — not the history-encoder choice. With -`generator = "flow"` it is worse still: ~10 ODE steps per token, so ~150 forwards -per step in the worst case. `generator = "wgan"` (one pass per token) is the only -configuration that plausibly meets the budget; flow AR is for quality comparison. - ---- - -## 7. Training loop - -Decision 2 (full mixed per-stage objectives) makes `giant/training/` 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 (v0.2 `train.py:262`) generalizes cleanly. -- `metrics.csv` and W&B metric names gain a stage prefix. **Implemented as** - `//` (`stage1/train/loss`, `stage2/train/d_loss`, - `stage1/lr`, `stage1/router/entropy`) plus an unprefixed run-level tail - (`val/loss`, `val/marginal_kl`, `grad_norm`, `gpu_mem_mb`, - `samples_per_sec`, `is_best`, `epoch_time_s`). Each name is declared once, - as a `MetricSpec` on the `StageTrainer` that computes it; the CSV header - and W&B payload are derived from those declarations - (`giant/training/metrics.py`). -- With `stage1_model.active = false`, stage 2 still needs its stage-1 context: it - comes from the ground-truth target already in the batch (`x1_s1`), which is - exactly what v0.2 does anyway. **Stage-2-only training is therefore a cheap - ablation, not new plumbing** — drop the stage-1 loss, skip building stage 1. - ---- - -## 8. Data and setup-cache changes - -Three config options need a "top 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/training/` | Per-stage trainers (§7), metric collection, checkpointing, the epoch loop. | -| `giant/cli.py` | Stage-prefixed flags for `train` and `new-run`; `build_models` now returns a dict (`:871`, `:1254`); `:1339` writes `model_config` into the rollout sidecar. | -| `giant/sample.py` | Sampler picked per stage from `stage*_model.generator`; new AR sampling loop with KV cache under `history = "attention"`. | -| `giant/rollout.py` | AR secondary generation; categorical class -> PDG decode; `other_policy` handling. | -| `giant/validate.py` | Stage-2 marginals gain a type-class marginal. | -| `giant/analysis/render.py` | `_router_summary(model_config)` at `:44` reads `model_config["router"]`. | -| `giant/analysis/router_gating.py` | Same, at `:76`. | -| `giant/constants.py` | `K_MAX` / `SEC_DIM` demoted to defaults (§8). | -| `configs/*.toml` | All eight shipped configs are v0.2-format; regenerate or rely on the shim. | -| `condor-gpu-train-rollout` branch | Submits `giant train` flags; needs rebasing onto the new flag surface. | - ---- - -## 11. Settled scope and open questions - -### 11.1 Settled — build as specified - -- **`other_policy` switch.** The three-way `"sample"` / `"modal"` / `"drop"` - design in §3.3 is approved as the config surface for turning a predicted "other" - class into a concrete PDG at rollout. `"sample"` (draw from the empirical - within-bucket distribution recorded at map-build time) stays the default as the - least-biased option. -- **`[stage*_model.flow]` and `[stage*_model.ddpm]` stay separate tables.** The - duplicated `time_dim` is accepted; a merged `[stage*_model.diffusion]` was - considered and rejected because the name fits flow matching poorly. - -### 11.2 Deferred — not implemented in v0.3.0 - -- **`n_sec.mode = "stop_token"`.** The value is **accepted by the schema** but - raises a clear "not implemented in v0.3.0" error if set. The key existing now - means landing the implicit-stop mechanism later is not a config break. `"head"` - is what v0.3.0 builds, per the meeting's §3 decision. -- **Charge conservation.** No key at all — `[stage2_model.conservation]` is absent - from v0.3.0 (decision 6), unlike `stop_token` above. **Deliberately left - undesigned:** the mechanism should be worked out on its own terms when it is - taken up, not pre-shaped by choices made for this refactor. Adding the block - later is a config addition, not a break. -- **`stage2_model.generator = "ddpm"`.** The value is **accepted by the - schema** (§3.3 lists `"flow" | "ddpm" | "wgan"` with no caveat) but - `FlowDDPMStageTrainer.__init__` (`giant/training/trainers.py`) raises - `NotImplementedError` for stage 2 — only `"flow"` and `"wgan"` have a - stage-2 secondary-decoder loss implemented. `stage1_model.generator = - "ddpm"` is unaffected; this restriction is stage-2-only. Landing stage-2 - ddpm later is a trainer addition, not a config break. - -### 11.3 Tracked as implementation work - -- **Log the L1 distance distribution at rollout** under - `particle_type.target = "embedding"`. Nearest-neighbour decode has no reject - option — an output far from every table row still snaps to its nearest neighbour, - with no "other" bin and no confidence signal. A heavy tail in that distribution - means the decoder is emitting vectors off the embedding manifold, which is the - direct analogue of the species-collapse symptom this redesign exists to fix. - Belongs with the `giant/rollout.py` decode work (step 6) and should surface as a - `giant analyze` diagnostic plot alongside `router_gating`. -- **`estimate_batch_size` recalibration** for AR stage 2 is the **last** step of - the implementation flow (§12 step 8), measured on real hardware using the example - configs — not guessed from the existing calibration constants. - -### 11.4 Accepted with a validation obligation: differentiability - -**Position taken: differentiability through the categorical type path is broken, -and that is accepted.** The expected contribution of the broken path to the total -gradient is small enough to ignore. **This is an assumption, not a result — it has -to be demonstrated later.** - -Recording it precisely, since "broken" covers three distinct things: - -| where | status | -|-------|--------| -| Per-token training loss under teacher forcing | **Fine.** Softmax cross-entropy needs no sampling; the loss is differentiable in the logits. | -| ST-Gumbel into the critic (`generator = "wgan"`, §2.1) | **Biased, not absent.** The forward pass is a hard one-hot; the backward pass pretends it was the soft sample. Gradient flows, but it is not the gradient of what was actually computed. | -| Full shower-rollout backprop | **Structurally broken regardless.** Already non-differentiable once secondaries spawn branches, independent of the type representation — so the categorical switch costs nothing that was not already lost. | - -The claim being accepted is about the middle row: the straight-through estimator's -bias, propagated into the shared trunk, is expected to be negligible against the -gradient from the continuous paths (stick-breaking energy, direction, and — when -stage 1 is active — the 9D primary target). The third row is the reason this is -tolerable at all: end-to-end differentiability was never available. - -**Validation obligation.** Do not treat this as settled until one of the following -has actually been run. Cheapest first: - -1. **Gradient-magnitude accounting.** Instrument a training run to log the norm of - the trunk gradient contributed through the type slice against the norm from the - continuous slices. "Negligible" should mean a stable, small ratio — not merely - small at initialization. This is the direct measurement of the claim and costs - almost nothing to add. -2. **Detached-type ablation.** Train with the type path detached from the shared - trunk entirely (type head still learns; no type gradient reaches the trunk) - against the ST-Gumbel default. Comparable species marginals and kinematics mean - the coupling was weak, which is the same conclusion by a different route. -3. **Estimator swap**, only if 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. **`giant/training/`** — per-stage trainers; `active = false` paths. At this point - Stage-2-only training works and the meeting's step 2 (one-shot WGAN baseline, - trained standalone) is runnable. -4. **Type 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. diff --git a/docs/v0.3.0-followups.md b/docs/v0.3.0-followups.md deleted file mode 100644 index 9a9caef..0000000 --- a/docs/v0.3.0-followups.md +++ /dev/null @@ -1,53 +0,0 @@ -# v0.3.0 — post-implementation audit: open discrepancies - -**Status:** steps 1–7 of `docs/v0.3.0-design.md` §12 are implemented (branch -`v0.3.0-stage2-autoregressive`, commits `eb6dd27`..`200c6d2`). This document -tracked discrepancies found between that implementation and the design contract -during a 2026-08-07 audit, as concrete work items. **All items (1-9) are now -resolved** — either implemented (1-5, 7-9) or explicitly deferred into -`docs/v0.3.0-design.md` §11.2 (6: `stage2_model.generator = "ddpm"`). Step 8 -(`estimate_batch_size` recalibration) was intentionally still outstanding per -§12 and was never tracked here. - ---- - -## Confirmed correct during the audit (no action needed) - -For reference — these were explicitly checked against the design doc and -match it, including two spots the doc itself flagged as likely stale that -turned out fine: - -- Config schema, `DEFAULT_CONFIG`, `migrate_config` table (§4), `save_config`/ - `merge_cli_overrides` recursion, `default_out_dir_name`, `Conditioning` enum, - `n_sec.mode = "stop_token"` error (§9, §11.2). -- `network.py`'s full class decomposition (§5.3), dict-returning - `build_models`/`build_critics` (§5.4), `ExpertTrunk` separate in/out dims, - ST-Gumbel wiring (§2.1), AR token layout (§6.1), Markov/Attention history - encoders (§6.2). -- Shared PDG top-N type map (one map, not two, per §8), `other_policy` - sample/modal/drop (§11.1), embedding L1-nearest decode, and the L1-distance - diagnostic surfaced in `giant analyze` (§11.3). -- `analysis/render.py`/`analysis/router_gating.py` correctly branch - old-flat vs new-nested `model_config["router"]` location — doc flagged this - as a likely stale spot (§10) but it's actually fine. -- `giant/training/`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence, - per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only - training via ground-truth `x1_s1` (§7). - -## Behavior change: WGAN best-checkpoint selection - -The `giant/training/` split fixed a dead guard in WGAN validation scoring. A -WGAN stage was *meant* to contribute its marginal KL to the `val_loss` that -drives `best.pt`, but the guard `if n not in val_loss_per_stage` could never -fire (every stage was pre-seeded to `0.0`), so the stage contributed a flat -`0.0` and the KL was written to `metrics.csv` without ever being used. -`WGANStageTrainer.val_objective` now returns the KL as intended. - -**Consequence:** any checkpoint selected before this commit under a config -with a WGAN stage — including the v0.3.0 default (`stage2_model.generator = -"wgan"`) — picked its best epoch on the non-adversarial stages alone. Measured -on the test harness's default flow+wgan config, `val_loss` went from `2.182` -(stage 1 only) to `15.137` (stage 1 + KL `12.954`), and which epoch won -changed. Do not compare `val/loss` or `best.pt` choice across this commit. -- `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init` - seeding, removed stale expert-size warning (§9, §10). diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index 8f5c751..a4a48d2 100644 --- a/giant/analysis/catalog.py +++ b/giant/analysis/catalog.py @@ -72,7 +72,7 @@ class Bundle: r_phys: pl.LazyFrame # rollout, physical steps only t_phys: pl.LazyFrame # reference, physical steps only checkpoint: str | None = None # from the rollout YAML; router_gating only - # §11.3 diagnostic pre-aggregated at rollout time (giant.rollout. + # Diagnostic pre-aggregated at rollout time (giant.rollout. # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance # only. Unlike checkpoint/router_gating, this needs no live model: it's # already a finished histogram, just passed through. diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index ec2d3e0..9cd1c64 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -83,7 +83,7 @@ _PLOT_META_KEYS = ( "best_val_loss", "training_config", "training_meta", - # §11.3 diagnostic — only present when giant rollout ran under + # Diagnostic — only present when giant rollout ran under # stage2_model.particle_type.target="embedding" (see giant/cli.py's # rollout command and giant.rollout.L1DistCollector); absent otherwise, # which the type_embedding_l1_distance PlotSpec (catalog.py) reads as diff --git a/giant/analysis/render.py b/giant/analysis/render.py index ec166ca..862e023 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -51,8 +51,8 @@ def _figure_params_v2(mc: dict, run_meta: dict) -> dict: """`_figure_params` for a new-shape (nested) `model_config` — has a `stage1_model` key. Reports stage 1's architecture (the headline generator); stage 2's generator is only added (`mode_s2`) when it - differs from stage 1's, since a mixed run (docs/v0.3.0-design.md's - `stage1=flow` + `stage2=wgan` case) is the interesting exception, not + differs from stage 1's, since a mixed run (the `stage1=flow` + + `stage2=wgan` case) is the interesting exception, not the common case.""" s1 = mc["stage1_model"] s2 = mc.get("stage2_model") or {} diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py index edfb826..491a433 100644 --- a/giant/analysis/router_gating.py +++ b/giant/analysis/router_gating.py @@ -71,8 +71,8 @@ def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str `giant.data.transforms.build_cond_features` — from either a v0.2 checkpoint's flat `model_config["conditioning"]` (one shared string, same for both axes) or a new-format one (independent - `model_config["conditioning"]["particle"/"material"]["type"]` — - docs/v0.3.0-design.md §3.1: the two axes may differ). Mirrors + `model_config["conditioning"]["particle"/"material"]["type"]` — the two + axes may differ). Mirrors `giant.cli._conditioning_axes`.""" raw = model_cfg.get("conditioning", default) if isinstance(raw, dict): diff --git a/giant/analysis/type_embedding_distance.py b/giant/analysis/type_embedding_distance.py index adb39c6..8304a2a 100644 --- a/giant/analysis/type_embedding_distance.py +++ b/giant/analysis/type_embedding_distance.py @@ -1,4 +1,4 @@ -"""Secondary-type embedding-distance diagnostic (docs/v0.3.0-design.md §11.3). +"""Secondary-type embedding-distance diagnostic. Unlike every other diagnostic in this package, the data isn't derivable from a rollout/reference parquet at all — it's the L1 distance between each diff --git a/giant/cli.py b/giant/cli.py index dcd2590..6a1170c 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -78,9 +78,8 @@ def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str `giant.data.transforms.build_cond_features`/`build_features` — from either a v0.2 checkpoint's flat `model_config["conditioning"]` (one shared string, same for both axes) or a new-format one (independent - `model_config["conditioning"]["particle"/"material"]["type"]` — - docs/v0.3.0-design.md §3.1: the two axes are configured independently and - may differ).""" + `model_config["conditioning"]["particle"/"material"]["type"]` — the two + axes are configured independently and may differ).""" raw = model_cfg.get("conditioning", default) if isinstance(raw, dict): return ( @@ -390,7 +389,7 @@ def train( typer.Option( "--stage1-generator", help="Stage 1's generative objective — overrides --mode for stage 1 " - "only (see docs/v0.3.0-design.md decision 7)", + "only", ), ] = None, stage1_hidden_dim: Annotated[ @@ -725,8 +724,7 @@ def train( # --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for # backward compatibility (they predate stage2_model having its own # flags); --stage1-*/--stage2-* below are the explicit, discoverable - # per-stage flags docs/v0.3.0-design.md decision 7 calls for, and take - # precedence when both are given. + # per-stage flags, and take precedence when both are given. cli_stage1_model: dict[str, object] = { k: v for k, v in { @@ -1360,7 +1358,7 @@ def predict( ) if coord == Coord.global_: - # A fresh v0.3.0 Stage1Model owns no n_sec_head (decision 1) — + # A fresh v0.3.0 Stage1Model owns no n_sec_head — # sample_stage1 returns n_sec_pred=None then, so ask stage 2. n_sec_pred = resolve_n_sec( model, sec_decoder, cc, ck, stage1_norm, n_sec_pred @@ -1764,7 +1762,7 @@ def rollout( writer = pq.ParquetWriter(out, table.schema) writer.write_table(table) - # Only meaningful under particle_type.target="embedding" (§11.3) — a + # Only meaningful under particle_type.target="embedding" — a # no-op collector otherwise, cheaper than branching the call itself. l1_dist_collector = L1DistCollector() @@ -1820,7 +1818,7 @@ def rollout( "rollout_seed": seed, "n_rows": summary["n_rows"], "termination_reason_counts": summary["termination_reason_counts"], - # §11.3 diagnostic — only present under + # Diagnostic — only present under # stage2_model.particle_type.target="embedding"; omitted (not # written as null) otherwise, so giant.analysis can tell "not # applicable to this checkpoint" apart from "collector empty". diff --git a/giant/config.py b/giant/config.py index 95fc1e6..c779969 100644 --- a/giant/config.py +++ b/giant/config.py @@ -325,8 +325,7 @@ def auto_device() -> torch.device: # batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this # is a rough estimate rather than a guaranteed bound. # NOTE: not yet recalibrated for the v0.3.0 autoregressive stage-2 trunk — -# see docs/v0.3.0-design.md §9/§11.3, deliberately last in the implementation -# order. +# deliberately last in the implementation order. _REF_BYTES = 7683 * 1024**2 _REF_BATCH_SIZE = 29696 _REF_HIDDEN_DIM = 1024 @@ -528,10 +527,9 @@ def migrate_config(cfg: dict) -> dict: Called on every config.toml load (see merge_cli_overrides) so old training configs on disk keep working under new code without hand- - editing (decision 3, docs/v0.3.0-design.md §4). `[meta].config_version == - CONFIG_VERSION` marks a dict as already-v0.3; its absence is read as - "this is v0.2" (the design doc's stated rule), so an already-migrated - dict is returned unchanged (deep-copied). + editing. `[meta].config_version == CONFIG_VERSION` marks a dict as + already-v0.3; its absence is read as "this is v0.2", so an + already-migrated dict is returned unchanged (deep-copied). Only keys actually present in `cfg` are translated — `cfg` may be a partial file (e.g. `[train]\\nepochs = 5\\n` with no [model] section at @@ -543,7 +541,7 @@ def migrate_config(cfg: dict) -> dict: to set. Operates on the config.toml shape. A checkpoint's `model_config` dict - (which additionally carries n_sec_head ownership, §4.1, and needs + (which additionally carries n_sec_head ownership and needs `network.build_models`'s cooperation) is a separate migration surface, deferred to the network.py refactor. """ @@ -729,7 +727,7 @@ def validate_config(cfg: dict) -> None: raise ValueError( "stage2_model.n_sec.mode = 'truth' is invalid for a " "rollout-capable checkpoint (both stage1_model.active and " - "stage2_model.active = true — see docs/v0.3.0-design.md §9): " + "stage2_model.active = true): " "'truth' takes n_sec from ground truth, which giant rollout " "doesn't have. 'truth' is for standalone stage-2 evaluation " "only — set stage1_model.active = false for that, or use " diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 6f5f216..13c78fc 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -44,7 +44,6 @@ class StreamingStepsDataset(IterableDataset): where: cond_cont: (B, COND_DIM) float32 cond_cat: (B, 2/3/4) int64 — width 2 unless conditioning="onehot" - (see docs/v0.3.0-design.md decision 4) target_s1: (B, 9) float32 — normalised Stage-1 primary target n_sec: (B,) int64 — true secondary count per step sec_cont: (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit, @@ -60,7 +59,7 @@ class StreamingStepsDataset(IterableDataset): ("onehot", "embedding"); zeros (unused) otherwise `k_max` (constructor arg, default the module constant) should match - `stage2_model.k_max` (docs/v0.3.0-design.md §8) — it sets the padded + `stage2_model.k_max` — it sets the padded width of `sec_cont`/`sec_type_idx` above. """ diff --git a/giant/data/loader.py b/giant/data/loader.py index e5a77bf..43a1a44 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -176,7 +176,7 @@ def iter_file_chunks( """Yield one parquet row-group at a time so a large file never fully loads. `k_max` sets the padded width of the sec_*_list columns (should match - `stage2_model.k_max` — see docs/v0.3.0-design.md §8); defaults to the + `stage2_model.k_max`); defaults to the module constant for callers that don't care (e.g. Stage-1-only reads).""" pf = pq.ParquetFile(path) for i in range(pf.num_row_groups): @@ -277,8 +277,7 @@ def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]: Returns `(class_map, other_members)` — `other_members` is `{key: count}` for every key bucketed into "other" (the empirical within-bucket - distribution, for `other_policy = "sample"` at rollout — see - docs/v0.3.0-design.md §8). + distribution, for `other_policy = "sample"` at rollout). """ ranked = sorted(counts, key=lambda k: counts[k], reverse=True) keep = ranked[: max(n_classes - 1, 0)] @@ -324,7 +323,7 @@ def build_topn_map_from_files( `_topn_plus_other_map`), generalized over the source column and key type. Used for the material axis (`column="material"`, `cast=str`, matching - `mat_map`'s key type) — see docs/v0.3.0-design.md §8. The PDG axis uses + `mat_map`'s key type). The PDG axis uses `build_pdg_topn_map_from_files` instead (it needs to pool two columns, which this single-column form can't express). Also records `other_members` (the empirical within-"other" distribution), needed @@ -341,8 +340,8 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap: plays in this dataset: a step's own primary particle (`pdg` column) and an emitted secondary's species (`sec_pdg_list`, exploded) — shared by `conditioning.particle.type = "onehot"` and - `stage2_model.particle_type.target = "onehot"` (docs/v0.3.0-design.md - §8). Pooling both is what keeps a species that's common as a secondary + `stage2_model.particle_type.target = "onehot"`. Pooling both is what + keeps a species that's common as a secondary but rare as a primary (or vice versa) from being pushed into "other" just because one role's count alone looks small — the meeting's failure mode (zero photon secondaries, hallucinated antineutrinos) was diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index ce38def..a0a3981 100644 --- a/giant/data/setup_cache.py +++ b/giant/data/setup_cache.py @@ -104,7 +104,7 @@ def normalizer_key( ) -> str: # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing # spurious cache misses between runs with the "same" val_fraction. The two - # conditioning axes are independent (docs/v0.3.0-design.md §3.1) and both + # conditioning axes are independent and both # affect which cond_cont columns are computed for real vs. zero-filled # (giant.data.transforms._physical_cond_columns), so both must be part of # the key or two mixed-axis runs could collide on the same cache entry. @@ -114,7 +114,7 @@ def normalizer_key( ) -# Top-N-map axes (docs/v0.3.0-design.md §8): "pdg" keys match pdg_map's int +# Top-N-map axes: "pdg" keys match pdg_map's int # keys (shared by conditioning.particle.type="onehot" and # stage2_model.particle_type.target="onehot" — one map for both), "material" # keys match mat_map's str keys. @@ -190,7 +190,7 @@ class SetupCache: proc_maps: dict[int, dict[str, int]] = field(default_factory=dict) normalizers: dict[str, NormalizerEntry] = field(default_factory=dict) topn_maps: dict[str, TopNMap] = field(default_factory=dict) - """Keyed by `topn_key(axis, n_classes)` — see docs/v0.3.0-design.md §8.""" + """Keyed by `topn_key(axis, n_classes)`.""" @classmethod def empty(cls, files: list[Path]) -> "SetupCache": diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 555ed22..bfe7234 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -561,9 +561,9 @@ def encode_secondary_type_idx( `class_map` is either a top-N-plus-other map's `class_map` (`stage2_model.particle_type.target = "onehot"`, see `giant.data.loader.build_pdg_topn_map_from_files`) or the dense `pdg_map` - (`target = "embedding"`). Not used at all for `target = "physical"` - (see docs/v0.3.0-design.md decision 1) — that target keeps using - `encode_secondaries`'s (log_mass, charge) columns unchanged. + (`target = "embedding"`). Not used at all for `target = "physical"` — + that target keeps using `encode_secondaries`'s (log_mass, charge) + columns unchanged. Padding slots get index 0 (their looked-up value is discarded downstream by the `sec_valid`/`n_sec` mask regardless, so any in-vocabulary dummy @@ -599,7 +599,7 @@ def decode_secondary_cont( stick-breaking energy split and local->world direction — generator/ `particle_type.target`-independent, since every target (`"physical"`, `"onehot"`, `"embedding"`) shares the same `CONT_SLOT_DIM`-wide - (stick_logit, dir) prefix (docs/v0.3.0-design.md §6.1) and differs only + (stick_logit, dir) prefix and differs only in what follows it. `decode_secondaries` (target="physical") is the original all-in-one form built on top of this; `target` in `("onehot", "embedding")` decodes their type slice separately via @@ -727,9 +727,8 @@ def _physical_cond_columns( ) -> np.ndarray: """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns. - The particle and material blocks are gated independently - (docs/v0.3.0-design.md §3.1: "configured independently and may mix - freely — e.g. material `physical` with particle `embedding`"), so e.g. + The particle and material blocks are gated independently and may mix + freely — e.g. material `physical` with particle `embedding` — so e.g. `particle_conditioning="embedding"` + `material_conditioning="physical"` zero-fills only the particle columns and computes the material ones for real. @@ -802,16 +801,15 @@ def build_cond_features( ) -> tuple[np.ndarray, np.ndarray]: """Build conditioning arrays only — no target, no post-step variables. - `particle_conditioning`/`material_conditioning` are independent - (docs/v0.3.0-design.md §3.1) — e.g. `particle_conditioning="embedding"` + + `particle_conditioning`/`material_conditioning` are independent — + e.g. `particle_conditioning="embedding"` + `material_conditioning="physical"` is a valid mix. `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat` - columns read by `ConditionEncoder`'s `"onehot"` mode - (docs/v0.3.0-design.md decision 4): pdg topN index at column 2 (iff - `pdg_topn_map` given), material topN index at column 3 (iff - `mat_topn_map` given, after column 2 if both are). Only ever given when + columns read by `ConditionEncoder`'s `"onehot"` mode: pdg topN index at + column 2 (iff `pdg_topn_map` given), material topN index at column 3 + (iff `mat_topn_map` given, after column 2 if both are). Only ever given when the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise. """ cond_cont = np.column_stack( @@ -936,9 +934,8 @@ def build_features( the secondary's real physical identity (from its ground-truth PDG code), a fixed regression target, not a learned/snapped one. Always computed the same way regardless of - `stage2_model.particle_type.target` (docs/v0.3.0-design.md - decision 1/3) — only actually used downstream under `target = - "physical"`. + `stage2_model.particle_type.target` — only actually used + downstream under `target = "physical"`. proc_idx: (N,) integer process-class label (ProcessRouter supervision only — never conditioning). Zeros when `proc_map` is None or the loaded data has no "process" column (e.g. pre-conversion parquet files). @@ -966,7 +963,7 @@ def build_features( dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself). `None` for `target = "physical"`. - k_max: should match `stage2_model.k_max` (docs/v0.3.0-design.md §8) — + k_max: should match `stage2_model.k_max` — overridden internally by `data["sec_E_list"]`'s own padded width when present (the loader already padded it to some k_max; that width is authoritative), so this only actually matters when secondary list @@ -1015,7 +1012,7 @@ def build_features( cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map)) if mat_topn_map is not None: cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map)) - cond_cat = np.column_stack(cat_cols) # (N, 2/3/4) — see decision 4 + cond_cat = np.column_stack(cat_cols) # (N, 2/3/4) n_sec_raw = data["n_sec"].astype( np.int64 @@ -1029,7 +1026,7 @@ def build_features( # The loader already padded sec_*_list to some k_max (see # giant.data.loader.iter_file_chunks); that padded width is # authoritative over whatever this call happened to pass in, so the - # two can never drift apart (docs/v0.3.0-design.md §8). + # two can never drift apart. k_max = sec_E_list.shape[1] # Clamp the classification label to k_max: the head only has k_max+1 diff --git a/giant/model/network.py b/giant/model/network.py index ce91890..8aeadfb 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -21,7 +21,7 @@ from giant.constants import ( ) # --------------------------------------------------------------------------- -# Building blocks (docs/v0.3.0-design.md §5.2/§5.3) +# Building blocks # --------------------------------------------------------------------------- @@ -47,7 +47,7 @@ def cat_col_layout( particle_type: str, material_type: str ) -> tuple[int | None, int | None]: """`cond_cat` column indices for each axis's top-N-onehot index, or - `None` if that axis isn't `"onehot"` (docs/v0.3.0-design.md decision 4). + `None` if that axis isn't `"onehot"`. Columns 0/1 are always the dense pdg/material vocab index. The particle top-N column (if any) comes next, then the material top-N column (if @@ -92,9 +92,9 @@ class ConditionEncoder(nn.Module): """Fuses continuous conditioning with particle/material identity. The particle and material axes are configured independently - (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}` — - see docs/v0.3.0-design.md §3.1) and may mix freely, e.g. material - "physical" with particle "embedding". Three modes per axis: + (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`) + and may mix freely, e.g. material "physical" with particle "embedding". + Three modes per axis: - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s dense training-vocab index. Memorizes the training menu. - "physical": an `n_layers`-deep MLP over the axis's raw physical @@ -104,8 +104,7 @@ class ConditionEncoder(nn.Module): - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other class map (`giant.data.loader.build_topn_map_from_files`/ `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra - top-N-index column(s) — see `_cat_col_layout` and - docs/v0.3.0-design.md decision 4. + top-N-index column(s) — see `_cat_col_layout`. """ def __init__( @@ -193,7 +192,7 @@ class ContextAdapter(nn.Module): fixed-width context vector for a downstream stage's conditioning — `stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj` (+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since - `SecondaryConditionEncoder` as a wrapper class disappears (design doc §5.2).""" + `SecondaryConditionEncoder` as a wrapper class disappears.""" def __init__(self, in_dim: int, context_dim: int) -> None: super().__init__() @@ -223,7 +222,7 @@ class ResBlock(nn.Module): # --------------------------------------------------------------------------- -# Routers — carried over unchanged from v0.2 (docs/v0.3.0-design.md §5.3) +# Routers — carried over unchanged from v0.2 # --------------------------------------------------------------------------- @@ -615,7 +614,7 @@ def _build_router_from_cfg( # --------------------------------------------------------------------------- -# Trunks (docs/v0.3.0-design.md §5.2 (b)) +# Trunks # --------------------------------------------------------------------------- @@ -783,7 +782,7 @@ def build_trunk( # --------------------------------------------------------------------------- -# History encoders — stage-2 autoregressive only (docs/v0.3.0-design.md §6.2) +# History encoders — stage-2 autoregressive only # --------------------------------------------------------------------------- @@ -791,8 +790,8 @@ class HistoryEncoder(nn.Module): """Interface for stage-2 autoregressive per-token history summaries: `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over a full (teacher-forced) token sequence — used by training. `MarkovHistory` - and `AttentionHistory` (docs/v0.3.0-design.md §6.2) are the two - implementations. Inference (`giant/sample.py`) generates one token at a + and `AttentionHistory` are the two implementations. Inference + (`giant/sample.py`) generates one token at a time and cannot afford `forward`'s per-step cost to be O(K) (attention would then be O(K^2) over a rollout's k_max loop); encoders that need incremental state for that path additionally implement `init_cache`/ @@ -806,13 +805,13 @@ class HistoryEncoder(nn.Module): class MarkovHistory(HistoryEncoder): """Summarizes the previous secondary's own `(energy_fraction, direction, - type_representation)` through one small MLP — the "markov" history - (docs/v0.3.0-design.md §6.2): token i+1 only ever sees token i plus the - running scalars (`remaining_frac`/`slot_idx`, fused in separately by + type_representation)` through one small MLP — the "markov" history: + token i+1 only ever sees token i plus the running scalars + (`remaining_frac`/`slot_idx`, fused in separately by `Stage2Autoregressive._token_cond`), not the full prefix. At slot 0 (`has_prev` False) substitutes a learned start vector rather - than zeros — a reasonable default, not specified by the design doc. + than zeros — a reasonable default. """ def __init__(self, in_dim: int, out_dim: int) -> None: @@ -879,9 +878,9 @@ class _CausalAttnBlock(nn.Module): class AttentionHistory(HistoryEncoder): - """Causal self-attention over the emitted-token prefix - (docs/v0.3.0-design.md §6.2) — the more expressive alternative to - `MarkovHistory`'s fixed previous-token-only summary. `feat`/`has_prev` + """Causal self-attention over the emitted-token prefix — the more + expressive alternative to `MarkovHistory`'s fixed previous-token-only + summary. `feat`/`has_prev` follow the same shifted-by-one convention `MarkovHistory` and `Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s own `(energy_fraction, direction, type_representation)`, with a learned @@ -893,9 +892,9 @@ class AttentionHistory(HistoryEncoder): `forward` is the parallel training path (one pass over the whole teacher-forced sequence); `init_cache`/`step` are the incremental inference path `giant/sample.py` uses, one new token per call, to avoid - re-encoding the whole prefix from scratch every slot (docs/v0.3.0-design.md - §10's "KV cache" note) — `step` must be called exactly once per slot (its - cache-extension is not idempotent), so a slot's output must be reused for + re-encoding the whole prefix from scratch every slot — `step` must be + called exactly once per slot (its cache-extension is not idempotent), + so a slot's output must be reused for every model call within that slot (`forward`'s ODE substeps, or a separate `predict_type` call) rather than re-derived — see `Stage2Autoregressive.history_step`. @@ -947,7 +946,7 @@ class AttentionHistory(HistoryEncoder): # --------------------------------------------------------------------------- -# Stage models (docs/v0.3.0-design.md §5.3) +# Stage models # --------------------------------------------------------------------------- @@ -955,7 +954,7 @@ def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int: """Width of a single secondary slot's type slice — `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are - `conditioning.particle.emb_dim` wide — docs/v0.3.0-design.md §3.3).""" + `conditioning.particle.emb_dim` wide).""" target = particle_type_cfg.get("target", "physical") return PARTICLE_PHYS_DIM if target == "physical" else emb_dim @@ -963,15 +962,15 @@ def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int: def stage2_trunk_sec_dim( particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int ) -> int: - """`Stage2OneShot`'s trunk output width (docs/v0.3.0-design.md decision 2). + """`Stage2OneShot`'s trunk output width. - `target = "physical"` is untouched from v0.2/today (decision 1): + `target = "physical"` is untouched from v0.2/today: `k_max * SEC_SLOT_DIM`, the type slice folded into the same flow-matched/WGAN vector as the continuous stick/dir slots. `target` in `("onehot", "embedding")`: under `generator == "wgan"` the type slice is still folded in (adversarial for onehot via ST-Gumbel, - already-continuous for embedding — §2.1), just `emb_dim` wide instead of + already-continuous for embedding), just `emb_dim` wide instead of `PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under `generator in ("flow", "ddpm")` the type slice isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk @@ -986,14 +985,14 @@ def stage2_trunk_sec_dim( class Stage1Model(nn.Module): - """Predicts the 9D primary post-step vector. No `n_sec_head` — decision 1 - (docs/v0.3.0-design.md §2) moves it to stage 2, except for a migrated - v0.2 checkpoint (`n_sec_head_k_max` given), where it stays attached here + """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs + move it to stage 2, except for a migrated v0.2 checkpoint + (`n_sec_head_k_max` given), where it stays attached here since that's where its weights live and what conditioning it was trained against (see `_migrate_legacy_model_config`). `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — `conditioning.share_stages = true` (§3.1): `build_models` + `ConditionEncoder` — `conditioning.share_stages = true`: `build_models` constructs one shared instance and passes it to both stages, halving the conditioning parameter count and forcing a common representation.""" @@ -1060,11 +1059,11 @@ class Stage1Model(nn.Module): ) -> torch.Tensor: """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0 - configs predict n_sec from Stage2OneShot instead (decision 1).""" + configs predict n_sec from Stage2OneShot instead.""" if self.n_sec_head is None: raise RuntimeError( "this Stage1Model has no n_sec_head — n_sec now lives on " - "stage 2 by default (decision 1); this method only exists " + "stage 2 by default; this method only exists " "for a migrated v0.2 checkpoint (legacy_owner='stage1')" ) c_emb = self.cond_enc(cond_cont, cond_cat) @@ -1073,17 +1072,16 @@ class Stage1Model(nn.Module): class Stage2OneShot(nn.Module): """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour, - reproduced exactly (see docs/v0.3.0-design.md §12 step 2 acceptance - criterion; `decoder = "autoregressive"` is `Stage2Autoregressive`, + reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`, step 4/5, not implemented yet). - Owns `n_sec_head` by default (decision 1) unless `build_n_sec_head=False` + Owns `n_sec_head` by default unless `build_n_sec_head=False` (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to Stage1Model — see `_migrate_legacy_model_config`). `particle_type_cfg["target"]` (default `"physical"`) selects the - secondary-type mechanism (docs/v0.3.0-design.md decision 2): - `"physical"` keeps the type slice folded into the trunk's own + secondary-type mechanism: `"physical"` keeps the type slice folded into + the trunk's own flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by the caller via `stage2_trunk_sec_dim` — already reflects this). Under `"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type @@ -1225,8 +1223,8 @@ class Stage2OneShot(nn.Module): class Stage2Autoregressive(nn.Module): - """Emits secondaries one at a time in descending-energy order - (docs/v0.3.0-design.md §6), instead of `Stage2OneShot`'s simultaneous + """Emits secondaries one at a time in descending-energy order, instead + of `Stage2OneShot`'s simultaneous k_max-slot prediction. `history` selects `MarkovHistory` or `AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only). `teacher_forcing` handling lives entirely in the trainer @@ -1555,7 +1553,7 @@ class CriticModel(nn.Module): # --------------------------------------------------------------------------- -# v0.2 -> v0.3 checkpoint migration (docs/v0.3.0-design.md §4.1, §4.3) +# v0.2 -> v0.3 checkpoint migration # --------------------------------------------------------------------------- @@ -1568,14 +1566,14 @@ def _migrate_legacy_model_config(model_config: dict) -> dict: Sets `stage2_model.n_sec.legacy_owner = "stage1"` so the n_sec_head weights a v0.2 checkpoint carries on its Stage-1 module keep loading - there (design doc §4.1) instead of the new default location + there instead of the new default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not just be labeled as such. Only the monolithic (non-routed) trunk shape is exercised by the step-2 - migration test (docs/v0.3.0-design.md §4.3); a routed v0.2 checkpoint - still builds correctly here (the router config passes through), but its + migration test; a routed v0.2 checkpoint still builds correctly here + (the router config passes through), but its state dict isn't covered by `migrate_legacy_state_dict` below. """ m = model_config @@ -1599,7 +1597,7 @@ def _migrate_legacy_model_config(model_config: dict) -> dict: "hidden_dim/n_res_blocks), so this checkpoint's routed experts " "have a different width/depth than the monolith — silently " "dropping these keys would resize the experts instead of " - "refusing (docs/v0.3.0-design.md §4.2). This checkpoint can " + "refusing. This checkpoint can " "only be loaded by v0.2 code." ) router_cfg.setdefault("enabled", False) @@ -1651,8 +1649,7 @@ def migrate_legacy_state_dict( `(Stage1Model, Stage2OneShot)` module structure produced by `build_models(_migrate_legacy_model_config(model_config))`. - Only the monolithic (non-routed) trunk shape is handled — see - docs/v0.3.0-design.md §4.3's migration test scope. + Only the monolithic (non-routed) trunk shape is handled. """ def _trunk_prefix(k: str) -> str: @@ -1682,7 +1679,7 @@ def migrate_legacy_state_dict( # --------------------------------------------------------------------------- -# Factories (docs/v0.3.0-design.md §5.4) +# Factories # --------------------------------------------------------------------------- @@ -1698,7 +1695,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: (v0.2's actual — probably accidental — behaviour: two routers built from one config with no semantic relationship between them). - `conditioning.share_stages = true` (§3.1) builds one `ConditionEncoder` + `conditioning.share_stages = true` builds one `ConditionEncoder` instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/ `Stage2Autoregressive`'s `cond_enc` param), instead of each stage building its own — halving the conditioning parameter count and forcing a diff --git a/giant/model/schedule.py b/giant/model/schedule.py index 27d8282..b8b3891 100644 --- a/giant/model/schedule.py +++ b/giant/model/schedule.py @@ -90,8 +90,8 @@ def flow_matching_loss_secondary( "physical"`'s width and the only case this function handled before v0.3.0 step 4. `0` means no type slice is in `x1` at all (`target` in `("onehot", "embedding")` under `generator in ("flow", "ddpm")` — - see docs/v0.3.0-design.md decision 2, `Stage2OneShot.type_head` - handles the type loss separately in that case). + `Stage2OneShot.type_head` handles the type loss separately in that + case). Only valid-slot dimensions contribute to the loss; padded slots are zeroed before averaging, so the loss is not diluted by empty slots. @@ -156,7 +156,7 @@ def flow_matching_loss_secondary_ar( model call signature differs enough (four extra per-token conditioning tensors) that merging would need an awkward shape-flag + closure. - Under teacher forcing (docs/v0.3.0-design.md §6.2 point 3) this is still a + Under teacher forcing this is still a single parallel pass over all K_MAX tokens — `x1`/`history_feat`/etc. are already built from ground truth for every slot by the caller (`giant.training.stage2_inputs._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`). diff --git a/giant/particles.py b/giant/particles.py index 7e2d5ac..c6be4e9 100644 --- a/giant/particles.py +++ b/giant/particles.py @@ -8,8 +8,8 @@ nuclide table doesn't cover — confirmed necessary for ~32% of the nuclear code actually present in the multi-material dataset (`0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet`). -Also holds the v0.3.0 stage-2 categorical-type rollout decode (§3.3/§8/§11.3 -of docs/v0.3.0-design.md): `decode_topn_class`/`decode_embedding_nearest` turn +Also holds the v0.3.0 stage-2 categorical-type rollout decode: +`decode_topn_class`/`decode_embedding_nearest` turn `Stage2Autoregressive`/`Stage2OneShot`'s `"onehot"`/`"embedding"` type predictions back into concrete PDG codes, the one place a secondary's categorical/continuous type representation is ever discretized (its @@ -145,8 +145,8 @@ def decode_topn_class( rng: np.random.Generator | None = None, ) -> np.ndarray: """`conditioning.particle.type` / `stage2_model.particle_type.target = - "onehot"` inference decode (docs/v0.3.0-design.md §3.3): per-row top-N - class index -> concrete PDG code. + "onehot"` inference decode: per-row top-N class index -> concrete PDG + code. class_idx: int array, any shape, values in `[0, n_classes)`. topn_map: the `TopNMap` (`giant.data.loader.build_pdg_topn_map_from_files`) @@ -205,9 +205,9 @@ def decode_embedding_nearest( emb_weight: np.ndarray, idx_to_pdg: dict[int, int], ) -> tuple[np.ndarray, np.ndarray]: - """`stage2_model.particle_type.target = "embedding"` inference decode - (docs/v0.3.0-design.md §3.3): L1-nearest row of the conditioning's own - particle embedding table, since a generative model's continuous output + """`stage2_model.particle_type.target = "embedding"` inference decode: + L1-nearest row of the conditioning's own particle embedding table, since + a generative model's continuous output essentially never lands within float tolerance of a table row (the exact- match form is only valid as a round-trip test assertion, never here). @@ -220,9 +220,9 @@ def decode_embedding_nearest( idx_to_pdg: `invert_dense_map(pdg_map)` — embedding row index -> PDG. Returns `(pdg, l1_dist)`, both shaped like `vectors.shape[:-1]`. `l1_dist` - is the §11.3 diagnostic: a heavy tail means the decoder is emitting - vectors off the embedding manifold, the direct analogue of the species- - collapse symptom this redesign exists to fix. + is a diagnostic: a heavy tail means the decoder is emitting vectors off + the embedding manifold, the direct analogue of the species-collapse + symptom this redesign exists to fix. """ emb_dim = vectors.shape[-1] flat = np.asarray(vectors, dtype=np.float64).reshape(-1, emb_dim) diff --git a/giant/pipeline.py b/giant/pipeline.py index c427a1b..843140d 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -68,7 +68,7 @@ def _seed_energy_router( """Mutate `router_cfg["centers_init"]` in place from real data quantiles, when this stage's router is an enabled EnergyRouter. Shared by both stages' router configs — each seeded independently, since v0.3.0 stages - may have entirely different router configs (see docs/v0.3.0-design.md).""" + may have entirely different router configs.""" active = router_cfg.get("enabled") and router_cfg.get("type") == "energy" if not active: return @@ -107,7 +107,7 @@ def run_setup_stage( `cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/ `stage2_model`), already passed through `giant.config.validate_config`. `conditioning.particle.type` and `conditioning.material.type` are - independent (docs/v0.3.0-design.md §3.1) and may differ. + independent and may differ. Reads from and writes to the `giant.data.setup_cache` sidecar when `cache_setup` is set (`rebuild_setup_cache` ignores — but still @@ -197,8 +197,8 @@ def run_setup_stage( if cache is not None: cache.proc_maps[n_experts] = proc_map - # Top-N-plus-other maps for onehot conditioning/type axes - # (docs/v0.3.0-design.md §8). The PDG axis is shared by + # Top-N-plus-other maps for onehot conditioning/type axes. + # The PDG axis is shared by # conditioning.particle.type="onehot" and # stage2_model.particle_type.target="onehot" (both key off # conditioning.particle.emb_dim), so at most one PDG scan is needed even @@ -403,11 +403,11 @@ def run_train_job( setup.n_train_steps, ) - # cond_cat's onehot columns (docs/v0.3.0-design.md decision 4) are - # present per-axis, independently, under that axis's own - # conditioning.{particle,material}.type == "onehot" (§3.1: the two axes - # may mix freely). run_setup_stage builds each map whenever its own axis - # is "onehot" (see its own particle_cfg["type"]/material_cfg["type"] + # cond_cat's onehot columns are present per-axis, independently, under + # that axis's own conditioning.{particle,material}.type == "onehot" + # (the two axes may mix freely). run_setup_stage builds each map + # whenever its own axis is "onehot" (see its own + # particle_cfg["type"]/material_cfg["type"] # checks), so they're guaranteed non-None here — asserted, not just # assumed, so a future wiring bug fails loudly instead of silently # dropping the onehot columns. @@ -422,7 +422,7 @@ def run_train_job( # The secondary type-index map depends on stage2_model.particle_type.target, # independently of conditioning's own onehot/embedding choice above - # (docs/v0.3.0-design.md §3.3 — physical stays untouched/None). + # (physical stays untouched/None). particle_type_target = ( cfg["stage2_model"].get("particle_type", {}).get("target", "physical") ) diff --git a/giant/rollout.py b/giant/rollout.py index c38edd4..bb6b9fd 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -53,14 +53,14 @@ if TYPE_CHECKING: class L1DistCollector: - """Accumulates the §11.3 L1-distance diagnostic across a whole rollout + """Accumulates the L1-distance diagnostic across a whole rollout run: the L1 distance between each emitted secondary's raw predicted embedding vector and the nearest table row it snapped to (only meaningful under `particle_type.target = "embedding"` — `giant.particles.decode_embedding_nearest`). A heavy tail means the decoder is emitting vectors off the embedding manifold — the direct analogue of the species-collapse symptom the v0.3.0 redesign exists to - fix (docs/v0.3.0-design.md §11.3). + fix. Not folded into `rollout()`'s own return value (which is shape-typed as step records, see `_RECORD_KEYS`/`RolloutSummary`) — passed in and read @@ -124,8 +124,7 @@ def decode_secondary_identity( np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None ]: """Decode Stage 2's raw (sec_cont, sec_type) output into physical - secondary attributes, branching on `sec_decoder.particle_type_cfg` - (docs/v0.3.0-design.md §3.3): + secondary attributes, branching on `sec_decoder.particle_type_cfg`: - `"physical"`: unchanged v0.2 path — `sec_type` already *is* (log_mass, charge), used as the secondary's identity as-is (no snapping). @@ -138,8 +137,8 @@ def decode_secondary_identity( - `"embedding"`: `sec_type` is a raw vector in the conditioning's own embedding space — `giant.particles.decode_embedding_nearest` L1-snaps it to the nearest table row for the PDG (+ physics via - `particle_phys_array`), and also returns the L1 distance (§11.3 - diagnostic — see `giant/rollout.py`'s L1-distance accumulator). + `particle_phys_array`), and also returns the L1 distance (see this + module's `L1DistCollector`). Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, sec_type_l1_dist) — the last is `None` except under `"embedding"`. @@ -480,24 +479,22 @@ def rollout( `n_events * max_steps * avg_tracks_per_event`. There is no `mode` parameter — each stage's generative objective is read - directly off the model instance's own `generator_kind` - (docs/v0.3.0-design.md decision 2: stage 1 and stage 2 objectives are - independent, e.g. `stage1_model.generator="flow"` + - `stage2_model.generator="wgan"`), and the decoder (one-shot vs + directly off the model instance's own `generator_kind` (stage 1 and + stage 2 objectives are independent, e.g. `stage1_model.generator="flow"` + + `stage2_model.generator="wgan"`), and the decoder (one-shot vs autoregressive) is inferred from `sec_decoder`'s own class — see `sample_stage1`/`sample_stage2` (giant.sample). `pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen - to share `pdg_topn_map` (docs/v0.3.0-design.md §8 — one PDG map, not - two): they're required whenever `particle_conditioning`/ - `material_conditioning` is `"onehot"` (feeds `build_cond_features`'s - extra `cond_cat` top-N columns — §3.1), and `pdg_topn_map`/`other_policy` - are additionally read under `stage2_model.particle_type.target = - "onehot"` (§3.3, secondary-species decode). `seed` seeds the - `other_policy = "sample"` draw only (torch/numpy sampling itself is - seeded by the caller, same as today). + to share `pdg_topn_map` (one PDG map, not two): they're required + whenever `particle_conditioning`/`material_conditioning` is `"onehot"` + (feeds `build_cond_features`'s extra `cond_cat` top-N columns), and + `pdg_topn_map`/`other_policy` are additionally read under + `stage2_model.particle_type.target = "onehot"` (secondary-species + decode). `seed` seeds the `other_policy = "sample"` draw only + (torch/numpy sampling itself is seeded by the caller, same as today). - `l1_dist_collector`, if given, accumulates the §11.3 embedding-distance + `l1_dist_collector`, if given, accumulates the embedding-distance diagnostic across the whole run — see `L1DistCollector`. Only populated under `particle_type.target = "embedding"`; a no-op otherwise. """ diff --git a/giant/sample.py b/giant/sample.py index e5fb65f..3683384 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -10,10 +10,9 @@ def _predict_n_sec_if_owned( model: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor ) -> torch.Tensor | None: """Stage-1 `n_sec_head` is only present on a migrated v0.2 checkpoint - (docs/v0.3.0-design.md decision 1 moves it to stage 2 for fresh runs — - see `Stage1Model`'s docstring). `None` here means "ask stage 2 instead", - which every caller (`giant/rollout.py`, `giant/cli.py`) must do for a - fresh checkpoint.""" + (fresh runs move it to stage 2 — see `Stage1Model`'s docstring). `None` + here means "ask stage 2 instead", which every caller (`giant/rollout.py`, + `giant/cli.py`) must do for a fresh checkpoint.""" if getattr(model, "n_sec_head", None) is None: return None logits = model.predict_n_sec(cond_cont, cond_cat) @@ -124,7 +123,7 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int: (continuous + type) under `particle_type.target = "physical"` or `generator = "wgan"`, continuous-only otherwise (the type slice then comes from `predict_type` instead — see `stage2_trunk_sec_dim`'s - docstring, docs/v0.3.0-design.md decision 2).""" + docstring).""" return stage2_trunk_sec_dim( sec_decoder.particle_type_cfg, sec_decoder.generator_kind, @@ -147,9 +146,9 @@ def _decode_stage2_flat( n_sec_pred: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Reshape a flat `(B, flat_width)` `Stage2OneShot` output into per-slot - tensors, generator/`particle_type.target`-agnostic (docs/v0.3.0-design.md - decision 2/3): shared by `sample_secondaries`/`sample_secondaries_wgan`, - which differ only in how `x` was produced. + tensors, generator/`particle_type.target`-agnostic: shared by + `sample_secondaries`/`sample_secondaries_wgan`, which differ only in how + `x` was produced. Returns (sec_cont, sec_type, sec_valid): sec_cont: (B, k_max, CONT_SLOT_DIM) — [stick_logit, local_dir] @@ -241,23 +240,22 @@ def sample_secondaries_ar( n_sec_pred: torch.Tensor, steps: int = 10, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """`Stage2Autoregressive` inference loop (docs/v0.3.0-design.md §6.4): - one token at a time, in descending-energy slot order, `k_max` sequential - calls. Unlike training (teacher forcing, §6.2 point 3 — a single - parallel pass over ground-truth tokens, see - `giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no ground truth at - inference: each token's conditioning is built free-running, from the - PREVIOUS TOKEN'S OWN just-generated output — the train/inference gap - §6.2 point 4 explicitly flags as the cost of markov history's + """`Stage2Autoregressive` inference loop: one token at a time, in + descending-energy slot order, `k_max` sequential calls. Unlike training + (teacher forcing — a single parallel pass over ground-truth tokens, see + `giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no + ground truth at inference: each token's conditioning is built + free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the + train/inference gap that is the cost of markov history's expressiveness. A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass — - §6.4's "K sequential forwards" cost note applies per-token here, not + the "K sequential forwards" cost applies per-token here, not once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per physics step. Under `history="attention"` the history encoding is computed once per - slot via `Stage2Autoregressive.history_step` (a KV-cache append, §10) + slot via `Stage2Autoregressive.history_step` (a KV-cache append) rather than re-derived by every model call inside that slot — so an ODE loop's `steps` substeps, and the separate `predict_type` call when the type slice isn't folded into the trunk output, all reuse the SAME `hist` @@ -389,9 +387,9 @@ def sample_secondaries_ar( # --------------------------------------------------------------------------- # Per-stage dispatch — shared by giant/rollout.py and giant/cli.py's # `predict` command, since both need "given a stage model, produce a -# sample" without hand-picking the sampler themselves (docs/v0.3.0-design.md -# decision 2: each stage's generative objective is independent, read off the -# model's own `generator_kind`, not a caller-supplied `mode` string). +# sample" without hand-picking the sampler themselves (each stage's +# generative objective is independent, read off the model's own +# `generator_kind`, not a caller-supplied `mode` string). # --------------------------------------------------------------------------- @@ -450,19 +448,17 @@ def resolve_n_sec( ) -> torch.Tensor: """`n_sec_pred` is already populated when `stage1_model` owns a legacy `n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s - docstring); otherwise ask stage 2, which owns it by default under - decision 1 (docs/v0.3.0-design.md §2). Raises if neither stage owns a - head at all — the only way that happens is `stage2_model.n_sec.mode` - other than `"head"` (`"truth"`/`"stop_token"`), neither of which is a - valid rollout-/predict-capable checkpoint (§3.3, §9).""" + docstring); otherwise ask stage 2, which owns it by default. Raises if + neither stage owns a head at all — the only way that happens is + `stage2_model.n_sec.mode` other than `"head"` (`"truth"`/`"stop_token"`), + neither of which is a valid rollout-/predict-capable checkpoint.""" if n_sec_pred is not None: return n_sec_pred if getattr(sec_decoder, "n_sec_head", None) is None: raise RuntimeError( "checkpoint has no n_sec_head on either stage — needs " "stage2_model.n_sec.mode = 'head' (the default); 'truth' is " - "standalone-evaluation-only and 'stop_token' isn't implemented " - "(docs/v0.3.0-design.md §3.3/§9)" + "standalone-evaluation-only and 'stop_token' isn't implemented" ) logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out) return logits.argmax(dim=-1) diff --git a/giant/training/loop.py b/giant/training/loop.py index 39719cc..b2cb192 100644 --- a/giant/training/loop.py +++ b/giant/training/loop.py @@ -72,9 +72,8 @@ class _GracefulShutdown: def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs): """Runs `validate_marginals` on `trainer`'s sampling model (EMA model if present, else the raw model). `validate_marginals` itself dispatches - through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec` - (docs/v0.3.0-design.md §10), so this is generator- and - one-shot-vs-autoregressive-agnostic.""" + through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec`, so + this is generator- and one-shot-vs-autoregressive-agnostic.""" model = trainer.sampling_model() return validate_marginals(model, val_loader, device=device, **kwargs) @@ -124,7 +123,7 @@ def train( wandb_log_every: int = 50, ) -> None: """Train whichever of stage1/stage2 are active, each through its own - `StageTrainer` (design doc §7). `models`/`critics` are the dicts + `StageTrainer`. `models`/`critics` are the dicts `giant.model.network.build_models`/`build_critics` return — a `None` entry means that stage is `active = false`. """ diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py index 2508f5d..b23ac5d 100644 --- a/giant/training/stage2_inputs.py +++ b/giant/training/stage2_inputs.py @@ -4,7 +4,7 @@ Pure functions, no optimizer/model state: they turn a batch's ground-truth secondary tensors into the per-token targets and autoregressive conditioning inputs `giant.training.trainers` feeds to `Stage2OneShot` / `Stage2Autoregressive`. Split out of the trainers so the (target, generator, -decoder) width rules — the fiddliest part of docs/v0.3.0-design.md §2.1/§6 — +decoder) width rules — the fiddliest part of this codebase — live in one place and stay unit-testable on their own. """ @@ -65,8 +65,7 @@ def _assemble_stage2_ar_target( analogue of `_assemble_stage2_real` (defined below in terms of this), matching whatever width `Stage2Autoregressive`'s (or `Stage2OneShot`'s) own trunk produces for this (target, generator) combination - (`giant.model.network.stage2_trunk_sec_dim`; docs/v0.3.0-design.md - decision 2/3): + (`giant.model.network.stage2_trunk_sec_dim`): - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit, dir, log_mass, charge) as-is. @@ -116,7 +115,7 @@ def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor: def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor: """(B, K_MAX) — fraction of the original e_sec budget unclaimed entering slot i: `1.0` at `i=0`, `prod_{j=1` - (docs/v0.3.0-design.md §6.3 — "no re-derivation needed": the existing + ("no re-derivation needed": the existing stick-breaking encoding is already scale-free, so this is derivable from the batch's ground-truth stick logits alone, no `e_sec` required).""" cumprod = torch.cumprod(1.0 - fraction, dim=1) @@ -166,7 +165,7 @@ def _assemble_stage2_ar_inputs( ) -> dict[str, torch.Tensor]: """Ground-truth per-token AR conditioning tensors — all `(B, K_MAX, ...)` or `(B, K_MAX)`, built in one vectorized pass (teacher forcing means - every token's input is ground truth, docs/v0.3.0-design.md §6.2 point 3). + every token's input is ground truth). Keys match `Stage2Autoregressive.forward`'s trailing kwargs.""" device = sec_cont.device B, K = sec_cont.shape[0], sec_cont.shape[1] @@ -187,9 +186,10 @@ def _stage2_tf_prob( mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int ) -> float: """P(condition slot k+1 on the TRUE token k rather than the model's own - prediction), for the current epoch (docs/v0.3.0-design.md §3.3 - `stage2_model.autoregressive.teacher_forcing`). `"always"`/`"never"` are - the two degenerate constants; `"scheduled"` linearly interpolates + prediction), for the current epoch + (`stage2_model.autoregressive.teacher_forcing`). + `"always"`/`"never"` are the two degenerate constants; `"scheduled"` + linearly interpolates `p_start` (epoch 0) to `p_end` (the final epoch) — standard scheduled sampling (Bengio et al. 2015).""" if mode == "always": @@ -240,7 +240,7 @@ def _assemble_stage2_ar_inputs_scheduled( sample_steps: int, ) -> dict[str, torch.Tensor]: """Scheduled-sampling counterpart of `_assemble_stage2_ar_inputs` - (docs/v0.3.0-design.md §3.3 `teacher_forcing` = "scheduled"/"never"): + (`teacher_forcing` = "scheduled"/"never"): each slot's history is the TRUE previous token with probability `p_tf` (an independent per-example, per-slot Bernoulli draw) and the model's own free-running prediction otherwise — closing the train/inference gap that @@ -252,8 +252,8 @@ def _assemble_stage2_ar_inputs_scheduled( `giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a cheap one-step proxy, so building it costs the same `k_max` (`* steps` for flow) sequential forwards `sample.py` pays at inference, EVERY batch - this is called on (§6.4's cost note, paid at train time too whenever - teacher_forcing != "always"). Fully detached: gradient only ever flows + this is called on (paid at train time too whenever teacher_forcing != + "always"). Fully detached: gradient only ever flows through the "real" target path each stage trainer already uses (`_assemble_stage2_ar_target`), never through this self-sample. """ @@ -302,7 +302,7 @@ def _relax_onehot_type_slice( ) -> torch.Tensor: """Straight-through Gumbel-softmax relaxation of the per-slot type slice inside a flattened `(B, k_max * (cont_dim + type_dim))` WGAN generator - output — decision 5 (docs/v0.3.0-design.md §2.1): the forward pass is a + output: the forward pass is a hard one-hot (matching what the critic sees from real data), the backward pass flows smooth gradient. Continuous slots (stick/dir, and the type slice itself under `target = "embedding"`, which never calls @@ -311,8 +311,8 @@ def _relax_onehot_type_slice( `grad_probe`, if given, gets `["cont"]`/`["type"]` populated with the L2 norm of the gradient reaching this split point during the next `.backward()` call that touches it — a backward hook, not a second - backward pass. This is the §11.4 differentiability validation-obligation - instrumentation (docs/v0.3.0-design.md): the trunk-gradient contribution + backward pass. This is the differentiability validation-obligation + instrumentation: the trunk-gradient contribution from the type slice vs. the continuous slices, for `particle_type.target="onehot"` + `generator="wgan"`. Only ever populated on a `did_g_step` batch — the critic step backprops through diff --git a/giant/training/trainers.py b/giant/training/trainers.py index bc597dd..e1ac249 100644 --- a/giant/training/trainers.py +++ b/giant/training/trainers.py @@ -166,7 +166,7 @@ class StageSpec: tf_p_end=ar_cfg.get("tf_p_end", 1.0), # AR self-sampling under scheduled/never teacher forcing reuses # train.validate_steps as its flow-matching ODE step count — no - # dedicated config key for this (docs/v0.3.0-design.md §3.3 lists + # dedicated config key for this (the autoregressive config lists # tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only). ar_sample_steps=t["validate_steps"], ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000), @@ -184,12 +184,11 @@ class StageTrainer: Reads only the shared batch tuple `(cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx, sec_type_idx)` — stage 2 always conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context = "truth"`, - stage-level teacher forcing; `"sampled"` is not implemented — see - docs/v0.3.0-design.md §3.3), so stage trainers never need each other's - output at train time. This means "stage-2-only training is a cheap - ablation, not new plumbing" (design doc §7) falls out for free: a trainer - only exists for active stages, and inactive stages are simply never - constructed. + stage-level teacher forcing; `"sampled"` is not implemented), so stage + trainers never need each other's output at train time. This means + "stage-2-only training is a cheap ablation, not new plumbing" falls out + for free: a trainer only exists for active stages, and inactive stages + are simply never constructed. Grad-norm clipping is per-stage here — v0.2's single shared optimizer clipped both stages' gradients jointly; splitting per stage is a small, @@ -381,7 +380,7 @@ class StageTrainer: non-adversarial and WGAN trainers. Gated on `n_sec_head is None`, not on `n_sec.mode`: a future - `mode="stop_token"` model (design doc §11.2, currently rejected in + `mode="stop_token"` model (currently rejected in `validate_config`) carries no head and would train its EOS signal in the generator/AR loss path instead, so this correctly stays zero. """ @@ -458,8 +457,7 @@ class FlowDDPMStageTrainer(StageTrainer): raise NotImplementedError( f"stage2_model.generator={spec.generator!r} is accepted by the " "schema but not implemented in v0.3.0 for stage 2 (only " - "'flow' and 'wgan' have a stage-2 secondary-decoder loss — " - "see docs/v0.3.0-design.md §11.2)" + "'flow' and 'wgan' have a stage-2 secondary-decoder loss)" ) super().__init__(spec, model, device) self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0) @@ -562,7 +560,7 @@ class FlowDDPMStageTrainer(StageTrainer): ): """CE (`target="onehot"`) or MSE (`target="embedding"`) loss for the stage-2 model's `type_head` — the non-adversarial counterpart to - WGANStageTrainer's ST-Gumbel-into-the-critic path (decision 2/5). + WGANStageTrainer's ST-Gumbel-into-the-critic path. Zero when this stage has no `type_head` (stage 1, or `particle_type.target = "physical"`).""" l_type = torch.zeros((), device=device) @@ -774,8 +772,8 @@ class WGANStageTrainer(StageTrainer): "grad_norm_g", ] if self.is_stage2 and self.particle_type_cfg.get("target") == "onehot": - # §11.4 differentiability instrumentation — only meaningful when - # the type slice is a straight-through Gumbel relaxation. + # Differentiability instrumentation — only meaningful when the + # type slice is a straight-through Gumbel relaxation. train_keys += ["grad_norm_type_slice", "grad_norm_cont_slice"] self.train_metrics = [train_metric(key) for key in train_keys] self.val_metrics = [] @@ -860,11 +858,11 @@ class WGANStageTrainer(StageTrainer): ) if self.particle_type_cfg.get("target", "physical") == "onehot": # Straight-through Gumbel-softmax relaxation of the type - # slice only (decision 5) — the critic must see a hard - # one-hot forward (matching what "real" data looks like) - # while gradient still flows smoothly to the generator. - # grad_probe captures the §11.4 gradient-magnitude - # instrumentation — see _relax_onehot_type_slice's docstring. + # slice only — the critic must see a hard one-hot forward + # (matching what "real" data looks like) while gradient + # still flows smoothly to the generator. grad_probe captures + # the gradient-magnitude instrumentation — see + # _relax_onehot_type_slice's docstring. tau = _gumbel_tau( global_step, self.total_steps, @@ -898,7 +896,7 @@ class WGANStageTrainer(StageTrainer): ) # On a non-generator-step batch with no n_sec_head on this stage - # (n_sec now defaults to stage 2, decision 1), there's nothing for + # (n_sec now defaults to stage 2), there's nothing for # the generator optimizer to do this batch — g_loss would otherwise # be a graph-less zero tensor, which .backward() rejects outright. skip_g_step = not did_g_step and self.model.n_sec_head is None diff --git a/giant/validate.py b/giant/validate.py index ea8a34e..c193f4f 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -85,7 +85,7 @@ def validate_marginals( When `sec_decoder` is given, also validates Stage 2 via `giant.sample.sample_stage2`/`resolve_n_sec` (generator- and - one-shot-vs-autoregressive-agnostic, docs/v0.3.0-design.md §10): n_sec + one-shot-vs-autoregressive-agnostic): n_sec distribution (+ classification accuracy), per-slot energy-fraction marginals, and a particle-type marginal whose shape depends on `sec_decoder.particle_type_cfg["target"]` — restricted to each side's own diff --git a/scripts/check_migration_v02_v03.py b/scripts/check_migration_v02_v03.py index 20ff32e..3136d11 100644 --- a/scripts/check_migration_v02_v03.py +++ b/scripts/check_migration_v02_v03.py @@ -1,6 +1,5 @@ -"""Portal-machine follow-up for v0.3.0 step 2 (docs/v0.3.0-design.md §4.3): -diff a real v0.2 checkpoint's outputs against the new `build_models` on the -same input batch. +"""Portal-machine follow-up for v0.3.0 step 2: diff a real v0.2 checkpoint's +outputs against the new `build_models` on the same input batch. `tests/test_migration_v02_v03.py` already proves this bit-identical with synthetic random weights, but that test can't run where it matters (no @@ -16,7 +15,7 @@ Usage (from the repo root, on a portal machine): uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1 Run it once against a flow (or ddpm) checkpoint and once against a wgan -checkpoint (design doc §4.3's "one flow checkpoint and one WGAN checkpoint"). +checkpoint ("one flow checkpoint and one WGAN checkpoint"). A routed checkpoint (`model_config["router"]["enabled"]`) is only checked for successful construction — `giant.model.network.migrate_legacy_state_dict` doesn't yet remap routed (Expert-per-router) state dicts, so the @@ -115,8 +114,7 @@ def main() -> int: print( " routed checkpoint: migrate_legacy_state_dict only handles the " "monolithic trunk shape — verifying construction only, skipping " - "the bit-identical weight/output comparison. See " - "docs/v0.3.0-design.md §2.4's scope note." + "the bit-identical weight/output comparison." ) print("PASS (construction only, routed checkpoint)") return 0 diff --git a/scripts/dwarf.py b/scripts/dwarf.py index 311e384..5e7cff5 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -536,7 +536,7 @@ def warm_cache( "--material-conditioning", help="Must match the `giant train` run(s)' conditioning.material.type " "to warm for — independent of --particle-conditioning " - "(docs/v0.3.0-design.md §3.1: the two axes may differ)", + "(the two axes may differ)", ), ] = Conditioning.physical, router: Annotated[ diff --git a/scripts/warm_setup_cache.py b/scripts/warm_setup_cache.py index 115463c..3168b9e 100644 --- a/scripts/warm_setup_cache.py +++ b/scripts/warm_setup_cache.py @@ -31,8 +31,8 @@ def run_warm_setup_cache( select the normalizer cache entry (`giant.data.setup_cache.normalizer_key`) — pass the same values a later `giant train` invocation will use so it hits this warmed entry. The two - conditioning axes are independent (docs/v0.3.0-design.md §3.1) and may - differ. `router_enabled`/`router_type`/`n_experts` only matter for + conditioning axes are independent and may differ. + `router_enabled`/`router_type`/`n_experts` only matter for `router_type == "process"` (warms that `n_experts`'s process map); the energy-router quantile summary is always collected regardless, so a later `--router-type energy` run never needs to rescan just to seed diff --git a/tests/legacy/network_v02_snapshot.py b/tests/legacy/network_v02_snapshot.py index 8d34eb7..ea9364c 100644 --- a/tests/legacy/network_v02_snapshot.py +++ b/tests/legacy/network_v02_snapshot.py @@ -1,13 +1,13 @@ """Frozen snapshot of `giant/model/network.py` as it stood at the v0.3.0 -"step 1" commit (eb6dd27), i.e. the last commit before the step-2 §5 -decomposition (see `docs/v0.3.0-design.md`). +"step 1" commit (eb6dd27), i.e. the last commit before the step-2 +composable-parts decomposition. This is a deliberate verbatim copy, not an import of the live module — the whole point is that this file's classes keep behaving exactly as v0.2 did even after `giant/model/network.py` itself is rewritten, so `tests/test_migration_v02_v03.py` has a stable "old" side to diff the new -`build_models`/`Stage1Model`/`Stage2OneShot` against (design doc §4.3's -bit-identical acceptance test). Do not edit this file to track future +`build_models`/`Stage1Model`/`Stage2OneShot` against (the bit-identical +acceptance test). Do not edit this file to track future `network.py` changes — it exists specifically to stop tracking them. """ diff --git a/tests/test_cli_train_overrides.py b/tests/test_cli_train_overrides.py index 2341b69..486ea4b 100644 --- a/tests/test_cli_train_overrides.py +++ b/tests/test_cli_train_overrides.py @@ -1,6 +1,5 @@ -"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md -decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must -independently override each stage's config block, and must take precedence +"""Tests for `giant train`'s stage-prefixed CLI flags: --stage1-*/--stage2-* +must independently override each stage's config block, and must take precedence over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still apply the same value to both stages for backward compatibility.""" diff --git a/tests/test_config.py b/tests/test_config.py index ca50de6..fb7b5c2 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -585,8 +585,8 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning(): def test_validate_config_mixed_particle_material_conditioning_is_valid(): - """docs/v0.3.0-design.md §3.1: the particle and material conditioning - axes are configured independently and may mix freely — e.g. material + """The particle and material conditioning axes are configured + independently and may mix freely — e.g. material "physical" with particle "embedding" — and the data pipeline (giant/data/transforms.py) now implements that end-to-end, so validate_config must not reject it.""" @@ -638,8 +638,8 @@ def test_validate_config_stop_token_not_implemented(): def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint(): - """docs/v0.3.0-design.md §9: 'n_sec.mode = "truth" is invalid for a - rollout-capable checkpoint' — both stages active means giant rollout + """'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint' — + both stages active means giant rollout could load this checkpoint, but 'truth' has no ground truth to draw n_sec from at rollout time.""" cfg = _cfg_with( diff --git a/tests/test_loader.py b/tests/test_loader.py index 23003a8..ad31f7a 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -210,8 +210,7 @@ def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path): def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path): """A species that's rare as a primary but common as a secondary must still rank by its pooled (primary + secondary) count, not just its - primary-role count alone — the whole point of pooling both roles - (docs/v0.3.0-design.md §8).""" + primary-role count alone — the whole point of pooling both roles.""" path = tmp_path / "a.parquet" # primary pdg: mostly 11 (electron), one lone 22 (photon) pdg = [11] * 5 + [22] * 1 diff --git a/tests/test_migration_v02_v03.py b/tests/test_migration_v02_v03.py index 6c5d10c..05829a3 100644 --- a/tests/test_migration_v02_v03.py +++ b/tests/test_migration_v02_v03.py @@ -1,13 +1,12 @@ -"""Migration acceptance test for v0.3.0 step 2 (docs/v0.3.0-design.md §4.3, -§12 step 2): "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." +"""Migration acceptance test for v0.3.0 step 2: "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." No `/ceph` access on this machine (see CLAUDE.md's Compute environment -section), so a real trained checkpoint can't be used here — see -docs/v0.3.0-design.md's plan for the separate portal-machine follow-up with a -real checkpoint. This test is the synthetic stand-in: build a v0.2-shaped +section), so a real trained checkpoint can't be used here — a separate +portal-machine follow-up with a real checkpoint is planned instead. This +test is the synthetic stand-in: build a v0.2-shaped model from the frozen `tests/legacy/network_v02_snapshot.py` classes with fixed-seed random weights (playing the role of "a v0.2 checkpoint"), migrate its config and remap its state dict onto the new `build_models` output, and @@ -134,7 +133,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None: assert isinstance(new_stage1, net.Stage1Model) assert isinstance(new_stage2, net.Stage2OneShot) # legacy_owner="stage1": n_sec lives on stage1, not stage2, for a - # migrated v0.2 checkpoint (design doc §4.1). + # migrated v0.2 checkpoint. assert new_stage1.n_sec_head is not None assert new_stage2.n_sec_head is None @@ -199,12 +198,11 @@ def test_migrate_legacy_model_config_shape(): def test_migrate_legacy_model_config_nonzero_expert_dims_raises(): - """docs/v0.3.0-followups.md item 8 regression: a v0.2 checkpoint's - model_config carrying a non-default expert_hidden_dim/expert_n_blocks - must fail loudly through this path too (§4.2) — not just - giant.config.migrate_config's parallel TOML-load path. Silently dropping - these keys (build_router's kwarg filtering) would resize the experts - instead of refusing.""" + """Regression: a v0.2 checkpoint's model_config carrying a non-default + expert_hidden_dim/expert_n_blocks must fail loudly through this path too + — not just giant.config.migrate_config's parallel TOML-load path. + Silently dropping these keys (build_router's kwarg filtering) would + resize the experts instead of refusing.""" legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical") legacy_cfg["router"] = { "enabled": True, diff --git a/tests/test_network.py b/tests/test_network.py index 0a254f6..bbf0552 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -84,7 +84,7 @@ def test_stage1_model_gradients_flow(): def test_stage1_model_no_n_sec_head_by_default(): """Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head — - decision 1 (docs/v0.3.0-design.md §2) moves it to stage 2.""" + it moves to stage 2.""" model = Stage1Model( pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG ) @@ -206,7 +206,7 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector(): assert torch.all(pdg_e.sum(dim=-1) == 1.0) -# --- Stage2OneShot particle_type architecture (docs/v0.3.0-design.md decision 2) -- +# --- Stage2OneShot particle_type architecture -------------------------------- def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot: @@ -300,7 +300,7 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type(): assert out.shape == (B, k_max * CONT_SLOT_DIM) -# --- MarkovHistory (docs/v0.3.0-design.md §6.2) ----------------------------- +# --- MarkovHistory ----------------------------------------------------------- def test_markov_history_shape(): @@ -314,8 +314,8 @@ def test_markov_history_shape(): def test_markov_history_uses_start_vector_when_no_prev(): """Slot 0's own raw feature must be ignored — a learned start vector is - substituted there instead (a reasonable default not specified by the - design doc, see Stage2Autoregressive's docstring).""" + substituted there instead (a reasonable default, see + Stage2Autoregressive's docstring).""" hist = MarkovHistory(in_dim=4, out_dim=6) B, K = 2, 3 has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) @@ -328,7 +328,7 @@ def test_markov_history_uses_start_vector_when_no_prev(): assert torch.allclose(out_a[:, 1:], out_b[:, 1:]) -# --- AttentionHistory (docs/v0.3.0-design.md §6.2, v0.3.0 step 7) ---------- +# --- AttentionHistory (v0.3.0 step 7) --------------------------------------- def test_attention_history_shape(): @@ -392,7 +392,7 @@ def test_attention_history_step_matches_forward(): assert torch.allclose(stepped, expected, atol=1e-5) -# --- Stage2Autoregressive (docs/v0.3.0-design.md §6, v0.3.0 step 5) --------- +# --- Stage2Autoregressive (v0.3.0 step 5) ----------------------------------- def _build_stage2_ar( @@ -671,8 +671,8 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete """The shared encoder's parameters must actually appear in both stages' own `.parameters()` — that's what makes each stage's independent optimizer include (and update) them, which is the actual mechanism behind - "shared weights, forced common representation" (docs/v0.3.0-design.md - §3.1), not just object identity on `.cond_enc`.""" + "shared weights, forced common representation", not just object identity + on `.cond_enc`.""" built = build_models(_minimal_model_config(share_stages=True)) stage1, stage2 = built["stage1"], built["stage2"] assert stage1 is not None and stage2 is not None diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index f202b9d..fae62c5 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -155,7 +155,7 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch): def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data): """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to - "onehot" (docs/v0.3.0-design.md §3.3/§8) — a plain _tiny_cfg() run must + "onehot" — a plain _tiny_cfg() run must build the shared pdg top-N map, cache it in the setup-cache sidecar, and persist it into the checkpoint, with no extra config needed.""" echo1 = _run(data, tmp_path / "out1") @@ -242,7 +242,7 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat def test_run_train_job_custom_k_max_end_to_end(tmp_path, data): - """docs/v0.3.0-followups.md item 3 regression: stage2_model.k_max other + """Regression: stage2_model.k_max other than the K_MAX module constant's default (15) must not produce a shape mismatch between the data pipeline (loader.py/transforms.py padding) and the model (network.py's trunks, sized from this same config value).""" @@ -254,9 +254,9 @@ def test_run_train_job_custom_k_max_end_to_end(tmp_path, data): def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data): - """docs/v0.3.0-followups.md item 4 regression: conditioning.particle.type + """Regression: conditioning.particle.type and conditioning.material.type are configured independently and may mix - freely (docs/v0.3.0-design.md §3.1) — e.g. particle "embedding" with + freely — e.g. particle "embedding" with material "physical" — end-to-end through the real data pipeline, not just accepted by validate_config.""" cfg = _tiny_cfg() @@ -286,7 +286,7 @@ def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, def test_run_train_job_share_stages_end_to_end(tmp_path, data): - """docs/v0.3.0-followups.md item 5 regression: conditioning.share_stages + """Regression: conditioning.share_stages = true must actually train (not raise NotImplementedError), and the resulting checkpoint's two stages must reload into a single shared ConditionEncoder instance rather than two independent ones.""" diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 7c230b2..da8d624 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -373,7 +373,7 @@ def _models_v3( particle_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1} material_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1} # A fresh v0.3.0 Stage1Model — no n_sec_head_k_max, unlike _models() above - # (decision 1 moves n_sec ownership to stage 2 by default). + # (n_sec ownership moves to stage 2 by default). s1 = Stage1Model( pdg_vocab=3, mat_vocab=2, @@ -470,9 +470,9 @@ def _run_v3( def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props): - """A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1) — n_sec - must come from Stage2's own head instead, and the run must still - complete and conserve energy.""" + """A fresh v0.3.0 Stage1Model has no n_sec_head — n_sec must come from + Stage2's own head instead, and the run must still complete and + conserve energy.""" s1, s2 = _models_v3() rec = _run_v3(s1, s2) assert len(rec["event_id"]) > 0 @@ -501,7 +501,7 @@ def test_rollout_physical_target_decoder_generator_matrix( ): """Every (decoder, stage2 generator) combination under particle_type.target="physical" must run to completion and conserve - energy — the matrix docs/v0.3.0-design.md §7 calls out for comparison.""" + energy.""" s1, s2 = _models_v3(decoder=decoder, generator2=generator2) rec = _run_v3(s1, s2) assert len(rec["event_id"]) > 0 @@ -555,10 +555,10 @@ def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props): _run_v3(s1, s2, pdg_topn_map=None) -# --- conditioning.{particle,material}.type = "onehot" (docs/v0.3.0-followups.md -# item 7) — a separate axis from stage2_model.particle_type.target above: this -# is what feeds cond_cat's extra top-N columns for ConditionEncoder's own -# "onehot" mode, not the secondary-species decode. --------------------------- +# --- conditioning.{particle,material}.type = "onehot" — a separate axis from +# stage2_model.particle_type.target above: this is what feeds cond_cat's +# extra top-N columns for ConditionEncoder's own "onehot" mode, not the +# secondary-species decode. --------------------------------------------- COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={}) COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={}) @@ -653,7 +653,7 @@ def test_rollout_embedding_target_end_to_end(decoder): def test_l1_dist_collector_populated_only_for_embedding_target(): - """§11.3: the L1-distance diagnostic only makes sense under + """The L1-distance diagnostic only makes sense under particle_type.target="embedding" — a physical-target run must leave the collector empty rather than silently accumulating garbage.""" s1, s2 = _models_v3(target="physical") diff --git a/tests/test_router.py b/tests/test_router.py index 8d791d7..01c0ded 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -896,8 +896,8 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): cond_cont, cond_cat = _cond(B, pdg=3, mat=2) stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) assert stage1_norm.shape == (B, X_DIM) - # A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1 moves it to - # stage 2) — sample_flow returns n_sec_pred=None here, and n_sec must be + # A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) — + # sample_flow returns n_sec_pred=None here, and n_sec must be # asked of stage2 instead, using the just-sampled stage1_norm as context. assert n_sec_pred is None n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1) @@ -1092,8 +1092,8 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow(): cond_cont, cond_cat = _cond(B, pdg=3, mat=2) stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) assert stage1_norm.shape == (B, X_DIM) - # A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1 moves it to - # stage 2) — sample_flow returns n_sec_pred=None here, and n_sec must be + # A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) — + # sample_flow returns n_sec_pred=None here, and n_sec must be # asked of stage2 instead, using the just-sampled stage1_norm as context. assert n_sec_pred is None n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1) diff --git a/tests/test_sample.py b/tests/test_sample.py index ff52610..f6b12e7 100644 --- a/tests/test_sample.py +++ b/tests/test_sample.py @@ -1,6 +1,6 @@ """Tests for giant/sample.py's v0.3.0 stage-model sampling — the AR loop (`sample_secondaries_ar`) and non-"physical" `particle_type.target` coverage -for the one-shot samplers (docs/v0.3.0-design.md step 6).""" +for the one-shot samplers.""" import pytest import torch @@ -38,7 +38,7 @@ def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tenso def _conditioning_for(target: str) -> str: # target="embedding" regresses against the conditioning's own embedding - # table (docs/v0.3.0-design.md §3.3) — only meaningful when the + # table — only meaningful when the # conditioning axis is itself "embedding". return "embedding" if target == "embedding" else "physical" @@ -103,7 +103,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int: return PARTICLE_PHYS_DIM if target == "physical" else emb_dim -# ── Stage-1 n_sec ownership (decision 1) ──────────────────────────────────── +# ── Stage-1 n_sec ownership ────────────────────────────────────────────────── def test_sample_flow_returns_none_n_sec_when_stage1_owns_no_head(): diff --git a/tests/test_train.py b/tests/test_train.py index 992adb8..ad14113 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -90,7 +90,7 @@ def test_wandb_run_config_handles_missing_model_config(): assert wcfg["model_config"] == {} -# --- AR helper functions (v0.3.0 step 5, docs/v0.3.0-design.md §6) --------- +# --- AR helper functions (v0.3.0 step 5) --------- def test_stick_fraction_matches_sigmoid_of_logit(): @@ -119,7 +119,7 @@ def test_ar_has_prev_false_only_at_slot_zero(): assert has_prev.tolist() == [[False, True, True, True, True]] -# --- _stage2_tf_prob (docs/v0.3.0-design.md §3.3, v0.3.0 step 7) ----------- +# --- _stage2_tf_prob (v0.3.0 step 7) ----------- def test_stage2_tf_prob_always_is_constant_one(): @@ -569,7 +569,7 @@ def test_metrics_csv_columns_are_stage_prefixed(): def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch(): """Regression test: on a non-generator-step batch, if this stage's model - has no n_sec_head (n_sec defaults to stage 2, decision 1), g_loss is a + has no n_sec_head (n_sec defaults to stage 2), g_loss is a graph-less zero — .backward() must not be called on it.""" cfg = _base_cfg() cfg["stage1_model"]["generator"] = "wgan" @@ -672,11 +672,11 @@ def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing( def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics(): - """§11.4 differentiability validation-obligation instrumentation: the + """Differentiability validation-obligation instrumentation: the trunk-gradient-norm-by-slice columns must appear and actually fire for generator='wgan' + particle_type.target='onehot' under decoder= - 'autoregressive' (added at v0.3.0 step 5 per the design doc's - instruction to accrue evidence during the architecture comparison).""" + 'autoregressive' (added at v0.3.0 step 5 to accrue evidence during the + architecture comparison).""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} diff --git a/tests/test_type_embedding_distance.py b/tests/test_type_embedding_distance.py index c17f965..6b3b8dd 100644 --- a/tests/test_type_embedding_distance.py +++ b/tests/test_type_embedding_distance.py @@ -1,5 +1,5 @@ """Tests for the secondary-type embedding-distance diagnostic -(giant.analysis.type_embedding_distance) — the §11.3 diagnostic.""" +(giant.analysis.type_embedding_distance).""" from __future__ import annotations diff --git a/tests/test_validate.py b/tests/test_validate.py index 533c933..1e29517 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -11,8 +11,8 @@ _K_MAX = 5 def _tiny_models(particle_type_cfg: dict | None = None): - """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head (decision 1), so - n_sec always comes from Stage2OneShot.""" + """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head, so n_sec always + comes from Stage2OneShot.""" s1 = Stage1Model( pdg_vocab=3, mat_vocab=2, diff --git a/uv.lock b/uv.lock index bd8cb35..fb16763 100644 --- a/uv.lock +++ b/uv.lock @@ -534,7 +534,7 @@ wheels = [ [[package]] name = "giant" -version = "0.2.0" +version = "0.3.0" source = { editable = "." } dependencies = [ { name = "numpy" },