LLM 증류 방식 개선, 확률값 대신 토큰 순서만 학습

핵심 요약

• 기존 온폴리시 증류(OPD)는 학생 모델을 교사 모델의 토큰 확률값에 맞춰 학습시킨다.
• 이 확률값에는 교사가 선호하는 토큰이 무엇인지와 얼마나 확신하는지 두 가지 정보가 섞여있는데, 후자는 교사의 용량과 보정을 반영하므로 작은 학생 모델이 재현하기 어렵다.
• 교사 온도 변화가 확률값을 바꾸지만 순서는 유지하므로 OPD는 온도 선택에 민감하다.
• 연구팀은 교사가 후보 토큰들을 순서 매긴 것만으로도 충분하다는 가정 아래 O³PD(순서 전용 온폴리시 증류)를 제안했다.
• O³PD는 학생이 생성한 각 접두사마다 교사의 상위-K 토큰의 순서를 재현하도록 학습하되, 확률값에는 맞추지 않는다.
• 이 방식은 교사 온도에 불변이므로 온도 튜닝이 필요 없다.
• 여섯 교사-학생 쌍을 대상으로 평가한 결과 O³PD는 HMMT25, AIME24, AIME25에서 모든 쌍에서 OPD를 앞섰으며, 최대 2.9점, AIME25에서는 5.6점까지 향상됐다.
• 다른 순위 손실함수도 유사한 성과를 내므로 이득은 순서 감독 자체에서 비롯된다.
• 분석 결과 O³PD는 OPD보다 교사 확률값에서 더 멀어지지만, 교사의 순위를 더 충실히 재현하고 더 강한 학생 모델을 만든다.
• 따라서 온폴리시 증류가 전달해야 할 지식은 확률값이 아니라 토큰 순위에 담겨있다.