[논문리뷰] TTRL: Test-Time Reinforcement Learning

🔗 https://arxiv.org/pdf/2504.16084.pdf

AI, 스스로 성장하다: 라벨 없는 테스트 데이터로 학습하는 혁신적인 RL 방법, TTRL

• 거대 언어 모델(LLM)은 점점 강력해지고 있지만, 새롭고 복잡한 문제에 부딪혔을 때 기존 학습 데이터만으로는 한계가 있습니다.
• 특히, 실시간으로 쏟아지는 라벨(정답)이 없는 새로운 데이터에 대해 모델 성능을 지속적으로 개선하는 것은 큰 도전 과제입니다.
• 기존의 강화 학습(RL)은 모델 성능 향상에 효과적이지만, 정확한 라벨이나 잘 만들어진 보상 모델이 필수적이라 새로운 데이터에 적용하기 어렵습니다.
• 이 논문은 라벨이 없는 테스트 데이터만으로도 LLM이 스스로 학습하여 성능을 높이는 TTRL이라는 방법을 제안하며, 이는 AI의 지속적인 발전 가능성을 보여줍니다.

서론

• 최근 LLM의 추론 능력을 향상시키는 테스트 시간 스케일링(Test-Time Scaling, TTS) 기법이 주목받고 있습니다.
• TTS는 사전 학습 시보다 추론 시에 더 많은 계산 자원을 투입하여 성능을 높이는 방식입니다.
• 강화 학습(RL)은 긴 추론 과정(Chain-of-Thought)을 개선하는 데 중요하지만, 테스트 시점에 라벨이 없는 데이터에 대해 어떻게 보상을 얻을 것인가가 문제입니다.
• 실제 세상의 복잡한 데이터를 대규모로 라벨링하는 것은 불가능하며, 이는 모델의 지속적인 학습에 큰 걸림돌입니다.
• TTRL은 이러한 문제를 해결하기 위해 라벨 없는 데이터에 RL을 적용하는 새로운 프레임워크입니다.

기존 연구

• 테스트 시간 스케일링(TTS): 추론 시점에 모델 성능을 높이는 기법들입니다.

• 여러 답변을 생성하고 가장 좋은 것을 선택하거나 (예: 다수결 투표), 탐색 알고리즘을 활용하는 방식입니다.
• 보통 정답률이나 보상 모델을 통해 가장 좋은 결과물을 고릅니다.

• 테스트 시간 학습(Test-Time Training, TTT): 테스트 데이터가 들어올 때 모델 파라미터를 업데이트하는 방식입니다.

• 주로 영상 인식 등 다른 분야에서 연구되었으며, LLM에 적용된 사례는 많지 않았습니다.

• 추론을 위한 RL: LLM의 추론 능력을 강화하기 위해 RL을 사용합니다.

• 인간 피드백 기반 RL(RLHF)이나 문제의 정답(수학 문제 답, 코드 실행 결과)을 통해 보상을 주는 방식 등이 있습니다.
• 하지만 대부분 대규모의 라벨링된 학습 데이터를 필요로 합니다.

• TTRL의 차별점: TTRL은 TTT와 RL을 결합하여, 라벨이 없는 테스트 데이터에서 RL 학습을 가능하게 한 최초의 시도입니다.

방법론

• TTRL은 라벨 없는 테스트 데이터로 모델을 RL 학습시킵니다.
• 핵심 아이디어: 테스트 시간 스케일링 기법 중 하나인 **다수결 투표(Majority Voting)**를 사용하여 보상을 추정합니다.

과정

• 모델에게 하나의 문제(입력)를 주고 여러 개의 답변을 생성하게 합니다 (예: 64개).
• 생성된 답변들에서 최종 답을 추출합니다.
• 추출된 답들 중 가장 많이 나온 답을 **'추정된 라벨'**로 삼습니다.
• 생성된 각 답변이 이 '추정된 라벨'과 일치하면 1점, 아니면 0점과 같은 보상을 줍니다.
• 이 추정된 보상을 사용하여 RL 알고리즘(논문에서는 GRPO 사용)으로 모델 파라미터를 업데이트합니다.

목표

• 모델이 '추정된 라벨'과 일치하는 답변을 더 자주 생성하도록 학습시키는 것입니다.

비유

• 여러 명의 학생이 각자 수학 문제를 풀고, 답을 맞춰봅니다. 가장 많은 학생이 쓴 답을 정답이라 가정하고, 그 답과 일치하는 학생은 칭찬(보상)을 받고, 다른 답을 쓴 학생은 그렇지 않습니다. 학생들은 이 과정을 통해 정답처럼 추정된 답을 맞출 확률을 높여갑니다. 실제 정답이 없어도 집단의 '지혜'로 학습하는 것과 비슷합니다.

실험 및 결과

사용 모델

• Qwen2.5-Math (1.5B, 7B) 기본/인스트럭트 모델, LLaMA-3.1-8B-Instruct 등 다양한 크기와 종류의 모델로 실험했습니다.

사용 데이터

• 라벨이 없는 수학 추론 벤치마크(AIME 2024, AMC, MATH-500)의 테스트 데이터셋을 사용했습니다.

주요 결과

• TTRL은 라벨 없는 데이터만으로도 상당한 성능 향상을 달성했습니다.
• 특히 어려운 AIME 2024 데이터셋에서 Qwen 7B 모델의 성능을 159% 향상시켰습니다 (16.7% -> 43.3%).
• 세 벤치마크 평균적으로 84.1%의 성능 향상을 보였습니다.
• TTRL의 성능 향상은 모델 크기에 비례하는 경향을 보였습니다 (1.5B보다 7B에서 더 큰 폭으로 상승).
• TTRL 학습은 학습에 사용된 데이터셋뿐만 아니라 다른 데이터셋에서도 성능을 향상시켜 일반화 능력이 있음을 보여주었습니다.
• 놀라운 점은, TTRL은 학습에 사용한 '초기 모델의 다수결 투표 성능'보다 최종 성능이 더 높게 나왔다는 것입니다. 즉, 자신의 한계를 뛰어넘어 스스로 발전했습니다.
• 심지어 라벨이 있는 테스트 데이터로 직접 학습시킨 결과(정보 유출 시나리오)에 가까운 성능을 달성했습니다.

결론 및 한계

결론

• TTRL은 라벨이 없는 테스트 데이터에 RL을 적용하여 LLM을 학습시키는 새로운 방법을 제시했으며, 다수결 투표 기반 보상 추정으로 효과적인 학습이 가능함을 보였습니다. 모델이 스스로 진화하며 성능을 향상시키는 가능성을 열었습니다.

한계

• TTRL은 초기 모델의 능력(사전 지식)에 크게 의존합니다. 모델이 너무 약하거나 데이터가 너무 어려우면 성능 향상이 미미할 수 있습니다.
• RL 학습에 사용되는 하이퍼파라미터(학습률, 배치 크기, 샘플링 온도 등)에 민감하며, 최적의 설정을 찾기 어렵습니다.
• 다수결 투표가 효과적이려면 '정답' 개념이 명확한 작업(수학 등)에 더 적합할 수 있습니다. 창의적인 글쓰기처럼 정답이 불분명한 작업에는 적용하기 어렵습니다.

향후 연구

• TTRL 학습의 이론적 분석, 실시간 스트리밍 데이터에 적용하는 온라인 학습, 대규모 비지도 RL 학습으로 확장, 더 복잡하고 개방적인 에이전트 작업에 적용 등을 계획하고 있습니다.

한 줄 요약

라벨이 없는 테스트 데이터만으로도 LLM이 스스로 여러 답변을 보고 '다수결 정답'을 추정해 RL 학습을 진행하며 성능을 비약적으로 향상시키는 놀라운 방법! 마치 AI가 집단 지성으로 스스로 공부하는 것 같아요.