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

Resumo Técnico: Repensando o Treinamento e a Inferência para Previsão

Definição do Problema
A previsão de trajetórias para condução autônoma avançou significamente, com muitos modelos de estado da arte gerando múltiplos candidatos para capturar a natureza multimodal do comportamento futuro de agentes. Esses modelos tipicamente produzem um conjunto de $K$ candidatos (frequentemente $K=64$) com probabilidades associadas, destinados a representar um Modelo de Mistura Gaussiana (GMM) condicional. No entanto, existe um gargalo crítico: as probabilidades dos modos previstos são frequentemente pouco informativas. Os modelos frequentemente atribuem verossimilhanças baixas a comportamentos corretos e comuns ou exibem probabilidades instáveis entre candidatos similares. Isso leva a uma poda de modos ineficaz, onde os métodos de seleção padrão (por exemplo, seleção gananciosa dos top-$K$ probabilidades) falham em recuperar as trajetórias futuras mais prováveis, podendo causar comportamentos de planejamento excessivamente conservadores ou inseguros.

Os autores identificam a causa raiz como um descompasso entre modelagem e treinamento. Embora esses modelos sejam projetados para representar GMMs, eles são treinados usando uma função de perda Winner-Take-All (WTA). O objetivo WTA atribui cada amostra de verdade fundamental (ground-truth) ao seu modo previsto mais próximo (atribuição rígida/hard assignment) para evitar o colapso de modo (onde todos os modos convergem para uma única previsão). Os autores argumentam que, embora o WTA incentive o sucesso na diversidade, essa atribuição rígida altera fundamentalmente o objetivo de aprendizagem de uma estimativa probabilística de GMM para um objetivo de agrupamento K-Means. Essa mudança causa dois problemas principais:

1. Sobre-segmentação: Um único modo comportamental verdadeiro (ex: "virar à direita") é fragmentado em múltiplos modos previstos distintos, cada um reivindicando uma parte dos dados.
2. Posteriores Pouco Informativos: Como a massa de probabilidade de um único comportamento verdadeiro é dividida entre muitos modos fragmentados, as probabilidades individuais dos modos parecem baixas, tornando-as não confiáveis para classificação ou seleção.

Metodologia
O artigo propõe uma perspectiva unificada visualizando os previsores treinados por WTA através da lente do agrupamento K-Means, em vez de GMMs puros. Com base nessa percepção, os autores introduzem dois tratamentos post-hoc leves que não exigem o retreinamento da arquitetura base:

1. Fusão de Posteriores Ponderada pelo Tempo de Teste (Test-Time Posterior-Weighted Merging):
   Reconhecendo que o WTA leva à sobre-segmentação, este método agrega trajetórias candidatas próximas. Ele emprega um processo iterativo semelhante ao agrupamento K-Means:

   • Atribui as trajetórias previstas iniciais aos centros de cluster mais próximos (baseado na distância euclidiana dos pontos finais/waypoints).
   • Recalcula os centros dos clusters usando uma fórmula de média ponderada, onde os pesos são as probabilidades de mistura previstas.
   • Repete até a convergência.
     Isso efetivamente funde modos fragmentados em clusters comportamentais coerentes, recuperando os componentes reais da distribuição.

2. Atualização de Expectativa-Maximização (EM) de Passo Único:
   Para abordar diretamente as probabilidades pouco informativas, os autores propõem substituir os rótulos rígidos usados durante o treinamento por responsabilidades suaves (soft responsibilities). Isso envolve um passo de EM único aplicado como um passo de pós-processamento ou ajuste fino (fine-tuning):

   • Passo E (Expectativa): Calcula responsabilidades suaves ($w_k$) para cada modo dada a verdade fundamental e os parâmetros atuais do modelo, calculando efetivamente a probabilidade posterior $p(k | \mathbf{y}^*)$.
   • Passo M (Maximização): Atualiza os pesos da mistura ($\hat{p}_k$) para maximizar a verossimilhança sob essas atribuições suaves.
     Este passo compartilha a massa de probabilidade entre modos vizinhos, corrigindo o viés de atribuição "rígida" e gerando classificações de modos mais fiéis.

Principais Contribuições

• Unificação Teórica: O artigo fornece uma explicação fundamentada ligando o objetivo de treinamento WTA ao agrupamento K-Means, esclarecendo por que modelos treinados por WTA exibem sobre-segmentação e probabilidades pouco informativas, apesar do sucesso na diversidade.
• Correções Práticas: Os autores introduzem duas estratégias post-hoc (fusão e EM de passo único) que alinham as suposições probabilísticas do tempo de inferência com os objetivos de treinamento. Esses métodos melhoram a qualidade dos posteriors dos modos sem retreinar as arquiteturas de aprendizado profundo subjacentes.
• Validação Empírica: A abordagem é validada em três modelos representativos treinados por WTA (Wayformer, MTR-e2e e EMP) nos conjuntos de dados NuScenes e Waymo Open Motion (WOMD).

Resultos
Experimentos demonstram que os tratamentos propostos melhoram significativamente o desempenho da previsão:

• Fusão (Merging): A fusão baseada em agrupamento supera a seleção gananciosa ingênua e a Supressão de Não-Máximo (NMS) em métricas de deslocamento (minADE, minFDE) e Brier-minFDE. Em alguns casos (ex: MTR-e2e), a fusão chega a superar levemente modelos treinados diretamente para prever o número final de modos ($M=5$).
• EM de Passo Único: Aplicar a atualização EM melhora a precisão dos métodos de seleção gananciosa e NMS. No NuScenes, o ajuste fino com EM de passo único reduziu o Brier-minFDE em até 29% para o Wayformer e 20% para o MTR-e2e em comparação com a seleção gananciosa original.
• Qualidade da Verossimilhança: A log-verossimilhança negativa (NLL) da verdade fundamental melhora significativamente após o passo de EM, indicando que as probabilidades corrigidas refletem melhor a distribuição real dos dados. A distribuição dos erros de minFDE desloca-se para previsões mais precisas, e a massa de probabilidade concentra-se de forma mais eficaz nos comportamentos futuros corretos.

Significância e Alegações
O artigo alega que sua análise oferece uma "correção fundamentada e prática" para um problema generalizado na previsão de movimento. Ao reenquadrar o objetivo WTA como uma variante de K-Means, os autores fornecem um embasamento teórico para o porquê de a seleção post-hoc padrão falhar e como corrigi-la. A significância reside em oferecer uma maneira de alinhar os objetivos de treinamento com as suposições de modelagem probabilística do tempo de inferência sem o custo computacional de um retreinamento total ou a complexidade de um treinamento EM iterativo durante o desenvolvimento inicial do modelo. Os autores sugerem que os profissionais devem usar a fusão baseada em agrupamento se a flexibilidade no número de modos for necessária, ou aplicar o ajuste fino de EM de passo único se a velocidade de inferência for priorizada, alcançando assim uma precisão maior do que os métodos de seleção gananciosa padrão.