기술 요약: 모드 붕괴를 넘어: 다양한 추론을 위한 분포 매칭

1. 문제 정의

본 논문은 온-정책 강화학습 (RL) 방법, 특히 검증 가능한 보상을 통한 강화학습 (RLVR) 으로 대형 추론 모델 (LRMs) 을 훈련시키는 데 사용되는 방법에서 **모드 붕괴 (mode collapse)**가 근본적인 한계임을 지적합니다. 그룹 상대적 정책 최적화 (GRPO) 와 같은 방법들은 수학 문제 해결과 코드 생성에서 성공을 거두었지만, 한 번 발견된 단일 고보상 궤적에 확률 질량을 집중시키는 경향을 보입니다. 이러한 행동은 대안적인 해결 전략을 억제하고, 탐색을 중단시키며, 일반화에 실패하는 취약한 추론으로 이어집니다.

저자들은 이 현상이 단순한 구현의 부산물이 아니라 최적화 목적 함수의 직접적인 결과라고 주장합니다. 표준 온-정책 RL 은 암묵적으로 역 Kullback-Leibler (KL) 발산($D_{KL}(\pi_\theta \| p^*)$) 을 최소화합니다. 역 KL 은 본질적으로 '모드 추구 (mode-seeking)' 성향을 가지며, 여러 개의 다양한 고보상 해결책에 대한 분포를 유지하기보다 처음 발견된 성공적인 궤적을 증폭시킵니다. 반면, 순방향 KL 발산($D_{KL}(p^* \| \pi_\theta)$) 은 '모드 커버링 (mode-covering)' 행동을 보여 정책이 목표 분포의 전체 지지 영역 (support) 을 매칭하도록 장려하지만, 전역 보상-가중 분포로부터의 샘플링이 필요하기 때문에 일반적으로 계산이 불가능합니다.

2. 방법론: DMPO

계산이 불가능한 전역 샘플링 없이 모드 붕괴를 해결하기 위해, 저자들은 **분포 매칭 정책 최적화 (Distribution-Matching Policy Optimization, DMPO)**를 제안합니다.

핵심 통찰

DMPO 는 전역적 수준이 아닌 그룹 수준에서 순방향 KL 최소화를 근사합니다. 전역 목표 분포로부터 샘플링하는 대신, 이 알고리즘은 주어진 프롬프트에 대해 샘플링된 궤적들의 특정 그룹에 대한 목표 분포를 구성합니다.

알고리즘 공식화

1. 그룹 수준 목표 분포: 현재 정책 $\pi_{\theta_{old}}$에서 샘플링된 $G$개의 궤적 집합 $\mathcal{O} = \{o_1, \dots, o_G\}$에 대해, 목표 분포 $p(o_i | \mathcal{O})$를 그들의 보상의 지수 함수에 비례하도록 구성합니다 (볼츠만 분포):
   $$p(o_i | \mathcal{O}) = \frac{\exp(r(o_i)/\alpha)}{\sum_{j=1}^G \exp(r(o_j)/\alpha)}$$
   이는 그룹 수준의 분할 함수만 필요로 하므로 계산적으로 실행 가능합니다.

2. 그룹 수준 정책 분포: 이 목표와 정책을 정렬하기 위해, 이 방법은 궤적 길이 편향을 보정하기 위해 길이로 정규화된 로그 가능도를 기반으로 정책 분포 $q_\theta(o_i | \mathcal{O})$를 정의합니다:
   $$q_\theta(o_i | \mathcal{O}) = \frac{\exp(\varphi(o_i))}{\sum_{j=1}^G \exp(\varphi(o_j))}$$
   여기서 $\varphi(o_i)$는 토큰당 평균 로그 확률입니다.

3. 분포 매칭 손실: 정책은 $q_\theta$와 $p$ 사이의 발산을 최소화함으로써 목표에 정렬됩니다. 저자들은 수치적 안정성을 위해 KL 발산 대신 **평균 제곱 오차 (MSE)**를 사용하여 정책 확률이 0 에 가까워질 때 그래디언트 폭발을 방지합니다:
   $$\mathcal{L}_{DM}(\theta) = \frac{1}{G} \sum_{i=1}^G (p(o_i | \mathcal{O}) - q_\theta(o_i | \mathcal{O}))^2$$

4. 통합 목적 함수: 최종 목적 함수는 표준 GRPO 보상 최대화 항과 분포 매칭 정규화 항을 결합합니다:
   $$\mathcal{L}_{DMPO}(\theta) = \mathcal{L}_{GRPO}(\theta) + \lambda \mathcal{L}_{DM}(\theta)$$
   여기서 $\mathcal{L}_{GRPO}$는 활용 (보상 최대화) 을 주도하고, $\mathcal{L}_{DM}$은 정책이 고보상 영역 내의 단일 모드로 붕괴하는 것을 방지함으로써 탐색을 강제합니다.

3. 주요 기여

A. DMPO 알고리즘

본 논문은 DMPO 를 모드 붕괴에 대한 원칙적이고 실용적인 해결책으로 제시합니다. 그룹 수준에서 순방향 KL 을 근사함으로써, 학습된 분할 함수나 전역 샘플링 없이도 모드 커버링 행동을 달성합니다. 이 방법은 기존 GRPO 프레임워크와 호환되며, 단일 정규화 항 추가만 필요합니다.

B. MM-NP-Bench

저자들은 추론 모델의 탐색을 엄격하게 평가하도록 설계된 멀티모달 벤치마크인 MM-NP-Bench를 소개합니다.

• 범위: 텍스트 기반 NP-Bench 를 확장하여 10 가지 NP-하드 조합 최적화 작업 (예: TSP, 그래프 색칠, 정점 덮개) 의 시각적 표현을 포함합니다.
• 이중 지표 평가: 모드 붕괴를 관찰 가능하게 만들기 위해 벤치마크는 두 가지 지표를 사용합니다:
  • 성공률 (SR): 제약 조건 충족 (실행 가능성) 을 측정합니다.
  • 품질 비율 (QR): 휴리스틱 기준에 대한 최적화 품질을 측정합니다.
  • 의의: 모드 붕괴는 높은 SR 이지만 낮은 QR (유효하지만 비최적의 해결책을 찾는 것) 로 나타납니다.
• 인프라: 준 최적 기준을 제공하기 위한 매개변수 생성기, 규칙 기반 검증기, 휴리스틱 솔버를 포함합니다.

C. 실증적 검증

본 논문은 DMPO 가 여러 도메인에서 강력한 기준선 (GRPO, GSPO, GPG, FlowRL 등) 보다 우수함을 보여주는 광범위한 실험을 제공합니다.

4. 결과

최적화 작업 수행

• MM-NP-Bench (멀티모달): DMPO 는 43.1% 의 품질 비율을 달성하여 GRPO(38.4%) 대비 12% 의 상대적 개선을 보였습니다. 또한 가장 높은 성공률 (61.9% 대 55.7%) 을 기록했습니다.
• NP-Bench (텍스트 전용): DMPO 는 43.9% 의 품질 비율을 달성하여 GRPO(40.1%) 대비 9% 의 상대적 개선을 보였습니다.
• 어블레이션: 이 방법은 하이퍼파라미터 선택 ($\lambda \in [0.3, 0.7]$ 및 $\alpha \in [0.8, 1.2]$) 에 대해 견고함을 보여주었습니다.

일반화 및 전이

• 수학적 추론: DMPO 를 사용하여 최적화 작업으로 훈련하는 것은 GRPO 대비 수학 추론 벤치마크 (AIME, AMC 등) 에서 **2.0%**만큼 성능을 향상시켰습니다.
• 도메인 외 작업: DMPO 는 도메인 외 시각 추론 벤치마크 (LogicVista, MathVista 등) 에서 **2.3%**의 개선을 보여주어, 다양성 보존 훈련이 일반 추론 능력을 향상시킨다는 것을 시사합니다.

훈련 역학

훈련 로그 분석 결과, GRPO 는 "길이 붕괴"(짧고 단순한 응답으로 퇴화) 와 조기 보상 정체에 시달리는 반면, DMPO 는 지속적인 개선을 유지하고 견고한 추론 체인 (~400 토큰) 을 보존하는 것으로 나타났습니다. DMPO 는 GRPO 보다 낮은 엔트로피를 유지하여 확률적 혼란이 아닌 "구조화된 다양성"(서로 다른 논리적 전략에 대한 자신감 있는 실행) 을 나타냅니다.

5. 의의 및 주장

본 논문은 온-정책 RL 에서 모드 붕괴를 방지하기 위한 실용적이고 원칙적인 접근법으로 분포 매칭을 확립한다고 주장합니다. 저자들은 다음과 같이 주장합니다:

1. 근본 원인 식별: RLVR 의 모드 붕괴는 근본적으로 역 KL 최소화의 모드 추구 성향에 의해 발생합니다.
2. 실행 가능한 해결책: 전통적으로 계산이 불가능한 것으로 간주되었던 순방향 KL 최소화는 해결책의 다양성을 보존하기 위해 그룹 수준에서 효과적으로 근사될 수 있습니다.
3. 일반적 이점: 특정 최적화 작업 훈련 중 다양성을 보존하는 것은 향상된 일반 추론 능력으로 이어지며, 이는 대형 추론 모델에서 견고한 추론을 위해 탐색이 핵심 구성 요소임을 시사합니다.

이 연구는 DMPO 를 단순한 다양성 정규화기가 아니라, 모드 추구 목적 함수가 놓치는 더 높은 품질의 해결책을 발견할 수 있게 하는 메커니즘으로 위치시킴으로써, 최적화 성능과 모달리티 전반의 추론 견고성 모두를 향상시킵니다.