https://arxiv.org/abs/2504.16084
TTRL: Test-Time Reinforcement Learning
This paper investigates Reinforcement Learning (RL) on data without explicit labels for reasoning tasks in Large Language Models (LLMs). The core challenge of the problem is reward estimation during inference while not having access to ground-truth informa
arxiv.org
최신 대형 추론 모델(LRM)들은 긴 CoT(Chain-of-Thought) 성능 향상을 위해 강화학습(RL)과 고비용의 인간 라벨링 데이터를 활용한다. 그러나 최근 등장한 ARC-AGI-2 같은 어려운 벤치마크에서 성능이 급격히 하락하며 라벨 없는 복잡한 문제들에 대한 새로운 접근이 필요하다. 이러한 문제들을 해결하기 위해, 라벨 없는 데이터에서 RL을 수행하는 Test-Time Training의 새로운 형태인 TTRL(Test-Time Reinforcement Learning)을 제안한다.
주요 아이디어
- 테스트 시점에 모델이 스스로 여러 출력을 샘플링 → 다수결을 통해 정답 유사도를 추정 → 이 결과를 보상 함수로 사용해 RL 수행
- 명시적 정답 없이도 모델이 성능을 자기 주도적으로 향상시킬 수 있도록 한다.
Test-Time Reinforcement Learning
TTRL은 기존 RL과 달리 ground-truth 없이, 테스트 데이터만 가지고 보상을 추정해 학습을 진행하는 방식으로 테스트 시점에서 모델을 강화학습 방식으로 적응시키는 것으로 label 없는 reasoning task에 특화되어 있다.
1. Methodology
- 입력 상태($x$): 프롬프트 (ex. 질문, 문장 등)
- 행동 (y): 정책 $\pi_\theta(y|x)$에서 샘플링한 출력
- 보상 구성 방식
- 동일한 프롬프트 $x$에 대해 모델이 여러 개의 후보 출력 {$\hat y_1,\hat y_2, \cdots, \hat y_M$}을 생성
- 다수결(Majority Vote) 또는 다른 집계 방식으로 합의 출력 $y^*$ 도출
- 보상 $R(y, y^*)$는 현재 행동 $y$가 $y^*$와 얼마나 유사한지를 평가
- 최적화 목표:
$$\max_\theta\mathbb{E_{y\sim \pi_\theta(\cdot|x)}}[r(y,y^*)]$$ - 파라미터 업데이트:
$$\theta\leftarrow\theta+\eta\nabla_\theta\mathbb E_{y\sim\pi_\theta(\cdot|x)}[r(y,y^*)]$$
→ 이를 통해 테스트 환경이 학습 환경과 분포가 다를 때도 라벨 없이 스스로 적응하며 성능을 향상시킬 수 있음.
2. Majority Voting Reward Function
- Majority Voting
$$y = {\arg\max}_{y'}\sum_{i=1}^N1[\hat y_i=y']$$ - 보상 함수
$$R(\hat y_i, y)=\begin{cases}1,\quad \mathrm{if}\ \hat y_i =y\\0, \quad \mathrm{otherwise}\end{cases}$$
→ unlabeled test data에서 모델이 자체적으로 생성한 출력들 중 가장 빈도가 높은 것을 pseudo-label로 간주하고, 이를 기준으로 Rule-based reward를 구성하는 방식
Experiments
- AIME 2024: 최소 73.6%, 최고 211.6% 수준의 성능 향상
- AMC: 최소 19.6%, 최고 91.3% 수준의 성능 향상
- MATH-500: 8.5%, 최고 123.2% 수준의 성능 향상
- GPQA: 최소 -4.8%, 최고 13.6% 수준의 성능 향상
→ 수학적인 계산 및 추론 기반 관련된 벤치마크에서는 모델의 성능 향상이 두드러지게 나타남
→ 과학적 개념 이해와 복잡한 논리 추론이 요구되는 GPQA의 경우 성능 향상의 폭이 작거나, 경우에 따라 성능 저하가 발생함
→ 모델 크기가 증가함에 따라 성능이 지속적으로 향상된다
다양한 모델 계열의 모델들(LLaMA, Mistral, DeepSeek)에 대한 실험 결과
고성능 추론 모델(LRM)에서도 우수한 성능을 보임
TTRL은 일반화 성능도 우수
Conclusion
TTRL은 정답 라벨 없이 테스트 데이터에서 직접 생성한 예측값을 기반으로 다수결 규칙에 의해 보상을 생성하고, 이 보상을 활용해 RL 학습을 진행하는 프레임워크이다.
- 라벨 없이도 학습 가능: ground-truth 없이도 학습 가능하도록 설계됨
- 규칙 기반 보상 설계: 모델이 생성한 여러 후보 중 다수결로 정답을 추정하며 보상함
- 범용성: 다양한 RL 알고리즘과 호환되며, 특정 태스크에 과적합되지 않고 다양한 태스크로 일반화됨
- 자기 개선: 다른 벤치마크에서도 향상된 성능을 보여, 지속적인 자기 개선이 가능함을 시사