은닉 상태(Hidden States) 정렬을 통한 LLM 증류(Distillation)의 안정화

은닉 표현(Hidden-representation) 정렬은 KL 분산(KL variance)을 정확히 0으로 유도하여, 온-폴리시(on-policy) LLM 증류(distillation)를 노이즈 없는 최적화 문제로 전환합니다. 롤아웃(rollouts)에 대한 결정론적 MSE 손실(MSE loss)은 정책 경사(policy-gradient) 증류 파이프라인을 오랫동안 괴롭혀온 확률성(stochasticity)을 제거합니다. \u200d[1]\u200d

OPRD 이전의 온-폴리시 증류(on-policy distillation, OPD)는 약 15만(150k) 개의 토큰 어휘(vocabularies)에 대해 KL 발산(KL divergence)을 샘플링하여, 몬테카를로 분산(Monte Carlo variance)을 팽창시키고 후기 단계의 정체(stagnation)를 유발했습니다. 교사(teacher) 모델은 블랙박스로 취급되었고 출력 로짓(output logits)만이 감독(supervision)을 받았기에, 모델 내부 역학(internal dynamics)의 대부분은 손대지 않은 채로 남겨졌습니다. \u200d[1]\u200d

OPRD의 결정론적 MSE 목적 함수(objective)는 상위-k OPD(top-k OPD) 대비 1.44배의 속도 향상을 제공하는 동시에, AIME 2024/2025 및 AIMO 벤치마크에서의 성능 격차를 해소합니다. 경험적으로 학생(student) 모델은 순수 OPD가 정체되는 지점에서 교사 수준의 점수에 도달하며, 이는 은닉 상태(hidden-state) 감독이 로짓(logits)만보다 더 풍부한 그래디언트(gradients)를 제공함을 확인시켜 줍니다. “OPRD의 MSE 목적 함수는 롤아웃(rollout)의 결정론적 함수입니다. 이 함수의 그래디언트는 추가적인 샘PLING 분산을 전혀 포함하지 않으므로, 구조적으로 OPD의 후기 단계 신호 대 잡음비(signal-to-noise) 붕괴를 제거합니다.” \u200d[1]\u200d

OPRD는 LM 헤드(LM head)를 폐기하고 선택된 은닉 레이어(hidden layers)만을 저장함으로써 OPD에 필요한 대규모 top-k 토큰 캐시(caches)를 피하기 때문에 메모리 소비가 54% 감소합니다. 이러한 감소를 통해 배치 크기(batch size)를 희생하지 않고도 일반 GPU에서 학습이 가능해지며, 이는 대규모 클러스터가 없는 팀에게 실질적인 이득이 됩니다. 코드베이스는 KL 항을 완전히 제거하는 REP_DISTILLATION_ONLY=True 플래그를 제공하며, 이를 통해 메모리 절약이 표현 전용(representation-only) 체제에서 비롯됨을 확인해 줍니다. \u200d[1]\u200d

이 논문의 분석은 또한 OPD가 왜 붕괴하는지를 밝혀냅니다: "(8)식의 OPD 분산(variance)은 $\alpha \to 0$일 때 사라지지 않으며, score-function 항을 통해 훈련 후반부에 정책 경사(policy gradient)를 지배합니다. 이것이 순수 OPD의 후기 단계 정체(stagnation) 뒤에 숨겨진 메커니즘입니다(Section 1). OPRD는 조건부 분산(conditional variance)을 0으로 추가하므로, 출력 분포(output distribution)가 거의 수렴한 이후에도 안정적인 최적화 신호(optimization signal)를 제공합니다." 이 논문은 OPRD가 정렬된 선택적 레이어들에 대해 샘플링 분산(sampling variance)을 제거한다는 것을 보여주지만, 더 깊은 레이어나 레이어 간 상호작용(cross-layer interactions)에 미치는 영향은 평가하지 않았습니다. 또한 이 방법은 여전히 각 학생(student) 업데이트마다 교사(teacher)의 롤아웃(rollouts)을 생성해야 합니다. 하나의 미해결 과제는 경량화된 teacher-free 프록시(proxy)가 대규모 환경에서도 zero-variance 특성을 유지할 수 있는지 여부입니다.

만약 은닉 상태(hidden-state) 정렬이 KL 분산(KL variance)을 진정으로 무효화한다면, 향후 압축 파이프라인(compression pipelines)은 기본 증류(distillation) 루틴으로서 일반적인 OPD를 OPRD로 대체해야 합니다. AIME 및 AIMO와 같은 벤치마크는 표현 전용 손실(representation-only loss)을 사용하여 재실행되어야 하며, 이를 통해 학생 모델의 정확도 향상과 계산 예산(compute budgets) 절감을 모두 기대할 수 있습니다.

References

1. OPRD: On-Policy Representation Distillation