Lookahead Tree-Based Rollouts for Enhanced Trajectory-Level Exploration in Reinforcement Learning with Verifiable Rewards

이 논문은 그룹 롤아웃 시 생성되는 경로의 다양성 부족 문제를 해결하기 위해, 불확실성이 높은 단계에서 분기하고 미리 보며 유사한 경로를 제거하는 'Lookahead Tree-Based Rollouts (LATR)' 전략을 제안하여 GRPO 및 DAPO 알고리즘의 학습 속도와 최종 성능을 크게 향상시켰음을 보여줍니다.

핵심

이 논문은 인공지능 (LLM) 이 수학이나 논리 문제를 더 잘 풀 수 있도록 돕는 새로운 훈련 방법인 LATR을 소개합니다. 복잡한 기술 용어 대신, 일상적인 비유를 통해 쉽게 설명해 드릴게요.

🧠 핵심 문제: "똑같은 생각만 반복하는 AI"

지금까지 AI 를 훈련시킬 때, 수학 문제를 풀게 하면 AI 는 같은 문제를 여러 번 풀게 했습니다. 하지만 AI 가 문제를 풀 때, 매번 거의 똑같은 생각의 흐름 (경로) 을 따라가는 경향이 있었습니다.

• 비유: 마치 10 명의 탐정 (AI) 이 같은 사건을 조사하라고 보냈는데, 10 명 모두 똑같은 길로 출발해서 똑같은 단서만 찾아내고 돌아오는 상황입니다.
• 결과: 10 명이 다 같은 결론만 내면, "어떤 방법이 더 좋은가?"를 가르쳐 줄 수 있는 비교 자료가 부족해집니다. AI 는 "아, 이 방법이 좋구나"라고 배울 기회를 잃게 되죠.

💡 해결책: LATR (앞을 내다보는 가지치기)

저자들은 이 문제를 해결하기 위해 LATR이라는 새로운 방법을 제안했습니다. 이는 마치 나뭇가지를 키우는 것과 같습니다.

1. 분기 (Branching - 갈라지기):
   AI 가 문제를 풀다가 "어? 이 단어 대신 저 단어를 써볼까?"라고 의심스럽거나 고민하는 순간 (불확실성이 높은 곳) 에, AI 가 두 가지 이상의 다른 길로 갈라지게 만듭니다.

   • 비유: 탐정들이 갈림길에서 "한 팀은 왼쪽 길로, 다른 팀은 오른쪽 길로 가보자!"라고 지시하는 것입니다.

2. 시뮬레이션 (Lookahead - 앞을 내다보기):
   갈라진 길로 조금만 더 나아가서 (예: 30 단어 정도) 가보게 합니다.

   • 비유: 두 팀이 조금 더 전진해 보니, 오른쪽 팀은 막다른 길로 들어가고 왼쪽 팀은 보물상자를 발견하는 것을 미리 확인하는 것입니다.

3. 가지치기 (Pruning - 다듬기):
   만약 갈라진 두 팀이 너무 비슷하게 움직여서 결국 똑같은 결론만 내린다면, 그중 하나는 잘라냅니다.

   • 비유: "오른쪽 팀이랑 왼쪽 팀이 똑같은 길만 걷고 있네? 그럼 한 팀은 퇴근시키고, 진짜 다른 길을 가는 팀만 남기자!"라고 정리하는 것입니다.

🚀 LATR 의 효과: 왜 더 잘할까?

이 방법을 쓰면 AI 는 매우 다양한 생각의 흐름을 경험하게 됩니다.

• 더 빠른 학습: "이건 안 되네, 저건 되네"를 빠르게 비교할 수 있어서, AI 가 정답을 찾는 속도가 약 2 배 (131% 향상) 빨라집니다.
• 더 높은 점수: 다양한 시도를 해보면서 더 효율적이고 정확한 해결책을 찾게 되어, 최종 정답률도 약 4% 이상 높아집니다.
• 간결한 답변: 불필요한 헛수고 (긴 설명) 를 줄이고 핵심만 짚는 답변을 내놓게 됩니다.

🎯 요약

기존의 AI 훈련은 "10 명이 똑같은 길로 가서 같은 답을 내놓게 하는 것"이었다면, LATR은 "10 명이 서로 다른 길을 탐험하게 하고, 진짜 의미 있는 길만 골라내어 AI 가 더 똑똑하게 배우게 하는 것"입니다.

이 방법은 AI 가 수학이나 논리 문제를 풀 때, 다양한 관점에서 문제를 바라보게 함으로써 더 빠르고 정확하게 성장하도록 돕는 혁신적인 기술입니다.

기술 요약

논문 요약: Lookahead Tree-Based Rollouts for Enhanced Trajectory-Level Exploration in RLVR

1. 연구 배경 및 문제 제기

최근 검증 가능한 보상 (Verifiable Rewards) 을 활용한 강화 학습 (RLVR) 은 대규모 언어 모델 (LLM) 의 추론 능력을 획기적으로 향상시켰습니다. 특히 그룹 상대 정책 최적화 (GRPO) 와 같은 알고리즘은 그룹 내 여러 응답을 비교하여 학습하는 방식으로 성공을 거두었습니다. 그러나 기존 RLVR 파이프라인에는 샘플링된 궤적 (trajectories) 의 다양성 부족이라는 심각한 병목 현상이 존재합니다.

• 핵심 문제: 현재의 토큰 수준 확률적 샘플링 (Token-level Stochastic Sampling) 은 각 시퀀스를 독립적으로 생성합니다. 이로 인해 로컬 수준의 작은 변동 (예: 'compute' 대신 'calculate' 사용) 이 발생하더라도, 전체적인 추론 경로 (reasoning path) 는 매우 유사하게 수렴하는 경향이 있습니다.
• 결과: 그룹 내 궤적들이 너무 비슷하면 (Homogeneous), 정책 업데이트에 필요한 상대적 이득 (relative advantage) 신호가 약화되어 학습 효율이 떨어지고, 최종 성능 향상에 한계가 발생합니다.

2. 제안 방법: LATR (Lookahead Tree-Based Rollouts)

저자들은 궤적 수준의 다양성을 명시적으로 증진시키기 위해 **Lookahead Tree-Based Rollouts (LATR)**라는 새로운 롤아웃 전략을 제안합니다. 이는 몬테카를로 트리 탐색 (MCTS) 에서 영감을 받아, 모델의 불확실성이 높은 단계에서 분기 (branching) 를 강제하고 유사한 경로를 제거하는 구조를 가집니다.

LATR 은 세 가지 주요 단계를 반복적으로 수행합니다:

1. 분기 (Branching):

   • 모델의 다음 토큰 예측 확률 분포에서 가장 확률이 높은 토큰을 기본 경로로 확장합니다.
   • 동시에, 절대적 확률 임계값 ($\tau_{abs}$) 을 넘고, 최상위 토큰과의 확률 차이가 상대적 임계값 ($\tau_{rel}$) 이내인 다른 후보 토큰들이 존재할 경우, 새로운 자식 분기를 생성합니다. 이는 모델이 의미적으로 다른 추론 경로를 고민하는 '교차로'에서 분기를 유도합니다.

2. 선도 시뮬레이션 (Lookahead Simulation):

   • 새로 생성된 분기는 고정된 길이 (r 토큰) 만큼 미리 생성 (시뮬레이션) 합니다.
   • 이를 통해 단순히 토큰 하나만 다른 것이 아니라, 실제로 의미 있는 다른 추론 경로로 이어지는지 확인합니다.

3. 가지치기 (Pruning):

   • 시뮬레이션된 분기들이 부모 분기나 다른 분기들과 너무 유사한지 (정규화된 편집 거리, Edit Distance 기준) 평가합니다.
   • 유사도가 임계값 ($\tau_{ed}$) 이하인 경우, 해당 분기와 그 하위 노드를 제거하여 계산 자원을 의미 있는 다양한 경로에 집중시킵니다.

이 과정이 목표하는 롤아웃 수 (k) 에 도달할 때까지 반복되며, 이후에는 표준 확률적 샘플링으로 전환됩니다. 또한, 학습 초기에는 LATR 을 많이 사용하고 후기에는 테스트 시간 행동과 유사하게 하기 위해 점진적으로 LATR 비율을 줄이는 하이브리드 롤아웃 전략을 사용합니다.

3. 주요 기여

1. 새로운 트리 기반 롤아웃 알고리즘: 토큰 수준의 샘플링을 넘어, 불확실성이 높은 지점에서 분기하고 시뮬레이션을 통해 유사한 경로를 제거함으로써 궤적 수준의 다양성을 최적화하는 LATR 을 도입했습니다.
2. 광범위한 실험적 검증: GRPO 와 DAPO 알고리즘에 LATR 을 적용하여 수학 및 논리 추론 태스크 (Countdown, DAPO-Math, AMC-2023 등) 에서 일관된 성능 향상을 입증했습니다.
3. 학습 효율성 및 최종 성능 동시 개선: 기존 방법 대비 학습 속도를 획기적으로 높이고 최종 정확도도 향상시켰습니다.

4. 실험 결과

저자들은 다양한 벤치마크에서 LATR 의 효과를 입증했습니다:

• 학습 가속화: LATR 은 정책 학습 속도를 평균 131% 가속화했습니다. 예를 들어, DAPO + LATR 은 DAPO + Stochastic Sampling 대비 Countdown 태스크에서 3 배, 수학 태스크에서 2 배 더 빠르게 수렴했습니다.
• 성능 향상: 최종 Pass@1 성능은 평균 4.2% 향상되었습니다.
  • Countdown 데이터셋: GRPO/DAPO 모두에서 정확도가 약 4.5% 상승.
  • 수학 데이터셋 (DAPO-Math, AMC-2023 등): 정확도 상승 및 Pass@8 성능 개선.
• 추론 길이 단축: LATR 을 사용한 모델은 더 간결하고 효율적인 추론 경로를 학습하여, 생성된 시퀀스의 평균 길이가 22% (Countdown) 까지 감소했습니다. 이는 불필요한 반복이나 장황한 설명을 줄인 결과입니다.
• 다양성 증대: LATR 그룹 내에서의 서로 다른 정답 수 (Distinct Answers) 가 확률적 샘플링 대비 유의미하게 증가하여, 더 넓은 해 공간 탐색이 가능함을 보였습니다.

5. 의의 및 결론

이 논문은 RLVR 에서 **궤적 수준의 다양성 (Trajectory-level Diversity)**이 모델 성능 확장의 핵심 요소임을 강조합니다. 기존의 토큰 단위 무작위 샘플링은 의미 있는 다양성을 보장하지 못한다는 한계를 지적하고, 이를 해결하기 위해 구조화된 탐색 (트리 기반 분기 및 가지치기) 을 도입했습니다.

LATR 은 모델 아키텍처 변경 없이도 기존 RL 알고리즘 (GRPO, DAPO) 에 쉽게 통합될 수 있으며, 더 적은 계산 비용으로 더 빠르고 정확한 추론 능력을 갖춘 모델을 학습시킬 수 있음을 증명했습니다. 이는 추론이 필요한 복잡한 태스크에서 LLM 의 성능을 극대화하기 위한 새로운 표준으로 자리 잡을 수 있는 중요한 기여입니다.