질의 한 개로 온폴리시 증류 성능 대부분을 회복한 실험
온폴리시 증류를 학습 질의 한 개로 돌려도 전체 데이터가 만든 향상분의 87%가 재현됐다는 arXiv 논문의 실험 구성과 수치, 한계를 정리했습니다.
왜 지금 이 주제인가
9월 3일 arXiv에 온폴리시 증류(on-policy distillation, 이하 OPD)를 다룬 논문이 올라왔습니다. OPD는 학생 모델이 스스로 만들어 낸 출력 위에서 교사 모델이 매 토큰의 확률 분포를 알려 주는 후속 학습 방식으로, 같은 연구 계열의 4월 논문은 이 기법이 대규모 언어 모델 후속 학습의 핵심 도구가 됐다고 적었습니다. 이번 두 번째 편은 그 전제를 뒤집는 쪽을 시험했습니다. 학습에 쓰는 질의를 하나로 줄여도 전체 데이터로 돌린 OPD가 얻는 향상의 상당 부분이 재현된다는 것입니다.
핵심 내용
실험 구성
OPD에서 학생은 자기 롤아웃을 샘플링하고, 교사는 학생이 지나간 모든 접두사에서 다음 토큰의 전체 분포를 제공합니다. 학습은 토큰 단위 KL 발산을 줄이는 방향으로 진행됩니다. 연구진은 수학, 코드 생성, 지시 따르기, 에이전트 도구 사용 네 영역에서 학생과 교사 쌍을 짰습니다. 수학에서는 DeepSeek-R1-Distill-Qwen-1.5B에 JustRL-DeepSeek-1.5B를, Llama-3.2-3B-Instruct에 GT-Llama-3.2-3B-Instruct-MATH를, OLMo-3-7B-Instruct-DPO에 OLMo-3-7B-Instruct를 교사로 붙였습니다. 코드 생성은 Nemotron-1.5B, 지시 따르기는 UltraData-IF-1.5B, 도구 사용은 Hammer2.1-1.5b가 교사를 맡았습니다. 모든 실험은 롤아웃 64개를 한 배치로 묶어 AdamW(학습률 1e-6)로 300스텝을 돌린 뒤 평가했습니다.
질의 하나로 얻은 결과
여기서 질의 하나란 학습 내내 같은 문제 하나만 학생에게 던지고, 매 스텝 새로 뽑은 롤아웃에 교사의 분포를 맞추는 설정을 말합니다. 수학에서 이렇게 돌린 OPD는 300스텝에서 68.5점을 기록했습니다. 같은 조건에서 전체 데이터로 돌린 OPD는 69.8점이었습니다. 논문은 이를 교사와 학생 사이 격차의 69%, 전체 데이터 OPD가 만든 향상분의 87%를 회복한 것으로 계산했습니다. 다른 영역의 격차 회복률은 코드 생성 73%, 지시 따르기 66%, 도구 사용 64%였습니다. 네 영역 모두에서 질의 하나짜리 학습이 수백 스텝 동안 계속 좋아졌다는 점이 공통 관찰이었습니다.
왜 하나로 충분한가
연구진은 상태 커버리지라는 지표를 도입했습니다. 교사 신호의 표현을 K-평균으로 200개 군집에 나눈 뒤, 학습 중 학생이 실제로 도달한 군집의 비율을 셉니다. 질의 하나짜리 OPD는 300스텝 시점에 71.5%에 이르렀고, 의미가 서로 다른 질의 16개를 쓰면 98.9%까지 올라갔습니다. 전체 데이터 기준값은 100%입니다. 교사를 여러 개 붙이는 다중 교사 OPD에서도 영역당 질의 16개면 전체 데이터 대비 향상분의 101%를 회복했습니다. 논문은 이 상황을 OPD가 데이터는 과잉 공급됐지만 알고리즘이 굶주린 상태라고 표현합니다. 롤아웃이 넓은 감독 신호를 금세 드러내는 데 비해, 학생이 그것을 흡수하는 속도가 점점 느려진다는 진단입니다.
의미와 한계
증류용 데이터를 크게 모으는 일은 지금까지 당연한 전제였습니다. 이 결과는 병목이 데이터 수집이 아니라 학생이 교사 신호를 소화하는 속도 쪽에 있다고 말합니다. 질의 열댓 개와 충분한 스텝만 확보하면 되는 구조라면, 대규모 데이터 파이프라인이 없는 팀도 같은 기법을 시도해 볼 여지가 생깁니다.
다만 격차가 사라지는 것은 아닙니다. 학습을 1,000스텝까지 늘리면 질의 하나짜리는 68.4점에 머문 반면 전체 데이터는 72.1점으로 올라가, 오래 돌릴수록 데이터의 몫이 다시 드러났습니다. 상태 커버리지도 저자들이 밝히듯 의미 수준의 대리 지표이고, 기준이 되는 공간 자체를 전체 데이터 롤아웃으로 만들었다는 점에서 순환적인 면이 있습니다. 다중 교사 실험은 영역 세 개에 교사 하나씩으로 한정돼, 교사를 늘렸을 때의 확장성은 확인되지 않았습니다. 실험에 쓰인 모델이 1.5B에서 7B 사이라는 점도 감안해야 합니다. 더 큰 모델에서 같은 비율이 유지될지는 이 논문만으로 판단하기 어렵고, 이 부분은 추측의 영역으로 남습니다.
한 줄 정리
질의 하나로도 온폴리시 증류 향상분의 대부분이 재현된다는 결과는, 증류의 병목이 데이터 양이 아닐 수 있음을 보여 줍니다.