기술 요약: 풍부한 피드백을 활용한 분포적 DAgger 기반 강화 학습

문제 정의

본 논문은 추론 모델 학습의 지배적인 패러다임인 검증 가능한 보상 기반 강화 학습(RLVR)의 한계를 다룹니다. RLVR에서 모델은 여러 개의 후보 솔루션을 생성하며, 각 솔루션은 최종 답변에 기초하여 단일 이진 보상(정답/오답)을 할당받습니다. 이 방식은 두 가지 결정적인 결함을 가집니다:

1. 부실한 신용 할당(Poor Credit Assignment): 최종 보상이 응답 내의 모든 토큰에 균등하게 배포되어, 어떤 특정 추론 단계가 성공이나 실패에 기여했는지를 구분하지 못합니다.
2. 검증 불가능한 도메인에 대한 적용 불가능성: RLVR은 최종 답변의 자동 검증에 의존하므로, 정답(ground-truth)을 알 수 없거나 확인하기 어려운 작업에는 적합하지 않습니다.

많은 설정이 "풍부한 피드백"(예: 실행 트레이스, 에러 로그, 전문가 교정 또는 자기 평가)을 제공하지만, 이를 온-폴리시 자기 증류(on-policy self-distillation, 예: SDPO, OPSD)를 통해 활용하려는 현재의 방법들은 이론적 및 실무적 장벽에 직면해 있습니다. 기존 접근 방식은 일반적으로 학생 정책(student policy)과 피드백이 조건화된 교사 정책(teacher policy) 사이의 $f$-발산(reverse KL 또는 Jensen-Shannon 등)을 최소화합니다. 저자들은 이러한 방법들이 단조적 정책 개선(monotonic policy improvement)에 대한 보장을 결여하고 있으며, 시퀀스 수준의 신용 할당을 수행하지 못해 종종 최적에 미치지 못하는 수렴으로 이어진다고 주장합니다.

방법론: DistIL

저자들은 풍부한 피드백이 있는 강화 학습에 적응시킨 고전적 DAgger 알고리즘의 분포적 변형인 DistIL(Distributional Imitation Learning)을 제안합니다.

핵심 목적 함수

DistIL은 학생이 방문한 상태(states)에서 학생 정책($\pi_\theta$)과 피드백이 조건화된 교사 정책($\pi_T$) 사이의 순방향 교차 엔트로피(forward cross-entropy) 손실을 최적화합니다:
$$\mathcal{L}_{\text{DistIL}}(\theta) := \mathbb{E}_{x \sim \rho, y \sim \pi_\theta(\cdot|x)} \left[ \sum_{t=1}^H H_\times(\pi_T(\cdot | s_t), \pi_\theta(\cdot | s_t)) \right]$$
여기서 $H_\times(p, q) = -\mathbb{E}_{y \sim p}[\log q(y)]$는 순방향 교차 엔트로피입니다. 결정적으로, 교사 $\pi_T$는 풍부한 피드백 $f$(예: 실행 로그 또는 정답)를 조건으로 하지만, 학생 업데이트 동안에는 고정된 상태로 유지됩니다.

그래디언트 및 신용 할당

국소적인 토큰 단위 불일치만을 사용하여 그래디언트를 근사하는(즉, 상태 분포를 고정하는) 기존의 자기 증류 방법들과 달리, DistIL은 **전체 그래디언트(full gradients)**를 활용합니다. 그래디언트는 두 가지 항으로 분해됩니다:

1. 국소 모방(Local Imitation): 현재 상태에서 학생이 교사의 분포를 따르도록 업데이트하는 교사 가중치 항입니다.
2. 미래 인지 신용 할당(Future-Aware Credit Assignment): 미래 시점에서 관찰된 교사-학생 간의 교차 엔트로피 불일치를 이전의 결정들로 역전파하는 항입니다. 이는 높은 교사-학생 불일치를 가진 궤적으로 이어지는 초기 토큰 선택에 적절한 신용이나 페널티를 부여하도록 보장합니다.

이 알고리즘은 교사의 로그 확률(log-probabilities)에 접근할 필요 없이 샘플이나 목적 함수를 추정할 수 있는 능력만 있으면 되므로 블랙박스 교사(외부 모델 또는 인간 전문가)를 지원합니다.

주요 기여

1. 기존 방법의 이론적 한계

본 논문은 기존의 $f$-발산(reverse KL 및 Jensen-Shannon 포함) 기반 온-폴리시 자기 증류 목적 함수가 두 가지 근본적인 문제를 겪고 있음을 입증하는 공식적인 증명을 제공합니다:

• 단조적 개선 보장 부재: 교사가 엄격하게 더 높은 보상을 가지고 있더라도, $f$-발산을 최소화하는 것이 학생의 기대 보상을 증가시킨다는 것을 보장하지 않습니다. 업데이트 방향이 로그 비율에 의한 가중치 부여로 인해 의도치 않게 차선의 행동 확률을 높일 수 있습니다.
• 국소 근사를 통한 최적 미달 수렴: 미래 상태 분포에 미치는 영향(future-credit term)을 무시함으로써, 국소 그래디언트 근사는 엄격하게 차선의 정책으로 수렴할 수 있습니다. 저자들은 단순한 2단계 MDP에서 국소 업데이트가 고보상 상태에 도달하는 데 중요한 초기 결정에 신용을 할당하는 데 실패할 수 있음을 증명합니다.

2. DistIL의 이론적 보장

DistIL은 최적화를 보상 최대화와 일치시키는 특성을 가짐이 입증되었습니다:

• 단조적 정책 개선: 표준적인 국소 실현 가능성(local realizability) 가정 하에서, 순방향 교차 엔트로피 목적 함수는 더 나은 교사를 향해 이동하는 것이 학생의 기대 보상을 증가시킨다는 것을 보장합니다.
• Regret Bounds: 이 알고리즘은 농축 계수(concentrability coefficients)와 교사의 분산에 따라 수렴 속도가 결정되는 아선형 레그렛(sublinear regret) 경계를 가집로 합니다.
• Pass@N 최적화: 이 목적 함수는 **교사 가중치 성공 가능성(teacher-weighted likelihood of success)**의 하한을 최대화함이 증명되었습니다. 이는 DistIL이 모든 $N$에 대해 Pass@N 지표(N개의 샘플 중 적어도 하나의 정답을 생성할 확률)를 개선하는 이유에 대한 원칙적인 설명을 제공합니다.

3. 실험적 성능

DistIL은 세 가지 도메인에서 평가되었습니다:

• 과학적 추론: SciKnowEval(화학, 물리학, 생물학, 재료과학)에서 DistIL은 SDPO, GRPO(온-폴리시 및 오프-폴리시)를 능가하였으며, 더 높은 Best@16 및 Maj@16 점수를 달성했습니다. 특히, SDPO는 특정 도메인에서 심각한 진동과 성능 붕괴를 보인 반면, DistIL은 안정적인 훈련을 보여주었습니다.
• 코딩: 실행 트레이스를 피드백으로 사용하는 LiveCodeBench에서, DistIL은 정확한 정답(exact correctness)과 부분 정확도(partial accuracy) 모두에서 SDPO 및 GRPO를 능가했습니다.
• 어려운 수학적 추론: 베이스 모델의 통과율(pass rate)이 0인 도전적인 문제(AIME, HMMT)에서, RLVR(GRPO) 방식은 어드밴티지 신호가 0이라서 개선에 실패했습니다. 반면 DistIL은 정답(ground-truth) 솔루션을 성공적으로 활용하여 성능을 개선했으며, 자기 증류 베이스라인(SDPO, OPSD) 및 SFT를 능가했습니다.

의의 및 주장

본 논문은 DistIL이 풍부한 피드백으로부터 학습하기 위한 이론적으로 근거 있는 프레임워크를 제공함으로써, 모방 학습과 강화 학습 사이의 간극을 메운다고 주장합니다. 그 의의는 다음과 같습니다:

1. 목적 함수의 교정: ($f$-발산 최소화와 같이 보상과 정렬되지 않을 수 있는 방식에서) 순방향 교차 엔트로피로 전환하여, 교사가 더 나을 때 단조적 개선을 보장합니다.
2. 시퀀스 수준 학습 가능: 전체 그래디언트를 활용함으로써, DistIL은 장기 호흡(long-horizon) 추론에서 발생하는 신용 할당 문제를 해결하고 초기 결정이 다운스트림 결과에 따라 평가되도록 보장합니다.
3. 실무적 견고성: DistIL은 블랙박스 교사와 다양한 피드백 유형(이진 정답부터 실행 로그까지)과 함께 작동하므로, 표준 RLVR보다 더 넓은 범위의 추론 작업에 적용 가능합니다.

저자들은 보상과 정렬된 목적 함수와 미래 인지 신용 할당의 결합이 풍부한 피드백을 추론 모델로 효과적으로 증류하는 데 필수적이며, 이를 통해 최신 베이스라인 대비 일관된 실험적 이득을 얻을 수 있다고 결론짓습니다.