Add class-balanced secondary particle-type loss (gitea #44)
CI / Lint (ruff check) (push) Successful in 36s
CI / Format (ruff format) (push) Successful in 38s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Lint (ruff check) (pull_request) Successful in 45s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 5m20s
CI / Tests (pull_request) Successful in 4m50s
CI / Lint (ruff check) (push) Successful in 36s
CI / Format (ruff format) (push) Successful in 38s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Lint (ruff check) (pull_request) Successful in 45s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 5m20s
CI / Tests (pull_request) Successful in 4m50s
The v0.3.0 pivot exists because the 2026-08-03 WGAN rollout benchmark produced zero photon secondaries and ~4M hallucinated antineutrinos — even with a correctly-sized top-N species vocabulary (gitea #29), plain cross-entropy over a class distribution spanning orders of magnitude still under-predicts rare-but-physical species. stage2_model.particle_type.class_weighting = "none" | "inverse_freq" (default "none", fully back-compat) weights the stage-2 type head's CE loss (FlowDDPMStageTrainer._type_loss) by inverse class frequency, normalized to mean 1 so switching it on doesn't rescale the type loss against particle_type.lambda / the generator loss it's summed with. The per-class counts the weighting needs don't already exist despite the issue's premise: _topn_plus_other_map (giant/data/loader.py) previously kept counts only for keys folded into "other", dropping the kept classes' counts on the floor. TopNMap now carries class_counts (index -> count), round-tripped through the setup-cache sidecar (format version bumped 3->4, since existing sidecars have none) and through checkpoints (tolerantly — a pre-#44 checkpoint decodes to {}, since only training-time loss weighting reads it, not inference). Decisions made during planning (with the user): dropped "effective_num" from the issue's proposed three-way enum (no beta hyperparameter to design around) — final domain is "none" | "inverse_freq". Weights are mean-1-normalized. validate_config rejects class_weighting != "none" combined with particle_type.target != "onehot" or stage2_model.generator == "wgan" (both have no class CE to weight), following the #28/#30 dead-key-must-not-go-silent convention. Branch fix/issue-44 off master. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -35,7 +35,10 @@ from giant.data.transforms import Normalizer, sorted_membership
|
||||
# v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
|
||||
# energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
|
||||
# v2 sidecar has no such grid to fall back on, so it must be recomputed.
|
||||
_CACHE_FORMAT_VERSION = 3
|
||||
# v4: TopNMap gained class_counts (gitea #44, stage2_model.particle_type.
|
||||
# class_weighting) — a v3 sidecar's cached topn_maps have no counts, so they
|
||||
# must be rebuilt rather than silently cached with class_counts={}.
|
||||
_CACHE_FORMAT_VERSION = 4
|
||||
|
||||
_DIMS = {
|
||||
"COND_DIM": COND_DIM,
|
||||
@@ -131,6 +134,7 @@ def topnmap_to_json(m: TopNMap) -> dict:
|
||||
return {
|
||||
"class_map": {str(k): v for k, v in m.class_map.items()},
|
||||
"other_members": {str(k): v for k, v in m.other_members.items()},
|
||||
"class_counts": {str(k): v for k, v in m.class_counts.items()},
|
||||
}
|
||||
|
||||
|
||||
@@ -139,6 +143,11 @@ def topnmap_from_json(d: dict, axis: str) -> TopNMap:
|
||||
return TopNMap(
|
||||
class_map={cast(k): v for k, v in d["class_map"].items()},
|
||||
other_members={cast(k): v for k, v in d["other_members"].items()},
|
||||
# Missing for a checkpoint's topn maps predating gitea #44 — {} is
|
||||
# the correct decode there (inference never reads class_counts; only
|
||||
# stage2_model.particle_type.class_weighting does, at train time, and
|
||||
# it raises loudly if it needs counts a checkpoint doesn't have).
|
||||
class_counts={int(k): v for k, v in d.get("class_counts", {}).items()},
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user