Beyond Token-Level Policy Gradients for Complex Reasoning with Large Language Models

1. 문제 제기 (Problem)

대형 언어 모델 (LLM) 의 복잡한 추론 작업 (수학 문제 해결, 코드 생성 등) 에서 기존 강화 학습 (Policy Gradient) 기반 미세 조정 방법론은 다음과 같은 근본적인 한계를 가지고 있습니다.

• 토큰 수준의 최적화와 의미론적 블록의 불일치: 기존의 PPO, GRPO, DAPO 등의 방법은 다음 토큰을 하나씩 선택하는 것을 '행동 (Action)'으로 간주합니다. 그러나 복잡한 추론 과정에서는 변수 정의, 방정식 구성, 함수 호출 등 하나의 의미론적 결정이 여러 개의 토큰에 걸쳐 이루어지는 경우가 많습니다.
• 의미론적 단위의 파편화: 토큰 단위로 최적화를 수행하면, 본래 하나의 논리적 흐름을 이루는 '블록'이 파편화되어, 모델이 장기적인 의존성 (Long-horizon dependency) 과 구조적 일관성을 유지하기 어렵게 됩니다.
• 그레인 (Granularity) 불일치: 추론의 정답은 토큰 단위가 아니라 의미 있는 문장이나 식 (Expression) 단위로 결정되는데, 토큰 수준의 최적화 신호는 이러한 고수준의 구조를 제대로 포착하지 못합니다.

2. 제안 방법: MPO (Multi-token Policy Gradient Optimization)

이러한 한계를 극복하기 위해 저자들은 MPO라는 새로운 프레임워크를 제안합니다. 이는 토큰 단위가 아닌 **연속된 K 개의 토큰을 하나의 통합된 의미론적 행동 (Unified Semantic Action)**으로 간주하여 정책 경사 (Policy Gradient) 를 최적화하는 방법입니다.

핵심 기술 요소

1. 블록 단위 행동 (Block-level Action):

   • 모델이 다음 토큰 하나만 예측하는 것이 아니라, K 개의 연속된 토큰 시퀀스 ($o_{t+1:t+K}$) 를 하나의 단위로 예측하고 평가합니다.
   • 이를 통해 변수 정의나 방정식 작성과 같은 추론 단계의 내부 구조를 보존합니다.

2. 통합 중요도 샘플링 비율 (United Importance Sampling Ratio):

   • 기존 PPO 의 중요도 비율 $r_t = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$을 블록 단위로 확장합니다.
   • 단순 곱셈 형태 ($\prod r_{t+n}$) 는 분산이 급격히 커질 수 있으므로, 가중 로그 - 합 (Weighted Log-sum) 방식을 도입하여 분산을 제어합니다.
   • 식 (9) 에 따르면, $K$개의 토큰에 대한 중요도 비율을 로그 공간에서 가중치 $\beta_n$을 부여하여 합산한 후 지수화합니다:
     $$\tilde{R}^{(K)}_{i,t}(\theta) = \exp \left( \sum_{n=1}^{K} \beta_n \log r_{i,t+n}(\theta) \right)$$
   • 여기서 $\beta_n$은 시간적 거리에 따라 감쇠하는 가중치 (Decaying weight) 를 사용하여, 가까운 토큰 예측에 더 큰 영향을 주면서도 장기적인 맥락을 고려합니다.

3. MTP (Multi-Token Prediction) 모듈 활용:

   • DeepSeek-V3 등에서 사용된 MTP 아키텍처를 기반으로, 백본 모델의 마지막 레이어를 복사하여 K 개의 추가 예측 헤드를 생성합니다.
   • Warm-up 전략: RL 학습 전, MTP 모듈이 여러 토큰을 정확히 예측할 수 있도록 지도 학습 (Supervised Fine-tuning) 으로 사전 훈련 (Warm-up) 을 수행합니다. 이는 RL 단계에서의 불안정성을 방지합니다.

4. 최적화 목표 (Objective):

   • PPO 의 클립된 손실 함수 (Clipped Surrogate Loss) 를 블록 단위 중요도 비율에 적용하여, 모델이 개별 토큰이 아닌 의미 있는 추론 세그먼트를 계획하도록 유도합니다.

3. 주요 기여 (Key Contributions)

1. 구조화된 추론 단위의 최적화: MPO 를 통해 정책이 고립된 토큰이 아닌 구조화된 추론 단위 (Block) 를 기반으로 최적화되도록 설계했습니다.
2. 후기 훈련 (Post-training) 단계의 통합: LLM 의 강화 학습 기반 후기 훈련 파이프라인에 구조적 다중 토큰 최적화를 성공적으로 통합하여, 추론의 일관성을 향상시키는 새로운 방향을 제시했습니다.
3. 실증적 유효성 증명: 수학 추론 (GSM8K, MATH) 및 코드 생성 (HumanEval) 벤치마크에서 기존 토큰 단위 정책 경사 방법 (PPO, GRPO, DAPO) 을 일관되게 능가하는 성능을 입증했습니다.

4. 실험 결과 (Experimental Results)

• 벤치마크: GSM8K (대학 수준 수학), MATH (경쟁 수준 수학), HumanEval (코드 생성) 에서 평가 수행.
• 모델: Llama3.2-1B, DeepSeek-Distilled-Qwen2.5-1.5B, 7B 등 다양한 크기의 모델에 적용.
• 성능:
  • 수학 추론: 모든 모델 크기와 벤치마크에서 MPO 가 PPO, GRPO, DAPO 보다 높은 정확도 (Pass@1) 를 기록했습니다. 특히 MATH 와 같은 고난도 문제에서 성능 격차가 두드러졌습니다.
  • 코드 생성: HumanEval 에서도 일관된 개선을 보여주어, 다중 단계의 의미적 의존성이 중요한 작업에서도 효과적임을 입증했습니다.
• 학습 안정성:
  • MPO 는 중요도 샘플링 비율의 분산 (Variance) 과 클립 비율 (Clip Fraction) 을 줄여 학습을 더 안정화시켰습니다.
  • Bias-Variance Trade-off: 블록 크기 (K) 와 감쇠율 ($\lambda$) 을 조절하여 최적의 성능을 얻었습니다. K=5, $\lambda=0.8$일 때 가장 안정적이고 효과적인 결과를 보였습니다.
  • 과도한 가중치 주의: MTP 모듈의 기여도가 너무 높으면 (예: 40% 이상) 다중 토큰 예측의 노이즈가 증가하여 성능이 저하됨을 확인했습니다 (최적은 약 10~20% 수준).

5. 의의 및 결론 (Significance)

• 패러다임 전환: LLM 의 추론 능력 향상을 위해 '토큰 단위' 최적화의 한계를 넘어, '의미론적 블록 단위' 최적화가 필요함을 강력하게 주장합니다.
• 효율성: MPO 는 추가적인 MTP 모듈을 도입하지만, 추론 시에는 이를 분리 (Detach) 하므로 추론 비용은 증가하지 않으며, 학습 시에도 기존 PPO 와 유사한 메모리 사용량을 보이며 약 30% 빠른 학습 속도를 달성했습니다.
• 미래 연구 방향: 추론 집약적 (Reasoning-intensive) 인 언어 작업에서 정책 학습의 세분화 (Granularity) 를 재고하고, 의미 구조에 기반한 의사결정 과정을 모델링하는 연구의 새로운 지평을 열었습니다.

요약하자면, 이 논문은 복잡한 추론 작업에서 토큰 단위의 국소적 최적화가 가지는 한계를 지적하고, 연속된 토큰을 하나의 의미 있는 행동 단위로 묶어 최적화하는 MPO를 제안함으로써, LLM 의 추론 일관성과 정확성을 획기적으로 개선하는 방법을 제시했습니다.