diff --git a/deeptab/distributions/negative_binomial.py b/deeptab/distributions/negative_binomial.py index 141cac89..bd0e02b6 100644 --- a/deeptab/distributions/negative_binomial.py +++ b/deeptab/distributions/negative_binomial.py @@ -37,9 +37,11 @@ def compute_loss(self, predictions, y_true): mean = self.mean_transform(predictions[:, self.param_names.index("mean")]) dispersion = self.dispersion_transform(predictions[:, self.param_names.index("dispersion")]) - # variance = mean + mean^2 / dispersion + # size r = 1 / dispersion, so variance = mean + dispersion * mean^2. + # torch's NegativeBinomial has mean = total_count * probs / (1 - probs), + # so probs must be mean / (r + mean) for the mean head to be the mean. r = torch.tensor(1.0) / dispersion # type: ignore[operator] - p = r / (r + mean) + p = mean / (r + mean) negative_binomial_dist = dist.NegativeBinomial(total_count=r, probs=p) diff --git a/tests/test_distributions.py b/tests/test_distributions.py index ada5ec68..8c2aa018 100644 --- a/tests/test_distributions.py +++ b/tests/test_distributions.py @@ -632,6 +632,40 @@ def test_loss_requires_grad(self): self.dist.compute_loss(preds, self.y).backward() assert preds.grad is not None + def test_mean_head_parameterizes_the_mean(self): + """Regression test: probs must follow torch's convention. + + The distribution built from a transformed mean of mu must actually + have mean mu (with the scipy-style probs it would be r^2/mu). + """ + import torch + import torch.distributions as dist + + mu = torch.tensor([0.5, 1.0, 10.0, 100.0]) + dispersion = torch.tensor([0.5, 1.0, 2.0, 4.0]) + r = 1.0 / dispersion + p = mu / (r + mu) + nb = dist.NegativeBinomial(total_count=r, probs=p) + assert torch.allclose(nb.mean, mu, rtol=1e-4) + + def test_nll_minimized_near_true_mean(self): + """The NLL of data with mean mu must be lower at mean=mu than far away.""" + import torch + + torch.manual_seed(0) + y = torch.poisson(torch.full((512,), 10.0)) + disp_raw = torch.zeros(512) + + def nll_at(mean_value): + # positive transform is softplus; invert it (in float64 to avoid + # expm1 overflow) to hit the target mean + raw = torch.log(torch.expm1(torch.tensor(mean_value, dtype=torch.float64))).float() + preds = torch.stack([raw.expand(512), disp_raw], dim=1) + return self.dist.compute_loss(preds, y).item() + + assert nll_at(10.0) < nll_at(1.0) + assert nll_at(10.0) < nll_at(100.0) + # --------------------------------------------------------------------------- # StudentTDistribution