Negative Binomial `Loss`
Neural Networks
TensorFlow
Keras
Machine Learning

Negative Binomial `Loss` in Neural Network using Tensorflow / Keras

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Negative binomial loss is useful when your target is count data and the variance is larger than the mean. In that setting, mean-squared error is often a poor fit, while a negative-binomial likelihood gives the model a loss function that matches the data-generating assumptions much more naturally.

Why Use Negative Binomial Loss

Count targets such as:

  • number of visits
  • number of clicks
  • number of claims

are often overdispersed. That means the variance is larger than the mean. A Poisson loss assumes mean and variance are equal, so it can struggle when the data is more variable than that.

The negative binomial distribution introduces an extra dispersion parameter, which makes it much more flexible for real-world count regression.

A Practical Parameterization

One common parameterization uses:

  • 'mu for the mean'
  • 'theta for the dispersion'

The model predicts mu, while theta may be fixed or learned separately. To keep mu positive, the output is usually transformed with softplus or exp.

A Custom Keras Loss

Here is a simple loss implementation with a fixed dispersion parameter:

python
1import tensorflow as tf
2
3
4def negative_binomial_loss(theta):
5    theta = tf.cast(theta, tf.float32)
6
7    def loss_fn(y_true, y_pred):
8        y_true_f = tf.cast(y_true, tf.float32)
9        mu = tf.nn.softplus(y_pred) + 1e-8
10
11        log_prob = (
12            tf.math.lgamma(y_true_f + theta)
13            - tf.math.lgamma(theta)
14            - tf.math.lgamma(y_true_f + 1.0)
15            + theta * tf.math.log(theta / (theta + mu))
16            + y_true_f * tf.math.log(mu / (theta + mu))
17        )
18
19        return -tf.reduce_mean(log_prob)
20
21    return loss_fn

This computes the negative log-likelihood, which is what the optimizer minimizes.

A Small Keras Model Example

python
1import numpy as np
2import tensorflow as tf
3
4x = np.random.rand(200, 4).astype("float32")
5y = np.random.negative_binomial(n=3, p=0.4, size=(200, 1)).astype("float32")
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Input(shape=(4,)),
9    tf.keras.layers.Dense(16, activation="relu"),
10    tf.keras.layers.Dense(1),
11])
12
13model.compile(
14    optimizer="adam",
15    loss=negative_binomial_loss(theta=2.0),
16)
17
18model.fit(x, y, epochs=5, batch_size=32, verbose=0)
19pred = tf.nn.softplus(model.predict(x[:3], verbose=0))
20print(pred.numpy())

The final layer returns an unconstrained value, and the loss converts it to a positive mean using softplus.

Fixed Dispersion Versus Learned Dispersion

The example above uses a fixed theta, which is simpler and often good enough for a first model. A more advanced setup predicts both:

  • the mean
  • the dispersion

That makes the output head and the loss more complex, because both parameters must remain in valid ranges. If you only need a working negative-binomial regression baseline, start with a fixed dispersion value and move to learned dispersion later.

Why Positivity Constraints Matter

A negative binomial mean cannot be negative. If the model outputs raw values directly into the likelihood without a positive transform, the loss becomes invalid or numerically unstable.

That is why transforms such as:

  • 'tf.nn.softplus'
  • 'tf.exp'

are so common in count-model implementations.

softplus is often nicer numerically because it grows more gently than exp.

Common Pitfalls

  • Using a count-distribution loss on targets that are not counts.
  • Forgetting to enforce positivity on the predicted mean.
  • Treating Poisson and negative binomial as interchangeable even when the data is clearly overdispersed.
  • Making the loss numerically unstable by allowing mu to reach zero exactly.
  • Jumping straight to a learned dispersion head before validating the simpler fixed-dispersion version.

Summary

  • Negative binomial loss is a strong choice for overdispersed count targets.
  • In Keras, you can implement it as a custom negative log-likelihood.
  • The predicted mean must stay positive, typically through softplus or exp.
  • A fixed dispersion parameter is the simplest place to start.
  • Use this loss when the count structure matters; otherwise a generic regression loss may be the wrong statistical model.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.