Secondary species vocabulary is accidentally sized by conditioning.particle.emb_dim #29

Closed
opened 2026-08-13 15:05:33 +02:00 by lars · 1 comment
Owner

Severity: medium-high (design defect, on the critical path of the v0.3 pivot).

giant/pipeline.py:189 and giant/model/models.py:19-25 both size the stage-2
particle-type head from conditioning.particle.emb_dim:

if particle_cfg["type"] == "onehot" or particle_type_target == "onehot":
    n_classes = particle_cfg["emb_dim"]

Under the default config that coupling is between two unrelated things:
conditioning.particle.type = "physical" (so emb_dim means "output width of
the physical-property MLP") and stage2_model.particle_type.target = "onehot"
(so emb_dim silently means "number of secondary species classes"). The
default emb_dim = 16 therefore fixes the secondary species vocabulary at
15 PDG codes + one other bucket, and the only way to widen it to 32
species is to widen the conditioning MLP's output to 32 as well.

The comment at pipeline.py:177-182 presents this as a deliberate optimisation
("both key off conditioning.particle.emb_dim, so at most one PDG scan is
needed"). Sharing the scan is fine; sharing the class count is what's wrong
— and the scan can be shared anyway by building one map at
max(cond_n_classes, type_n_classes) and slicing.

The reason this is more than cosmetic: the entire v0.3.0 redesign exists because
the 2026-08-03 WGAN benchmark produced zero photon secondaries and ~4M
hallucinated -14s. The species vocabulary is the thing being fixed, and it is
currently not independently settable.

Fix: add stage2_model.particle_type.n_classes (default 0 → inherit
conditioning.particle.emb_dim, preserving today's behaviour and every existing
checkpoint), thread it through stage2_type_dim / stage2_trunk_sec_dim /
pipeline.py, and store the resolved value in the checkpoint. See also Issue 19
(class-balanced type loss) for the class-imbalance half of the same problem.


Migrated from issues.md (v0.3.0 branch review, 2026-08-13), Issue 3.

**Severity: medium-high (design defect, on the critical path of the v0.3 pivot).** `giant/pipeline.py:189` and `giant/model/models.py:19-25` both size the stage-2 particle-type head from `conditioning.particle.emb_dim`: ```python if particle_cfg["type"] == "onehot" or particle_type_target == "onehot": n_classes = particle_cfg["emb_dim"] ``` Under the **default** config that coupling is between two unrelated things: `conditioning.particle.type = "physical"` (so `emb_dim` means "output width of the physical-property MLP") and `stage2_model.particle_type.target = "onehot"` (so `emb_dim` silently means "number of secondary species classes"). The default `emb_dim = 16` therefore fixes the secondary species vocabulary at **15 PDG codes + one `other` bucket**, and the only way to widen it to 32 species is to widen the conditioning MLP's output to 32 as well. The comment at `pipeline.py:177-182` presents this as a deliberate optimisation ("both key off `conditioning.particle.emb_dim`, so at most one PDG scan is needed"). Sharing the *scan* is fine; sharing the *class count* is what's wrong — and the scan can be shared anyway by building one map at `max(cond_n_classes, type_n_classes)` and slicing. The reason this is more than cosmetic: the entire v0.3.0 redesign exists because the 2026-08-03 WGAN benchmark produced zero photon secondaries and ~4M hallucinated `-14`s. The species vocabulary is the thing being fixed, and it is currently not independently settable. **Fix:** add `stage2_model.particle_type.n_classes` (default `0` → inherit `conditioning.particle.emb_dim`, preserving today's behaviour and every existing checkpoint), thread it through `stage2_type_dim` / `stage2_trunk_sec_dim` / `pipeline.py`, and store the resolved value in the checkpoint. See also Issue 19 (class-balanced type loss) for the class-imbalance half of the same problem. --- Migrated from `issues.md` (v0.3.0 branch review, 2026-08-13), Issue 3.
lars added the architecturebug labels 2026-08-13 15:05:33 +02:00
Author
Owner

Fixed in 32aa5a5 on v0.3.0-stage2-autoregressive.

Added stage2_model.particle_type.n_classes (default 0 = inherit conditioning.particle.emb_dim, fully back-compat) and a single resolve_type_n_classes helper (giant/model/models.py) now used everywhere the coupling used to be implicit: pipeline.py's PDG top-N map build, Stage2OneShot/Stage2Autoregressive's type head sizing, build_models/build_critics, and StageSpec's training loss width.

The checkpoint's single shared pdg_topn_map is now split: pdg_topn_map stays conditioning-only, and a new sec_type_topn_map key covers the secondary-species decode side, built independently through the existing (axis, n_classes)-keyed setup cache — no extra parquet scan in the common case where the two still resolve to the same N. Threaded through giant predict/giant rollout's decode path (giant/rollout.py, giant/cli.py). A checkpoint with no sec_type_topn_map key (pre-#29) falls back to reusing pdg_topn_map in checkpoint_io.py, reproducing the old shared behavior exactly — no existing checkpoint becomes unloadable.

Decisions made during planning (with the user): commit directly on this branch rather than a new one; represent the split as an additive sec_type_topn_map checkpoint key rather than conditionally reusing pdg_topn_map; build the two top-N maps independently rather than the issue's proposed build-at-max-and-slice, since giant/data/setup_cache.py's topn_maps cache is already keyed by (axis, n_classes) and already avoids redundant scans across runs.

All four checks (pytest, ruff format, ruff check, ty check) clean. Gitea #44 (class-balanced secondary particle-type loss), referenced by this issue for context, is a separate open issue and was left untouched.

Fixed in 32aa5a5 on v0.3.0-stage2-autoregressive. Added `stage2_model.particle_type.n_classes` (default 0 = inherit `conditioning.particle.emb_dim`, fully back-compat) and a single `resolve_type_n_classes` helper (giant/model/models.py) now used everywhere the coupling used to be implicit: pipeline.py's PDG top-N map build, Stage2OneShot/Stage2Autoregressive's type head sizing, build_models/build_critics, and StageSpec's training loss width. The checkpoint's single shared `pdg_topn_map` is now split: `pdg_topn_map` stays conditioning-only, and a new `sec_type_topn_map` key covers the secondary-species decode side, built independently through the existing (axis, n_classes)-keyed setup cache — no extra parquet scan in the common case where the two still resolve to the same N. Threaded through giant predict/giant rollout's decode path (giant/rollout.py, giant/cli.py). A checkpoint with no sec_type_topn_map key (pre-#29) falls back to reusing pdg_topn_map in checkpoint_io.py, reproducing the old shared behavior exactly — no existing checkpoint becomes unloadable. Decisions made during planning (with the user): commit directly on this branch rather than a new one; represent the split as an additive sec_type_topn_map checkpoint key rather than conditionally reusing pdg_topn_map; build the two top-N maps independently rather than the issue's proposed build-at-max-and-slice, since giant/data/setup_cache.py's topn_maps cache is already keyed by (axis, n_classes) and already avoids redundant scans across runs. All four checks (pytest, ruff format, ruff check, ty check) clean. Gitea #44 (class-balanced secondary particle-type loss), referenced by this issue for context, is a separate open issue and was left untouched.
lars closed this issue 2026-08-13 16:11:31 +02:00
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: lars/giant#29