Complementing reinforcement learning with SFT through logit averaging in the post training of LLMs

기술 요약: 로그이트 평균화를 통한 강화학습의 SFT 보완

1. 문제 제기

대형 언어 모델 (LLM) 의 사후 학습은 일반적으로 두 가지 패러다임에 의존합니다: 지도 미세 조정 (SFT) 과 검증 가능한 보상을 활용한 강화학습 (RLVR). SFT 는 포맷팅과 스타일을 학습하기 위해 토큰 단위의 밀집된 감독을 제공하지만, 학습 데모의 분포에 의해 제한을 받습니다. 반면, Proximal Policy Optimization (PPO) 과 Group Relative Policy Optimization (GRPO) 과 같은 RL 방법은 이러한 분포적 한계를 초월할 수 있지만, 희소하고 노이즈가 많은 보상 신호로 고통받습니다.

표준 접근법은 이를 순차적으로 결합합니다: 먼저 SFT 를 적용한 후, 학습 가능한 정책 $\pi_\theta$를 SFT 참조 정책 $\pi_{\text{sft}}$에 고정하기 위해 Kullback-Leibler (KL) 정규화 항 $\beta \text{KL}(\pi_\theta \| \pi_{\text{sft}})$을 포함한 RL 을 수행합니다. 저자들은 이러한 "KL 고정" 레시피에 근본적인 긴장 관계가 있음을 지적합니다: KL 패널티가 $\pi_\theta$를 $\pi_{\text{sft}}$ 쪽으로 끌어당겨, 정책이 SFT 체크포인트의 한계 (예: 대수적 오류, 반복 루프) 에 구속되고 RL 동안 개발된 추론 능력을 완전히 활용하지 못하게 할 수 있습니다. 실무자들은 종종 이를 완화하기 위해 작은 $\beta$를 선택하지만, 저자들은 정규화에 대한 구조적 대안을 제안합니다.

2. 방법론

이 논문은 KL 정규화를 로그이트 평균화로 대체하는 방법을 소개합니다. 발산을 패널티하는 대신, 학습 가능한 정책 $\pi_\theta$와 고정된 참조 정책 $\pi_{\text{ref}}$(예: SFT) 의 로그이트를 평균화하여 혼합 정책 $\pi^{\text{mix}}_{\theta, \alpha}$를 형성함으로써 두 정책을 결합합니다.

핵심 메커니즘:
가중치 $\alpha \in [0, 1]$와 상태 $s$가 주어졌을 때, 혼합 로그이트 $z^{\text{mix}}_{\theta, \alpha}(s)$는 다음과 같이 정의됩니다:
$$z^{\text{mix}}_{\theta, \alpha}(s) = (1 - \alpha) z_\theta(s) + \alpha z_{\text{ref}}(s)$$
그런 다음 혼합 정책은 $\pi^{\text{mix}}_{\theta, \alpha}(a | s) = \text{softmax}(z^{\text{mix}}_{\theta, \alpha}(s))$가 됩니다.

학습 목적:
저자들은 GRPO 목적 함수를 이 혼합 정책에 맞게 조정합니다:

• 롤아웃: 궤적은 순수 학습 가능한 정책 $\pi_\theta$가 아닌 혼합 정책 $\pi^{\text{mix}}_{\theta, \alpha}$에서 샘플링됩니다.
• 어드밴티지 추정: 그룹 상대적 어드밴티지는 이러한 혼합 롤아웃에서 얻은 보상을 기반으로 계산됩니다.
• 최적화: 학습 가능한 매개변수 $\theta$는 혼합 정책 하에서 보상을 최대화하도록 업데이트됩니다. 참조 로그이트 $z_{\text{ref}}$는 전체 과정에서 고정된 상태로 유지됩니다.
• 가중치 방식:
  1. 고정 가중치: $\alpha$는 학습 내내 일정하게 유지됩니다 (예: $\alpha=0.5$).
  2. 적응형 가중치: $\alpha$는 홀드아웃 검증 세트에 기반하여 동적으로 업데이트됩니다. 업데이트 규칙은 순수 정책과 혼합 정책의 정확성을 비교하는 "순이익 (net-gain)" 지표를 사용하며, 혼합이 순이익을 낼 경우 $\alpha$를 1 로, 그렇지 않을 경우 0 으로 이동시킵니다.

RLVR 과의 비교:
크리틱과 KL 정규화를 사용하는 표준 RLVR 방법 (PPO, GRPO, DAPO) 과 달리, 이 접근법은 KL 항이나 크리틱이 없습니다. 참조 정책은 로그이트 평균화를 통해 탐색에 참여하며, 손실 함수 내의 중요도 비율은 현재 및 이전 혼합 정책 간에 계산됩니다.

3. 주요 기여 및 발견

상호 보완적 기술 (패턴 1 및 2):
MATH, cn-k12, MMLU 데이터셋에 대한 실증 분석은 SFT 와 RL 학습 정책이 상호 보완적인 기술을 보유하고 있음을 보여줍니다:

• SFT는 포맷팅 (출력 템플릿 준수) 에 뛰어나지만 복잡한 추론에서는 실패할 수 있습니다.
• **RL 정책 ($\pi_\theta$)**은 추론 (올바른 수치 답 도출) 에 뛰어나지만 종종 포맷팅 규율을 잃습니다.
• 상호 보완성: 문제의 상당 부분 (22%–38%) 은 두 정책 중 정확히 하나에 의해 올바르게 해결됩니다. 로그이트 평균화는 "전문가들의 곱 (Product of Experts)" 역할을 하여, SFT 의 포맷팅 강점과 RL 정책의 추론 강점을 결합합니다.

성능 결과:
세 가지 데이터셋에 걸쳐 Qwen2.5-Instruct 모델 (1.5B, 3B, 7B) 을 평가한 결과:

• 정확도: 로그이트 평균화 방법은 모든 9 가지 (데이터셋 $\times$ 모델 크기) 구성에서 표준 KL 정규화 GRPO 대비 더 높거나 동등한 답변 정확도를 달성합니다.
• 어블레이션: SFT 앵커를 기본 Instruct 모델로 대체하면 성능이 현저히 저하되어, 성취된 이득이 포맷 인식 참조와의 혼합에서 비롯됨을 확인시켜 줍니다.
• 로그이트 대 확률 평균화: 로그이트 평균화 (분포의 기하평균) 는 확률 평균화 (산술평균) 보다 우수합니다. 로그이트 평균화는 더 작은 혼합 가중치로 더 높은 해결률을 달성하며 $\alpha=0.5$에서 정점을 찍는 반면, 확률 평균화는 더 높은 가중치를 필요로 하며 $\alpha=0.6$에서 정점을 찍습니다.

이론적 연결:
이 논문은 로그이트 평균화가 수학적으로 **전문가들의 곱 (Product of Experts, POE)**과 동등함을 입증합니다. 결과적인 정책은 RL 정책과 참조 정책으로의 KL 발산의 가중 합을 최소화합니다. 이 구조는 정책이 두 전문가가 선호하는 토큰에 확률을 집중하게 하여, SFT 의 대수적 오류를 효과적으로 수정하면서도 포맷팅 제약을 유지할 수 있게 합니다.

4. 중요성 및 주장

저자들은 이 방법이 "포맷팅 앵커"를 "그래디언트 제약"에서 분리함으로써 표준 KL 정규화 레시피보다 뚜렷한 이점을 제공한다고 주장합니다.

• 한계 극복: KL 정규화가 $\pi_\theta$를 $\pi_{\text{sft}}$에 가깝게 유지하게 함으로써 (잠재적으로 그 추론 결함을 물려받을 수 있음) 로그이트 평균화는 $\pi_\theta$가 참조의 한계를 극복하면서도 혼합 정책의 구조적 결합을 통해 포맷팅 이점을 유지할 수 있게 합니다.
• 암시적 신뢰 영역: 혼합은 암시적 신뢰 영역으로 작용하여, 대규모 순수 RL 에서 종종 발생하는 포맷팅 스타일에서의 분포적 이탈을 방지합니다.
• 밀집된 사전 지식: 최종 보상 토큰에서만 신호를 전파하는 RL 과 달리, 로그이트 평균화는 모든 디코딩 단계에서 SFT 앵커로부터의 밀집된 토큰 단위 사전 지식을 통합합니다.
• 확장성: 이 접근법은 소규모 언어 모델 (SLM) 앙상블로 자연스럽게 확장되어, 시스템이 거대한 단일 정책을 학습하는 비용 없이 전문화된 전문가들 (예: 코드, 안전, 도메인 용어) 을 활용할 수 있게 합니다.

이 논문은 로그이트 평균화가 KL 정규화에 대한 견고한 대안을 제공하며, SFT 와 RL 간의 상호 보완성을 성공적으로 활용하여 포맷팅 품질을 희생하지 않고 추론 정확도를 향상시킨다고 결론지었습니다.