기술 요약: 선택적 지식 증류(Selective Knowledge Distillation)의 재고

문제 정의

대규모 언어 모델(LLM)은 최첨단 성능을 달-성하지만, 서빙 비용이 많이 들고 적응시키기 어렵다. 지식 증류(Knowledge Distillation, KD)는 더 작은 학생(student) 모델이 더 큰 교사(teacher) 모델을 모방하도록 학습함으로써 이를 해결한다. 그러나 표준적인 KD는 훈련 시퀀스의 모든 토큰 위치에 대해 조밀하고 균일한 감독(supervision)을 적용하며, 이는 종종 최적의 방법이 아니다. 최근 연구들은 특정 토큰 위치, 어휘 클래스 또는 훈련 샘플을 선택하는 것이 성능을 향상시킬 수 있음을 시사하지만, 어떤 중요도 신호와 선택 정책이 가장 효과적인지, 그리고 이 축들이 어떻게 상호작용하는지는 여전히 불분명하다. 특히, 어떤 신호가 가치 높은 토큰을 신뢰성 있게 식별하는지, 그리고 서로 다른 선택 정책이 효과적인 증류 커리큘럼을 어떻게 형성하는지에 대한 체계적인 비교가 부족하다.

방법론

저자들은 선택적 지식 증류(Selective Knowledge Distillation, SKD)를 위한 통합 프레임워크를 제안하며, 증류 과정을 정렬 기준(alignment criterion), 위치(positions), 클래스(classes), 샘플(samples), 특징(features)이라는 다섯 가지 직교 설계 축으로 분리한다. 본 연구는 이 중 처음 세 가지 축에 집중한다.

1. 선택 축의 체계적 분석

본 논문은 다음을 체계적으로 평가한다:

• 위치 중요도 신호(Position Importance Signals): 학생의 교차 엔트로피(student cross-entropy), 교사의 불확실성(예: $1 - p_t(y_t)$), 교사-학생 불일치(KL divergence), 학생 엔트로피를 포함한 다양한 메트릭을 비교한다.
• 선택 정책(Selection Policies): top-$k\%$ 선택, 전역 임계값 설정(GLS), 커리큘럼 학습, 확률적 할당(RS-KD)과 같은 방법들을 평가한다.
• 다중 축 상호작용(Multi-Axis Interaction): 위치 선택이 클래스(어휘) 및 샘플 축을 따른 희소화(sparsification)와 어떻게 상호작용하는지 조사한다.

2. 제안된 방법: SE-KD 및 SE-KD3X

분석을 바탕으로 저자들은 두 가지 새로운 방법을 도입한다:

• SE-KD (Student Entropy-guided Position-selective KD): 이 방법은 학생 엔트로피($H(q_t)$)를 위치 중요도를 식별하기 위한 주요 신호로 사용한다. 불확실성이 가장 높은 상위 $k\%$의 위치를 선택하여 감독한다.
• SE-KD3X: 세 가지 축 모두에서 동시에 선택을 적용하는 확장 모델이다:
  • 위치(Position): 학생 엔트로피를 기반으로 한 top-$k\%$ 선택.
  • 클래스(Class): 희소한 교사 타겟을 생성하기 위해 교사 로짓(teacher logits)을 무작위 샘플링(RS-KD)함.
  • 샘플(Sample): (동결된 학생의 단일 순전파를 통해 계산된) 평균 학생 엔트로피를 기준으로 상위 $\ell\%$의 훈련 샘플을 선택함.

3. 효율성 최적화

이러한 선택 전략을 지원하기 위해 저자들은 두 가지 기술적 최적화를 구현한다:

• 선택적 LM 헤드(Selective LM Head): 선택된 위치에 대해서만 로짓을 계산하고 그래디언트를 역전파하여, 전체 $[B, L, V]$ 로짓 텐서의 실체화(materialization)를 방지한다.
• 청크 단위 엔트로피 계산(Chunked Entropy Computation): 전체 로짓을 저장하지 않고 작은 청크 단위로 위치별 엔트로피를 계산하여 메모리 오버헤드를 줄인다.
• 오프라인 교사 캐싱(Offline Teacher Caching): SE-KD3X를 위해 선택된 샘플 인덱스를 캐싱함으로써, 메인 증류 단계 동안 교사 모델을 메모리에서 제거할 수 있게 하여 저장 요구 사항을 크게 줄인다.

주요 결과

일반 목적 증류에서의 성능

Qwen3-1.7B(학생)와 Qwen3-8B(교사)를 사용하여 FineWeb-Edu의 80M 토큰에 대해 평가하였다:

• 정확도(Accuracy): SE-KD (Top-20% 학생 엔트로피)는 평균 정확도 **64.8%**를 달성하여, Full KD(64.4%)와 무작위 위치 선택(64.2%)을 능가했다.
• 퍼플렉서티(Perplexity): SE-KD는 퍼플렉서티를 Full KD의 7.3에서 6.9로 낮추었다.
• 지시 이행(Instruction Following): SE-KD는 IFEval 점수를 20.5(Full KD)에서 21.4로 향상시켰다.
• 보정(Calibration): SE-KD는 일부 베이스라인에 비해 기대 보정 오차(ECE)가 약간 증가(27.6 vs 27.3)했으나, 그 차이는 미미했으며, 이는 이득이 신뢰도 변화보다는 더 나은 감독 할당에서 기인함을 시사한다.

다중 축 및 효율성 이득

• SE-KD3X: 경쟁력 있는 성능(정확도 64.4%)을 달성하는 동시에 막대한 효율성 이득을 가능하게 했다.
• 저장 공간(Storage): 클래스 샘플링(RS-KD)과 샘플 선택을 결합함으로써, SE-KD3X는 Full KD 대비 저장 요구 사항을 99.96% 감소시켰으며(100B 토큰 기준 10,000 TB에서 3.84 TB로), RS-KD 단독 사용 대비 **80%**를 절감했다.
• 런타임(Runtime): 샘플 선택은 오프라인 캐시 재사용을 통해 총 실행 시간을 70% 단축시켰다(약 23시간에서 약 3.5시간으로).
• 메모리(Memory): 선택적 LM 헤드와 청크 단위 엔트로피 계산을 통해 피크 GPU 메모리를 18.3% 감소시켰다 (33.18 GB $\to$ 27.10 GB).

태스크 특화 증류 (GSM8K)

태스크 특화 환경(GSM8K)에서의 결과는 더 미묘했다:

• 오프-폴리시(Off-Policy): SE-KD 단독으로는 Full KD(69.5% vs 71.6%)보다 낮은 성능을 보였다. 그러나 SE-KD와 샘플 선택을 결합한 방식(SE-KD + TopSmp)은 훨씬 적은 감독만으로도 Full KD(70.9%)에 근접한 성능을 유지했다.
• 온-폴리시(On-Policy): SE-KD + TopSmp는 Full KD(70.6%)를 능가하는 가장 높은 정확도(71.2%)를 달성했다.

의의 및 기여

본 논문은 다음과 같은 기여를 주장한다:

1. 통합 프레임워크: 기존 방법들을 정리하고, 이전 연구에서 미흡했던 다중 축 선택을 강조하는 선택적 KD를 위한 일반적인 이론적 프레임워크를 제안한다.
2. 신호로서의 학생 엔트로피: 학생 엔트로피가 위치 선택을 위한 매우 신뢰할 수 있는 신호임을 밝혀냈으며, 이는 교사 기반의 불확실성이나 불일치 메트릭보다 종종 더 우수한 성능을 보인다.
3. 희생 없는 효율성: 위치, 클래스, 샘플 선택을 결합(SE-KD3X)함으로써, 성능 저하 없이 런타임, 피크 메모리, 저장 공간을 대폭 절감(기존 희소 방법 대비 저장 공간 최대 80% 절감)할 수 있음을 입증한다.
4. 오프라인 캐싱의 타당성: 본 연구는 이러한 방법들을 통해 오프라인 교사 캐싱이 실질적으로 가능하다는 것을 보여주며, 이를 통해 자원 발자국을 크게 줄인 증류를 가능하게 한다.

저자들은 조밀하고 균일한 로짓 감독이 종종 불필요하다고 결론짓는다. 특히 학생 엔트로피에 의해 유도되고 샘플/클래스 희소화와 결합될 때, 높은 불확실성을 가진 소수의 위치에 감독을 집중하는 것이 자기회귀(autoregressive) LLM에 있어 유리한 정확도-효율성 트레이드오프를 제공한다.