Rethinking Training and Inference for Trajectory Forecasting: Linking Winner-Take-All back to GMMs
Abstract
Trajectory forecasting for autonomous driving has advancedrapidly, yet representative models often produce uninformative posteri-ors over forecast modes, causing problems for mode pruning. We tracethis to a modeling-training mismatch: forecasters are typically modeledas conditional Gaussian mixture models (GMMs) but trained with awinner-take-all (WTA) loss that assigns each sample to its nearest mode.We argue that this K-means-like hard assignment (one-hot), while pre-venting mode collapse, is the source of uninformative mode probabil-ities: it over-segments the trajectory space, ignores relatedness amongnearby modes, and yields assignment instability under small perturba-tions. Guided by this lens, we introduce two post-hoc treatments: (1)test-time posterior-weighted merging that aggregates nearby candidatetrajectories; and (2) a one-step expectation-maximization (EM) updatethat replaces hard labels with soft responsibilities, sharing probabilitymass across neighboring modes. Across several WTA-trained architec-tures, these lightweight steps produce more informative, faithfully rankedmode posteriors and strengthen final forecasts on popular displacementmetrics—without retraining. Our analysis unifies recent design choicesthrough a GMM-vs-K-means perspective and offers principled, practicalcorrections that better align training objectives with inference.