🤖 machine learning

Rethinking Training & Inference for Forecasting: Linking Winner-Take-All back to GMMs

This paper identifies a modeling-training mismatch in autonomous driving trajectory forecasting, where winner-take-all losses applied to Gaussian mixture models cause uninformative posteriors, and proposes lightweight post-hoc inference corrections—specifically posterior-weighted merging and a one-step EM update—to recover soft mode assignments and improve forecast accuracy without retraining.

Original authors: Qiyuan Wu, Katie Z Luo, Bharath Hariharan, Wei-Lun Chao, Mark Campbell

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Qiyuan Wu, Katie Z Luo, Bharath Hariharan, Wei-Lun Chao, Mark Campbell

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict where a pedestrian or a car will go next at a busy intersection. You know they might turn left, turn right, or keep going straight. To be safe, a self-driving car needs to prepare for all these possibilities.

This paper tackles a specific problem with how computers currently learn to make these predictions. It argues that while the computers are good at finding the right paths, they are terrible at ranking them by how likely they are to happen.

Here is the breakdown of the paper's story, using simple analogies:

The Problem: The "Best Guess" Trap

Most modern self-driving models try to predict 64 different possible paths (modes) for a car. They are supposed to act like a weather forecaster saying, "There's a 60% chance of rain, 30% chance of sun, and 10% chance of snow."

However, the paper found that these models are trained using a method called "Winner-Take-All" (WTA).

The Analogy: Imagine a teacher grading a student who draws 64 different maps of a city. The teacher only looks at the one map that is closest to the real city and says, "Good job, this one is right!" The teacher ignores the other 63 maps completely.
The Result: The student (the AI) learns to draw 64 different maps that are all slightly different from each other, but it has no idea which one is actually the most likely. It treats all 64 maps as if they are equally important, or it gets confused about which one is the "winner."

The Root Cause: Clustering vs. Probability

The authors discovered that this "Winner-Take-All" training method accidentally turns the AI into a K-Means Clustering machine instead of a Probability machine.

The Analogy: Imagine you have a jar of red, blue, and green marbles.
- What we want (Probability): The AI should say, "This jar is 50% red, 30% blue, 20% green."
- What the AI is doing (K-Means): The AI is told to sort the marbles into piles. Because it's forced to pick a "winner" for every marble, it ends up splitting the "Red" pile into 20 tiny, separate piles of red marbles just to make sure every single marble has a pile.
- The Consequence: Now, instead of one big "Red" pile with a high probability, you have 20 tiny "Red" piles, each with a tiny, confusing probability. If you ask the AI, "What is the chance of Red?", it looks at one tiny pile and says, "Very low!" even though Red is actually the most common color.

This is called Over-segmentation. The AI breaks one big, likely future (like "turning right") into many tiny, unlikely fragments.

The Solution: Two "Post-It Note" Fixes

The authors propose two ways to fix this without having to re-teach the AI from scratch (which would be expensive and slow). Think of these as "post-processing" steps you do after the AI has already done its work.

Fix 1: The "Group Hug" (Test-Time Merging)

Since the AI has broken one big idea into many tiny fragments, we can just glue them back together.

How it works: When the AI outputs 64 paths, we look at them and say, "Hey, these 10 paths are all trying to go to the same place." We merge them into one single, strong path.
The Result: Instead of 10 weak signals, you get 1 strong, clear signal. This makes the final prediction much more accurate.

Fix 2: The "Reality Check" (One-Step EM Update)

The AI is bad at assigning probabilities because it was trained to only care about the "winner." This fix gives the AI a quick "reality check" right before it makes its final decision.

How it works: We take the AI's 64 paths and run a quick mathematical calculation (called Expectation-Maximization) that asks, "If we look at all 64 paths together, how much probability mass should each one actually get?"
The Result: It redistributes the "votes." If 10 paths are all trying to go right, the math realizes they are all part of the "Turn Right" group and gives that group a high probability score, rather than spreading the score thin across 10 tiny groups.

The Bottom Line

The paper shows that by using these two simple tricks (merging similar paths or running a quick probability check), we can make self-driving cars much better at understanding which future is most likely, without needing to retrain the complex AI models.

Before: The AI sees 64 paths, gets confused about which is best, and might pick a wrong one because its probability scores are messy.
After: The AI still sees 64 paths, but we clean up the mess, group the similar ones, and get a clear, accurate ranking of the most likely futures.

The authors tested this on real-world driving data (NuScenes and Waymo datasets) and found that these simple fixes significantly improved the accuracy of the predictions.

Technical Summary: Rethinking Training & Inference for Forecasting

Problem Statement
Trajectory forecasting for autonomous driving has advanced significantly, with many state-of-the-art models generating multiple candidate trajectories to capture the multi-modal nature of future agent behavior. These models typically output a set of $K$ candidates (often $K=64$ ) with associated probabilities, intended to represent a conditional Gaussian Mixture Model (GMM). However, a critical bottleneck exists: the predicted mode probabilities are often uninformative. Models frequently assign low likelihoods to correct, common behaviors or exhibit unstable probabilities among similar candidates. This leads to ineffective mode pruning, where standard selection methods (e.g., greedy selection of top- $K$ probabilities) fail to retrieve the most probable future trajectories, potentially causing overly conservative or unsafe planning behaviors.

The authors identify the root cause as a modeling-training mismatch. While these models are designed to represent GMMs, they are trained using a Winner-Take-All (WTA) loss function. The WTA objective assigns each ground-truth sample to its single nearest predicted mode (hard assignment) to prevent mode collapse (where all modes converge to a single prediction). The authors argue that while WTA successfully encourages diversity, this hard assignment fundamentally shifts the learning objective from a probabilistic GMM estimation to a K-Means clustering objective. This shift causes two primary issues:

Over-segmentation: A single true behavioral mode (e.g., "turn right") is fragmented into multiple distinct predicted modes, each claiming a subset of the data.
Uninformative Posteriors: Because the probability mass of a single true behavior is split across many fragmented modes, individual mode probabilities appear low, making them unreliable for ranking or selection.

Methodology
The paper proposes a unified perspective viewing WTA-trained forecasters through the lens of K-Means clustering rather than pure GMMs. Based on this insight, the authors introduce two lightweight, post-hoc treatments that do not require retraining the base architecture:

Test-Time Posterior-Weighted Merging:
Recognizing that WTA leads to over-segmentation, this method aggregates nearby candidate trajectories. It employs an iterative process similar to K-Means clustering:
- Assign initial predicted trajectories to the closest cluster centers (based on Euclidean distance of final waypoints).
- Recompute cluster centers using a weighted mean formula, where weights are the predicted mixture probabilities.
- Repeat until convergence.
  This effectively merges fragmented modes into coherent behavioral clusters, recovering the true distribution components.
One-Step Expectation-Maximization (EM) Update:
To address the uninformative probabilities directly, the authors propose replacing the hard labels used during training with soft responsibilities. This involves a single EM step applied as a post-processing or fine-tuning step:
- E-step: Compute soft responsibilities ( $w_k$ ) for each mode given the ground truth and current model parameters, effectively calculating the posterior probability $p(k | \mathbf{y}^*)$ .
- M-step: Update the mixture weights ( $\hat{p}_k$ ) to maximize the likelihood under these soft assignments.
  This step shares probability mass across neighboring modes, correcting the "hard" assignment bias and yielding more faithful mode rankings.

Key Contributions

Theoretical Unification: The paper provides a principled explanation linking the WTA training objective to K-Means clustering, clarifying why WTA-trained models exhibit over-segmentation and uninformative probabilities despite successful diversity.
Practical Corrections: The authors introduce two post-hoc strategies (merging and one-step EM) that align the inference-time probabilistic assumptions with the training objectives. These methods improve the quality of mode posteriors without retraining the underlying deep learning architectures.
Empirical Validation: The approach is validated across three representative WTA-trained models (Wayformer, MTR-e2e, and EMP) on the NuScenes and Waymo Open Motion (WOMD) datasets.

Results
Experiments demonstrate that the proposed treatments significantly improve forecasting performance:

Merging: Clustering-based merging outperforms naive greedy selection and Non-Maximum Suppression (NMS) on displacement metrics (minADE, minFDE) and Brier-minFDE. In some cases (e.g., MTR-e2e), merging even slightly outperforms models trained directly to predict the final number of modes ( $M=5$ ).
One-Step EM: Applying the EM update improves the accuracy of greedy and NMS selection methods. On NuScenes, fine-tuning with one-step EM reduced Brier-minFDE by up to 29% for Wayformer and 20% for MTR-e2e compared to the original greedy selection.
Likelihood Quality: The negative log-likelihood (NLL) of the ground truth improves significantly after the EM step, indicating that the corrected probabilities better reflect the true distribution of the data. The distribution of minFDE errors shifts toward more accurate predictions, and probability mass concentrates more effectively on correct future behaviors.

Significance and Claims
The paper claims that its analysis offers a "principled, practical correction" to a widespread issue in motion forecasting. By reframing the WTA objective as a K-Means variant, the authors provide a theoretical grounding for why standard post-hoc selection fails and how to fix it. The significance lies in offering a way to align training objectives with inference-time probabilistic modeling assumptions without the computational cost of full retraining or the complexity of iterative EM training during the initial model development. The authors suggest that practitioners should use clustering-based merging if flexibility in mode count is required, or apply one-step EM fine-tuning if inference speed is prioritized, thereby achieving higher accuracy than standard greedy selection methods.