SAKI: Maximal-Coupling-Routed Teacher Supervision for On-Policy Distillation
On-policy 증류 시 발생하는 학생 모델의 경로 이탈 문제를 해결하기 위해, 확률적 결합(Maximal Coupling)을 이용해 정교하게 감독 신호를 배분하는 SAKI 프레임워크를 제안합니다.
논문이 다루는 내용
On-policy distillation(OPD)은 학생 모델이 생성한 궤적을 사용하여 학습함으로써 훈련-테스트 간의 상태 불일치를 줄이지만, 성능이 낮은 학생 모델은 교사 모델과 정렬되지 않은 경로를 탐색하는 문제가 있습니다. 이를 해결하기 위해 SAKI는 KL 발산 제약 하의 롤아웃과 Maximal-Coupling을 결합하여 토큰 단위의 감독 신호를 라우팅합니다. 수락된 위치에는 샘플링된 토큰에 대한 역-KL(reverse-KL) 감독을, 수정이 필요한 위치에는 교사 모델의 최고 확률 토큰에 대한 직접 감독을 적용합니다. 이 방식은 롤아웃 편차와 개입 빈도를 동일한 신뢰 영역 반경 내에서 제어할 수 있게 합니다. 실험 결과, SAKI는 수학적 추론 벤치마크에서 학생 모델의 성능을 유의미하게 향상시켰으며, 엔진 레지던트 검증기를 통해 처리량 또한 크게 개선했습니다.
핵심 결과
-
Maximal-Coupling을 활용하여 학생의 궤적 이탈 시점에 맞춰 감독 신호를 동적으로 배분하는 SAKI 알고리즘 제안
-
수락된 위치(Accepted)와 수정 위치(Correction)를 구분하여 서로 다른 형태의 최적화 신호를 제공하는 라우팅 전략
-
정확한 궤적 분포를 유지하면서도 롤아웃 처리량을 4.22배 향상시킨 엔진 레지던트 투기적 검증기 구현
실무에서 볼 만한 점
작은 모델(Student)을 더 큰 모델(Teacher)의 성능에 가깝게 학습시킬 때, 잘못된 생성 경로로 인한 학습 효율 저하를 방지하는 실전적인 기법입니다.
읽을 때 확인할 점
-
작은 규모의 LLM 증류 작업에 SAKI의 라우팅 로직(Accept/Correction 구분) 적용 실험
-
KL-divergence 제약 조건 변화에 따른 학생 모델의 수렴 안정성 테스트
-
기존 On-policy 학습 방식과 SAKI 기반 학습 방식의 추론 효율성 비교