TEMPO: 대형 추론 모델을 위한 테스트 시간 학습 확장 프레임워크

기존의 테스트 시간 학습(TTT)은 추론 과정에서 모델을 업데이트할 때 정답 레이블이 없어 성능이 정체되거나 답변의 다양성이 급격히 떨어지는 한계가 있었다. TEMPO는 레이블이 있는 데이터로 비판 모델(Critic)을 주기적으로 재보정하는 방식을 도입해 추론 단계에서도 모델 성능이 지속적으로 우상향할 수 있음을 증명했다.

기존의 대형 언어 모델은 학습이 끝난 후 가중치가 고정된 상태로 추론을 수행한다. 최근에는 추론 시점에 입력된 문제를 바탕으로 모델을 미세 조정하는 테스트 시간 학습(TTT)이 시도되었으나, 정답(Ground-truth)이 없는 상태에서 모델이 스스로 내놓은 답변에 점수를 매기다 보니 특정 답변 패턴에만 매몰되어 성능이 정체되고 답변의 다양성이 사라지는 '다양성 붕괴' 현상이 발생했다.

TEMPO는 이 문제를 해결하기 위해 통계학의 EM 알고리즘을 도입했다. 추론 중인 문제의 정답 여부를 알 수 없는 '숨겨진 변수'로 설정하고, 이를 예측하는 비판 모델(Critic)을 별도로 운영한다. 핵심은 비판 모델이 길을 잃지 않도록, 이미 정답을 알고 있는 별도의 데이터를 사용해 비판 모델의 기준을 주기적으로 바로잡아주는(Recalibration) 것이다.

이 과정을 통해 모델은 자신이 생성한 답변 중 어떤 것이 실제로 정답에 가까운지 더 정확하게 판단할 수 있게 된다. 결과적으로 모델은 추론 단계에서 더 많은 계산 자원을 투입할수록 정답률이 계속해서 높아지는 선순환 구조를 갖추게 되며, 수학이나 논리 퍼즐처럼 복잡한 사고가 필요한 영역에서 비약적인 성능 향상을 달성한다.

TEMPO는 정책 모델(Actor)과 비판 모델(Critic)로 구성된 액터-크리틱 구조를 기반으로 하며, 전체 과정을 EM 알고리즘의 관점에서 정형화했다.

E-step(비판 모델 재보정)에서는 레이블이 있는 데이터셋 $D_L$을 사용하여 비판 모델 $V_\phi$를 업데이트한다. 토큰 수준의 가치 추정치와 실제 정답 여부 사이의 평균 제곱 오차(MSE)를 계산하여 이를 최소화하는 방향으로 가중치를 갱신한다. [정답 레이블과 모델의 예측 점수 차이를 입력으로] → [L2 Loss 연산을 수행해] → [오차 값을 얻고] → [이 값이 작아지도록 비판 모델을 학습시켜 보상 신호의 정확도를 높인다.]

M-step(정책 최적화)에서는 레이블이 없는 테스트 문제 $D_u$에 대해 정책 모델이 생성한 궤적을 비판 모델로 평가한다. 비판 모델이 부여한 점수를 보상으로 활용하여 정책 경사(Policy Gradient) 방법으로 모델 파라미터 $\theta$를 최적화한다. [비판 모델의 가치 예측값과 실제 보상의 차이(Advantage)를 입력으로] → [로그 확률에 곱하는 연산을 수행해] → [기대 보상을 최대화하는 방향의 그래디언트를 얻고] → [정책 모델이 더 나은 추론 경로를 선택하도록 가중치를 갱신한다.]

AIME 2024 벤치마크에서 OLMO3-7B 모델의 성능을 기존 33.0%에서 51.1%로, Qwen3-14B 모델은 42.3%에서 65.8%로 크게 향상시켰다. 특히 기존 TTT 방식인 TTRL이나 EMPO가 학습 초기 이후 성능이 정체되거나 하락하는 것과 대조적으로, TEMPO는 학습 단계가 진행됨에 따라 성능이 꾸준히 상승하는 우수한 확장성을 보였다.

답변의 다양성 측면에서도 pass@k 지표를 통해 분석한 결과, 기존 방식들은 특정 답변으로 수렴하며 다양성이 붕괴되었으나 TEMPO는 높은 정확도를 유지하면서도 다양한 해결 경로를 탐색하는 능력을 보존했다. 수학 외에도 BigBenchHard(BBH), ZebraLogic 등 일반 논리 추론 과제에서도 OLMO3-7B 기준 각각 +21.4%p, +12.9%p의 성능 향상을 기록하며 범용성을 입증했다.

TEMPO는 테스트 시간 학습을 증거 하한(ELBO) 최적화 문제로 정의한다. 추론 시 정답 여부를 관찰되지 않은 latent variable로 취급하여, 사후 분포 $q(y|x)$를 추정하는 과정(E-step)과 이를 바탕으로 하한을 최대화하는 과정(M-step)을 반복한다. 기존 TTT 방법론들이 E-step을 생략하거나 단순 다수결(Majority Voting)로 대체하여 사후 분포가 실제 정답 분포에서 멀어지는 문제를 해결한 것이 핵심 기술적 차별점이다. 구현상으로는 PPO 알고리즘을 활용하며, 오프-폴리시(Off-policy) 학습의 안정을 위해 시퀀스 클리핑 메커니즘을 적용했다.

액터와 크리틱 모델을 동시에 유지해야 하므로 단일 모델 기반 TTT 방식보다 GPU 메모리 점유율과 계산 오버헤드가 높다. 또한, 비판 모델의 재보정을 위해 소량의 레이블된 데이터셋($D_L$)에 대한 접근이 필수적이며, 이 데이터의 분포가 테스트 도메인과 다를 경우 성능에 영향을 미칠 수 있다.