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 information. While this setting appears elusive, we find that common practices in Test-Time Scaling (TTS), such as majority voting, yield surprisingly effective rewards suitable for driving RL training. In this work, we introduce Test-Time Reinforcement Learning (TTRL), a novel method for training LLMs using RL on unlabeled data. TTRL enables self-evolution of LLMs by utilizing the priors in the pre-trained models. Our experiments demonstrate that TTRL consistently improves performance across a variety of tasks and models. Notably, TTRL boosts the pass@1 performance of Qwen-2.5-Math-7B by approximately 211% on the AIME 2024 with only unlabeled test data. Furthermore, although TTRL is only supervised by the maj@n metric, TTRL has demonstrated performance to consistently surpass the upper limit of the initial model maj@n, and approach the performance of models trained directly on test data with ground-truth labels. Our experimental findings validate the general effectiveness of TTRL across various tasks and highlight TTRL's potential for broader tasks and domains. GitHub: https://github.com/PRIME-RL/TTRL

한 줄 요약

TTRL은 라벨 없는 데이터에서 강화학습을 통해 대형 언어 모델을 테스트 시간에 자가 진화시키는 새로운 방법이다.

핵심 기여도

TTRL은 라벨 없는 데이터에서 강화학습을 가능하게 하며, 다수결 투표를 기반으로 보상을 생성함.
Qwen-2.5-Math-7B 모델에서 AIME 2024 데이터셋에서 pass@1 성능이 211% 향상됨.
TTRL은 기존 모델의 maj@n 성능 상한을 초과하며, 라벨이 있는 데이터로 학습한 모델에 근접한 성능을 보임.
TTRL은 다양한 규모와 유형의 모델에 적용 가능하며, 기존 RL 알고리즘과 통합 가능함.

핵심 아이디어

TTRL은 테스트 시간에 강화학습을 통해 모델이 자가 진화하도록 유도하는 새로운 접근법이다. 기존 강화학습은 라벨 데이터에 의존하지만, TTRL은 라벨 없는 데이터에서도 모델이 스스로 경험을 생성하고 보상을 추정하도록 설계되었다. 핵심 아이디어는 **다수결 투표**(majority voting)를 기반으로 **rule-based reward**를 생성하는 방식이다. 이는 모델이 여러 샘플링된 출력 중 가장 높은 공감도를 가진 출력을 선택하도록 유도하며, 이를 기반으로 보상을 계산한다.

이러한 접근은 기존 테스트 시간 스케일링(TTS) 방법과 유사하지만, TTRL은 이를 강화학습 프레임워크로 확장하여 모델이 **자체적으로 학습**하고 **성능을 개선**하도록 유도한다. 학습 목표는 샘플링된 출력과 다수결 출력 간 일치도를 최대화하는 것으로, 이는 다음과 같은 기대 보상을 최대화하는 방식으로 수행된다:

$$
\max_{\theta} \mathbb{E}_{y \sim \pi_\theta(y|x)}[r(y, y^*)]
$$

기술적 접근법

**TTRL**(Test-Time Reinforcement Learning)은 테스트 시간에 강화학습을 수행하는 알고리즘.
**Rollout phase**에서 모델은 여러 후보 출력을 샘플링하여 다수결 투표로 최적 출력을 추정.
**Reward function**은 샘플링된 출력과 다수결 출력 간 일치도를 기반으로 생성됨.
**Policy update**는 기대 보상을 최대화하는 방향으로 수행되며, 학습률 $\eta$를 사용한 경사 상승법을 적용.
**Hyperparameters**: 최대 생성 길이 3072 토큰, 샘플 수 16개, temperature 0.6, top-p 0.95.
**Evaluation protocol**: pass@k, 특히 pass@1을 사용하여 정확도를 측정.

주요 결과

**Qwen-2.5-Math-7B** 모델에서 **AIME 2024** 데이터셋에서 **pass@1 성능이 12.9% → 40.2%**로 **211% 향상**됨.
**AIME 2024, AMC, MATH-500, GPQA** 데이터셋에서 평균 **76% 성능 향상**.
**maj@n** 지표로만 학습된 TTRL은 기존 모델의 maj@n 상한을 초과하며, 라벨 데이터로 학습한 모델에 근접한 성능을 보임.
TTRL은 다양한 모델 규모와 유형에 적용 가능하며, 기존 RL 알고리즘과 통합 가능함.

의의 및 한계

TTRL은 강화학습을 라벨 없는 데이터에서 수행할 수 있는 새로운 가능성을 제시하며, **자가 라벨링**(self-labeling) 기반 학습의 첫 단계로 기능한다. 특히, **ARC-AGI-2**와 같은 라벨이 부족한 복잡한 문제에 적용 가능성이 높다. 또한, **대규모 비감독 학습**으로의 확장 가능성과 **지속 학습**(continual learning)을 위한 기반을 제공한다.

그러나 TTRL은 **테스트 데이터의 분포에 의존**하며, **샘플링 기반 추정**이 오차를 유발할 수 있다는 한계가 있다. 또한, **복잡한 다단계 추론**이 필요한 문제에서는 다수결 투표가 충분한 보상을 제공하지 못할 수 있다.

실용적 활용

TTRL은 **수학 문제 해결**, **복잡한 추론**, **자연어 이해** 등 라벨이 부족한 분야에서 유용하게 활용될 수 있다. 특히, **AI 모델의 지속적 진화**와 **비용 효율적인 학습**을 필요로 하는 산업 및 연구 환경에서 적용 가능하다. 예를 들어, **교육 AI**, **자율 시스템**, **대규모 언어 모델 학습**에 활용할 수 있다.