[Paper] LAD: 추론을 위한 Learning Advantage Distribution

Source: arXiv - 2602.20132v1

개요

논문 “LAD: Learning Advantage Distribution for Reasoning” 은 복잡한 추론 작업을 위해 대형 언어 모델(LLM)을 훈련시키는 새로운 방식을 제안한다. 일반적인 강화학습(RL) 목표가 모델이 단일 기대 보상을 최대화하도록 하는 것과 달리, 저자들은 Learning Advantage Distributions (LAD) 를 도입하여 모델이 이점(advantage)의 분포와 일치하도록 학습시킨다. 이는 보다 다양하고 신뢰할 수 있는 추론 결과를 제공하며, RL‑미세조정된 LLM에서 자주 발생하는 “모드 붕괴” 현상을 피한다.

주요 기여

• Advantage‑distribution objective: 고전적인 advantage 최대화를 대체하여 정책의 출력 분포와 advantage‑유도 목표 분포 사이의 f‑divergence에 기반한 분포 매칭 손실을 사용합니다.
• Theoretical equivalence proof: RL에서 최적 정책 업데이트가 이 발산을 최소화하는 것과 수학적으로 동등함을 보여주어, 방법론을 견고한 이론에 기반합니다.
• Entropy‑free regularization: LAD 손실은 과도하게 확신하는 확률 스파이크를 자연스럽게 억제하여, 다른 RL 기반 파인튜닝 방법에서 흔히 사용되는 추가 엔트로피 보너스가 필요 없게 합니다.
• Zero extra compute: 이 알고리즘은 최첨단 GRPO (Generalized Reward‑Based Policy Optimization)와 비교했을 때 추가 연산 오버헤드가 없으며, 표준 LLM 사전 학습 직후 바로 적용할 수 있습니다.
• Empirical validation: LAD가 합성 밴딧 실험에서 다중모드 advantage 분포를 복원하고, 여러 LLM 백본에 걸쳐 수학 및 코드 추론 벤치마크에서 정확도와 출력 다양성을 일관되게 향상시킴을 입증합니다.

Source: …

Methodology

1. Advantage‑induced distribution:

   • 프롬프트에 대한 가능한 각 응답 (y)에 대해, advantage (A(y) = r(y) - V) (보상에서 기준값을 뺀 값)를 계산한다.
   • 이 advantage들을 목표 확률 분포 (p_A(y) \propto \exp(A(y))) 로 변환한다. advantage가 높은 응답은 더 높은 확률을 갖지만, 모든 유리한 대안들은 일정량의 질량을 유지한다.

2. Policy‑induced distribution:

   • 현재 LLM은 소프트맥스 로짓을 통해 응답에 대한 확률 분포 (p_\theta(y)) 를 정의한다.

3. LAD objective:

   • f‑divergence (D_f(p_A ,|, p_\theta)) 를 최소화한다. 실제 구현에서는 KL‑divergence를 사용하여 손실을 다음과 같이 정의한다:

[ \mathcal{L}{\text{LAD}} = \mathbb{E}{y \sim p_A}!\big[ \log p_A(y) - \log p_\theta(y) \big]. ]

• 이 손실에 대한 경사 하강은 높은 advantage를 가진 답변의 가능성을 높이고 낮은 advantage를 가진 답변의 가능성을 낮추면서, 분포가 과도하게 뾰족해지는 것을 방지한다.

4. Training pipeline:
   • 후보 완성 집합을 생성한다(예: nucleus sampling 사용).
   • 각 후보를 작업‑특화 보상 모델로 평가한다(예: 수학 풀이의 정답 여부).
   • advantage를 계산하고 (p_A) 를 만든 뒤, LAD 손실을 이용해 LLM을 업데이트한다.

손실이 보상을 얻기 위한 순전파와 KL 항에 대한 표준 역전파만을 필요로 하기 때문에, 이 방법은 기존의 RL‑from‑human‑feedback 혹은 RL‑fine‑tuning 루프에 자연스럽게 통합될 수 있다.

결과 및 발견

* 다양성은 평균 토큰 수준 엔트로피와 프롬프트당 서로 다른 유효 솔루션 비율로 측정됩니다.

핵심 요약

• 작업 전반에 걸쳐 2–4 %의 정확도 향상이 나타났으며, 엔트로피 정규화 RL 방법과 비교해도 동등하거나 더 좋습니다.
• 생성 다양성이 눈에 띄게 향상되어 모델이 동일한 “안전한” 답변을 반복해서 출력할 가능성이 낮아졌음을 나타냅니다.
• 제어된 밴딧 설정에서 LAD는 이론적인 이점 분포와 완벽히 일치하여 공식의 정확성을 확인했습니다.

실용적 함의

• 보다 견고한 LLM 어시스턴트: 챗봇, 튜터링 시스템, 코드 어시스턴트를 개발하는 개발자들은 LAD를 채택하여 정답이면서도 다양성을 갖춘 답변을 얻을 수 있어, 반복적이거나 과도하게 보수적인 응답 위험을 줄일 수 있다.

• 무 비용 파인튜닝: LAD는 일반적인 보상 평가 외에 추가적인 forward pass를 요구하지 않으므로, 기존 RL‑HF 파이프라인에 추가 GPU 예산 없이 바로 삽입할 수 있다.

• 안전‑중요 분야에서의 향상된 탐색: 여러 고‑이득 추론 경로를 보존함으로써 단일 목표 RL이 놓칠 수 있는 새로운 솔루션을 발견할 수 있다(예: 자동 정리 증명, 과학적 발견).

• 간소화된 하이퍼파라미터 튜닝: 이 방법은 엔트로피 계수를 조정할 필요성을 없애며, PPO‑스타일 RL을 LLM에 적용할 때 흔히 겪는 어려움을 해소한다.

제한 사항 및 향후 연구

• 보상 모델 의존성: LAD의 성능은 기본 보상 추정기의 품질에 좌우됩니다; 편향되거나 노이즈가 많은 보상은 학습된 이점 분포에 직접적인 영향을 미칩니다.
• 후보 생성의 확장성: 이 접근법은 프롬프트당 적당한 수의 샘플링된 완성을 필요로 하며, 매우 큰 모델의 경우 이 단계를 실현 가능하게 유지하기 위해 신중한 예산 관리가 필요할 수 있습니다.
• KL 발산에 대한 이론적 초점: 논문에서는 일반적인 f‑divergence에 대한 동등성을 증명했지만, 실험에서는 KL만을 탐구했습니다. 다른 발산(예: 역 KL, α‑divergence)을 살펴보면 탐색과 활용 사이의 다른 트레이드‑오프를 발견할 수 있습니다.
• 보다 넓은 작업 범위: 현재 평가가 수학 및 코드 추론에 집중되어 있어, LAD를 개방형 생성(예: 스토리 작성)에 적용하는 것은 아직 미해결 과제입니다.

전반적으로, LAD는 개념적으로 단순하면서도 강력한 RL 기반 LLM 파인‑튜닝 조정 방법을 제공하여 정확성과 창의성을 동시에 향상시킬 수 있습니다—추가 연산 비용 없이 언어 모델에서 더 많은 가치를 끌어내고자 하는 개발자에게 매력적인 제안입니다.

저자

• Wendi Li
• Sharon Li

논문 정보

• arXiv ID: 2602.20132v1
• 분류: cs.LG
• 출판일: 2026년 2월 23일
• PDF: PDF 다운로드