기술 요약: LLM 증류에서 KL 과 궤적의 분리

문제 제기

지식 증류는 특히 추론 작업을 위해 강화 학습 (RL) 과 병행하여 수행되는 대규모 언어 모델 (LLM) 후속 학습의 핵심 기법이 되었습니다. 그러나 증류의 설계 공간은 여전히 잘 이해되지 않고 있습니다. 현재 파이프라인은 일반적으로 두 가지 패러다임 중 하나를 따릅니다: 오프-폴리시 증류(교사 생성 트레이스를 기반으로 한 지도 미세 조정 (SFT) 을 통한 학습) 또는 온-폴리시 증류 (OPD)(학생 생성 롤아웃에 대한 교사 피드백을 통한 학습).

이 논문은 이러한 지배적인 패러다임들이 두 가지 직교하는 설계 선택을 암묵적으로 결합하고 있다고 주장합니다:

1. 접두사 소스: 감독이 교사 (오프-폴리시) 가 생성한 접두사에 적용되는지, 아니면 학생 (온-폴리시) 이 생성한 접두사에 적용되는지 여부.
2. 토큰 수준 KL 방향: 목적 함수가 순방향 KL($KL(p_{teacher} \parallel q_{student})$) 을 최소화하는지, 아니면 역방향 KL($KL(q_{student} \parallel p_{teacher})$) 을 최소화하는지 여부.

기존 관행은 교사 접두사를 순방향 KL 과, 학생 접두사를 역방향 KL 과 결합합니다. 저자들은 이러한 결합이 토큰 수준에서 본질적인 것이 아니며, 이러한 축을 분리하면 네 가지 유효한 목적 함수를 갖는 더 넓은 설계 공간이 드러난다고 주장합니다. 각 목적 함수는 고유한 학습 체제에 해당합니다.

방법론

이론적 프레임워크: 축의 분리

저자들은 자기회귀 연쇄 법칙을 사용하여 시퀀스 수준의 KL 발산을 분해합니다. 그들은 다음과 같이 증명합니다:

• 순방향 KL($KL(p_T \parallel q_S)$) 은 기대값이 교사의 궤적 분포에서 취해지므로 교사 접두사(오프-폴리시) 와 자연스럽게 짝을 이룹니다.
• 역방향 KL($KL(q_S \parallel p_T)$) 은 기대값이 학생의 궤적 분포에서 취해지므로 학생 접두사(온-폴리시) 와 자연스럽게 짝을 이룹니다.

이들을 독립적인 축으로 취급함으로써 저자들은 네 가지 목적 함수를 정의합니다:

1. 교사 접두사 + 순방향 KL: 오프-폴리시 SFT(표준 증류) 에 해당합니다.
2. 학생 접두사 + 순방향 KL: DAgger 스타일의 온-폴리시 SFT 에 해당합니다.
3. 교사 접두사 + 역방향 KL: 오프라인 RL 스타일 증류 목적 함수에 해당합니다.
4. 학생 접두사 + 역방향 KL: 밀도 보상 온-폴리시 RL 로 간주되는 온-폴리시 증류 (OPD) 에 해당합니다.

이 논문은 다음과 같은 그래디언트 수준의 항등식을 확립합니다:

• 순방향 KL은 교사 소프트 타겟을 갖는 SFT 스타일 교차 엔트로피 그래디언트를 생성합니다.
• 역방향 KL은 교사 - 학생 로그 비율로 정의된 밀도 보상 ($r = \log p_T - \log q_S$) 을 갖는 RL 스타일 정책 그래디언트를 생성합니다.

실험 설정

저자들은 Qwen3 모델 (0.6B 학생, 4B/8B 교사) 을 사용하여 수학 추론에 대한 통제된 연구를 수행합니다. 그들은 두 가지 체제에서 네 가지 목적 함수를 평가합니다:

1. 단독 증류: 각 목적 함수로만 학습된 모델의 최종 성능을 비교합니다.
2. 증류 후 RL: 증류된 체크포인트를 그룹 상대 정책 최적화 (GRPO) 의 초기화로 사용합니다.

그들은 다양한 시퀀스 길이 (128 대 4096 토큰) 와 벤치마크 (MATH500, GSM8K, AIME24, AMC23) 에서 정확도, 예측 엔트로피, 응답 길이를 추적합니다.

주요 기여

1. 세 가지 트레이드오프의 식별

실증적 분석은 증류 설계를 지배하는 세 가지 근본적인 트레이드오프를 드러냅니다:

• 정확도 - 엔트로피 트레이드오프 (KL 방향): 역방향 KL 은 일관되게 더 높은 단독 정확도 (Avg@k) 를 달성하지만, 예측 엔트로피를 붕괴시켜 다양성을 감소시키고 (낮은 Pass@k), 후속 RL 에 대한 모델의 안정성을 떨어뜨립니다. 순방향 KL 은 엔트로피를 보존하여 더 안정적인 RL 개선을 지원하지만, 초기 단독 정확도는 더 낮습니다.
• 품질 - 계산 트레이드오프 (접두사 소스): 학생 접두사 (온-폴리시) 는 학생이 실제로 방문하는 상태를 감독함으로써 매칭된 학습 단계에서 더 높은 정확도를 제공합니다. 그러나 교사 접두사 (오프-폴리시) 는 캐시된 교사 로짓과 오프라인 궤적을 재사용하고 온라인 학생 생성 비용을 피할 수 있으므로 매칭된 FLOP 하에서 더 계산 효율적입니다.
• 정확도 - 안정성 트레이드오프 (학습 길이): 긴 시퀀스 증류 (4096 토큰) 는 추론 정확도를 향상시키지만, 특히 역방향 KL 하에서 심각한 엔트로피 붕괴와 "길이 과잉"(생성 한계에 도달하는 응답) 을 초래합니다. 짧은 시퀀스 학습은 더 안정적이지만 정확도는 낮습니다.

2. 트레이드오프 균형을 위한 제안된 방법

이러한 발견에 영감을 받아 저자들은 두 가지 실용적인 방법을 제안합니다:

• KL 혼합: 각 접두사에서 순방향 및 역방향 KL 의 가중치 조합입니다. 저자들은 긴 시퀀스 증류에서 순방향 우세 혼합이 중요하다고 발견합니다. 이는 엔트로피 붕괴와 길이 과잉을 방지하면서 역방향 KL 의 정확도 이점을 유지하여 학습 역학을 안정화합니다.
• 엔트로피 게이트 길이 커리큘럼: 예측 엔트로피가 안정성 임계값 위에 있는 동안만 시퀀스 길이를 증가시키는 짧은 지평선으로 시작하는 학습 일정입니다. 이 접근법은 고정된 긴 지평선 학습과 관련된 엔트로피 붕괴를 방지합니다.

결과

단독 증류

• 역방향 KL은 평균 +2.45 점의 이득으로 일반적으로 단독 정확도 (Avg@k) 에서 순방향 KL 을 능가합니다.
• 학생 접두사는 특히 긴 시퀀스 설정에서 (+3.55 Avg@k 이득) 교사 접두사보다 정확도에서 더 우수합니다.
• 엔트로피 붕괴: 긴 시퀀스 설정에서 순수 역방향 KL 은 엔트로피를 0 에 가깝게 밀어붙이고 응답 길이를 최대 한도 (예: 8192 토큰) 로 팽창시켜 Pass@k 를 저하시킵니다.

증류 후 RL (GRPO)

• 초기화 안정성: 순방향 KL로 초기화된 모델 (단독 정확도가 낮더라도) 은 더 높은 엔트로피를 유지하기 때문에 GRPO 동안 더 신뢰할 수 있는 개선을 보입니다.
• 역방향 KL 함정: 역방향 KL 로 초기화된 모델은 종종 더 높은 정확도로 시작하지만 낮은 엔트로피와 감소된 탐색 능력으로 인해 RL 동안 성능이 저하되거나 정체됩니다.
• 커리큘럼 성공: 엔트로피 게이트 길이 커리큘럼은 성능을 크게 향상시킵니다. 고정된 4096 토큰 학습과 비교하여 Avg@k 를 3.6 점, Pass@k 를 최대 5.8 점 향상시켰으며, 평균 응답 길이는 약 3 배 감소시켰습니다 (예: ~8185 에서 ~2688 토큰으로).

KL 혼합 결과

긴 시퀀스 설정에서 높은 순방향 KL 가중치 (예: 0.8 순방향 / 0.2 역방향) 를 가진 혼합은 엔트로피를 안정화하고 길이 과잉을 방지하면서 역방향 KL 의 정확도 대부분을 유지했습니다.

중요성과 주장

이 논문은 SFT, DAgger, 오프라인 RL, OPD 와 같은 다양한 학습 레시피를 두 가지 직교하는 설계 선택의 특정 인스턴스로 연결하는 통합 이론적 프레임워크를 제공한다고 주장합니다. 접두사 소스와 KL 방향을 분리함으로써 저자들은 실무자들이 경직된 패러다임을 넘어 다음을 명시적으로 균형 잡는 추론 증류 목적 함수를 설계할 수 있다고 주장합니다:

• 정확도(역방향 KL 및 긴 시퀀스를 통해),
• 다양성/탐색(순방향 KL 및 엔트로피 보존을 통해),
• 계산 효율성(교사 접두사/캐싱을 통해),
• 하류 RL 학습 가능성(안정적인 엔트로피를 통해).

제안된 방법 (KL 혼합 및 엔트로피 게이트 커리큘럼) 은 식별된 트레이드오프를 탐색하기 위한 실용적인 도구로 제시되며, 정확하고 추가 정책 최적화를 위해 안정적인 더 작은 추론 모델을 학습할 수 있게 합니다. 저자들은 "최고의"단독 목적 함수가 반드시 RL 파이프라인을 위한 최고의 초기화가 아니라고 강조하며, 후속 학습 과정에 대한 전체적인 관점의 필요성을 강조합니다.