diff --git a/docs/v0.3.0-design.md b/docs/v0.3.0-design.md index 5b4408d..f1b35c4 100644 --- a/docs/v0.3.0-design.md +++ b/docs/v0.3.0-design.md @@ -24,7 +24,7 @@ to build it. Read it before touching `giant/config.py` or `giant/model/network.p 8. [Data and setup-cache changes](#8-data-and-setup-cache-changes) 9. [Config machinery changes](#9-config-machinery-changes) 10. [Callers that need updating](#10-callers-that-need-updating) -11. [Open questions](#11-open-questions) +11. [Settled scope and open questions](#11-settled-scope-and-open-questions) 12. [Implementation order](#12-implementation-order) --- @@ -201,7 +201,7 @@ active = true generator = "flow" hidden_dim = 256 n_res_blocks = 6 -dropout = 0.1 +dropout = 0.0 lambda = 1.0 ``` @@ -211,7 +211,7 @@ lambda = 1.0 | `generator` | `"flow"` \| `"ddpm"` \| `"wgan"` | `"flow"` | The generative objective. `"flow"`: conditional flow matching (Lipman et al. 2022), ~10 ODE steps at inference. `"ddpm"`: cosine-schedule diffusion baseline. `"wgan"`: WGAN-GP, single forward pass at inference. Replaces the global `train.mode`. Objective-specific knobs live in the matching sub-table below. | | `hidden_dim` | int | `256` | Trunk width — also the width of **every expert** under a routed trunk. | | `n_res_blocks` | int | `6` | Number of `ResBlock`s in the trunk, and in every expert under a routed trunk. Was `model.n_blocks`; renamed for clarity since v0.3.0 also has attention layers in stage 2. | -| `dropout` | float | `0.1` | Dropout inside each `ResBlock`. | +| `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"` @@ -304,7 +304,7 @@ decoder = "autoregressive" generator = "wgan" hidden_dim = 256 n_res_blocks = 6 -dropout = 0.1 +dropout = 0.0 lambda = 1.0 k_max = 15 context_dim = 64 @@ -318,7 +318,7 @@ stage1_context = "truth" | `generator` | `"flow"` \| `"ddpm"` \| `"wgan"` | `"wgan"` | As stage 1. Under `"autoregressive"` this is the objective for **each token**: a WGAN token costs one forward pass, a flow token costs ~10 ODE steps. See the cost note in §6.4. | | `hidden_dim` | int | `256` | Trunk width. | | `n_res_blocks` | int | `6` | Trunk depth. | -| `dropout` | float | `0.1` | Dropout inside each `ResBlock`. | +| `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`. | @@ -328,7 +328,7 @@ stage1_context = "truth" | Key | Type | Default | Meaning | |-----|------|---------|---------| -| `mode` | `"head"` \| `"stop_token"` \| `"truth"` | `"head"` | `"head"`: a classifier over `{0..k_max}` on the condition encoding alone (no diffusion noise), so it is callable independently at inference — v0.2 behaviour, and what the meeting's §3 explicitly decided to keep. `"stop_token"`: an EOS-style implicit stop, documented in §3 as a later possibility, not a decision. `"truth"`: take `n_sec` from ground truth — only valid for standalone stage-2 evaluation, never for rollout. | +| `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]` @@ -456,7 +456,7 @@ num_workers = 4 seed = 0 validate_every = 10 validate_steps = 10 -wandb = false +wandb = true wandb_project = "giant" wandb_run_name = "" wandb_log_every = 50 @@ -476,7 +476,7 @@ wandb_log_every = 50 | `seed` | int | `0` | Seeds Python/numpy/torch and the event split. | | `validate_every` | int | `10` | Epochs between full marginal/KL validation passes. | | `validate_steps` | int | `10` | Sampler steps used during those passes. | -| `wandb` | bool | `false` | Opt-in W&B logging. | +| `wandb` | 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. | @@ -845,8 +845,8 @@ All in `giant/config.py`: `n_sec.mode = "truth"` is invalid for a rollout-capable checkpoint. - **`estimate_batch_size`** — its calibration constants assume the v0.2 architecture ("post-Phase-2, including the Stage-2 secondary decoder and n_sec - head"). AR stage 2 changes the activation-memory profile; re-measure and note - the new calibration point. + head"). AR stage 2 changes the activation-memory profile. **Recalibrate last** + (§12 step 8), measured on real hardware with the example configs. --- @@ -868,41 +868,55 @@ All in `giant/config.py`: --- -## 11. Open questions +## 11. Settled scope and open questions -1. **"Other" bucket at rollout.** A predicted "other" class has no concrete PDG, - so no mass/charge for that secondary's own downstream conditioning. - `other_policy` is the config surface; the mechanism is a judgement call - (`"sample"` from the empirical within-bucket distribution seems least biased). - **Not decided in the meeting.** -2. **Charge conservation.** Out of scope for v0.3.0 (decision 6), but worth - recording *why* it interacts with decision 5: a categorical head makes a hard - mechanism tractable — mask classes whose charge cannot fit the remaining charge - budget before the softmax, analogous to the energy simplex. This is **only** - possible with a categorical type; the continuous `(mass, charge)` target admits - no such mask. That is a second, independent argument for the §5 switch beyond - the species-collapse one. -3. **Does `"stop_token"` deserve building in v0.3.0?** The meeting kept the `n_sec` - head and recorded the implicit stop as a later possibility. The config key - exists; the implementation can be deferred. -4. **`estimate_batch_size` calibration** for AR stage 2 (§9). -5. **Nearest-neighbour decode has no reject option.** Under - `particle_type.target = "embedding"`, an output far from every table row still - snaps to its nearest neighbour — there is no equivalent of `"onehot"`'s "other" - bin, and no confidence signal. Worth logging the L1 distance distribution at - rollout: a heavy tail means the decoder is emitting vectors off the embedding - manifold, which would be the direct analogue of the species-collapse symptom - this redesign is chasing. -6. **Should `flow` and `ddpm` share one sub-table?** Both need `time_dim`; only - ddpm needs `n_steps`. Currently duplicated as `[stage*_model.flow]` and - `[stage*_model.ddpm]`. A single `[stage*_model.diffusion]` would avoid the - duplication at the cost of a name that fits flow matching poorly. -5. **Differentiability sanity check.** The meeting's §5 notes that the original - argument for the continuous type target rested on avoiding a non-differentiable - categorical *sampling* step — but per-token training loss is differentiable - either way under teacher forcing, and full shower-rollout backprop is already - structurally non-differentiable once secondaries spawn branches. The note flags - this as "worth confirming explicitly with Jan rather than assuming it". +### 11.1 Settled — build as specified + +- **`other_policy` switch.** The three-way `"sample"` / `"modal"` / `"drop"` + design in §3.3 is approved as the config surface for turning a predicted "other" + class into a concrete PDG at rollout. `"sample"` (draw from the empirical + within-bucket distribution recorded at map-build time) stays the default as the + least-biased option. +- **`[stage*_model.flow]` and `[stage*_model.ddpm]` stay separate tables.** The + duplicated `time_dim` is accepted; a merged `[stage*_model.diffusion]` was + considered and rejected because the name fits flow matching poorly. + +### 11.2 Deferred — not implemented in v0.3.0 + +- **`n_sec.mode = "stop_token"`.** The value is **accepted by the schema** but + raises a clear "not implemented in v0.3.0" error if set. The key existing now + means landing the implicit-stop mechanism later is not a config break. `"head"` + is what v0.3.0 builds, per the meeting's §3 decision. +- **Charge conservation.** No key at all — `[stage2_model.conservation]` is absent + from v0.3.0 (decision 6), unlike `stop_token` above. **Deliberately left + undesigned:** the mechanism should be worked out on its own terms when it is + taken up, not pre-shaped by choices made for this refactor. Adding the block + later is a config addition, not a break. + +### 11.3 Tracked as implementation work + +- **Log the L1 distance distribution at rollout** under + `particle_type.target = "embedding"`. Nearest-neighbour decode has no reject + option — an output far from every table row still snaps to its nearest neighbour, + with no "other" bin and no confidence signal. A heavy tail in that distribution + means the decoder is emitting vectors off the embedding manifold, which is the + direct analogue of the species-collapse symptom this redesign exists to fix. + Belongs with the `giant/rollout.py` decode work (step 6) and should surface as a + `giant analyze` diagnostic plot alongside `router_gating`. +- **`estimate_batch_size` recalibration** for AR stage 2 is the **last** step of + the implementation flow (§12 step 8), measured on real hardware using the example + configs — not guessed from the existing calibration constants. + +### 11.4 Still open + +- **Differentiability sanity check with Jan.** The meeting's §5 notes that the + original argument for the continuous type target rested on avoiding a + non-differentiable categorical *sampling* step — but per-token training loss is + differentiable either way under teacher forcing, and full shower-rollout backprop + is already structurally non-differentiable once secondaries spawn branches. The + note flags this as "worth confirming explicitly with Jan rather than assuming + it". Decision 5 rests on this reasoning, so it is worth raising before step 4 + rather than after. --- @@ -923,10 +937,16 @@ All in `giant/config.py`: 5. **`Stage2Autoregressive`** with `history = "markov"`, `teacher_forcing = "always"`. The meeting's step 3. 6. **`sample.py` / `rollout.py`** — AR generation and class -> PDG decode, so an AR - model can actually be rolled out and put through `giant analyze`. + 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. +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.