From 424f2aaf1b484fc495404e7d284c19315c182590 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Mon, 27 Jul 2026 21:24:04 +0200
Subject: [PATCH] fix(distributions): use torch probs convention in
 NegativeBinomial

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>
---
 deeptab/distributions/negative_binomial.py |  6 ++--
 tests/test_distributions.py                | 34 ++++++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

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