강화학습 모델을 학생 모델에 가르치는 새로운 방식, 표현층에서 직접 학습
원제 The Teacher Is a Direction, Not a Destination: Extrapolating RL-Induced Representation Residuals in On-Policy Distillation
추천 273댓글 2
Key Point
모델 증류에서 강화학습의 효과를 더 효율적으로 전달하는 새로운 방법이 제시되어, 더 작은 모델로도 큰 모델 수준의 성능을 얻을 수 있는 경로를 열어준다.
핵심 요약
- 기존 온-폴리시 디스틸레이션(OPD)은 학생 모델이 자신의 궤적 위에서 교사의 다음 토큰 분포를 모방하도록 훈련하며, 출력 공간에서 암묵적 보상을 외삽하여 교사를 능가하게 할 수 있다.
- 그러나 언어모델 헤드는 이 변화를 비등방적으로 감쇠시키며, 교사 은닉층에 인코딩된 대부분의 변화가 로짓에는 작은 가중치로만 도달하고, 출력 공간 외삽이 의존하는 샘플 토큰 로그확률 비율은 훈련 불안정성을 초래하는 노이즈를 증폭시킨다.
- 연구진은 강화학습이 모델의 내부 표현을 기본 체크포인트에서 변화시키고, 이 변화의 방향을 모든 층에서 측정할 수 있음을 관찰했다.
- 이를 바탕으로 RIDE(RL-Induced Direction Extrapolation)를 제안하는데, 이는 표현 공간에서 RL 유도 변화를 직접 외삽하며, 모든 층과 토큰 위치에서 교사와 RL 이전 체크포인트 간 잔차를 계산하고 학생 은닉층을 이 잔차를 따라 교사를 넘어 변위된 목표로 회귀시킨다.
- 샘플 궤적에 조건부인 이 회귀는 교사를 중심으로 한 이차 페널티 아래 잔차로 정의된 선형 방향 보상을 최대화하는 것과 동치이며, 학생이 RL 유도 방향을 따라 이동하면서 교사로부터의 편차를 제한하는 방식을 명확히 한다.
- 네 가지 기본/RL 교사 쌍(서로 다른 규모, 아키텍처, 사전학습 계보)에서 RIDE는 모든 쌍에서 RL 훈련 교사에 접근하거나 이를 초과하며, 유일하게 평균적으로 그렇게 하는 방법이다.
- RIDE는 출력 공간 외삽을 일관되게 능가하며, 특히 교사가 기본 모델에 가까울 때 출력 공간 외삽이 학생 성능을 저하시키는 문제를 해결한다.