기술 요약: 대규모 언어 모델 증류에서 온도의 역할에 대한 재고

문제 정의

지식 증류(Knowledge Distillation, KD)는 교사(Teacher) 모델로부터 학생(Student) 모델로 지식을 전이함으로써 대규모 언어 모델(LLM)을 압축하는 표준적인 패러다임이다. 고전적 KD의 핵심 구성 요소는 온도 파라미터($\tau$)로, 이는 교사의 확률 분포를 부드럽게 만들어 "다크 노리지(dark knowledge)"(최상위 예측 이외의 유익한 신호)를 드러내는 역할을 한다. 그러나 최근의 LLM 증류 문헌에서는 온도를 생략하거나 $\tau=1$로 고정하여 사용하는 경우가 빈번하다.

이러한 관행은 역방향 쿨백-라이블러(Reverse Kullback–Leibler, RKL) 발산이 LLM 증류에서 본질적으로 순방향 쿨백-라이블러(Forward KL, FKL)보다 우월하다는 지배적인 경험적 결론을 낳았다. 저자들은 이러한 결론이 교사의 분포가 충분히 부드럽게 만들어지지 않은 설정($\tau=1$) 하에서 이러한 목적 함수들을 평가한 데서 기인한 인위적인 결과일 가능성이 높다고 주장한다. $\tau=1$일 때 교사의 소프트 타겟(soft targets)은 매우 날카로우며, 이는 FKL이 의존하는 비타겟(non-target) 확률 질량을 억제하여 FKL을 열등하게 보이게 만들 수 있다. 즉, 이는 FKL의 내재적 한계 때문이 아니라 평가 설정 때문일 수 있다. 본 논문은 FKL이 체계적으로 과소평가되었는지, 그리고 온도가 FKL과 RKL의 비교를 어떻게 근본적으로 변화시키는지 규명하고자 한다.

방법론

저자들은 KD 온도와 KL 기반 목적 함수 간의 상호작용을 조사하기 위해 이론적 분석과 경험적 평가를 결합하여 사용한다.

이론적 분석

본 논문은 학생 로짓(student logits)에 대한 FKL($L_{FKL} = D_{KL}(p_\tau \| q_\tau)$)과 RKL($L_{RKL} = D_{KL}(q_\tau \| p_\tau)$)의 그래디언트를 분석한다.

• 고온 영역(High-Temperature Regime): 저자들은 $\tau \to \infty$로 갈 때, FKL과 RKL의 그래디언트가 모두 동일한 로짓 매칭 동작($\nabla L \propto \frac{1}{V\tau^2}(z_s - z_t)$)으로 수렴함을 증명한다. 즉, 두 목적 함수는 동등해진다.
• 저온 영역($\tau \approx 1$): 두 목적 함수는 크게 갈라진다.
  • FKL: 그래디언트는 부드러워진 교사 분포($p_\tau$)와 학생 분포($q_\tau$) 사이의 차이에 직접적으로 의존한다. $\tau$를 높이는 것은 수퍼비전 타겟(supervision target) 자체를 변화시켜, 확률 질량을 비지배적 토큰으로 재분배하고 FKL을 위한 신호를 풍부하게 만든다.
  • RKL: 그래디언트는 학생 분포($q_\tau$)에 의해 가중치가 부여된 교사-학생 간의 로짓 격차에 의해 구동된다. 여기서 온도는 FKL과 달리 타겟 분포를 근본적으로 바꾸기보다는, 주로 그래디언트를 *재조정(rescale)*하고 학생 측의 가중치를 부드럽게 만드는 역할을 한다.

경험적 평가

GPT-2 및 OPT 모델 제품군(예: GPT-2 XL $\to$ GPT-2 Base, OPT 6.7B $\to$ OPT 1.3B)을 사용하여 명령 수행 벤치마크(Dolly Eval, Self-Instruct, Vicuna Eval, Super-NI, UnNI)에서 실험을 수행하였다.

• 설정: 학생이 고정된 교사 생성 응답을 학습하는 오프-폴리시(off-policy) 증류 방식을 채택했다.
• 변수: 연구는 $\tau=1$에서의 표준 목적 함수(FKL, RKL, Sym-KL, JS, SFKL, SRKL, AKL, AB, DRKL)와 온도 튜닝된 변형 모델($\tau \in \{1.5, \dots, 4.0\}$에 대한 그리드 탐색)을 비교한다.
• 지표: 다섯 개의 무작위 시드에 대해 평균화된 ROUGE-L 점수를 사용한다.

주요 기여

1. 온도의 비대칭적 효과

본 논문은 온도가 FKL과 RKL에 비대칭적으로 영향을 미친다는 점을 입증한다.

• FKL: 온도는 교사의 비지배적 토큰 정보를 노출함으로써 학습 신호를 실질적으로 풍부하게 한다. 즉, 온도는 FKL이 무엇을 배우는지를 변화시킨다.
• RKL: 온도는 확률 타겟으로서의 부드러운 교사 분포에 직접 접근하는 대신, 주로 그래디언트를 재조정하고 가중치를 부드럽게 만든다. 즉, 온도는 RKL이 어떻게 배우는지를 변화시킨다.

2. 성능 순위의 역전

저자들은 RKL의 우위가 온도에 매우 의존적임을 보여준다.

• $\tau=1$일 때, RKL은 기존 문헌과 일치하게 FKL보다 일관되게 우수한 성능을 보인다.
• 높은 온도(통상적으로 $\tau \in [2.5, 3.5]$)에서는 FKL이 모든 모델 규모와 벤치마크에서 RKL을 일관되게 앞선다. 예를 들어, GPT-2 Base의 경우 $\tau=3.0$에서 FKL은 ROUGE-L이 19.41에서 20.97로 향상된 반면, RKL은 거의 정체되어 있었다.

3. KL 기반 목적 함수의 광범위한 개선

온도는 단순히 FKL을 위한 보정 도구가 아니라, 적절한 온도 튜닝을 통해 Sym-KL, JS, SFKL, AKL을 포함한 광범위한 KL 기반 목적 함수 군의 성능을 개선하는 일반적인 요인이다.

• 적절한 온도 튜닝을 거친 단순 KL 기반 방식들은 최근의 SOTA 방식인 AB-KD 및 DRKL과 경쟁할 만한 성능을 달 achieves 한다.
• 예시로, GPT-2 Base에서 $JS+\tau$는 AB(23.91)와 DRKL(25.16)을 능가하는 평균 ROUGE-L 25.63을 달성했다.

결과

• 성능 향상: 온도 튜닝은 FKL 기반 목적 함수(예: GPT-2 Base에서 FKL +5.83, AKL +6.62)에 상당한 이득을 가져왔다. RKL 기반 목적 함수의 이득은 미미했다(예: RKL +0.49).
• 질적 개선: Unnatural Instructions 벤치마크에 대한 사례 연구에 따르면, 온도 튜닝된 모델은 $\tau=1$인 모델에 비해 출력 형식 및 엔티티 추출에 관한 지시사항을 더 잘 따랐으며, 환각(hallucination)과 누락을 줄였다.
• 순위 변화: 온도의 도입은 증류 목적 함수의 상대적 순위를 변화시켰으며, 이는 $\tau=1$에서 수행된 이전의 비교들이 특정 목적 함수(RKL 등)의 내재적 우월성에 대한 편향된 결론을 낼 수 있음을 시사한다.

의의 및 주장

본 논문은 KD 온도가 단순한 구현 세부 사항이 아니라, LLM 증류에서 공정한 비교를 위해 반드시 고려되어야 하는 핵심적인 목적 함수 의존적 요소라고 주장한다.

• FKL의 재평가: 저자들은 온도의 누락으로 인해 LLM 증류 커뮤니티에서 FKL이 체계적으로 과소평가되었다고 주장한다. 적절히 부드럽게 만들어질 때, FKL은 RKL보다 뛰어난 성능을 낼 수 있다.
• 단순성 vs 복잡성: 연구 결과는 적절한 온도 스케일링과 결합된 단순한 KL 기반 목적 함수가 최근의 복잡한 SOTA 증류 프레임워크와 경쟁할 수 있음을 시사하며, 이는 향후 연구를 위한 설계 공간을 단순화할 수 있다.
• 공정한 벤치마킹: 본 논문은 향후 증류 목적 함수를 비교할 때 RKL과 같은 특정 목적 함수의 내재적 우월성에 대한 편향된 결론을 피하기 위해 반드시 온도 튜닝을 포함해야 한다고 강조한다.

저자들은 실험이 최대 7B 파라미터의 교사 모델로 제한되었고, 온-폴리시 샘플링이나 더 강력한 학생 초기화를 제외했다는 점을 한계로 인정하며, 이를 향후 연구 과제로 남겨두었다.