기술 요약: DISA - 분포 매칭 LLM-RL 을 위한 오프라인 중요도 샘플링

1. 문제 정의

현대 추론 에이전트들은 단일 정답 생성 능력보다는 다양한 유효한 해결 경로 (예: 다양한 사고 연쇄, 계획, 또는 도구 사용 흔적) 를 생성하는 능력으로 평가받는 경우가 increasingly 늘고 있습니다. Proximal Policy Optimization(PPO) 또는 Group Relative Policy Optimization(GRPO) 과 같은 표준 강화학습 (RL) 목적 함수는 보상 최대화에 중점을 둡니다. 단일 정답 정확도에는 효과적이지만, 이러한 방법들은 **모드 붕괴 (mode collapse)**를 겪는 경향이 있어, 정책이 단일 고보상 궤적에 확률 질량을 집중시킴으로써 견고한 평가 (예: pass@k 지표) 에 필요한 유효한 해결책의 다양성을 제거합니다.

분포 매칭 RL은 궤적 확률이 지수화된 보상에 비례하는 목표 분포에 정책을 적합하도록 훈련함으로써 대안을 제시합니다:
$$\tilde{\pi}(o | q) = \frac{1}{Z(q)} \pi_{\text{ref}}(o | q) \exp(\beta \tilde{r}(q, o))$$
여기서 $Z(q)$는 프롬프트 $q$에 의존하는 **정규화 상수 (partition function)**입니다.

기존 분포 매칭 방법들 (예: FlowRL, GFlowNets) 의 핵심 병목 현상은 정책 ($\pi_\theta$) 과 정규화 상수 추정기 ($Z_\phi$) 간의 온라인 결합입니다. 이러한 방법들은 공동 손실 함수 (예: 궤적 균형, Trajectory Balance) 를 사용하여 공유된 온라인 롤아웃에서 두 구성 요소를 함께 훈련합니다. 이 결합은 두 가지 치명적인 문제를 야기합니다:

1. 오차 전파: 정규화 상수 추정기 $Z_\phi$의 보정 오차가 정책 경사에 직접 흘러들어 업데이트를 왜곡합니다.
2. 진단 불가: 정규화 상수에 대한 독립적인 품질 신호가 없습니다. 그 성능은 정책의 손실과 얽혀 있어, 실패가 정책에서 비롯된 것인지 정규화 상수 추정에서 비롯된 것인지 진단할 수 없습니다.

2. 방법론: 분해된 중요도 샘플링 앵커링 (DISA)

저자들은 정규화 상수 추정을 정책 최적화 루프에서 분리하는 3 단계 파이프라인인 DISA를 제안합니다. 핵심 통찰은 정규화 상수 $Z(q)$가 훈련 가능한 정책 $\pi_\theta$가 아닌 참조 정책 $\pi_{\text{ref}}$, 보상 함수, 그리고 역온도 $\beta$에만 의존한다는 점입니다. 따라서 RL 훈련 시작 전에 오프라인으로 추정할 수 있습니다.

1 단계: 오프라인 중요도 샘플링 (IS) 추정

오프라인 프롬프트 집합 $\mathcal{D}_Z$에 대해, 이 방법은 강력한 제안 모델 $p_T(\cdot|q)$ (예: 더 큰 교사 모델) 로부터 $N$개의 제안 궤적을 추출합니다.

• 중요도 가중치: 각 궤적 $o_i$에 대해 중요도 가중치를 계산합니다: $w(q, o_i) = \frac{\pi_{\text{ref}}(o_i|q)}{p_T(o_i|q)} \exp(\beta \tilde{r}(q, o_i))$.
• 로그 공간 집계: 수치적 안정성을 보장하고 편향을 줄이기 위해, 프롬프트 $q$에 대한 정규화 상수 추정은 logsumexp 집계기를 사용하여 로그 공간에서 계산됩니다:
  $$\log \hat{Z}_{\text{IS}}(q) = \text{logsumexp}_i \left( \log \pi_{\text{ref}}(o_i|q) - \log p_T(o_i|q) + \beta \tilde{r}(q, o_i) \right) - \log N$$
  이는 선형 스케일에서 $Z(q)$의 편향 없는 추정을 제공하며, 로그 스케일에서는 통제된 $O(1/N)$ 편향을 가집니다.

2 단계: 회귀를 통한 상계 (Amortization)

프롬프트별 추정치 $\log \hat{Z}_{\text{IS}}(q)$는 경량 프롬프트 조건부 회귀기 $g_\psi(q)$ (예: 작은 MLP) 를 훈련하는 라벨로 사용됩니다.

• 회귀기는 예측값과 오프라인 IS 라벨 사이의 제곱 오차를 최소화하도록 최소제곱 회귀를 통해 훈련됩니다.
• 훈련이 완료되면 회귀기 매개변수 $\psi^*$는 동결됩니다. 이는 초기 추정 집합 외부의 프롬프트에도 일반화되는 재사용 가능한 오프라인 정규화 상수 추정을 생성합니다.

3 단계: 앵커링된 온라인 정책 학습

온라인 RL 단계 동안, 동결된 회귀기 $g_{\psi^*}(q)$는 궤적 균형 (TB) 손실에서 학습 가능한 정규화 상수 $Z_\phi(q)$를 대체합니다.

• 목적 함수: 정책 $\pi_\theta$는 다음 잔차를 최소화하도록 최적화됩니다:
  $$\mathcal{L}_{\text{DISA}}(\theta) = \mathbb{E} \left[ \sum_i \omega_i \left( g_{\psi^*}(q) + \log \pi_\theta(o_i|q) - \log \pi_{\text{ref}}(o_i|q) - \beta \tilde{r}(q, o_i) \right)^2 \right]$$
• 분해: 핵심적으로, 기울기는 $g_{\psi^*}$를 통해 흐르지 않도록 차단됩니다. 정규화 상수는 정책 업데이트 동안 고정된 프롬프트별 상수 ("앵커") 로 작용합니다.

3. 주요 기여

1. 결합을 핵심 장애물로 규명: 이 논문은 정규화 상수 추정과 정책 학습의 얽힘을 분포 매칭 RL 에서 보정 오차와 진단 무지의 주요 원인으로 규명합니다.
2. DISA 파이프라인: 중요도 샘플링을 통해 정규화 상수 추정을 오프라인으로 이동시키고, 회귀를 통해 상계하며, 온라인 RL 을 위해 이를 동결하는 새로운 3 단계 접근법입니다.
3. 이론적 보장:
   • 정확한 보존: 오프라인 정규화 상수가 정확하다면, DISA 목적 함수의 전역 최소값은 이론적 보상 기울기 목표 분포와 일치합니다.
   • 편향과 함께 하는 정상성: 프롬프트만의 편향 (오차가 궤적이 아닌 $q$에만 의존하는 경우) 이 있더라도, 목표 분포는 온-정책 경사 흐름의 정상점으로 남습니다.
   • 유계 섭동: 잔차 추정 오차는 코시 - 슈바르츠를 통해 유도된 유계 섭동 포락선을 통해서만 정책 목적 함수에 진입하여 안정성을 보장합니다.
4. 실증적 검증: 수학 및 코드 벤치마크에 대한 포괄적인 평가로, DISA 가 온라인 결합의 함정을 피하면서도 분포 매칭의 이점을 보존함을 입증했습니다.

4. 실험 결과

저자들은 두 가지 오픈 가중 백본 (Qwen2.5-7B 및 Qwen3-4B-Base) 을 대상으로 6 개의 수학 벤치마크 (AIME 2024/2025, AMC 2023, MATH-500 포함) 와 3 개의 코드 벤치마크 (LiveCodeBench, CodeForces, HumanEval+) 에서 DISA 를 평가했습니다.

• 수학적 추론:
  • DISA 는 온라인 결합 기반인 FlowRL과 일치하거나 약간 상회했습니다 (예: Qwen3-4B-Base 에서 Mean@8 65.5 대 64.8).
  • 보상 최대화 기반인 GRPO 및 GSPO보다 현저히 우수한 성능을 보였습니다 (예: Qwen3-4B-Base 에서 GRPO 대비 +16.1 점).
• 코드 생성:
  • 해결책 다양성에 민감한 지표인 pass@16에서 DISA 는 모든 방법 중 가장 높은 점수를 기록했습니다 (예: Qwen3-4B-Base 에서 51.5), FlowRL, GSPO, GRPO 를 능가했습니다.
  • pass@1에서는 보상 최대화 방법 대비 유의미한 후퇴 없이 경쟁력 있는 성능을 유지했습니다.
• 다양성 분석:
  • 해결 전략 다양성에 대한 LLM-as-judge 평가 (1-5 리커트 척도 사용) 에서 DISA 는 백본의 원래 다양성 중 가장 큰 비율 (3.72/5.0) 을 유지한 것으로 나타났으며, FlowRL(3.40) 보다 높고 GRPO/GSPO(~3.0) 보다 훨씬 우수했습니다.
• 절차 연구 (Ablation Studies):
  • 제안 강도: 약한 제안 모델 (Qwen3-4B 대 235B) 을 사용하면 오프라인 라벨의 분산이 증가하여 성능이 하락했으나, 회귀기는 여전히 수렴하여 이론에서 예측한 편향 - 분산 트레이드오프를 확인했습니다.
  • 역온도 ($\beta$): 민감도 연구는 과소 기울기 ($\beta=10$) 가 약한 신호 손실을 초래하는 반면, 과다 기울기 ($\beta=20$) 는 "분산 절벽"과 모드 붕괴를 유발하여 중요도 샘플링 가중치에 대한 이론적 분석을 검증했습니다.
• 증류 vs RL: DISA 의 오프라인 추정에 사용된 동일한 고보상 궤적에 대한 미세 조정 (SFT) 과 비교했을 때, DISA 는 최대 13.8 Mean@8 점의 추가 향상을 제공했습니다. 이는 SFT 의 하드 타겟으로 버려진 지형 정보를 부드러운 보상 기울기 분포에 적합함으로써 회복함을 보여줍니다.

5. 의의 및 주장

이 논문은 DISA 가 추론을 위한 LLM 의 사후 훈련에서 중요한 아키텍처 전환을 나타낸다고 주장합니다. 정규화 상수 추정을 정책 루프에서 분해함으로써 DISA 는 다음과 같은 이점을 제공합니다:

• 보정 오차 제거: 정규화 상수 추정 오차가 정책 경사를 왜곡하는 것을 방지합니다.
• 독립적 진단 가능: 회귀기의 검증 손실을 통해 정책 훈련과 별도로 정규화 상수의 품질을 모니터링할 수 있습니다.
• 다양성 보존: 여러 유효한 경로가 존재하는 복잡한 추론 작업에 필수적인 해결 공간의 다중 모드 구조를 성공적으로 유지합니다.
• 이론적 건전성: 추정 오차가 유계라면 오프라인 추정 과정이 분포 매칭 목적 함수의 이론적 고정점을 변경하지 않는다는 공식적 보장을 제공합니다.

저자들은 DISA 가 온라인 결합 분포 매칭에 대한 견고하고 이론적으로 근거 있는 대안을 제공하며, 정규화 상수와 정책을 함께 훈련하는 것과 관련된 불안정성 없이 정확성과 다양성 모두에서 최첨단 성능을 달성한다고 결론지었습니다.