기술 요약: ARKD – 적응형 강화 학습 가이드 기반 양방향 KL 발산 증류 (Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation)

1. 문제 정의

지식 증류(Knowledge Distillation, KD)는 대규모 언어 모델(LLM)의 성능을 유지하면서 추론 비용을 줄이기 위해 모델을 압축하는 주요 기술입니다. 그러나 단일 쿨백-라이블러(Kullback-Leibler, KL) 발산 목적 함수에 의존하는 전통적인 KD 방식은 개방형 텍스트 생성에서 근본적인 트레이드오프에 직면합니다:

• 순방향 KL (Forward KL, FKL): "모드 커버링(mode-covering)" 동작을 장려하며, 학생 모델이 교사 분포의 모든 모드를 커버하도록 유도합니다. 이는 종종 신뢰도가 낮은 롱테일(long-tail) 영역에서 확률을 과대평가하게 만들어 생성 품질을 저하시킵니다.
• 역방향 KL (Reverse KL, RKL): "모드 시킹(mode-seeking)" 동작을 장려하며, 지배적인 모드에 집중하고 확률이 낮은 영역은 무시합니다. 이는 롱테일 과대평가는 방지하지만, 다양성과 커버리지를 희생시켜 모드 붕괴(mode collapse)를 초래할 수 있습니다.

기존의 정적 보간법(예: 고정된 0.5:0.5 가중치)은 훈련이 진행됨에 따라 변화하는 학생-교사 간의 격차를 포착하지 못합니다. 정적 계수는 분포 정렬(distribution alignment)을 위한 변화하는 요구사항, 즉 훈련 단계마다 달라지는 FKL과 RKL의 균형을 맞추는 데 적합하지 않습니다.

2. 방법론: ARKD

저자들은 강화 학습(RL) 에이전트를 사용하여 FKL과 RKL에 대한 가중치를 동적으로 할당하는 프레임워크인 ARKD(Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation)를 제안합니다.

이론적 토대

논문은 먼저 FKL과 RKL의 이론적 일관성을 확립합니다. 분석 결과, 두 목적 함수 모두 그래디언트가 0일 때 동일한 해($q_\theta = p$)로 수렴함을 보여줍니다. 이는 두 발산 모두 분포 매칭을 위한 유효한 방법임을 확인하며, 이들의 결합 사용이 정당함을 입증합니다. 즉, 두 방식은 최적화 경로 상에서 서로 보완적인 동작을 보입니다.

RL 기반 프레임워크

ARKD는 가중치 할당 문제를 순차적 의사결정 과정(Markov Decision Process)으로 재정의합니다:

• 상태 (State, $s$): 교사/학생 엔트로피($H_T, H_S$), 분산($\sigma^2_T, \sigma^2_S$), 그리고 현재 FKL/RKL 손실 값($L_{FKL}, L_{RKL}$)을 포함하여 현재 훈련 역학을 포착하는 6차원 벡터입니다.
• 행동 (Action, $\alpha$): 정책 네트워크($\pi_\phi$)는 FKL 항에 가중치를 두는 $\alpha$와 RKL 항에 가중치를 두는 $(1-\alpha)$를 결정하는 스칼라 가중치 $\alpha \in (0, 1)$를 출력합니다.
• 보상 (Reward, $r$): 현재 배치의 음의 증류 손실($r = -L_{distil}$)로 정의됩니다. 목표는 훈련 궤적 전체에 걸쳐 누적 보상을 최대화(총 손실 최소화)하는 것입니다.
• 최적화:
  • 학생 업데이트: 가중치가 적용된 증류 손실과 언어 모델링 손실을 최소화합니다.
  • 정책 업데이트: 훈련을 안정화하고 경계 포화(boundary saturation)를 방지하기 위해 지수 이동 평균(EMA) 베이스라인과 엔트로피 정규화를 사용하는 REINFORCE 알고리즘을 사용합니다.

훈련 역학

정책 네트워크는 다음과 같이 트레이드오프를 적응적으로 탐색하도록 학습됩니다:

1. 탐색 단계 (Exploration Phase): 훈련 초기에는 정책이 더 넓은 범위의 $\alpha$(예: 0.3–0.6)를 탐색하여, 광범위한 커버리지를 확보하고 모드 붕p를 방지하기 위해 FKL을 강조합니다.
2. 수렴 단계 (Convergence Phase): 학생 분포가 교사 분포에 가까워짐에 따라, 정책은 롱테일 과대평가를 피하고 정밀한 정렬에 집중하기 위해 더 낮은 $\alpha$ 값(RKL 우세)으로 전환합니다.
3. 안정화 단계 (Stability Phase): 정책은 최종 미세 조정을 위해 특정 $\alpha$(예: $\approx 0.18$)에서 안정화됩니다.

3. 주요 기여

1. 이론적 분석: 본 논문은 FKL과 RKL의 상보성 및 수렴 일관성을 확립하는 이론적 분석을 제공하여 적응형 증류의 기초를 마련했습니다.
2. ARKD 프레임워크: 정적 휴리스틱을 넘어, 상태 기반 특성에 따라 FKL과 RKL의 가중치를 동적으로 조절하는 새로운 RL 기반 공동 최적화 프레임워크를 도입했습니다.
3. 경험적 검증: 광범적인 실험을 통해 ARKD가 다양한 벤치마크와 모델 규모에 걸쳐 생성 품질, 분포 적합도 및 일반화 성능을 일관되게 향상시킴을 입증했습니다.

4. 실험 결과

저자들은 GPT-2 (1.5B $\to$ 120M/340M) 및 LLaMA (13B $\to$ 7B) 구성을 사용하여 인스트럭션 튜닝 작업에 대해 ARKD를 평가했습니다.

• 내부 분포 (In-Distribution, ID) 성능: DollyEval 벤치마크에서 ARKD는 SFT, SeqKD, FKL, RKL, FKL+RKL, 그리고 MiniLLM과 같은 베이스라인들과 비교하여 가장 높은 BertScore와 경쟁력 있거나 우수한 Rouge-L 점수를 달성했습니다. 예를 들어, GPT-2 1.5B $\to$ 120M 설정에서 ARKD는 Rouge-L 24.5, BertScore 0.815를 기록하며 강력한 베이스라인인 MiniLLM을 능가했습니다.
• 외부 분포 (Out-of-Distribution, OOD) 일반화: ARKD는 미학습 데이터셋(SelfInst, VicunaEval, S-NI, UnNI)에서 탁월한 일반화 능력을 보여주었으며, 두 모델 제품군 모두에서 모든 베이스라인을 일관되게 앞질렀습니다.
• 절제 연구 (Ablation Studies):
  • 정적/선형 스케줄링 대비: ARKD는 정적 가중치(고정 $\alpha$) 및 선형 스케줄링 전략(시간 기반 감소/증가)보다 유의미하게 우수한 성능을 보였는데, 이러한 방식들은 복잡한 정렬 역학을 포착하지 못했습니다.
  • 그리디 휴리스틱 대비: ARKD는 각 단계에서 손실이 더 낮은 발산을 선택하는 그리디 휴리스틱보다 0.4~0.6 포인트 더 높은 성능을 보였습니다. 그리디 방식은 근시안적이어서 종종 순수 RKL로 퇴화하는 경향이 있는 반면, ARKD는 지속적인 혼합을 통해 유익한 작은 FKL 성분을 유지했습니다.
  • 궤적 분석: 학습된 $\alpha$ 궤적은 지능적인 전략을 확인시켜 주었습니다: 커버리지를 위해 초기에 FKL을 강조한 후 정밀도를 위해 RKL로 전환하는 패턴을 보였으며, 이는 수동 설계로는 달성할 수 없는 패턴입니다.

5. 의의 및 주장

본 논문은 ARKD가 정적 KL 기반 방식의 한계를 해결함으로써 적응형 지식 증류를 발전시킨다고 주장합니다. 강화 학습을 활용하여 모드 커버링과 모드 시킹 동작의 균형을 동적으로 조절함으로써, ARKD는 생성 품질과 일반화 능력을 모두 향상시키는 더 나은 분포 정렬을 달상합니다.

저자들은 ARKD를 고정된 휴리스틱보다 동적이고 상태 인식적인 가중치 부여가 필요함을 입증하는 원칙적인 접근 방식으로 규정합니다. 현재의 정책이 수작업으로 설계된 상태 벡터와 즉각적인 배치 보상에 의존하고 있지만, 이 프레임워크는 더 넓은 $f$-발산 제품군 및 온폴리시(on-policy) 패러다임과 호환 가능합니다. 이 연구는 적응형 RL 가이드 최적화가 다양한 작업과 모델 규모에서 정적 및 그리디 베이스라인을 일관되게 능가할 수 있음을 보여줌으로써 효율적인 LLM 압축을 위한 견고한 토대를 제공합니다.