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

Resumen Técnico: Replanteando el Entrenamiento y la Inferencia para la Predicción

Planteamiento del Problema
La predicción de trayectorias para la conducción autónoma ha avanzado significamente, con muchos modelos de vanguardia generando múltiples candidatos para capturar la naturaleza multimodal del comportamiento de los agentes en el futuro. Estos modelos suelen generar un conjunto de $K$ candidatos (frecuentemente $K=64$) con probabilidades asociadas, destinados a representar un Modelo de Mezcla Gaussiana (GMM) condicional. Sin embargo, existe un cuello de botella crítico: las probabilidades de los modos predichos suelen ser poco informativas. Los modelos frecuentemente asignan verosimilitudes bajas a comportamientos correctos y comunes, o exhiben probabilidades inestables entre candidatos similares. Esto conduce a una poda de modos ineficaz, donde los métodos de selección estándar (por ejemplo, la selección codiciosa de los top-$K$ de probabilidad) fallan al recuperar las trayectorias futuras más probables, lo que puede causar comportamientos de planificación excesivamente conservadores o inseguros.

Los autores identifican la causa raíz como un desajuste entre el modelado y el entrenamiento. Aunque estos modelos están diseñados para representar GMM, se entrenan utilizando una función de pérdida de Ganador-Se-Lleva-Todo (WTA, por sus siglas en inglés). El objetivo WTA asigna cada muestra de verdad fundamental (ground-truth) a su modo predicho más cercano (asignación dura) para evitar el colapso de modos (donde todos los modos convergen a una única predicción). Los autores argumentan que, si bien el WTA fomenta con éxito la diversidad, esta asignación dura desplaza fundamentalmente el objetivo de aprendizaje de una estimación probabilística de GMM hacia un objetivo de agrupamiento K-Means. Este cambio provoca dos problemas principales:

1. Sobresegmentación: Un único modo de comportamiento real (por ejemplo, "girar a la derecha") se fragmenta en múltiples modos predichos distintos, cada uno de los cuales reclama una parte de los datos.
2. Posteriores Poco Informativos: Debido a que la masa de probabilidad de un único comportamiento real se divide entre muchos modos fragmentados, las probabilidades individuales de los modos parecen bajas, lo que las hace poco fiables para el ordenamiento o la selección.

Metodología
El artículo propone una perspectiva unificada que observa los pronosticadores entrenados con WTA a través de la lente del agrupamiento K-Means en lugar de los GMM puros. Basándose en este conocimiento, los autores introducen dos tratamientos post-hoc ligeros que no requieren el reentrenamiento de la arquitectura base:

1. Fusión Ponderada por Posteriores en Tiempo de Prueba:
   Reconociendo que el WTA conduce a la sobresegmentación, este método agrega trayectorias candidatas cercanas. Emplea un proceso iterativo similar al agrupamiento K-Means:

   • Asigna las trayectorias predichas iniciales a los centros de clúster más cercanos (basándose en la distancia euclidiana de los puntos finales de la trayectoria).
   • Recalcula los centros de los clústeres utilizando una fórmula de media ponderada, donde los pesos son las probabilidades de la mezcla predichas.
   • Repite hasta la convergencia.
     Esto fusiona eficazmente los modos fragmentados en clústeres de comportamiento coherentes, recuperando los componentes reales de la distribución.

2. Actualización de Esperanza-Maximización (EM) de un Solo Paso:
   Para abordar directamente las probabilidades poco informativas, los autores proponen reemplazar las etiquetas duras utilizadas durante el entrenamiento con responsabilidades suaves. Esto implica un único paso de EM aplicado como un paso de post-procesamiento o de ajuste fino (fine-tuning):

   • Paso E: Calcula las responsabilidades suaves ($w_k$) para cada modo dada la verdad fundamental y los parámetros actuales del modelo, calculando efectivamente la probabilidad posterior $p(k | \mathbf{y}^*)$.
   • Paso M: Actualiza los pesos de la mezcla ($\hat{p}_k$) para maximizar la verosimilitud bajo estas asignaciones suaves.
     Este paso comparte la masa de probabilidad entre los modos vecinos, corrigiendo el sesgo de la asignación "dura" y produciendo clasificaciones de modos más fieles.

Contribuciones Clave

• Unificación Teórica: El artículo proporciona una explicación fundamentada que vincula el objetivo de entrenamiento WTA con el agrupamiento K-Means, aclarando por qué los modelos entrenados con WTA exhiben sobresegmentación y probabilidades poco informativas a pesar de tener una diversidad exitosa.
• Correcciones Prácticas: Los autores introducen dos estrategias post-hoc (fusión y EM de un solo paso) que alinean las suposiciones probabilísticas en el tiempo de inferencia con los objetivos de entrenamiento. Estos métodos mejoran la calidad de los posteriores de los modos sin reentrenar las arquitecturas de aprendizaje profundo subyacentes.
• Validación Empírica: El enfoque se valida a través de tres modelos entrenados con WTA (Wayformer, MTR-e2e y EMP) en los conjuntos de datos NuScenes y Waymo Open Motion (WOMD).

Resultados
Los experimentos demuestran que los tratamientos propuestos mejoran significativamente el rendimiento de la predicción:

• Fusión: La fusión basada en el agrupamiento supera a la selección codiciosa (greedy) ingenua y a la Supresión de No Máximos (NMS) en métricas de desplazamiento (minADE, minFDE) y Brier-minFDE. En algunos casos (por ejemplo, MTR-e2e), la fusión incluso supera ligeramente a los modelos entrenados para predecir directamente el número final de modos ($M=5$).
• EM de un Solo Paso: Aplicar la actualización EM mejora la precisión de los métodos de selección codiciosa y NMS. En NuScenes, el ajuste fino con la actualización EM de un solo paso redujo el Brier-minFDE hasta en un 29% para Wayformer y un 20% para MTR-e2e en comparación con la selección codiciosa original.
• Calidad de la Verosimilitud: La log-verosimilitud negativa (NLL) de la verdad fundamental mejora significativamente después del paso EM, lo que indica que las probabilidades corregidas reflejan mejor la verdadera distribución de los datos. La distribución de los errores minFDE se desplaza hacia predicciones más precisas, y la masa de probabilidad se concentra de manera más efectiva en los comportamientos futuros correctos.

Significancia y Reivindicaciones
El artículo afirma que su análisis ofrece una "corrección fundamentada y práctica" a un problema generalizado en la predicción de movimiento. Al reformular el objetivo WTA como una variante de K-Means, los autores proporcionan una base teórica de por qué la selección post-hoc estándar falla y cómo solucionarlo. La importancia radica en ofrecer una forma de alinear los objetivos de entrenamiento con las suposiciones de modelado probabilístico en el tiempo de inferencia sin el costo computacional de un reentrenamiento completo o la complejidad de un entrenamiento EM iterativo durante el desarrollo inicial del modelo. Los autores sugieren que los profesionales deben utilizar la fusión basada en el agrupamiento si se requiere flexibilidad en el número de modos, o aplicar el ajuste fino de un solo paso EM si se prioriza la velocidad de inferencia, logrando así una mayor precisión que los métodos de selección codiciosa estándar.