기술 요약: 대규모 언어 모델을 위한 분산 보정 오프라인 데이터 증류

문제 정의

본 논문은 복잡한 추론 능력을 대규모의 강력한 "교사" 대규모 언어 모델 (LLM) 에서 더 작고 자원이 제한된 "학생" 모델로 전이하는 과제를 다룹니다. 오프라인 증류(교사가 생성한 고정된 추론 궤적을 기반으로 학생을 학습시키는 방식) 는 샘플 효율성이 높고 고품질의 감독을 제공하지만, 근본적인 분산 드리프트로 고통받습니다.

학습 중에는 학생 모델이 교사 생성 접두어에 조건부로 학습되지만, 추론 중에는 학생이 스스로의 접두어를 autoregressive 방식으로 생성해야 합니다. 학생이 생성한 토큰이 교사의 궤적에서 벗어나면, 학습 데이터에서 대표성이 부족한 토큰 공간 영역으로 진입하게 됩니다. 이는 누적 오차를 초래하고 긴 추론 궤적 전반에 걸쳐 신뢰할 수 없는 행동을 유발하며, 기존에 정제된 증류 목적 함수 (예: 양방향 KL 발산 최소화) 는 여전히 교사가 제공하는 분산에 대해 최적화하기 때문에 이 한계를 완전히 해결하지 못합니다.

반면, 온-정책 방법(예: GRPO, RLHF) 은 자기 생성 데이터를 기반으로 학습하여 이러한 드리프트를 완화하지만, 온라인 샘플링에 따른 높은 계산 비용이 발생하며 초기 학습 단계에서 저품질의 궤적으로 인해 어려움을 겪는 경우가 많습니다.

방법론: DISCORD

저자들은 DISCORD(분산 보정 증류) 를 제안합니다. 이는 교사 생성 데이터의 효율성을 유지하면서 분산 불일치를 보정하는 원칙적인 오프라인 프레임워크입니다. 핵심 메커니즘은 학생의 현재 지지 (해당 토큰을 생성할 가능성) 에 기반하여 교사 토큰 감독을 적응적으로 재가중치하는 것입니다.

주요 기술 구성 요소:

1. 학생 지지 인식 재가중치:
   이 방법은 학습 궤적 $\{x, y \sim \pi_T(x; c)\}$를 재가중치하기 위해 밀도 비율을 도입합니다:
   $$\frac{\bar{\pi}_\theta(y|x)}{\bar{\pi}_T(y|x;c)}$$
   여기서 $\pi_\theta$는 학생, $\pi_T$는 교사이며, 막대는 stop-gradient 연산을 나타냅니다. 이 비율은 추론 중 학생이 마주할 가능성이 높은 교사 토큰을 우선시합니다.

2. 안정화된 변환:
   밀도 비율을 직접 적용하면 무제한 값과 긴 시퀀스에서의 높은 분산으로 인해 학습 불안정성이 발생합니다. 저자들은 두 가지 안정화 기법을 적용합니다:

   • 시그모이드 변환: 로그 밀도 비율을 시그모이드 함수 $\sigma(\tau \cdot \log(\dots))$에 통과시켜 보정을 유계 범위 $(0, 1)$로 매핑합니다. 이는 비율의 순서를 유지하면서 폭발을 방지합니다.
   • 토큰 수준 근사: 길이 편향 (긴 시퀀스가 더 큰 집계 가중치를 누적하는 현상) 을 완화하기 위해 보정을 토큰 수준에서 적용합니다. 궤적 내의 각 토큰 $y_t$에 대해 가중치는 다음과 같습니다:
     $$r_t = \sigma\left(\tau \log \frac{\bar{\pi}_\theta(y_t | x; y_{<t})}{\bar{\pi}_T(y_t | x; c; y_{<t})}\right)$$
     이 가중치는 기본 목적 함수 (예: SFT 또는 KL 발산) 학습 중 기울기 승수로 적용됩니다.

3. 해석:
   보정 가중치 $r_t$ ($\tau=1$) 는 균일한 사전 확률을 가정할 때, 해당 토큰이 교사 분산이 아닌 학생 분산에서 추출되었을 확률 (후확률) 로 해석됩니다. 이는 대조적 디코딩 원리와 연결되지만, 학생이 안정적으로 존재할 수 있는 영역에서의 학습을 우선시하기 위해 학습 중에 적용됩니다.

주요 기여

• 원칙적인 오프라인 프레임워크: DISCORD는 비용이 많이 드는 온라인 학생 롤아웃 없이 유일하게 오프라인 교사 제공 데이터만을 사용하여 교사 - 학생 분산 드리프트를 효과적으로 보정하는 최초의 방법입니다.
• 유연한 통합: 토큰 수준 재가중치 메커니즘은 표준 감독 미세 조정 (SFT) 과 KL 발산 기반 목적 함수를 포함한 다양한 오프라인 증류 백본과 호환됩니다.
• 안정화 기법: 시그모이드 변환과 토큰 수준 근사의 도입은 긴 시퀀스에서 밀도 비율 기반 보정에 내재된 분산 및 불안정성 문제를 해결합니다.

실험 결과

이 방법은 Qwen2.5 및 Llama-3.2 모델을 사용하여 수학 추론 벤치마크 (GSM8K, MATH, MATH500) 와 더 어려운 홀드아웃 경쟁 스타일 작업 (AMC, AIME, OlympiadBench, Omni-Math) 에서 평가되었습니다.

• 성능 향상: DISCORD는 도메인 내 및 도메인 외 (OOD) 작업에서 강력한 오프라인 베이스라인 (SFT, ToDi, 대칭 KL) 보다 일관되게 우수한 성과를 냅니다. 예를 들어, Qwen2.5-7B $\to$ 3B 증류에서 DISCORD는 표준 SFT 대비 MATH500 정확도를 4.60 점, 대칭 KL 대비 2.60 점 향상시켰습니다.
• 접두어 드리프트에 대한 강건성: ExAccErr(분산 발산) 를 사용한 평가에 따르면, 접두어 길이가 증가함에 따라 (학생 생성 접두어가 교사로부터 멀어짐) 표준 SFT 성능이 크게 저하되는 반면, DISCORD는 안정적인 성능을 유지하여 누적 오차가 감소했음을 나타냅니다.
• 온라인 방법 대비 효율성: DISCORD는 온라인 학생 롤아웃 방법 (GKD 등) 과 유사한 성능을 달성하지만 계산 비용은 훨씬 낮습니다. 7B $\to$ 3B 설정에서 DISCORD는 온라인 학생 토큰을 사용하지 않으면서 GKD 의 49.56 GPU 시간 대비 3.12 GPU 시간만 필요했습니다.
• 추론 품질: 정성적 분석에 따르면 DISCORD 는 더 깔끔하고 간결하며 논리적으로 일관된 궤적을 생성합니다. 반복, 정답 이후 계속 작성, 여러 개의 서로 다른 최종 답변과 같은 일반적인 실패 모드를 줄입니다. 특히 "실패 복구" 사례에서 DISCORD 는 교사의 표면적 포맷을 모방하지 않더라도 교사의 솔루션에 내재된 수학적 구조를 유지합니다.

의의 및 주장

본 논문은 가벼운 분산 보정 인식 학습이 온라인 상호작용 없이 오프라인 추론 증류를 크게 강화할 수 있다고 주장합니다.

• 격차 해소: 이 연구는 새로운 데이터를 생성하는 것이 아니라 기존 오프라인 데이터를 지능적으로 재가중치함으로써 오프라인 및 온라인 증류 방법 간의 성능 격차를 좁힐 수 있음을 보여줍니다.
• 일반화 가능성: 실험은 수학 추론에 초점을 맞추었지만, 저자들은 이 방법이 코드 생성 및 과학적 추론을 포함한 기타 장기 생성 설정에도 적용 가능하다고 주장합니다.
• 실제 배포: 온라인 샘플링의 높은 계산 비용을 피함으로써 DISCORD 는 자원이 제한된 환경에서 능력 있는 추론 모델을 배포할 수 있는 실현 가능한 경로를 제공합니다.

저자들은 겸손하게도, 그들의 방법이 기존 오프라인 베이스라인을 개선하지만, 학생의 롤아웃 분산을 직접 관찰하기 때문에 GKD 와 같은 온라인 방법이 일부 특정 도메인 내 벤치마크에서는 여전히 우위를 점한다고 지적합니다. 그러나 DISCORD 는 성능과 계산 효율성 간의 더 나은 균형을 제공합니다.