기술적 요약: 할인된 강화학습을 통한 효율적 추론 학습

문제 제기

대규모 추론 모델 (LRMs) 은 최종 답변을 emit 하기 전에 중간 추론 토큰을 생성함으로써 복잡한 작업 (예: 수학, 코드) 을 해결합니다. 강화학습 (RL) 후처리는 정확도를 향상시키지만, 종종 지나치게 긴 사고의 사슬을 초래하여 계산 비용, 메모리 사용량 (증가하는 KV 캐시를 통해), 그리고 추론 지연 시간을 증가시킵니다.

이 분야의 지배적인 가정은 더 긴 응답이 본질적으로 정확도를 향상시킨다는 것으로, 응답 길이와 정확성 간의 트레이드오프를 시사합니다. 본 논문은 이러한 가정에 도전합니다. 모델 클래스와 문제 인스턴스에 의해 결정되는 영역까지, 본질적인 트레이드오프는 존재하지 않는다고 주장합니다. 즉, 정확도를 희생하지 않고도 응답 길이를 특정 인스턴스 의존적 임계값까지 줄일 수 있습니다. 목표는 확률적 최단 경로 (SSP) 문제에서 가장 짧은 성공적인 궤적을 찾는 것과 유사하게, 올바른 답변에 도달하는 LRMs 을 효과적이고 효율적으로 훈련하는 것입니다.

방법론

이론적 프레임워크: 블랙웰 최적성

저자들은 검증자 기반 추론을 이진 종료 보상 (정답은 1, 오답은 0) 을 갖는 유한 시간 마르코프 결정 과정 (MDP) 으로 모델링합니다. 정확도를 극대화하면서 경로 길이를 최소화하는 목표를 공식화하기 위해 **블랙웰 최적성 (Blackwell optimality)**을 활용합니다.

• 정의: 정책이 충분히 1 에 가까운 모든 할인 인자 $\gamma$에 대해 최적일 때, 그 정책은 블랙웰 최적입니다.
• 메커니즘: 이진 종료 보상을 갖는 목표 달성 MDP 에서, $\gamma \to 1$로 갈 때 할인된 반환 $J_\gamma(\pi)$를 극대화하는 것은 먼저 성공 확률 (정확도) 을 극대화합니다. 최대 성공 확률을 달성하는 정책들 중에서 할인 메커니즘은 더 짧은 기대 경로 길이를 가진 정책을 선호함으로써 동점을 깨뜨립니다.
• 핵심 통찰: 유한 정책 클래스 $\Pi$에 대해, 임계 할인 인자 $\gamma_{bw} < 1$이 존재하여 모든 $\gamma \in (\gamma_{bw}, 1)$에 대해 $\gamma$-최적 정책의 집합은 블랙웰 최적 정책의 집합과 일치합니다. 이러한 정책들은 정확도 극대화자들 중에서 동시에 정확도를 극대화하고 기대 경로 길이를 최소화합니다.

실용적 훈련 레시피

이론적 결과에 guided 되어, 저자들은 그룹 상대적 정책 최적화 (GRPO) 를 사용한 구체적인 훈련 방법론을 제안합니다:

1. 할인된 환경 보상만 적용: 할인 인자 $\gamma \in (0, 1)$은 외적 환경 보상 (정확도) 에만 배타적으로 적용됩니다. 내재적 보상 (포맷팅, 구조적 형성) 은 모델이 잘 구조화된 출력을 생성하도록 보장하기 위해 할인되지 않습니다.
2. 추론 토큰 마스킹: 할인은 "추론 구간" (예: <reasoning> 과 </reasoning> 태그 사이의 토큰) 으로 식별된 토큰에만 적용됩니다. 프롬프트 준수, 포맷팅, 그리고 최종 답변에 필요한 토큰은 할인되지 않습니다.
3. KL 정규화: 정책 붕괴 (모델이 너무 일찍 추론을 중단하는 경우) 를 방지하고 학습을 안정화하기 위해 이동하는 참조 정책 (주기적으로 업데이트됨) 에 대해 KL 페널티를 적용합니다.
4. 비교 가능한 토큰 예산: 공정한 비교를 위해, 할인된 방법과 할인되지 않은 방법 간의 훈련 중 처리된 총 토큰 수가 비교 가능하도록 롤아웃 횟수를 조정합니다.

최적화되는 목적 함수는 다음과 같습니다:
$$\mathbb{E}_{S_1 \sim \mu, \tau \sim \pi(S_1)} [\gamma^{K(\tau)} r_e(\tau) + r_f(\tau)] - \beta \text{KL}(\pi || \pi')$$
여기서 $K(\tau)$는 추론 토큰의 수, $r_e$는 환경 보상, $r_f$는 포맷팅 보상입니다.

주요 기여

1. 트레이드오프 부재에 대한 이론적 보장: 본 논문은 고정된 (제한적일 수 있는) 정책 클래스 내에서 블랙웰 최적 정책이 할인되지 않은 성공 확률을 극대화하고, 이러한 극대화자들 중에서 기대 궤적 길이를 최소화함을 증명합니다. 이는 정확도와 경로 길이가 인스턴스 의존적 임계값까지 본질적으로 트레이드오프되지 않음을 확립합니다.
2. 블랙웰 인자의 존재: 유한 정책 클래스에 대해, 저자들은 블랙웰 할인 인자 $\gamma_{bw} < 1$의 존재를 증명합니다. 그들은 응답을 단축하면서 정확도를 유지하기 위해 $\gamma$가 1 에 얼마나 가까워야 하는지에 대한 경계를 제공하여, 제한된 배포 클래스에 대한 하이퍼파라미터 선택을 명확히 합니다.
3. 실증적 검증: 할인된 목적 함수를 사용한 GRPO 를 통해, 저자들은 수학 벤치마크 (GSM8K, MATH, AMC, AIME 등) 에서 모델이 평균 응답 길이를 현저히 줄이면서도 할인되지 않은 베이스라인의 Pass@1 정확도와 일치함을 입증했습니다.

실험 결과

저자들은 다양한 수학 데이터셋에서 네 가지 지시 튜닝 모델 (Qwen2.5 7B/14B, Llama 3 8B, Phi-4) 을 평가했습니다.

• GSM8K 및 MATH: 할인된 모델은 토큰 수를 현저히 줄이면서 할인되지 않은 베이스라인의 Pass@1 정확도와 일치하거나 약간 초과했습니다.
  • Qwen2.5 7B (GSM8K): 정확도는 안정적으로 유지되었으며 (91.06% 대 91.07%), 평균 응답 길이는 217.60 에서 170.08 토큰으로 감소했습니다 (~22% 감소).
  • Llama 3 8B (GSM8K): 정확도가 약간 향상되었으며 (80.87% 대 81.07%), 길이는 125.43 에서 108.67 토큰으로 감소했습니다 (~13% 감소).
• 더 어려운 벤치마크 (AMC, AIME, MINERVA, OLYMPIAD): DeepScaleR 에서 훈련된 더 큰 모델 (Qwen2.5 14B, Phi-4) 에도 이 경향이 유지되었습니다. 할인된 모델은 현저히 짧은 흔적을 달성하면서도 동등하거나 더 높은 정확도를 달성했습니다 (예: AMC 2023 의 Phi-4: 정확도 51.00% $\to$ 61.00%, 토큰 1134 $\to$ 716).
• 할인 인자 민감도: $\gamma$를 변화시킨 실험은 이론적 트레이드오프를 확인했습니다: $\gamma$가 1 에서 멀어질수록 응답 길이는 신뢰성 있게 단축되지만, 할인이 너무 공격적이면 결국 정확도가 떨어집니다.

중요성과 주장

본 논문은 큐레이션된 데이터나 프롬프트 엔지니어링에 의존하지 않는 효율적 추론을 위한 원칙적이고 인스턴스 인식 메커니즘을 제공한다고 주장합니다. 블랙웰 최적성을 활용함으로써, 이 접근법은 "효율적 추론"의 문제를 특정 할인 전략을 가진 표준 RL 목적 함수로 재구성합니다.

저자들은 그들의 결과가 다음과 같다고 주장합니다:

• 본질적인 길이 - 정확도 트레이드오프 반박: 모델의 용량에 의해 결정되는 특정 임계값까지 정확도 손실 없이 더 짧은 추론이 가능함을 보여줍니다.
• 길이 페널티 통합: 이 작업은 작은 단계별 음수 보상을 사용하는 방법들이 이 설정에서 할인과 동등함을 시사하며, 올바르게 구현될 때 여러 기존 길이 페널티 방법들이 동일한 정확도 - 그다음 - 길이 순서를 회복할 수 있음을 암시합니다.
• 실용적 레시피 제공: 제안된 방법은 표준 정책 최적화 (GRPO 등) 에 알고리즘적 변경이 필요하지 않으며, 보상 계산 (추론 토큰 할인) 의 수정과 이동하는 참조 정책의 사용만 필요합니다.

본 논문은 더 짧은 흔적이 효율성을 향상시키지만, 압축된 추론이 일반화를 개선할 수 있는 가능성은 향후 연구의 열린 질문으로 남는다고 결론지었습니다.