🤖 machine learning

OPD+: Rethinking the Advantage Design for On-Policy Distillation

이 논문은 어드밴티지 추정 시 표준 스톱 그래디언트(stop-gradient) 연산에 의해 발생하는 편향을 수학적으로 증명하고, 추론 및 도구 사용 벤치마크에서 성능을 향상시키는 일반적인 f-다이버전스 기반 최적화 방법을 제안하는 교정된 온-폴리시 증류 프레임워크인 OPD+를 소개한다.

원저자: Hanyang Zhao, Haoxian Chen, Han Lin, Genta Indra Winata, David Yao, Wenpin Tang

게시일 2026-06-02

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Hanyang Zhao, Haoxian Chen, Han Lin, Genta Indra Winata, David Yao, Wenpin Tang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 숙련된 요리사(교사 모델)를 관찰하며 복잡한 퍼즐을 푸는 법을 배우려는 어린 견습생(학생 모델)을 가르치고 있다고 상상해 보세요. 견습생은 요리사의 움직임을 관찰하고, 그 동작을 따라 하려 노력하며, 자신의 움직임이 마스터의 레시피와 얼마나 유사했는지에 대한 피드백을 받습니다. 이 과정을 **온폴리시 증류(On-Policy Distillation, OPD)**라고 부릅니다.

오랫동안 연구자들은 피드백을 주는 데 있어 가장 효과적인 특정 "레시피"가 하나 존재한다고 믿었습니다. 바로 **역 KL(Reverse KL)**이라 불리는 방법입니다. 이는 마치 "학습하는 유일한 방법은 정확히 이 방식으로 차이를 측정하는 것뿐이다"라고 말하는 것과 같았습니다.

하지만 이 논문의 저자들은 **OPD+**를 통해, 피드백이 계산되는 방식에 숨겨진 결함이 있다는 것을 발견했습니다. 그들은 피드백을 전달하는 방식이 실제로 고장 나 있으며, 이로 인해 특히 다양한 유형의 피드백 레시피를 시도할 때 혼란과 학습 저하를 초래한다는 사실을 찾아냈습니다.

다음은 이들의 발견을 쉬운 비유를 통해 정리한 내용입니다.

1. 고장 난 피드백 루프 ("스톱 그래디언트" 문제)

기존 방식에서는 견습생이 학습하려고 할 때, 교사는 "네 움직임의 점수는 이렇다"라고 말한 뒤, 그 점수가 변하지 않도록 즉시 고정해 버렸습니다. 이는 마치 교사가 정적인 차트를 가리키며 "이것이 목표치이니, 이 차트가 어떻게 그려졌는지는 신경 쓰지 마라"라고 말하며 안정성을 유지하려는 것과 같습니다.

논문은 이것이 수학적으로 잘못되었다고 주장합니다.

비유: 당신이 화살 쏘기를 배우고 있다고 상상해 보세요. 선생님은 "화살이 왼쪽으로 5인치 지점에 떨어졌다"라고 말합니다. 하지만 곧이어 선생님은 "그 5인치를 어떻게 계산했는지는 중요하지 않으니, 그냥 고정된 사실로 받아들여라"라고 말합니다.
문제점: 그런데 이 "5인치"라는 수치는 사실 당신이 활을 어떻게 잡고 있는지(학생의 현재 상태)에 따라 달라집니다. 만약 당신이 움켜쥐는 법을 바꾼다면, 그 거리도 변합니다. 하지만 계산을 고정해 버림으로써, 선생님은 당신에게 거짓말을 하고 있는 것입니다. 당신은 자신의 움켜쥐는 법을 조절하여 개선하려고 노력하지만, 새로운 움켜쥐기에 맞춰 업데이트되기를 거부하는 숫자에 기반하여 움직이고 있습니다. 이는 **편향된 추정치(biased estimates)**로 이어집니다. 즉, 당신은 자신이 발전하고 있다고 생각하지만, 실제로는 잘못된 방향으로 움직이고 있는 것입니다.

2. 새로운 해결책: OPD+

저자들은 피드백이 학생이 학습함에 따라 실제로 업데이트되도록 "수학을 바로잡는" 방법인 **OPD+**를 제안합니다.

해결책: 점수를 고정하는 대신, OPD+는 "여기 점수가 있고, 만약 네가 손을 움직인다면 이 점수가 정확히 어떻게 변하는지도 알려주겠다"라고 말합니다. 이는 피드백에 작은 보정 항(correction term)을 추가하는 것입니다.
결과: 이제 선생님은 "화살이 왼쪽으로 5인치 떨어졌다. 하지만 기억해라, 네가 움켜쥐는 힘을 조절하면 그 거리가 1인치 줄어들 것이다"라고 말하는 것과 같습니다. 이 작은 보정 덕분에 학습 과정이 정직하고 정확해집니다.

3. 놀라운 사실: 다른 레시피들도 효과가 있다!

수년 동안 사람들은 역 KL만이 학생과 교사 사이의 차이를 측정하는 유일하게 좋은 방법이라고 생각했습니다. 그들은 **순방향 KL(Forward KL)**이나 JSD 같은 다른 방법들은 쓸모없으며, 학생을 "붕괴(collapse)"(학습을 완전히 멈추게 함)시킬 것이라고 믿었습니다.

논문은 이러한 다른 방법들이 실제로 나빴던 것이 아니라, 고장 난 피드백 루프("스톱 그래디언트" 문제) 때문에 제대로 사용되지 못했을 뿐이라는 것을 보여줍니다.

비유: 마치 특정 종류의 망치가 못을 박는 유일한 도구라고 모두가 생각했던 것과 같습니다. 그들은 드라이버를 사용해 보았고, 드라이버가 나무를 망가뜨리는 것을 보고 드라이버는 쓸모없다고 결론 내렸습니다.
발견: 저자들은 드라이버를 잡는 방식(수학)을 고쳤습니다. 그러자 갑자기 드라이버가 완벽하게 작동하기 시작했습니다! 실제로 어떤 경우(JSD 방식의 경우)에는 드라이버가 망치보다 더 잘 작동하여, 학생이 더 어려운 수학 문제를 풀고 도구를 더 효과적으로 사용할 수 있게 해주었습니다.

4. 테스트 대상

그들은 두 가지 매우 어려운 과제를 테스트했습니다:

수학적 추론: 높은 수준의 경시대회 수학 문제(AIME 및 HMMT 등) 해결.
도구 사용: AI가 외부 도구(계산기나 검색 엔진 등)를 사용하여 문제를 해결하도록 가르치기.

결과:

기존 방식(고장 난 수학 기반)은 일부 학습 스타일을 완전히 실패(정확도 0%)하게 만들었습니다.
새로운 방식(OPD+)은 이러한 실패를 해결했습니다.
"표준" 방식(역 KL)에 대해서도, 새로운 수학을 적용했을 때 학생이 더 빠르게 학습하고 더 높은 정점의 성능에 도달했습니다.

요약

이 논문은 우리가 현재 AI 모델이 서로를 모방하도록 가르치는 방식에 근본적인 수학적 오류가 있다고 주장합니다. 피드백을 "고정"하는 코드 한 줄을 수정함으로써 우리는 다음을 달uç할 수 있습니다:

학습 과정을 수학적으로 정직하게 만듭니다.
이전에는 고장 났다고 여겨졌던 다양한 학습 전략의 잠재력을 끌어냅니다.
우리가 이미 보유한 모델을 가지고도 수학 및 도구 사용과 같은 어려운 작업에서 더 나은 결과를 얻을 수 있습니다.

요컨대: 피드백을 고정하지 마세요. 수학이 업데이트되도록 두면, AI는 더 잘 학습합니다.

기술 요약: OPD+: 온-폴리시 증류를 위한 어드밴티지 설계의 재고

1. 문제 정의

온-폴리시 증류(On-policy distillation, OPD)는 유능한 교사 언어 모델로부터 기초 학생 모델로 능력을 전이하기 위해 사용되는 기법으로, 학생이 생성한 롤아웃(rollout)에서 유도된 조밀한 토큰 레벨 감독 신호를 최적화합니다. 기존의 OPD 구현체(예: Tinker)는 주로 쿨백-라이블러(Kullback-Leibler, KL) 발산을 주요 목적 함수로 사용하지만, 다음과 같은 특정 설계 선택을 채택합니다: **보상 항(학생과 교사의 우도 비율의 로그 차이)에 대한 그래디언트 정지(stopping the gradient)**를 적용하는 것입니다.

이 설계는 딥 Q-네트워크(DQN)의 타겟 네트워크와 유사하게 수치적 안정성을 확보하려는 의도에서 비롯되었습니다. 그러나 저자들은 이 접근 방식에서 치명적인 이론적 결함을 식별했습니다:

편향된 그래디언트 추정: 일반적인 $f$ -발산의 경우, 보상 함수는 학생 정책 파라미터 $\theta$ 에 의존합니다. 보상 항에 그래디언트 정지를 적용하면 $\theta$ 에 대한 보상의 직접적인 의존성을 무시하게 되어 수학적으로 잘못된 그래디언트 추정치를 초래합니다.
성능 저하: 이러한 편향은 Reverse KL 이외의 발산 함수(예: Forward KL 및 Jensen-Shannon Divergence)를 사용할 때 파괴적인 학습 붕괴를 일으켜 학습을 불가능하게 만듭니다. Reverse KL의 경우에도, 스코어 항등식(score identity) 덕분에 편향이 이론적으로는 사라지지만, 실험 결과는 여전히 그래디언트 정지 설계가 최적의 성능을 내지 못함을 시사합니다.

2. 방법론: OPD+ 프레임워크

본 논문은 $f$ -발산에 기반하여 온-폴리시 증류를 위한 어드밴티지 설계를 재평가하는 수정된 최적화 프레임워크인 **OPD+**를 제안합니다.

2.1 이론적 기초

저자들은 OPD 목적 함수를 학생 정책 $p_\theta$ 와 교사 정책 $q$ 사이의 $f$ -발산을 최소화하는 것으로 공식화합니다. 목적 함수는 다음과 같이 정의됩니다:
$J_f(\theta) = \mathbb{E}_{x \sim X, y \sim p_\theta(\cdot|x)} \left[ \sum_{n=1}^{L_y} f(u_{\theta,n}) \right]$
여기서 $u_{\theta,n} = q(y_n | y_{<n}, x) / p_\theta(y_n | y_{<n}, x)$ 는 토큰 레벨 밀도 비율이며, $f$ 는 $f(1)=0$ 을 만족하는 볼록 함수입니다.

2.2 수정 사항

논문은 이 목적 함수의 진정한 그래디언트를 도출합니다. 표준 OPD는 보상 $r = -f(u_{\theta,n})$ 를 고정된 것으로 취급하는 "그래디언트 정지" 정책 그래디언트 추정기를 사용합니다. 저자들은 진정한 그래디언트에 두 가지 항이 포함되어 있음을 보여줍니다:

표준 스코어 함수 항: $\mathbb{E}[r \nabla_\theta \log p_\theta]$ .
직접 보상-그래디언트 항: $\mathbb{E}[\nabla_\theta r]$ .

일반적인 $f$ -발산의 경우, 직접 보상-그래디언트 항은 사라지지 않습니다. 구체적으로, $\nabla_\theta r = u_{\theta,n} f'(u_{\theta,n}) \nabla_\theta \log p_\theta$ 입니다.
따라서 올바른 어드밴티지 함수는 단순히 $-f(u)$ 가 아니라, 다음과 같은 가중치 항이 됩니다:
$w_f(u) = -f(u) + u f'(u)$

2.3 특수 사례: Reverse KL

Reverse KL의 경우, $f(u) = -\log u$ 입니다. 수정 항은 다음과 같습니다:
$w(u) = -(-\log u) + u(-1/u) = \log u + 1$
추가된 $+1$ 항은 상수 베이스라인입니다. $\mathbb{E}[\nabla_\theta \log p_\theta] = 0$ (스코어 항등식)이므로, 이 항은 이론적으로 기대값에서 사라집니다. 그러나 저자들은 실제 상황에서 이 상수 베이스라인이 분산 감소 메커니즘으로 작용하여 수렴 속도와 안정성을 개선한다고 주장합니다.

2.4 구현

OPD+는 기존 OPD 구현에서 생으로 된 음의 발산 보상 $-f(u)$ 를 수정된 가중치 $w_f(u)$ 로 교체하는 단 한 줄의 코드 변경만을 요구합니다. 이 프레임워크는 Forward KL, Reverse KL, Jensen-Shannon Divergence(JSD)를 포함한 다양한 $f$ -발산을 지원합니다.

3. 주요 기여

범용 최적화 프레임워크: 본 논문은 일반적인 $f$ -발산에 기반한 온-폴리시 증류를 위한 통합된 프레임워크를 제공하며, 보상의 파라미터 의존성을 고려한 정확한 그래디언트와 손실 함수를 도출합니다.
편향에 대한 이론적 증명: 표준 그래디언트 정지 연산이 일반적인 $f$ -발산에 대해 편향된 그래디언트 추정치를 생성하며, 보상을 고정된 것으로 취급하는 일반적인 관행을 무효화한다는 것을 수학적으로 증명합니다.
OPD+ 알고리즘: 저자들은 Forward KL 및 JSD에서 관찰되는 학습 붕괴를 제거하고, Reverse KL의 성능을 개선하는 수정된 알고리즘인 OPD+를 도입합니다.
실험적 검증: 본 연구는 다른 $f$ -발산들이 (올바르게 최적화될 경우) Reverse KL보다 더 우수한 성능을 낼 수 있음을 입증함으로써, 온-폴리시 증류에서 Reverse KL이 유일하게 실행 가능한 거리 척위라는 기존의 믿음에 도전합니다.

4. 실험 결과

저자들은 QWEN3 모델 제품군(1.7B 및 8B 파라미터)을 대상으로 두 가지 별도의 도메인인 경쟁 수학(DeepMath, AIME 24/25, HMMT 25) 및 에이전트 도구 사용(Agentic Tool Use)에서 OPD+를 평가했습니다.

4.1 "학습 불가능한" 발산의 구제

표준 그래디언트 정지 설계(OPD) 하에서 Forward KL과 JSD는 도구 사용 작업에서 정확도가 0%로 떨어지고, 수학 작업에서 SFT 초기화 베이스라인을 따라잡지 못하는 등 파괴적인 실패를 보였습니다.

OPD+ 수정: 수정된 어드밴티지 $w(u)$ 를 적용함으로써, OPD+는 이러한 실패 모드를 완전히 제거했습니다.
JSD 성능: JSD+는 안정화되었을 뿐만 아니라 표준 Reverse KL 베이스라인을 넘어섰습니다. 까다로운 AIME 25 벤치마크에서 JSD+는 Reverse KL의 48.54% 대비 **50.42%**의 정확도를 달성했습니다. 도구 사용 작업에서 JSD+는 **60.29%**에 도달하여 Reverse KL 베이스라인인 **59.97%**를 앞질렀습니다.

4.2 Reverse KL에서의 개선

전통적으로 Reverse KL이 지배적인 영역에서도 OPD+는 엄격한 실증적 개선을 보여주었습니다:

AIME 24: Reverse KL+는 64.27%에서 **64.69%**로 향상되었습니다.
도구 사용: Reverse KL+는 59.97%에서 **61.44%**로 향상되었습니다.
수렴: 학습 역학(그림 1)은 OPD+ 변형들이 초기 수렴을 가속화하고 표준 베이스라인에서 보이는 성능 정체 또는 저하를 방지함을 보여줍니다.

4.3 자기 증류 (OPSD)

이 방법은 모델이 학생이자 교사 역할을 하는 온-폴리시 자기 증류(On-Policy Self-Distillation, OPSD) 환경에서도 검증되었습니다. OPD+는 모든 발산 유형에 대해 모든 경우에 대해 바닐라 그래디언트 정지 버전을 지속적으로 능가했으며, 수정된 Reverse KL이 모든 발산 중 가장 좋은 성능을 보였습니다.

5. 의의 및 주장

본 논문은 다른 $f$ -발산(Forward KL, JSD)의 역사적인 저조한 성능이 발산 자체의 내재적 한계가 아니라, 그래디언트 정지 설계로 인한 잘못 지정된 그래디ient 계수의 부작용이라는 점을 주장합니다.

$f$ -발산 목적 함수에 수학적으로 충실하도록 어드挽티지 함수를 수정함으로써, OPD+는 다음과 같은 효과를 얻습니다:

설계 공간 확장: 산업 표준인 Reverse KL을 넘어 더 다양한 통계적 거리를 증류에 사용할 수 있게 합니다.
성능 향 향상: 추론 및 도구 사용 벤치마크에서 새로운 최고 성능(SOTA)을 달립합니다.
관습에 대한 도전: Reverse KL이 안정성 때문에 본질적으로 우월하다는 일반적인 믿음에 의문을 제기하며, 그 안정성이 발산의 선택이라기보다는 최적화 방법의 산물이었음을 시사합니다.

저자들은 자신들의 연구가 더 수학적으로 정확하고 경험적으로 효과적인 알고리즘을 제공하며, 이는 온-폴리시 증류 방법의 어드밴티지 함수 설계에 대한 추가적인 연구를 촉발할 수 있다고 결론짓습니다. 또한 모델 스케일링 분석 및 단일 토큰 우도와 로짓 증류 간의 비교에 관한 제한 사항을 언급하며 이를 향후 연구 과제로 남겨두었습니다.