Reimplement dropout_weights() correctly for log-space weights

[MaxGhenis]

Summary

Finding #2 (HIGH) in the bug-hunt report. dropout_weights held weights in log space but set masked entries to literal 0, which is exp(0) = 1 in linear space — the opposite of dropping. It then normalised by dividing masked log-weights by their sum and multiplying by total_weight, which is not a meaningful operation on logs. On realistic survey-weight scales (hundreds to thousands, log ~6–8), masked_weights.sum() could cross zero, producing Inf/NaN that crashed the forward pass via the NaN guard in loss().

Rewritten as standard inverted dropout performed in log space: dropped entries go to -inf (exp = 0) and surviving entries are shifted by -log(1-p) so the expected linear-space sum is preserved. Hoisted to module scope so it is directly testable, and invalid probabilities now raise.

Test plan

• Add tests/test_dropout_regression.py: p=0 identity, p=1 zeroes all linear-space weights, expected-sum preservation on realistic-scale weights (Monte Carlo tolerance), approximate drop fraction, input validation, and end-to-end smoke test that training with dropout no longer poisons the loss.
• All existing tests pass (uv run pytest tests -x -q -> 21 passed).

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
microcalibrate	Ready	Preview, Comment	Apr 17, 2026 0:40am

[MaxGhenis]

LGTM (cannot self-approve; posting as comment).

Correct inverted-dropout semantics in log space:

• p=0 returns input unchanged (identity).
• p=1 returns -inf tensor → exp = 0 everywhere (all dropped).
• Masked entries go to -inf (not 0, which was the original bug).
• Survivors shifted by -log(1-p) so E[exp(dropout(logw))] = exp(logw) in linear space, matching standard inverted dropout.
• p validated to [0, 1] with a clear ValueError.

Tests cover identity, full dropout, approximate drop fraction, expected-sum preservation on realistic (hundreds-to-thousands) weight scales, input validation, and end-to-end reweight() with dropout_rate > 0. Good coverage of the original bug's exact failure mode (NaN/Inf on realistic scales).

Hoisting dropout_weights to module scope is also a small structural win — now directly unit-testable.