기술 요약: REFT – 첫 번째 토큰 다양화를 통한 롤아웃 탐색

문제 정의

검증 가능한 보상을 활용한 강화학습 (RLVR) 은 라벨링된 추론 궤적에 의존하지 않고 대규모 언어 모델 (LLM) 의 추론 능력을 향상시키는 주요 방법으로 부상했습니다. GRPO 및 DAPO 와 같은 그룹화 된 RLVR 방법에서는 정책이 단일 프롬프트에 대해 유한한 수의 롤아웃 (추론 궤적) 을 생성하고, 이를 자동 검증기가 점수화합니다. 학습 신호는 이 그룹 내 상대적 보상에서 도출됩니다.

이 패러다임의 결정적인 병목 현상은 롤아웃 다양성입니다. 그룹 내 모든 롤아웃이 유사한 추론 경로를 따르거나 동일한 보상을 받는다면, 정책 업데이트를 위한 대조적 신호는 거의 제공되지 않습니다. 기존 접근법들은 샘플링 온도를 높이거나, 고엔트로피 추론 분기점에서 분기하거나, 다양한 최종 답변을 선택함으로써 탐색을 확장하려고 시도합니다. 그러나 이러한 방법들은 종종 초기 접두사 토큰이 고정된 발판이거나, 가치 있는 탐색이 고불확실성인 실질적인 추론 단계에만 존재한다고 가정합니다.

저자들은 구조적 비효율성을 식별했습니다: 추론 마커 (<thought>) 이후 생성되는 첫 번째 토큰은 일반적으로 낮은 콘텐츠의 담화 시작어 (예: "First," "Let," "To") 입니다. 표준 샘플링 정책은 이러한 토큰에 대해 급격히 뾰족한 분포를 보이며, 롤아웃 그룹을 단일 시작어로 크게 편향시킵니다. 결정적으로, 저자들은 이러한 첫 번째 토큰의 확률이 매우 집중되어 있는 반면, 상위 20 개 후보에 걸친 생성된 롤아웃의 정확도는 거의 평탄하다고 관찰했습니다. 이는 "낮은 부하, 높은 레버리지"의 기회를 창출합니다: 첫 번째 토큰은 의미론적 부하가 최소화되었지만 높은 분포적 레버리지를 발휘하여 이후 전체 연장을 서로 다른 의미 영역으로 라우팅합니다.

방법론: REFT

이 논문은 모델 아키텍처, 검증기, 보상 함수, 또는 RL 목적 함수를 변경하지 않고 RLVR 파이프라인에서 첫 번째 토큰을 다양화하는 최소한의 개입인 **REFT(첫 번째 토큰 다양화를 통한 롤아웃 탐색)**를 소개합니다.

핵심 메커니즘

REFT 는 첫 번째 토큰의 표준 확률적 샘플링을 계층적 샘플링 전략으로 대체합니다:

1. 후보 선택: 주어진 프롬프트 $x$에 대해 첫 번째 토큰 위치에서 정책의 분포를 평가합니다. 공백 및 특수 토큰을 제외한 유효한 상위-$N$개 의미 토큰을 후보 집합 $F_N(x)$로 식별합니다.
2. 균등 샘플링: $F_N(x)$에서 $K$개의 서로 다른 첫 번째 토큰을 복원 없이 균등하게 샘플링하여 집합 $S_K(x)$를 형성합니다.
3. 예산 할당: 총 롤아웃 그룹 크기 $G$를 선택된 $K$개의 토큰 사이에 균등하게 나눕니다. 선택된 각 첫 번째 토큰 $f \in S_K(x)$에 대해, 모델은 시퀀스의 나머지 부분에 대해 변경되지 않은 기준 디코더를 사용하여 $f$에 조건부인 $G/K$개의 연장을 생성합니다.
4. 학습: 생성된 다양화된 그룹은 첫 번째 토큰 샘플링 불일치에 대한 중요도 가중치 보정 없이 표준 RLVR 알고리즘 (예: GRPO 또는 DAPO) 에 의해 처리됩니다.

주요 설계 선택

• 낮은 부하, 높은 레버리지: 이 방법은 모델이 "인위적으로 확실한" (특정 시작어에 대한 높은 확률) 위치를 대상으로 하지만, 대안들 사이에서 정확성이 안정적인 위치를 타겟팅합니다.
• 전역 노이즈 부재: 전체 궤적에 걸쳐 분포를 평탄하게 만들어 (나중에 중요한 추론 단계를 손상시킬 수 있음) 온도를 높이는 것과 달리, REFT 는 첫 번째 토큰에서만 다양성을 주입하여 이후 추론 사슬의 무결성을 유지합니다.
• 드롭인 교체: REFT 는 기본 RL 목적 함수, 이점 추정, 또는 검증기 로직에 변경이 필요하지 않습니다.

주요 기여

1. 진단적 통찰: 저자들은 추론 모델의 첫 번째 토큰 분포가 급격히 집중되어 있지만 정확성과는 분리되어 있음을 입증했습니다. 그들은 정상적으로 샘플링된 연장이 제공된다면, 상위 20 개 내의 낮은 확률 첫 번째 토큰이 가장 확률 높은 토큰만큼이나 올바른 롤아웃으로 이어질 수 있음을 보여줍니다.
2. 라우팅 효과: 실험은 서로 다른 첫 번째 토큰이 모델을 서로 다른 의미적 연장 영역으로 라우팅함을 확인했습니다. 표준 샘플링에 비해 상위 20 개 첫 번째 토큰의 균등 샘플링은 (첫 번째 토큰을 제거한 후) 나머지 연장에서 훨씬 더 높은 의미적 다양성을 산출합니다.
3. REFT 알고리즘: 타당한 첫 번째 토큰 시작어에 고정된 롤아웃 예산을 재할당함으로써 롤아웃 그룹을 다양화하는 간단하고 계산 효율적인 방법입니다.
4. 과도한 신용 부여 완화: 이 논문은 표준 GRPO 업데이트가 검증기가 해당 선호도를 강력하게 지지하지 않더라도 훈련 단계에 걸쳐 상위 1 위 첫 번째 토큰에 대한 선호도를 강화하는 경향이 있음을 보여줍니다. REFT 는 여러 타당한 시작어에 대한 커버리지를 보장함으로써 정책이 특정 담화 선호도에 과도하게 신용을 부여하는 것을 방지합니다.

실험 결과

저자들은 세 가지 데이터셋 (GSM8K, BigMath-Easy, BigMath) 에서 네 가지 기본 모델 (Qwen2.5-0.5B, 3B, 7B 및 Llama3.2-3B) 과 두 가지 RLVR 목적 함수 (GRPO 및 DAPO) 에 대해 REFT 를 평가했습니다.

• 성능 향상: REFT 는 Pass@1, Pass@8, Pass@64 지표를 통해 일관되게 집계 성능을 향상시킵니다.
  • GSM8K 에서 REFT 가 강화된 GRPO 와 DAPO 는 기준선보다 우수한 성능을 보입니다.
  • "Math Avg"(MATH-500, AIME 2024/2025, AMC 2023 의 평균) 에서 향상은 특히 두드러지는데, 이는 더 넓은 탐색이 중요한 경우 REFT 가 효과적임을 시사합니다.
  • 가장 작은 규모 (0.5B 파라미터) 에서도 개선이 관찰되어, 이 방법이 대규모 모델 용량에 의존하지 않음을 나타냅니다.
• 다양성 분석:
  • 훈련 다양성: REFT 는 훈련 중에 가장 의미적으로 다양한 롤아웃 그룹을 생성합니다.
  • 추론 집중: 역설적으로, REFT 로 훈련된 모델은 추론 시 더 집중된 출력을 생성하는데, 이는 이 방법이 정책이 다양한 경로를 탐색한 후 성공적인 패턴에 더 효과적으로 수렴하도록 학습하는 데 도움이 됨을 시사합니다.
  • 제로-분산 분해: REFT 는 "모두 틀린" 그룹 (올바른 궤적이 발견되지 않는 제로-분산 그룹) 의 비율을 줄이지만 "모두 맞는" 포화 상태를 증가시키지 않아, 더 유용한 탐색을 나타냅니다.
• 효율성: REFT 는 거의 무시할 수 있는 런타임 오버헤드를 도입합니다. 일부 경우, 더 높은 정확도로 인한 평균 완료 길이 단축으로 인해 GPU 시간이 약간 감소하기도 합니다.

중요성 및 주장

이 논문은 REFT 가 일반적인 탐색 직관의 "단순한 반전"을 제공한다고 주장합니다. 고엔트로피이고 불확실한 추론 분기점에서 다양성을 찾는 대신, 유용한 다양성은 급격히 편향되어 있지만 정확성과는 약하게 연결된 낮은 부하의 접두사 선택에서 회복될 수 있습니다.

저자들은 REFT 가 새로운 RL 알고리즘이 아니라 롤아웃 샘플러의 첫 번째 결정을 대상으로 한 교체임을 강조합니다. 그 중요성은 다음과 같습니다:

• 효율성: 롤아웃 예산이나 계산 비용을 증가시키지 않고 훈련 신호를 확장합니다.
• 강건성: 단일 샘플 정확도 (Pass@1) 와 유한 예산 내에서 올바른 솔루션을 찾을 확률 (Pass@8/64) 을 모두 향상시켜, RLVR 에서 흔히 볼 수 있는 Pass@1 향상이 접근 가능한 분포를 좁히는 트레이드오프를 해결합니다.
• 일반화 가능성: 이 방법은 다양한 모델 규모와 RLVR 목적 함수에서 작동하여, "첫 번째 토큰 라우팅" 현상이 더 나은 훈련을 위해 활용할 수 있는 현재 추론 모델의 근본적인 특성임을 시사합니다.

이 논문은 추론 마커 이후의 첫 번째 토큰이 RLVR 탐색을 위한 예상치 못하게 저렴하고 효과적인 개입 지점이며, 모델이 추론 궤적의 안정성을 유지하면서 더 넓은 연장 영역을 탐색할 수 있게 한다고 결론지었습니다.