💬 NLP

KL for a KL: On-Policy Distillation with Control Variate Baseline

본 논문은 폐형식 토큰 단위 음의 역 KL 발산을 제어 변수 기준값으로 활용하여 학습 분산을 줄이는 안정적인 온-정책 증류 방법인 vOPD를 제안하며, 이는 추가적인 추론 오버헤드 없이 고비용 전체 어휘 기준값과 비교 가능한 성능을 달성합니다.

원저자: Minjae Oh, Sangjun Song, Gyubin Choi, Yunho Choi, Yohan Jo

게시일 2026-05-11

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Minjae Oh, Sangjun Song, Gyubin Choi, Yunho Choi, Yohan Jo

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

한 학생 (AI 모델) 이 뛰어난 교사 (더 강력한 AI 모델) 를 연구하도록 하여 복잡한 수학이나 과학 문제를 해결하는 법을 가르친다고 상상해 보세요.

목표는 학생이 교사의 사고 과정을 완벽하게 모방하여 스스로 문제를 해결할 수 있도록 하는 것입니다. 이 과정은 **온-정책 증류 (On-Policy Distillation, OPD)**라고 불립니다.

문제: 학습의 "롤러코스터"

이 작업을 수행하는 표준 방식에서 학생은 단어 단위로 답변을 생성합니다. 매 단어마다 시스템은 "교사가 이 단어를 말했는가?"를 확인합니다.

교사가 말했다면 학생은 작은 "잘했다"는 칭찬을 받습니다.
교사가 말하지 않았다면 학생은 "잘못했다"는 평가를 받습니다.

문제는 이 피드백이 노이즈가 많고 불안정하다는 점입니다. 학생이 줄타기를 하고 있다고 상상해 보세요. 때로는 아주 작은 밀림이 "잘못했다"는 신호가 너무 가혹하고 무작위적이어서 줄에서 날아오르게 만듭니다. 이 논문은 이를 **높은 기울기 분산 (high gradient variance)**이라고 부릅니다. 이는 누군가가 계속 무작위로 당신을 밀어 떨어뜨리는 사이 자전거 타기를 배우려는 것처럼 학습을 느리고 예측 불가능하게 만듭니다.

이 문제를 해결하려는 이전 시도들은 다음과 같은 방식으로 흔들림을 멈추려 했습니다:

완벽한 점수를 계산하기 위해 사전에 있는 모든 가능한 단어를 한 번에 살펴보는 것. (너무 느립니다. 도서관의 모든 책을 한 권씩 확인하여 한 문장을 찾는 것과 같습니다.)
학생이 말할 가능성이 높은 상위 20 개 단어만 살펴보는 것. (더 빠르지만, 나머지 사전 전체를 무시하여 편향을 도입합니다. 이는 시끄러운 목소리만 듣고 조용하지만 중요한 목소리는 무시하는 것과 같습니다.)

해결책: vOPD ("안정화 장치")

저자들은 vOPD라는 새로운 방법을 제안합니다. 그들은 학습 과정을 강화 학습 (Reinforcement Learning) 게임처럼 취급하며 **"제어 변수 기준선 (Control Variate Baseline)"**이라는 교묘한 트릭을 사용합니다.

다음은 비유입니다:
말 경주에 베팅한다고 상상해 보세요.

보상: 당신의 말이 이기면 돈을 따게 됩니다.
문제: 보상은 엄청나고 예측 불가능합니다. 어느 날은 크게 이기고, 다음 날은 크게 잃습니다. 전략을 배우기 어렵습니다.
기준선 트릭: 경주 전에 일반적인 경주의 평균 보상을 계산합니다.
- 당신의 말이 이기면 단순히 "이겼다!"라고 말하지 않습니다. 대신 "평균 보다 더 많이 이겼다"라고 말합니다.
- 당신의 말이 지면 단순히 "졌다"라고 말하지 않습니다. 대신 "평균 보다 덜 잃었다"라고 말합니다.

이 "평균"(기준선) 을 결과에서 빼면 급격한 요동을 부드럽게 만듭니다. 학습의 방향을 바꾸지는 않습니다 (어떤 말이 더 좋은지 여전히 알 수 있습니다). 하지만 롤러코스터 효과를 일으키는 노이즈를 제거합니다.

마법의 재료: 무료 점심

대부분의 AI 시스템에서 이 "평균"을 계산하려면 학생과 함께 실행되는 두 번째 고비용 AI 모델 ("비평가") 이 필요합니다. 이는 모든 것을 느리게 만듭니다.

vOPD 의 획기적인 점은 이 특정 유형의 학습에서는 "평균"(기준선) 을 학생이 이미 생성하는 정확한 데이터를 사용하여 수학적으로 실시간으로 계산할 수 있다는 것을 깨달았다는 것입니다.

마치 두 번째 사람이 계산을 할 필요 없이 뇌에 내장된 계산기가 즉시 "평균" 점수를 알려주는 것과 같습니다.
이 기준선은 단순히 학생이 생각하는 것과 교사가 생각하는 것 사이의 차이일 뿐입니다.

왜 작동하는가

핫스팟을 식히다: 학생과 교사가 극단적으로 불일치할 때 ("높은 불일치"), 표준 방법은 거대하고 노이즈가 많은 신호로 "오류!"라고 소리칩니다. vOPD 는 이 거대한 불일치를 감지하고 기준선을 계산한 후, "좋아, 이건 큰 차이이지만 신호가 너무 무섭지 않도록 조정하자"라고 말합니다. 이는 학습 과정을 위한 쇼크 업소버처럼 작용합니다.
편향 없음: 속이지 않습니다. 목표를 바꾸지 않으며, 단지 목표에 가는 길을 부드럽게 만들 뿐입니다.
빠름: 두 번째 모델이 필요하거나 전체 사전을 확인하지 않기 때문에 원래의 불안정한 방법과 거의 같은 속도로 실행됩니다.

결과

저자들은 수학 및 과학 문제 (방정식 풀이나 화학 문제 등) 에 대해 이를 테스트했습니다.

성능: vOPD 는 표준 방법보다 더 빠르게 학습하고 더 좋은 점수를 받았습니다. "매우 느리고 매우 정확한" 방법 (전체 사전 확인) 의 성능과 일치했지만 훨씬 빠르게 수행했습니다.
안정성: 학습 과정이 훨씬 더 매끄러웠습니다. 시스템에 가해지는 "충격"이 10 배에서 100 배까지 감소하여 AI 가 뛰어다니는 대신 꾸준히 학습하게 되었습니다.
효율성: 그들은 기준선의 "대략적인 추측"(상위 20 개 단어만 확인) 이도 완벽한 계산과 거의 비슷하게 작동하여 실행 비용을 극도로 낮췄다는 것을 발견했습니다.

요약

vOPD는 AI 모델을 가르치는 더 지적인 방법입니다. AI 가 노이즈가 많고 무작위적인 피드백에 압도되도록 내버려 두는 대신, 요동을 부드럽게 만드는 내장된 "현실 점검"(기준선) 을 추가합니다. 이를 통해 AI 는 추가적인 컴퓨팅 파워 없이도 더 빠르고 안정적으로 복잡한 추론 기술을 학습할 수 있습니다.

기술 요약: vOPD (제어 변수 기준선을 활용한 온-정책 증류)

1. 문제 제기

온-정책 증류 (OPD) 는 특히 추론 분야에서 대규모 언어 모델 (LLM) 을 위한 주류 사후 학습 패러다임으로 부상했습니다. 희소 종단 보상에 의존하는 검증 가능 보상 강화 학습 (RLVR) 과 달리, OPD 는 밀집된 토큰 수준 신호를 사용하여 학생 모델과 교사 모델 간의 역 Kullback-Leibler (KL) 발산을 최소화합니다. 이를 통해 더 빠른 학습 속도와 긴 사고의 사슬 (CoT) 작업에 대한 효과적인 처리가 가능해집니다.

그러나 OPD 는 실제 적용에서 중요한 안정성 문제에 직면해 있습니다. 표준 구현은 그래디언트를 근사하기 위해 단일 샘플 몬테카를로 추정기를 사용합니다. 이 접근법은 높은 그래디언트 분산으로 인해 불안정한 학습 동역학을 초래합니다. 기존 안정화 방법들은 이를 완화하려 시도하지만 상당한 트레이드오프를 도입합니다:

전체 어휘 OPD ( $OPD_{full-V}$ ): 전체 어휘에 걸쳐 KL 발산을 계산합니다. 이는 분산을 제거하지만, 매 단계에서 전체 어휘 (~15 만 토큰 등) 에 대한 역전파를 요구하는 금지적인 계산 오버헤드를 초래합니다.
Top-k OPD ( $OPD_{top-k}$ ): KL 계산을 가장 확률이 높은 토큰 상위 $k$ 개로 제한합니다. 이는 더 가볍지만, 지지를 절단함으로써 그래디언트에 편향을 도입하며, 실험 결과 불안정한 기본 OPD 대비 성능 향상이 미미한 것으로 나타났습니다.

핵심 과제는 단일 샘플 추정기의 계산 효율성을 희생하거나 목적 함수에 편향을 도입하지 않고 OPD 학습을 안정화하는 것입니다.

2. 방법론: vOPD

저자들은 vOPD(제어 변수 기준선을 활용한 온-정책 증류) 를 제안하며, OPD 를 정책 그래디언트 강화 학습 (RL) 문제로 재구성하고 표준 분산 감소 기법인 제어 변수 기준선 차감을 적용합니다.

2.1 이론적 기반

정책 그래디언트 RL 에서 보상 $r_t$ 에서 기준선 $b_t(c_t)$ 를 차감하면 이득 함수 $a_t = r_t - b_t$ 가 도출됩니다. 중요한 점은 기준선이 샘플링된 행동과 독립적일 경우, 그래디언트는 편향되지 않은 상태로 유지되지만 분산은 감소한다는 것입니다. 기준선의 표준적인 선택은 가치 함수 $V^{\pi}(c_t)$ 입니다.

vOPD 의 핵심 통찰은 OPD 의 경우 가치 함수가 추가적인 크리틱 모델이나 추론 없이 **폐형 해 (closed-form solution)**를 가진다는 점입니다:
$V^{\pi_\theta}(c_t) = \mathbb{E}_{y_t \sim \pi_\theta} [r_t(c_t, y_t)] = -D_{KL}(\pi_\theta(\cdot|c_t) \parallel \pi_T(\cdot|c_t))$
여기서 보상은 $r_t = \log \pi_T - \log \pi_\theta$ 로 정의됩니다. 학생 분포 하에서의 기대 보상은 해당 컨텍스트에서 학생과 교사 간의 음의 역 KL 발산과 동일합니다.

2.2 알고리즘 구현

vOPD 그래디언트 추정기는 다음과 같이 공식화됩니다:
$\nabla_\theta J_{vOPD} = \mathbb{E} \left[ \sum_{t} \left( r_t(c_t, y_t) + D_{KL}(\pi_\theta \parallel \pi_T) \right) \nabla_\theta \log \pi_\theta(y_t|c_t) \right]$

편향 없음: 기준선 (KL 항) 은 분포 $\pi_\theta$ 와 $\pi_T$ 를 기반으로 계산되지만, 샘플링된 토큰 $y_t$ 에 대한 그래디언트 흐름과는 연결이 끊겨 (detached) 있습니다. 따라서 그래디언트는 편향되지 않은 상태로 유지됩니다.
효율성: KL 기준선은 순전파 (이미 보상 계산을 위해 필요함) 중에 계산됩니다. 역전파는 여전히 샘플링된 단일 토큰에만 영향을 미치므로 기본 OPD 의 계산 효율성이 유지됩니다.
분산 감소 메커니즘: 기준선은 정규화제 역할을 합니다. 학생과 교사 분포가 강력하게 불일치할 때 (높은 KL), 보상 $r_t$ 는 크고 음수입니다. 기준선 $D_{KL}$ 은 크고 양수이므로, 음수 보상을 효과적으로 상쇄하여 불안정을 유발하는 두꺼운 꼬리 (heavy-tailed) 그래디언트를 감쇠시킵니다.

2.3 Top-k 근사 ( $vOPD_{top-k}$ )

전체 어휘 KL 계산의 순전파 비용을 추가로 줄이기 위해, 저자들은 학생 모델의 상위 $k$ 개 토큰만을 사용하여 기준선을 근사하는 것을 제안합니다.

손실 목적 함수를 대체하여 편향을 도입하는 $OPD_{top-k}$ 와 달리, $vOPD_{top-k}$ 는 기준선 에만 상위 $k$ 개 근사를 사용합니다.
기준선은 샘플링된 토큰과 연결이 끊겨 있으므로, $k$ 값에 관계없이 그래디언트는 편향되지 않은 상태로 유지됩니다.
실험적으로, 작은 $k$ (예: $k=5$ ) 만으로도 학습을 안정화하기에 충분한 분산 감소를 제공합니다.

3. 주요 기여

폐형 해 가치 기준선: 본 논문은 OPD 가치 함수가 정확히 음의 역 KL 발산임을 증명하여, 추가 모델 없이 순전파에서 직접 이를 얻을 수 있음을 보여줍니다.
편향 없는 안정화: vOPD 는 편향을 도입하는 Top-k 변형과 달리 단일 샘플 몬테카를로 추정기의 편향 없는 특성을 유지하면서 그래디언트 분산을 감소시켜 학습을 안정화합니다.
계산 효율성: vOPD 는 기본 OPD 의 단일 토큰 역전파를 유지합니다. $vOPD_{top-k}$ 변형은 순전파 비용을 추가로 감소시켜 전체 어휘 기반 방법 대비 우수한 정확도 - 계산 트레이드오프를 제공합니다.
이론적 분석: 저자들은 분산 감소가 KL 발산의 제곱에 비례함을 보여주는 유도를 제공하여, 이 방법이 불안정을 유발하는 고불일치 토큰을 구체적으로 표적한다는 것을 확인합니다.

4. 실험 결과

저자들은 네 가지 모델 (Qwen3-1.7B, 4B, Olmo-3-7B) 에 대해 여섯 가지 추론 벤치마크 (MATH500, Minerva Math, AMC23, AIME24/25, SciKnowEval, GPQA-Diamond) 에서 vOPD 를 평가했습니다.

성능 향상: vOPD 는 기본 OPD 를 일관되게 능가했습니다.
- MATH500 에서 vOPD 는 기본 OPD 대비 최대 **+6.2%**의 절대적 개선을 달성했습니다.
- 모든 벤치마크에서 vOPD 변형들은 기본 OPD 대비 평균 **+3%**의 개선을 보였습니다.
- vOPD 성능은 가장 비용이 많이 드는 방법인 $OPD_{full-V}$ 와 일치하거나 약간 상회하면서도 훨씬 더 빨랐습니다.
안정화 변형과의 비교:
- $OPD_{top-k}$ 는 기본 OPD 대비 평균 +0.4% 정도의 미미한 개선만 보여 편향 문제를 확인시켰습니다.
- vOPD 변형들은 $OPD_{top-k}$ 를 크게 능가했으며 $OPD_{full-V}$ 와 일치했습니다.
효율성:
- vOPD 는 $OPD_{full-V}$ 대비 최대 **57.7%**의 실제 소요 시간을 단축했습니다.
- $vOPD_{top-k}$ 가 가장 효율적이었으며, 가장 우수한 정확도 - 계산 트레이드오프를 제공했습니다.
안정성 지표:
- 그래디언트 노름: vOPD 는 기본 OPD 대비 1~2 차수 낮은 그래디언트 노름을 유지하여 노이즈가 많고 불안정을 유발하는 그래디언트의 억제를 확인시켰습니다.
- 하이퍼파라미터 민감도: 성능은 $k$ 값 (5 에서 100) 전반에 걸쳐 안정적이었으며, 심지어 거친 근사 ( $k=5$ ) 만으로도 충분한 안정화를 제공했습니다.

5. 의의 및 주장

본 논문은 vOPD 가 원칙적이고 효율적인 안정적 온-정책 증류 접근법을 제공함으로써 강화 학습과 지식 증류 간의 간극을 메운다고 주장합니다.

메커니즘: 단일 샘플 추정기에서 비롯된 높은 그래디언트 분산을 OPD 불안정의 근본 원인으로 규명하고, 문제의 구조 (KL 발산) 에서 도출된 제어 변수 기준선을 사용하여 이를 해결합니다.
실용적 영향: 추가 모델 (크리틱) 이나 비싼 역전파 없이 기본 OPD 를 바로 대체할 수 있는 솔루션을 제공하여, 계산 효율성이 최우선인 산업 규모 사후 학습에 적합합니다.
일반화: 이 방법은 다양한 모델 규모, 아키텍처 (Qwen, Olmo), 그리고 추론 영역 (수학 및 과학) 에서 견고함이 입증되었습니다.

저자들은 vOPD 가 불안정을 유발하는 보상 토큰에 대한 원칙적인 정규화제 역할을 하여, OPD 를 처음부터 매력적이게 만들었던 효율성을 훼손하지 않으면서 불안정한 휴리스틱에서 신뢰할 수 있는 학습 패러다임으로 변환한다고 결론지었습니다.