기술 요약: LLM 증류용 다중 세분성 궤적 정렬 (MTA)

문제 제기

대형 언어 모델 (LLM) 을 위한 기존 지식 증류 (KD) 방법들은 주로 고정된 정적 레이어에서 출력 확률 분포 (logits) 나 중간 은닉 상태를 정렬하는 데 중점을 둡니다. 이러한 접근 방식은 국소 신호 전달에는 효과적이지만, 네트워크 깊이에 걸쳐 균일한 토큰 수준 정렬 전략을 종종 사용합니다.

본 논문은 이러한 균일한 처리가 인간 언어의 계층적 본성과 트랜스포머 모델의 기능적 조직과 부합하지 않는다고 주장합니다. 언어학적 연구와 트랜스포머 해석 가능성 연구는 다음과 같은 사실을 시사합니다:

1. 하위 레이어는 표면적, 어휘적, 사실적 검색 특징을 강조합니다.
2. 상위 레이어는 점진적으로 추상적 의미 정보와 구성적 추론을 인코딩합니다.

모든 토큰을 고립된 상태로 취급하고 동일한 정렬 목표를 균일하게 적용함으로써, 현재 방법론들은 어휘 단위에서 추상적 의미 구조로 변환되는 표현의 점진적 진화 궤적을 포착하지 못합니다. 이러한 한계는 학생 모델의 하류 작업 일반화 능력을 제한합니다.

방법론: 다중 세분성 궤적 정렬 (MTA)

이러한 한계를 해결하기 위해 저자들은 MTA를 제안합니다. 이는 레이어 적응형 전략을 사용하여 교사와 학생 모델의 표현 궤적을 정렬하는 모듈식 프레임워크입니다. MTA 는 증류가 트랜스포머 레이어의 깊이 의존적 의미 역할을 존중해야 한다는 가설에 기반합니다.

핵심 구성 요소

1. 레이어 적응형 다중 세분성 스패인 (Spans):

   • 하위 레이어: 어휘적 기반과 세밀한 의존성을 보존하기 위해 단어 수준 (Word Spans) 에서 정렬이 수행됩니다.
   • 상위 레이어: 구성적 의미와 고차원 구조적 관계를 포착하기 위해 정렬은 구 수준 (명사구 및 동사구) 으로 전환됩니다.
   • 스팬 구성: 의미 단위는 구문 분석기 (예: spaCy) 를 사용하여 추출됩니다. 토큰 중요도 가중치는 자기 루프를 제외한 양방향 자기 주의 메커니즘을 통해 계산되어 토큰 표현을 스패인 표현으로 집계합니다.

2. 동적 구조 정렬 손실 ($L_{DSA}$):

   • 이 목적 함수는 교사와 학생 간의 기하학적 일관성을 강제합니다.
   • 개별 토큰 벡터를 매칭하는 대신, $L_{DSA}$는 특정 레이어 내 의미 스패인 (단어 또는 구) 간의 쌍별 관계 거리의 불일치를 최소화합니다.
   • 구성 토큰의 집계된 중요도에 기반하여 정렬을 가중함으로써 의미적으로 중요한 스패인을 우선시합니다.
   • 수학적으로, 표현 공간의 위상 구조를 정렬합니다:
     $$L_{DSA} = \frac{1}{\|L_{key}\|} \sum_{l \in L_{key}} \sum_{i,j} w_{ij} \left( d(U^S_{i,l}, U^S_{j,l}) - d(U^T_{i,\phi(l)}, U^T_{j,\phi(l)}) \right)^2$$
     여기서 $d(\cdot)$는 코사인 거리이며, $U$는 가중치 부여된 스패인 표현을 나타냅니다.

3. 은닉 표현 정렬 손실 ($L_{Hid}$):

   • 구조적 제약을 보완하기 위해, 특징 수준 손실이 특정 은닉 상태를 정렬합니다.
   • 학생의 차원이 일반적으로 교사보다 작기 때문에, 학습 가능한 선형 프로젝터가 학생 상태를 교사의 공간으로 매핑합니다.
   • 이 손실은 가장 정보량이 많은 토큰에 초점을 맞추어 투영된 학생 상태와 교사 상태 간의 가중치 부여된 코사인 거리를 최소화합니다.

4. 최적화:

   • MTA 는 플러그인 모듈로 설계되었습니다. 총 손실은 기존 프레임워크 (예: DistiLLM, FDD) 의 기본 증류 손실 ($L_{Base}$) 에 제안된 구조 및 은닉 손실을 결합합니다:
     $$L_{Total} = L_{Base} + \lambda_{DSA}L_{DSA} + \lambda_{Hid}L_{Hid}$$

주요 기여

1. 한계 식별: 저자들은 균일한 중간 정렬이 표현의 계층적 진화를 무시하여 비최적의 지식 전달로 이어진다고 식별했습니다.
2. 레이어 적응형 프레임워크: 저자들은 하위 레이어의 단어 수준 정렬에서 상위 레이어의 구 수준 구조 정렬로 동적으로 전환되는 MTA 를 도입하여, 언어의 하향식 구성 본성을 반영합니다.
3. 일반화된 모듈: MTA 는 토큰화를 공유하는 최첨단 증류 방법 (DistiLLM, DistiLLM-2, FDD) 과 호환되는 일반화된 최적화 레이어로 입증되었습니다.
4. 실증적 검증: 광범위한 실험을 통해 MTA 가 다양한 모델 아키텍처 (GPT-2, Qwen1.5, OPT) 와 벤치마크에서 일관되게 성능을 향상시키는 것을 보여주었습니다.

실험 결과

저자들은 세 가지 학생 - 교사 쌍을 사용하여 지시 따르기 작업에서 MTA 를 평가했습니다:

• GPT-2 1.5B $\to$ 120M
• Qwen1.5 1.8B $\to$ 0.5B
• OPT 6.7B $\to$ 1.3B

성능 지표:

• ROUGE-L: MTA 는 모든 데이터셋 (Dolly, SelfInst, Vicuna, S-NI) 에서 베이스라인 (FDD, DistiLLM, DistiLLM-2) 을 일관되게 능가했습니다. 예를 들어, GPT-2 1.5B $\to$ 120M 설정에서 DistiLLM 에 MTA 를 추가하면 평균 ROUGE-L 이 20.21 에서 21.45 로 향상되었습니다.
• LLM-as-a-Judge: GPT-4o-mini 를 사용하여 사실적 정확성, 일관성, 유용성을 평가한 결과, MTA 가 적용된 모델은 베이스라인 대비 일관된 점수 향상을 보였습니다.

어블레이션 연구:

• 손실 구성 요소: 구조적 ($L_{DSA}$) 과 특징적 ($L_{Hid}$) 인 두 구성 요소 모두 긍정적으로 기여했으며, 이들의 조합이 최상의 결과를 낳았습니다.
• 세분성: "전체 수준" (레이어 적응형) 전략이 정적 전략 (단어 전용 또는 구 전용) 보다 우수하여, 하위 레이어에서는 어휘적 기반을 정렬하고 상위 레이어에서는 구성적 의미를 정렬하는 것이 중요함을 확인했습니다.
• 레이어 선택: 이 방법은 다양한 레이어 스케줄에 대해 견고하지만, 테스트된 구성에서 1 개의 단어 수준 레이어 : 2 개의 구 수준 레이어라는 하이브리드 할당이 최적임이 발견되었습니다.

효율성:

• MTA 는 구문 스패인 추출 (최적화를 통해 약 0.27 초/스텝으로 감소) 과 기하학적 매칭으로 인해 주로 훈련 오버헤드를 도입합니다.
• 추론 비용은 테스트 시간에 스패인 추출을 사용하지 않으므로 변경되지 않습니다.
• 시간 매칭 비교는 성능 향상이 증가된 훈련 예산이 아닌 정렬 목적 함수로 인한 것임을 확인했습니다.

중요성과 주장

본 논문은 MTA 가 LLM 의 계층적 표현 궤적을 명시적으로 모델링함으로써 지식 증류에 의미론적으로 충실한 접근 방식을 제공한다고 주장합니다. 정적, 토큰 수준 매칭을 넘어 동적, 다중 세분성 구조 정렬로 이동함으로써, MTA 는 학생 모델이 교사의 내부 정보 처리 역학을 더 잘 복제할 수 있게 합니다.

저자들은 MTA 를 기존 증류 프레임워크를 대체하는 것이 아니라, 압축된 LLM 의 일반화와 견고성을 강화하는 보조적이고 플러그 앤 플레이 가능한 모듈로 위치시킵니다. 이 연구는 향후 모델 압축 기술 설계에서 언어적 계층과 깊이 의존적 특징 진화를 존중하는 것의 중요성을 강조합니다.