Skip to content

Reject all-zero weights in MixtureSampler#54

Merged
mmshad merged 1 commit intomainfrom
mixture-sampler-zero-weights
Apr 26, 2026
Merged

Reject all-zero weights in MixtureSampler#54
mmshad merged 1 commit intomainfrom
mixture-sampler-zero-weights

Conversation

@mmshad
Copy link
Copy Markdown
Collaborator

@mmshad mmshad commented Apr 22, 2026

Summary

  • Add _validate_weights(weights, context) in kempnerforge/data/sampler.py; call it from both MixtureSampler.__init__ and MixtureSampler.update_weights.
  • Reject empty, negative, and all-zero weight lists with a ValueError. Previously temperature == 1.0 crashed with ZeroDivisionError at the normalization divide and temperature != 1.0 silently produced uniform sampling via the max(w, 1e-12) clamp.
  • TrainingPhase.dataset_weights schema validator stays >= 0 so individual datasets can still be zeroed out; the sampler-level gate catches the all-zero case that the schema cannot see at config-parse time.

Closes #53

Test plan

  • uv run pytest tests/unit/test_mixing.py::TestMixtureSamplerZeroWeights -v (8 tests).
  • Full unit suite clean.

The two normalization branches of MixtureSampler disagreed on all-zero
input: temperature == 1.0 crashed with ZeroDivisionError at the divide,
temperature != 1.0 silently degraded to uniform sampling via the
max(w, 1e-12) clamp. Both are reachable from the annealing phase path
in scripts/train.py, where TrainingPhase.dataset_weights validates >= 0
(not > 0) to intentionally allow dropping individual datasets.

A user who zeroes out every dataset in a phase transition would get a
mid-run crash on one branch and silent training-quality drift on the
other. Add a single _validate_weights() helper that rejects empty,
negative, or all-zero weight lists with a clear ValueError, and call
it from both __init__ and update_weights. The TrainingPhase schema-
level validator stays permissive because it can't know the effective
mixture at config-parse time; sampler-level is where the hard gate
belongs.
@mmshad mmshad requested a review from Naeemkh April 22, 2026 01:02
@mmshad mmshad self-assigned this Apr 22, 2026
@mmshad mmshad merged commit b84d70e into main Apr 26, 2026
5 checks passed
@mmshad mmshad deleted the mixture-sampler-zero-weights branch April 26, 2026 15:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

MixtureSampler accepts all-zero weights inconsistently

2 participants