fix(distributions): use torch probs convention in NegativeBinomial - #429
Open
ChrisW09 wants to merge 1 commit into
Open
fix(distributions): use torch probs convention in NegativeBinomial#429ChrisW09 wants to merge 1 commit into
ChrisW09 wants to merge 1 commit into
Conversation
probs was computed as r/(r+mean), the scipy convention, but torch's NegativeBinomial defines mean = total_count * probs/(1-probs). The constructed distribution therefore had mean r^2/mu instead of mu: the network learned to output a value that is not the mean, so predict(), _extract_mean-based RMSE/MAE, and NegativeBinomialDeviance (which reads column 0 as mu) all received garbage. probs is now mean/(r+mean). Fixes #418 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #418
Problem
NegativeBinomialDistribution.compute_lossbuilt the distribution with the scipy success-probability convention:torch defines
mean = total_count * probs / (1 - probs), soprobsmust bemean / (r + mean). As written, the distribution mean is r²/µ instead of µ (verified: µ=10, r=1 → distribution mean 0.1).Training still converges — the network just learns whatever head output makes
softplus(head)equalr²/µ— but the column trained as "mean" is not the mean, sopredict(), mean-based metrics, andNegativeBinomialDevianceall silently consume garbage.Fix
p = mean / (r + mean), plus a comment recording the convention and the resulting variance form.Tests
test_mean_head_parameterizes_the_mean— the distribution built from a transformed mean of µ actually has mean µ.test_nll_minimized_near_true_mean— NLL of data with mean 10 is lower at mean=10 than at 1 or 100.Full
tests/test_distributions.pypasses.🤖 Generated with Claude Code