대형 언어 모델을 위한 온폴리시 증류(On-Policy Distillation) 서베이

기존의 정적인 데이터 학습 방식(Off-policy)은 모델이 스스로 생성한 오류에 대처하지 못하는 '노출 편향' 문제를 야기한다. 이 논문은 모델이 직접 생성한 결과물에 대해 피드백을 받으며 학습하는 '온폴리시 증류' 기술을 체계적으로 정리하여, 더 작고 효율적이면서도 강력한 추론 능력을 갖춘 모델 제작의 이정표를 제시한다.

기존의 지식 증류는 스승 모델이 미리 만들어둔 정답지(Static dataset)를 학생 모델이 그대로 외우는 방식(Off-policy)이었다. 하지만 실제 서비스 시점(Inference)에서 학생 모델은 자신이 방금 내뱉은 토큰을 바탕으로 다음 토큰을 생성하는 자기회귀(Autoregressive) 방식을 취한다. 이때 학습 데이터에 없던 미세한 실수를 한 번이라도 하면, 그 실수가 다음 생성에 연쇄적으로 영향을 미쳐 전체 결과가 망가지는 '노출 편향(Exposure Bias)'이 발생한다.

온폴리시 증류(OPD)는 학생 모델이 직접 문장을 생성해보고(On-policy sampling), 그 과정에서 발생한 자신의 실수에 대해 스승 모델로부터 실시간 피드백을 받는 방식이다. 이는 모방 학습(Imitation Learning)의 DAgger 알고리즘 원리와 유사하며, 학생 모델이 겪을 수 있는 다양한 오류 상태(Out-of-distribution states)에서 어떻게 올바른 궤적으로 복귀해야 하는지를 직접 학습하게 한다.

이 과정을 통해 오류 누적 속도를 시퀀스 길이 T의 제곱(T²)에서 선형(T) 수준으로 억제할 수 있다. 결과적으로 학생 모델은 단순히 스승의 말투를 흉내 내는 수준을 넘어, 복잡한 추론 과정에서 발생하는 예외 상황에 강건하게 대응하는 능력을 갖추게 된다.

통합 OPD 목적 함수는 샘플링 정책(π_mix)과 디버전스 지표(D_f)를 결합한 형태이다. [학생 모델의 현재 정책이나 스승 모델과의 혼합 정책을 입력으로] → [시퀀스 궤적을 샘플링하고 각 토큰 위치에서 스승과 학생의 분포 차이를 계산하여] → [전체 손실 값을 얻고] → [이 값을 줄이는 방향으로 학생 모델의 파라미터를 업데이트한다].

화이트박스 OPD는 스승 모델의 전체 로짓(Logit)에 접근하여 토큰별 분포를 맞춘다. Forward KL은 스승의 모든 모드(Mode)를 덮으려 하여 환각을 유발할 수 있고, Reverse KL은 특정 모드에 집중하는 성질이 있어 추론 작업에 유리하다. [스승의 확률 p_T와 학생의 확률 p_θ를 입력으로] → [KL Divergence 연산을 수행해] → [분포 간 거리 숫자를 얻고] → [학생이 스승의 판단 경계를 정확히 학습하게 한다].

블랙박스 OPD는 로짓 대신 스승 모델의 최종 출력이나 점수(Reward)를 활용한다. GAD와 같은 방식은 [학생과 스승의 출력물을 판별기(Discriminator)에 입력으로] → [어느 것이 스승의 것인지 구분하는 연산을 수행해] → [보상 신호를 생성하고] → [학생 모델이 스승과 유사한 고품질 문장을 생성하도록 유도한다].

토큰 수준의 OPD 기법인 DistiLLM-2는 지시 이행, 코드 생성, 수학 추론 등 다양한 벤치마크에서 기존 오프폴리시 SFT 대비 월등한 성능을 보였다. 특히 수학 문제(GSM8K, MATH)에서 Reverse KL 기반의 온폴리시 학습이 스승의 추론 경로를 더 정확히 복제함을 확인했다.

DeepSeek-R1의 경우 예외적으로 오프폴리시 증류를 사용했음에도 높은 성능을 냈는데, 이는 데이터 자체가 자가 수정(Self-correction) 과정을 포함하고 있어 노출 편향 문제를 데이터 수준에서 완화했기 때문으로 분석됐다. 하지만 R1 증류 모델 위에 다시 온폴리시 RL을 적용했을 때 추가적인 성능 향상이 관찰되어 온폴리시 방식의 우월성을 입증했다.

효율성 측면에서 Fast OPD는 샘플링 구간을 앞부분으로 제한함으로써 성능 저하 없이 학습 연산량(FLOPs)을 2배에서 최대 47배까지 절감할 수 있음을 보여주었다.

OPD의 핵심은 f-디버전스 프레임워크를 통한 분포 매칭이다. GKD, MiniLLM, DistiLLM 등은 모두 이 프레임워크 내에서 f 함수의 선택과 인자 순서(Forward vs Reverse)를 달리한 변형들이다. 이론적으로 온폴리시 샘플링은 모방 학습의 성능 하한(Performance bound)을 O(εT²)에서 O(εT)로 개선한다. 이는 자기회귀 생성 모델의 고질적인 문제인 공변량 변화(Covariate shift)를 직접적으로 해결하는 수학적 근거가 된다.

최근 연구는 적응형 디버전스(Adaptive Divergence)로 진화하고 있다. ToDi나 Entropy-Aware OPD는 토큰의 위치나 스승 모델의 확신도(Entropy)에 따라 KL의 방향을 동적으로 전환하여, 정답이 명확한 추론 토큰과 창의성이 필요한 일반 토큰을 다르게 처리한다. 구현 상의 난제는 메모리 오버헤드다. 70B 스승 모델과 학생 모델을 동시에 메모리에 올리고 KV 캐시를 관리해야 하므로, FP8 양자화나 로짓 오프로딩(Logit Offloading) 같은 엔지니어링 기법이 필수적으로 동반된다.

온폴리시 증류는 오프폴리시 대비 3~8배의 연산 비용이 발생하며, 스승 모델의 로짓을 실시간으로 계산해야 하므로 막대한 GPU 메모리가 요구된다. 또한 스승 모델 자체가 잘못된 정보를 줄 경우 학생 모델이 이를 맹목적으로 학습하는 '에코 챔버' 현상이 발생할 수 있다.