wenmin-wu/timeseries-k-mode-gaussian-nll-loss

Overview

Single-trajectory regression (MSE) assumes one future is correct. Real driving, motion, and forecasting problems are multimodal — a vehicle might turn left or go straight, and both are plausible. The standard fix is to predict K trajectories plus K confidence weights, and optimize the negative log-likelihood of the ground truth under a mixture of K isotropic Gaussians (variance = 1). The math simplifies to -logsumexp(log(conf) - 0.5 * Σ error²) along the K axis, which rewards any mode that nails the ground truth while letting the confidences learn the mixture weights. Used in Lyft Motion Prediction, Waymo Open, Argoverse, and most AV trajectory benchmarks.

Quick Start

import torch
import numpy as np

def neg_multi_log_likelihood(gt, pred, confidences, avails):
    """
    gt:          (B, T, 2)    ground-truth trajectory
    pred:        (B, K, T, 2) K candidate trajectories
    confidences: (B, K)       softmax over modes
    avails:      (B, T)       1/0 mask for valid future steps
    """
    gt = gt.unsqueeze(1)                    # (B, 1, T, 2)
    avails = avails[:, None, :, None]       # (B, 1, T, 1)
    err = torch.sum(((gt - pred) * avails) ** 2, dim=(2, 3))  # (B, K)

    with np.errstate(divide='ignore'):
        err = torch.log(confidences) - 0.5 * err               # (B, K)

    # logsumexp trick for numerical stability
    max_val, _ = err.max(dim=1, keepdim=True)
    err = -(torch.log(torch.sum(torch.exp(err - max_val), dim=1, keepdim=True))
            + max_val).squeeze(1)
    return err.mean()

Workflow

1. Model outputs (pred, confidences) where pred has shape (B, K, T, 2) and confidences is softmax-normalized over K
2. Multiply per-timestep squared error by the availability mask to ignore missing steps
3. Compute per-mode log-likelihood log(conf) - 0.5 * sum(err²)
4. Apply the log-sum-exp trick: subtract the max before exp, add it back after the log
5. Negate and average over the batch

Key Decisions

• Isotropic unit variance: simpler than learning per-mode covariance, and almost always enough. Covariance regression is brittle without a lot of data.
• Softmax confidences, not sigmoid: the mixture weights must sum to 1.
• Logsumexp is mandatory: without it, raw exp underflows to 0 and gradients vanish on well-fit batches.
• vs. winner-take-all MSE: WTA only updates the best mode, freezing the others. NLL trains all modes jointly and produces calibrated confidences.

References

• Lyft: Complete train and prediction pipeline
• Pytorch Baseline - Train