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

기술 요약: 예측을 위한 학습 및 추론의 재고

문제 정의
자율 주행을 위한 궤적 예측 기술은 크게 발전하였으며, 많은 최첨단 모델들은 에이전트의 미래 행동이 갖는 다중 모드(multi-modal) 특성을 포착하기 위해 여러 후보 궤적을 생성합니다. 이러한 모델들은 일반적으로 가우시안 혼합 모델(GMM)을 표현하기 위해 의도된 $K$개의 후보(흔히 $K=64$)와 연관된 확률들을 출력합니다. 그러나 결정적인 병목 현상이 존재합니다: 예측된 모드 확률들이 정보력이 없는 경우가 많다는 점입니다. 모델들은 종종 올바르고 일반적인 행동에 낮은 가능성(likelihood)을 할당하거나, 유사한 후보들 사이에서 불안정한 확률을 보입니다. 이는 표준적인 선택 방식(예: 상위-$K$ 확률에 대한 그리디 선택)이 가장 확률이 높은 미래 궤적을 검색하는 데 실패하게 만들어, 결과적으로 지나치게 보수적이거나 안전하지 못한 계획 행동을 초래할 수 있습니다.

저자들은 그 근본 원인을 **모델링-학습 불일치(modeling-training mismatch)**로 지목합니다. 이러한 모델들은 GMM을 표현하도록 설계되었지만, 실제로는 승자 독식(Winner-Take-All, WTA) 손실 함수를 사용하여 학습됩니다. WTA 목적 함수는 모든 모드가 하나의 예측값으로 수렴하는 모드 붕괴(mode collapse)를 방지하기 위해, 각 정답 샘플을 단 하나의 가장 가까운 예측 모드에 할당(hard assignment)합니다. 저자들은 WTA가 다양성을 장려하는 데는 성공하지만, 이 하드 할당이 학습 목표를 순수한 GMM 추정에서 K-평균(K-Means) 클러스터링 목표로 근본적으로 변화시킨다고 주장합니다. 이러한 변화는 두 가지 주요 문제를 야기합니다:

1. 과분할(Over-segmentation): 하나의 진정한 행동 모드(예: "우회")가 여러 개의 별개로 예측된 모드로 파편화되어, 각 모드가 데이터의 일부를 점유하게 됩니다.
2. 정보력이 없는 사후 확률(Uninformative Posteriors): 단일 진정한 행동의 확률 질량이 여러 파편화된 모드에 분산됨에 따라, 개별 모드의 확률이 낮아져 순위 지정이나 선택에 신뢰할 수 없게 됩니다.

방법론
본 논문은 WTA로 학습된 예측기를 순수 GMM이 아닌 K-평균 클러스터링의 관점에서 바라보는 통합된 시각을 제안합니다. 이러한 통찰을 바탕으로, 저자들은 기존 베이스 아키텍처를 재학습할 필요가 없는 두 가지 경량화된 사후 처리(post-hoc) 기법을 도입합니다:

1. 테스트 시점 사후 확률 가중 병합(Test-Time Posterior-Weighted Merging):
   WTA가 과분할을 초래한다는 점을 인식하여, 이 방법은 인접한 후보 궤적들을 집계합니다. 이는 다음과 같은 반복적인 과정을 수행합니다:

   • 초기 예측된 궤적들을 가장 가까운 클러스터 중심점에 할당합니다(최종 웨이포인트의 유클리드 거리 기준).
   • 예측된 혼합 확률을 가중치로 사용하는 가중 평균 공식을 사용하여 클러스터 중심점을 재계산합니다.
   • 수렴할 때까지 이 과정을 반복합니다.
     이는 효과적으로 파편화된 모드들을 일관된 행동 클러스터로 병합하여, 진정한 분포 성분을 복구합니다.

2. 단일 단계 EM(Expectation-Maximization) 업데이트:
   예측 확률의 정보력을 직접 해결하기 위해, 저자들은 학습 시 사용된 하드 레이블을 소프트 책임(soft responsibilities)으로 대체할 것을 제안합니다. 이는 다음과 같은 단일 EM 단계를 사후 처리 또는 미세 조정 단계로 적용하는 것을 포함합니다:

   • E-단계(E-step): 정답 및 현재 모델 파라미터가 주어졌을 때 각 모드에 대한 소프트 책임($w_k$)을 계산하여, 실질적으로 사후 확률 $p(k | \mathbf{y}^*)$를 구합니다.
   • M-단계(M-step): 이러한 소프트 할당 하에서 가능도(likelihood)를 최대화하도록 혼합 가중치($\hat{p}_k$)를 업데이트합니다.
     이 단계는 인접한 모드들 사이에 확률 질량을 공유함으로써 "하드" 할당 편향을 교정하고, 더 충실한 모드 순위를 생성합니다.

주요 기여

• 이론적 통합: 본 논문은 WTA 학습 목적 함수를 K-평균 클러스터링의 변형으로 연결함으로써, 왜 WTA로 학습된 모델이 다양성 확보에는 성공하면서도 과분할 및 정보력 없는 확률을 보이는지에 대한 원리적인 설명을 제공합니다.
• 실용적 교정: 저자들은 추론 시의 확률적 가정과 학습 목적을 일치시키는 두 가지 사후 처리 전략(병합 및 단일 단계 EM)을 도입합니다. 이 방법들은 기존 딥러닝 아키텍처를 재학습하지 않고도 모드의 사후 확률 품질을 개선합니다.
• 경험적 검증: 제안된 방식은 NuScenes 및 Waymo Open Motion (WOMD) 데이터셋에서 세 가지 대표적인 WTA 학습 모델(Wayformer, MTR-e2e, EMP)을 통해 검증되었습니다.

결과
실험 결과, 제안된 처리가 예측 성능을 크게 향상시키는 것으로 나타났습니다:

• 병합(Merging): 클러스터링 기반 병합은 변위 지표(minADE, minFDE) 및 Brier-minFDE에서 단순 그리디 선택 및 비최대 억제(NMS)보다 우수한 성능을 보였습니다. 일부 사례(예: MTR-e2e)에서는 병합 방식이 최종 모드 수($M=5$)를 직접 예측하도록 학습된 모델보다 약간 더 나은 성능을 보이기도 했습니다.
• 단일 단계 EM: EM 업데이트를 적용하면 그리디 및 NMS 선택 방식의 정확도가 향상됩니다. NuScenes에서 단일 단계 EM을 통한 미세 조정은 기존 그리디 선택과 비교했을 때 Wayformer의 경우 Brier-minFDE를 최대 29%, MTR-e2e의 경우 20%까지 감소시켰습니다.
• 가능도 품질: EM 단계 이후 정답에 대한 음의 로그 가능도(NLL)가 크게 개선되었으며, 이는 교정된 확률이 실제 데이터 분포를 더 잘 반영함을 나타냅니다. minFDE 오차 분포는 더 정확한 예측 쪽으로 이동하며, 확률 질량이 올바른 미래 행동에 더욱 효과적으로 집중됩니다.

의의 및 주장
본 논문은 자신의 분석이 모션 예측 분야의 광범위한 문제에 대해 "원리적이고 실용적인 교정"을 제공한다고 주장합니다. WTA 목적 함수를 K-평균의 변형으로 재정의함으로써, 저자들은 왜 표준적인 사후 선택이 실패하는지, 그리고 이를 어떻게 해결할 수 있는지에 대한 이론적 근거를 제시합니다. 이 연구의 의의는 초기 모델 개발 단계에서 전체 재학습의 비용이나 반복적인 EM 학습의 복잡성을 들이지 않고도, 학습 목적과 추론 시의 확률적 모델링 가정을 일치시키는 방법을 제공한다는 점에 있습니다. 저자들은 모드 수의 유연성이 필요한 경우 클러스터링 기반 병합을 사용하고, 추론 속도를 우선시하는 경우 단일 단계 EM 미세 조정을 적용하여 표준 그리디 선택 방식보다 더 높은 정확도를 달성할 수 있다고 제안합니다.