LSPD가 OPD를 KL 정규화 강화학습으로 재해석해 롤아웃 효율을 높인다
An RL View of OPD: Least Square Policy Distillation for Sample-Efficient LLM Reasoning
무엇인가
온폴리시 증류(OPD)는 학생 모델이 직접 생성한 궤적 위에서 교사 모델이 토큰 단위 감독을 주는 방식으로, 교사가 학생이 실제로 마주친 프리픽스에서 지도를 할 수 있다는 장점이 있다. 문제는 구현이다. 최근 OPD는 PPO나 GRPO 같은 정책 기반 강화학습 프레임워크로 구현되는데, 이들은 본질적으로 온폴리시라서 학습 내내 새 롤아웃을 뽑아야 하고 과거 데이터 재사용과 탐색 지원이 제한적이다. 이 논문은 OPD를 강화학습 관점에서 다시 읽으면 이 두 가지 한계를 값 기반(value-based) KL 정규화 RL의 도구로 풀 수 있다고 주장한다.
어떻게 동작하나
출발점은 정확한 재해석이다. 저자들은 고정된 참조 정책 π_ref에 대해 보상 R(x_<t, y_t) = log(π_E(y_t|x_<t) / π_ref(y_t|x_<t))를 정의하면, OPD의 토큰 단위 reverse-KL 목적함수가 η=1인 KL 정규화 정책 최적화 문제와 정확히 같아진다는 것을 보인다(식 4.1). 추가 정규화 항을 붙인 근사가 아니라 등가 변환이라는 점이 핵심이다. 이 관점에서 보상은 학생에 대해 고정되어 있고, KL 페널티는 각 프리픽스에서 참조 정책으로부터의 이탈을 억제한다.
무엇과 다른가
제안 방법 LSPD(Least Square Policy Distillation)는 여기에 값 기반 RL의 낙관적 탐색을 이식한다. k번째 반복까지 모은 프리픽스-토큰 쌍으로 보상을 최소제곱 추정하고(식 4.2), 누적 데이터와 통계적으로 일관된 보상 함수들의 신뢰집합 C_k를 만든 뒤(식 4.3), 그 안에서 각 토큰에 대해 가장 유리한 낙관적 보상 r_k^+를 취한다. 이 낙관적 보상은 π_{k+1}(y_t|x_<t) ∝ π_ref(y_t|x_<t)·exp(η·r_k^+(x_<t, y_t)) 형태의 닫힌 해 토큰 업데이트로 이어진다. 이 점별 최적화를 라그랑주 완화로 풀면 결국 학생과 교사의 로그 확률을 제곱으로 맞추는 목적함수가 나온다(식 4.5). 실전에서는 여기에 Huber형 로버스트 페널티와 엔트로피 보너스를 결합한 식 4.6을 쓰는데, 이 목적함수는 리플레이 버퍼 B 위에서 정의되므로 자연스럽게 오프폴리시 최적화를 지원한다. 최신 롤아웃으로 한 번 또는 여러 번 업데이트하는 LSPD와, 과거 정책이 만든 응답까지 재사용하는 완전 오프폴리시 변형 LSPD-RB로 나뉜다.
어떻게 쓰나
실험은 Qwen3 계열 세 가지 교사-학생 조합(Qwen3-8B→Qwen3-4B-Base, Qwen3-4B→Qwen3-1.7B-Base, Qwen3-1.7B→Qwen3-0.6B-Base)에서 수행했고, 학습 데이터는 DAPO-Math-17K, 평가는 MATH-500, Minerva, Olympiad-Bench, AMC23, AIME24, AIME25 여섯 개 벤치마크다. 베이스라인은 KD, OPD, EOPD다. 18개 모델-벤치마크 조합 평균에서 LSPD는 Avg@16 31.60, Pass@16 53.30으로 가장 강한 베이스라인 EOPD를 +0.91/+0.85점, 표준 OPD를 +1.99/+2.27점 앞섰다. LSPD-RB는 Avg@16 31.51로 비슷하면서 Pass@16을 54.86까지 올려 EOPD 대비 +0.82/+2.42점, OPD 대비 +1.89/+3.84점을 기록했다. LSPD는 18개 설정 중 Avg@16에서 11개, Pass@16에서 9개 최고 또는 공동 최고를 냈고, LSPD와 LSPD-RB를 합치면 36개 지표-설정 조합 전부에서 둘 중 하나가 상위 2위 안에 든다.
전제와 한계
데이터 효율이 이 논문의 실질적 승부처다. 롤아웃 배치당 업데이트 횟수를 N ∈ {1,4,16,64}로 늘리면 포화 성능에 도달하는 데 필요한 배치 수가 줄어들었고, AIME24·AIME25에서는 N≥4일 때 약 30 스텝에서 포화한 반면 N=1은 50 스텝 이상이 필요했다. 다만 N=4를 넘으면 이득이 급격히 줄어든다. 완전 오프폴리시인 LSPD-RB는 약 10개 롤아웃 배치에서 포화 성능에 도달했는데, 이는 배치당 1회 업데이트하는 LSPD의 40개 이상과 비교된다. 결론부 표현대로 LSPD-RB는 전체 롤아웃 배치의 처음 25%만으로 기존 OPD에 필적하는 성능을 낸다. 다양성 측면에서는 Pass@k를 k=64까지 측정했을 때 LSPD가 샘플 예산이 커질수록 더 강해져 AMC23·AIME24·AIME25 세 벤치마크 모두에서 Pass@32와 Pass@64 최고 또는 공동 최고를 기록했고, 세 벤치마크 평균 Pass@32/64는 58.51/61.94로 EOPD의 57.00/60.43을 +1.51/+1.52점 앞섰다. 학습 중 학생 정책 엔트로피도 LSPD와 EOPD가 약 0.5를 유지해 표준 OPD보다 높았다. 엔트로피 정규화를 제거하는 절제 실험에서는 Avg@16은 평균 0.46점만 떨어졌지만 Pass@16은 1.95점 떨어졌고, 12개 비교 전부에서 Pass@16이 하락했다.
이론적으로 저자들은 이상화된 낙관적 정식화에 대해 온라인 탐색 하에서 Õ(log K) 리그렛을 보인다(정리 6.2). 이는 혼합 정책이 최적 정책에 대해 reverse-KL ε 이내가 되는 데 Õ(ε^{-1}) 샘플 복잡도를 함의한다. 흥미로운 비교는 forward-KL 기반 SFT다. 저자들은 MLE 기반 모방이 전문가 결정성에 의존해 빠른 속도를 얻는 반면, OPD·LSPD 같은 reverse-KL 최소화의 빠른 속도는 학생과 교사 사이의 커버리지(가정 6.1, 로그 비율 보상의 유계성)에 의존한다고 정리한다. 이는 실무에서 SFT로 초기 정책을 세운 뒤 reverse-KL 기반 정제를 하는 순서가 합리적이라는 설명으로 이어진다.
개발자 입장에서 이 논문의 실용적 요점은 목적함수 교체만으로 롤아웃 비용을 크게 줄일 수 있다는 것이다. 기존 OPD 파이프라인에서 reverse-KL 정책 그래디언트를 로그 확률 제곱 매칭(Huber형) + 엔트로피 보너스로 바꾸고 리플레이 버퍼를 붙이면, 같은 성능에 도달하는 데 필요한 생성량을 대략 4분의 1 수준으로 줄일 수 있다는 관측이다. 다만 이는 Qwen3 계열과 수학 추론 데이터(DAPO-Math-17K)에서 얻은 결과이므로, 다른 도메인이나 교사-학생 격차가 큰 조합에서는 커버리지 가정이 깨질 수 있다는 점을 먼저 확인해야 한다. 코드는 공개되어 있다.
한계도 저자들이 명시한다. 이론 분석은 알고리즘 2로 요약된 이상화된 정식화에 대한 것이고, 실제 구현(식 4.6)은 로버스트 페널티와 엔트로피 항을 쓰는 대리 목적함수라 이론과의 간극이 있다. 또한 노이즈가 있는 교사 피드백(가정 C.1)과 학생-교사 간 커버리지, 로그 비율 보상의 유계성을 전제하며, reverse-KL 계열의 빠른 수렴 속도가 이 커버리지에 의존한다는 점을 인정한다. 실험적으로도 롤아웃 배치당 업데이트를 N=4 이상으로 늘리면 샘플 효율 이득이 포화되고, 엔트로피 정규화 제거 시 평균 정확도 손실은 작지만 다양성 지표 손실은 일관되게 크다는 점이 남는다.