기술 요약: SAGE – LLM 의 RLVR 에서의 안내 탐색을 위한 앵커 형성

1. 문제 제기

최근 검증 가능한 보상 (Verifiable Rewards) 을 활용한 강화 학습 (RLVR) 의 발전은 추론 과제에서 대규모 언어 모델 (LLM) 의 pass@1 정확도를 향상시킬 수 있음을 입증했습니다. 그러나 중요한 한계가 드러났습니다. RLVR 은 종종 단일 샘플 성공률을 높이지만, $k$개 샘플 중 정답을 찾을 확률인 pass@k 에서 동등한 향상을 이끌어내지 못하는 경우가 많습니다.

이전 분석들은 이 불일치가 표준 RLVR 이 모드 붕괴 (또는 "지지 축소") 를 유발하기 때문이라고 시사합니다. 정책이 새로운 추론 전략을 발견하는 대신, 기본 참조 모델에 존재하는 기존 고밀도 추론 모드의 좁은 하위 집합으로 수렴하는 것입니다. 이 행동은 표준 역 Kullback-Leibler (reverse-KL) 정규화 항의 구조적 제약에 기인합니다. 역-KL 은 정책을 참조 분포 ($\pi_{ref}$) 에 고정함으로써 필요한 학습 안정성을 제공하지만, 본질적으로 저밀도 경로를 억제하여 참조 모델의 즉각적인 지지 범위 밖에 있는 대체적이고 유효한 추론 경로를 모델이 탐색하는 것을 효과적으로 막습니다.

KL 항을 제거하거나 forward-KL 로 대체하여 이를 해결하려는 시도는 최적의 결과를 내지 못했습니다. KL 항을 제거하면 종종 보상 해킹과 불안정성이 초래되며, forward-KL 은 지지를 확장하더라도 보상과 무관한 영역에 확률 질량을 할당하여 pass@1 을 개선하지 못합니다. 핵심적인 과제는 역-KL 의 안정화 이점을 유지하면서 제어된 경험적 지지 확장을 가능하게 하는 것입니다.

2. 방법론: SAGE 프레임워크

저자들은 SAGE (Shaping Anchors for Guided Exploration) 를 제안합니다. 이는 역-KL 정규화를 약화시켜야 할 정적 제약이 아니라, 탐색을 안내하기 위해 목표 분포를 재형성할 수 있는 동적 앵커로 재해석하는 프레임워크입니다.

핵심 메커니즘

표준 역-KL RLVR 은 다음을 최적화합니다:
$$J(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta D_{KL}(\pi_\theta \parallel \pi_{ref})]$$
이는 정책이 참조의 지수적 기울기 (exponential tilt) 가 되는 정상 해를 유도합니다: $\pi_\theta \propto \pi_{ref} \exp(r/\beta)$. 이 곱셈적 업데이트는 $\pi_{ref}$의 고밀도 영역을 증폭시키는 반면 저밀도 영역은 억제합니다.

SAGE 는 가이드 함수 $q(x, y) > 0$을 도입하여 앵커 분포 자체를 재형성함으로써 이를 수정합니다. 새로운 목적 함수는 다음과 같습니다:
$$J'(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta \tilde{D}_{KL}(\pi_\theta \parallel q \cdot \pi_{ref})]$$
여기서 $\tilde{D}_{KL}$은 정규화되지 않은 함수를 위해 정의된 의사-KL 발산 (pseudo-KL divergence) 입니다. 정상 해는 다음과 같습니다:
$$\pi_\theta(y | x) \propto q(x, y) \pi_{ref}(y | x) \exp\left(\frac{1}{\beta} r(x, y)\right)$$
$q(x, y)$를 미탐색되었으나 잠재적으로 유효한 추론 모드에 대해 크게 설계함으로써, SAGE 는 앵커 $\pi_{ref}$가 제공하는 안정성을 깨뜨리지 않고 이러한 모드의 확률 질량을 증폭시킵니다.

내재적 신호를 통한 구체화

가이드 구축 시 명시적인 보상 정보가 필요하지 않도록 $q(x, y)$를 실용적으로 만들기 위해, 저자들은 추론 과정 자체에서 도출된 내재적 신호를 사용할 것을 제안합니다:

1. 무작위 탐색: 확률적 교란을 사용하는 기준선.
2. 토큰 수준 탐색: 놀라움 (surprisal, $-\log p$) 을 사용하여 희귀한 어휘 선택을 편향시킵니다. 저자들은 이것이 표면 수준의 편차를 포착하지만 구조적 추론 차이를 반영하지는 않을 수 있다고 지적합니다.
3. 분기 수준 탐색 (SAGE-Branch): 주요 제안 변형입니다. 엔트로피 ($H_t$) 를 추론 분기 신호로 사용합니다. 가이드 함수는 다음과 같이 정의됩니다:
   $$q_{Branch}(y_t | x, y_{<t}) = 1 + \gamma (H_t - \tau)_+$$
   여기서 $H_t$는 토큰 수준의 엔트로피, $\tau$는 활성화 임계값, $\gamma$는 크기를 조절합니다. 이는 결정적 기호 단계에서의 노이즈는 억제하면서 고불확실성 의사 결정 지점 (분기 노드) 에서의 탐색을 선택적으로 증폭시킵니다.

이론적 보장

이 논문은 표준 RLVR 이 참조 모델의 경험적 지지로 제한되는 반면, SAGE 는 이론적으로 경험적 지지를 확장할 수 있음을 보여주는 이론적 분석 (Theorem 4.3) 을 제공합니다. 구체적으로, 가이드 함수 $q$가 현재 참조 정책의 경험적 지지 밖에 있는 보상 유효 해 $y^*$에 대해 충분히 크다면, SAGE 는 이를 무시할 수 없는 확률을 할당하여 "숨겨진" 추론 모드를 효과적으로 복원할 수 있습니다.

3. 주요 기여

1. RLVR 한계 분석: 저자들은 역-KL 앵커가 안정화시키지만 RLVR 에서의 모드 붕괴의 주요 구조적 원인임을 규명했습니다. 단순히 KL 항을 제거하거나 대체하는 것만으로는 효율성 - 커버리지 트레이드오프를 방해하므로 불충분하다고 주장합니다.
2. SAGE 프레임워크: 앵커를 제거하는 대신 가이드 함수 $q(x, y)$를 통해 KL 앵커 분포를 재형성함으로써 제어 가능한 경험적 지지 확장을 가능하게 하는 원칙적인 방법을 제안합니다.
3. 내재적 신호 통합: 추가적인 순방향 통과나 외부 데이터 없이 효과적인 가이드 함수를 구축하는 데 경량의 내재적 신호 (특히 엔트로피와 놀라움) 를 사용할 수 있음을 입증합니다.
4. 실증적 검증: 광범위한 실험을 통해 SAGE 변형, 특히 분기 수준 탐색이 다양한 수학 추론 벤치마크에서 pass@1 과 pass@k 를 일관되게 향상시킴을 보여줍니다.

4. 실험 결과

저자들은 Qwen2.5-Math-7B-Base 모델과 GRPO 알고리즘을 사용하여 AIME 2024-2025, AMC23, MATH-500 데이터셋에서 SAGE 를 평가했습니다.

• 성능 향상: GRPO + Branch 변형이 가장 강력한 평균 성능을 달성하여, pass@1 을 0.266(바닐라 GRPO) 에서 0.284로, pass@256 을 0.714 에서 0.729로 향상시켰습니다. 세 가지 데이터셋 전반에서 기준선보다 우수한 성과를 보였습니다.
• 대안과의 비교:
  • KL 없는 GRPO: 더 높은 학습 보상을 보였지만 테스트 시간 성능 (pass@1 및 pass@256) 은 현저히 낮아, KL 앵커 없이 보상 해킹의 위험이 있음을 확인했습니다.
  • Forward-KL: pass@256 을 약간 향상시켰지만 SAGE 의 pass@1 향상과 일치하지 않아, 오프-타겟 영역에 질량을 할당함을 나타냅니다.
  • 기타 RL 알고리즘: SAGE 는 BNPO 및 DAPO와 성공적으로 통합되어, 이러한 방법들이 일반적으로 pass@1 을 향상시키는 대신 pass@k 를 희생하는 트레이드오프를 일관되게 해결했습니다.
• 추론 패턴 분석: 정성적 분석에 따르면 표준 RLVR 은 "제약 설정"과 같은 지배적 패턴을 강화한 반면, SAGE 는 드물지만 유효한 추론 모드인 모순에 의한 증명의 빈도를 GRPO 의 0.8% 에서 GRPO+Branch 의 3.3% 로 크게 증가시켰습니다.
• 강건성: 분포 외 평가 (Knights and Knaves 벤치마크) 에서 SAGE 로 훈련된 모델은 분포 이동 하에서도 성능을 유지한 반면, 표준 GRPO 모델은 현저한 저하를 보였으며, 이는 SAGE 가 모드 붕괴를 완화함을 시사합니다.
• 일반화: 이 방법은 장추론 모델 (DeepSeek-R1-Distill-Qwen-7B) 과 소형 모델 (Qwen2.5-1.5B/3B) 에서도 효과적이었으며, 다양한 모델 규모와 추론 패러다임 전반에 걸쳐 강건함을 입증했습니다.

5. 중요성과 주장

이 논문은 SAGE 가 역-KL RLVR 에서의 모드 sharpening 을 완화하는 실용적이고 원칙적인 접근법을 제공한다고 주장합니다. KL 정규화를 경직된 제약이 아닌 제어 가능한 도구로 취급함으로써, SAGE 는 LLM 이 기본 모델에서 소수이지만 유효한 다양한 추론 궤적을 발견하고 활용하도록 합니다.

저자들은 SAGE 가 기존 기술 (예: 어드밴티지 형성 또는 동적 샘플링) 과 모듈화되어 있으며 직교한다고 강조합니다. 이는 최적화 구조가 아니라 앵커 분포 자체를 수정하기 때문입니다. 그들은 SAGE 를 안정적이고 KL 호환적인 경험적 지지 확장을 위한 일반 프레임워크로 위치시킵니다.

한계와 범위:
저자들은 겸손하게 SAGE 가 고정된 참조 정책의 지지 내에서 작동함을 인정합니다. 이는 $\pi_{ref}$하에서 이미 0 이 아닌 확률을 갖는 경로들 사이에서 확률 질량을 재분배할 뿐, 기본 모델이 생성할 수 없는 완전히 새로운 추론 패턴을 생성할 수는 없습니다. 그들은 더 강력한 교사 모델로부터의 증류가 추론 경계를 확장할 수 있지만, SAGE 는 그 경계 내부에서의 안정적 탐색을 가능하게 함으로써 이를 보완한다고 지적합니다. 현재 구체화는 경량의 내재적 신호에 의존하며, 저자들은 더 강력한 모델로부터 더 표현력 있는 가이드 함수를 학습하는 것이 향후 연구의 유망한 방향이라고 제안합니다.