Poisson Regression by Maximum Likelihood

Fit a model for count data by maximum likelihood, with torch.distributions.

Counts such as clicks, arrivals or defects are whole numbers that are zero or more. Their spread grows with their mean. A Poisson regression models each count as a draw from a Poisson distribution whose rate depends on the features through a log link:

rate_i = exp(X_i @ w + b)                 # always positive
y_i    ~ Poisson(rate_i)
NLL    = -mean(log P(y_i | rate_i))

Write two functions:

• poisson_nll(X, y, w, b) returns the mean negative log-likelihood as a scalar tensor. Compute it with torch.distributions.Poisson(rate).log_prob(y). That includes the log(y!) term, so the value is a true log-likelihood.
• fit_poisson_regression(X, y) minimises poisson_nll with a torch optimiser and returns (w, b) as detached tensors.

Every loss you have used is a negative log-likelihood. Mean squared error comes from a Gaussian, and cross-entropy comes from a categorical distribution. Here you choose the distribution and get the loss from it.

Function signature

def poisson_nll(X, y, w, b):
def fit_poisson_regression(X, y):

Example

Input

X = torch.tensor([[0.0], [1.0]])
y = torch.tensor([1.0, 3.0])
poisson_nll(X, y, w=torch.tensor([1.0]), b=torch.tensor(0.0))

Output

tensor(1.2550)

Reasoning

The rates are exp(0) = 1 and exp(1) = e. log P(1 | 1) = -1, and log P(3 | e) = 3 * 1 - e - log(3!) = -1.5100. The mean of the negated values is (1 + 1.5100) / 2 = 1.2550.

Constraints and notes

• X has shape (n, d) and y has shape (n,), both float32. y holds whole numbers that are zero or more, stored as floats.
• w has shape (d,) and b has shape (). There is no regularisation.
• The fitted coefficients must match the maximum likelihood estimate to 0.01.