torch.distributions.NegativeBinomial

class NegativeBinomial extends Distribution

new NegativeBinomial(options: {
total_count: number | Tensor;
probs?: number | Tensor;
logits?: number | Tensor;
} & DistributionOptions)

readonlytotal_count(Tensor): – Number of successes (shape parameter).
readonlyarg_constraints(unknown)
readonlysupport(unknown)
readonlyprobs(Tensor)
readonlylogits(Tensor)
readonlymean(Tensor)
readonlymode(Tensor)
readonlyvariance(Tensor)

Negative Binomial distribution: flexible distribution for overdispersed count data.

Parameterized by total_count (r, number of successes) and probs (p, success probability). Represents the number of failures that occur BEFORE achieving total_count successes in a sequence of independent Bernoulli trials. Unlike Poisson, it allows mean ≠ variance, crucial for real-world count data which is almost always overdispersed. Essential for:

Count data regression (overdispersed counts more common than Poisson)
Modeling sequences with variance > mean (distinguishing feature)
RNA-seq and genomics (standard model for gene expression counts)
Failure counting (how many failures before r successes?)
Zero-inflated count models (NB is mixture baseline)
Ranking and recommendation systems (pairwise comparisons)
Network analysis (degree distributions often follow NB)
Gamma-Poisson mixture representation

Key Advantage over Poisson: Poisson assumes mean = variance (restrictive). NegativeBinomial allows variance > mean (overdispersion), matching real data. For counts with dispersion parameter > 1, use NegativeBinomial instead of Poisson.

Mixture Interpretation: NegativeBinomial(r, p) = ∫ Poisson(λ) dP(λ) where P(λ) = Gamma(r, (1-p)/p). This Gamma-Poisson mixture naturally accommodates heterogeneity across observations.

\begin{aligned} \text{PMF: } P(X = k) = \binom{k+r-1}{k} p^r (1-p)^k \quad k = 0,1,2,\ldots \\ \text{Mean: } \mathbb{E}[X] = \frac{r(1-p)}{p} \\ \text{Variance: } \text{Var}(X) = \frac{r(1-p)}{p^2} = \text{Mean}^2/r + \text{Mean} \\ \text{Dispersion: } \frac{\text{Var}(X)}{\text{Mean}} = \frac{1-p}{p} \quad \text{(>1 for p<0.5)} \\ \text{Mode: } \lfloor \frac{(r-1)(1-p)}{p} \rfloor \quad \text{for } r > 1 \end{aligned}

Variance = Mean × (1 + Mean/r): Overdispersion increases as r decreases
Poisson special case: As r → ∞, NegativeBinomial → Poisson(mean)
Relationship to Binomial: NB is NOT generalization of Binomial (different support)
Mixture interpretation: Gamma-Poisson mixture models heterogeneity
Dispersion parameter: r controls how much variance exceeds mean
Mode is 0: For many practical cases, P(X=0) is substantial
Overdispersion metric: Var/Mean = 1 is Poisson; 1 means NB needed

r must be positive: r ≤ 0 causes errors
Overdispersion required: If variance ≈ mean, use Poisson (simpler)
Small r gives extreme variance: Very small r produces heavy right tail
Probs near 1: When p → 1, mean → 0 but variance doesn't vanish

Examples

// Simple case: count failures before 5 successes with 30% success rate
const nb = new torch.distributions.NegativeBinomial({
  total_count: 5,
  probs: 0.3  // 30% success rate per trial
});
const sample = nb.sample();  // typical: 10-20 (mean = 11.67)
const samples = nb.sample([1000]);  // 1000 failure counts

// Gene expression modeling: RNA-seq counts are overdispersed
// Size parameter (r) inversely controls dispersion
const gene_size = 0.5;  // size parameter (smaller = more overdispersion)
const gene_probs = 0.1;  // gene expression probability
const gene_dist = new torch.distributions.NegativeBinomial({
  total_count: gene_size,
  probs: gene_probs
});
const gene_counts = gene_dist.sample([10000, 20000]);  // 10k genes, 20k samples
// Variance >> mean, as observed in RNA-seq
n *
// Batched distributions with different parameters
const total_counts = torch.tensor([1, 5, 10, 20]);
const probs = torch.tensor([0.3, 0.3, 0.3, 0.3]);
const dist = new torch.distributions.NegativeBinomial({ total_count: total_counts, probs });
const samples = dist.sample();  // [4] shaped samples
const means = dist.mean;  // [2.33, 11.67, 23.33, 46.67]
const vars = dist.variance;  // [7.78, 38.89, 77.78, 155.56]
n *
// Comparing to Poisson: overdispersion
const poisson = new torch.distributions.Poisson(5);
const poisson_var = poisson.variance;  // 5 (mean = variance)\n * const nb = new torch.distributions.NegativeBinomial({
  total_count: 2,
  probs: torch.tensor([2.0/7.0])
});
const nb_mean = nb.mean;  // 5
const nb_var = nb.variance;  // 12.5 (variance > mean!)

// Failure analysis: how many defects before 10 products pass QC?
// If 90% pass QC (p=0.9), expecting failures before 10 passes
const pass_prob = 0.9;
const passing_required = 10;
const qa_dist = new torch.distributions.NegativeBinomial({
  total_count: passing_required,
  probs: pass_prob
});
const expected_defects = qa_dist.mean.item();  // (10 * 0.1) / 0.9 ≈ 1.1