Making Knowledge Distillation Cheap Enough to Run at Scale

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

교사의 위치별 상위 100개 로짓을 한 번만 저장하고 출력 투영까지 청크 단위로 결합한 KL 손실을 사용해, 지식 증류의 품질을 유지하면서 장문맥 훈련의 메모리와 계산 비용을 크게 줄인 연구다.

📌 핵심 요약

• 기존 온라인 지식 증류는 매 훈련 단계마다 교사를 다시 실행하고 교사와 학생의 전체 어휘 분포를 동시에 보관하므로, 모델과 문맥이 커질수록 막대한 연산량과 GPU 메모리를 요구한다.
• 연구진은 각 토큰 위치에서 교사의 확률이 높은 상위 100개 로짓만 미리 계산해 저장하는 오프라인 증류를 도입함으로써, 학생 훈련 중 교사를 메모리에 올리거나 반복 실행할 필요를 없앴다.
• 핵심 기법인 융합 청크형 KL 손실은 학생의 전체 로짓 행렬을 만들지 않고 출력 투영과 손실 계산을 청크별로 수행하며, 역전파 때 청크를 재계산하는 대신 최대 메모리를 크게 낮춘다.
• 단일 H200의 8K 문맥 실험에서는 온라인 증류와 세 가지 오프라인 방식이 거의 같은 훈련 손실을 보였고, 융합 청크형 방식은 최대 메모리를 온라인 방식의 102.8GB에서 58.3GB로 줄였다.
• 32K 손실 커널 실험에서는 최대 메모리가 85.2GiB에서 5.45GiB로 15.6배 감소했으며, GPT-OSS 20B의 32,768토큰 증류에서는 구성을 GPU 노드 4대에서 1대로 줄이고 단계 시간을 57.0초에서 12.23초로 단축했다.

🧩 주요 포인트

1. 교사의 출력을 반복 계산하는 온라인 구조를 재사용 가능한 상위 100개 로짓 캐시로 전환해, 교사 추론 비용과 학생 훈련 비용을 분리하고 반복 실험의 부담을 낮췄다.
2. 학생의 전체 어휘×시퀀스 로짓을 보관하지 않는 융합 청크형 계산은 짧은 문맥에서 일부 재계산 비용이 있지만, 문맥이 길어질수록 메모리 확장을 억제하고 실행 가능한 훈련 범위를 넓힌다.
3. K에서의 손실 곡선 일치와 약 32억 매개변수 학생의 평가 결과는 이 실험 구성에서 비용 절감이 단순한 실행 최적화에 그치지 않고 증류 품질 보존과 함께 달성됐음을 보여준다.

🧠 상세 정리

1. 지식 증류가 다시 중요해진 배경

지식 증류는 큰 교사 모델의 성능을 더 작은 학생 모델이 모방하도록 훈련하는 기법으로, 최근 gpt-oss·Qwen·GLM·Kimi 같은 오픈소스 대규모 언어 모델이 늘면서 다시 주요 연구 주제가 됐다. 초대형 모델은 배포 자체가 비싸며, 본문이 예로 든 Kimi-K3는 2조 8천억 개의 매개변수를 가지고 로드하는 데만 약 3TB의 VRAM이 필요하다. 이에 따라 큰 모델을 압축한 뒤 증류로 원래 능력을 회복하는 방식이 일반적인 접근으로 자리 잡았고, Nvidia의 Nemotron 3 Puzzle 75B와 Multiverse Computing의 Hypernova 60B 같은 압축 모델도 공개됐다. 그러나 최종 품질을 크게 좌우하는 증류 단계가 압축 파이프라인에서 가장 비싼 부분인 경우가 많아, 대규모로 반복 적용하려면 증류 시스템 자체의 비용 구조를 바꿀 필요가 있었다.

2. 온라인 증류의 메모리와 계산 병목

일반적인 온라인 증류는 KL 발산 손실을 사용하면서 교사와 학생 모델을 동시에 메모리에 올리고, 매 훈련 단계마다 교사의 전체 순전파를 다시 수행한다. 각 토큰 위치에서 교사와 학생의 전체 어휘 확률 분포를 비교해야 하므로, 모델 가중치뿐 아니라 어휘 크기와 시퀀스 길이를 곱한 거대한 텐서도 함께 유지해야 한다. gpt-oss-120b는 어휘가 201,088개이기 때문에 시퀀스 길이 32K와 배치 크기 4에서 교사 확률 텐서 하나만 4×201,088×32,768 크기가 되며, bfloat16 기준 약 50GB의 VRAM을 차지한다. 여기에 그래디언트, 활성값, 모델 가중치, 옵티마이저 상태를 더하면 한 번의 증류 반복이 약 250GB까지 치솟아 141GB 용량의 H200 한 대로는 처리할 수 없다. 교사의 동작이 훈련 중 변하지 않는데도 같은 출력을 계속 재계산한다는 점 역시 온라인 방식의 구조적인 비효율이다.

3. 교사 출력을 재사용하는 오프라인 증류

첫 번째 변경은 교사를 매 단계 실행하는 대신, 데이터의 각 토큰 위치에서 교사가 예측한 확률 상위 100개 토큰의 로짓을 한 번 계산해 캐시에 저장하는 것이다. 캐시가 만들어진 뒤에는 학생 훈련 중 교사 모델을 메모리에 유지할 필요가 없고, 동일한 데이터에 대해서는 교사를 다시 실행할 필요도 없다. 이 캐시는 여러 절제 실험과 훈련 설정에서 반복 사용할 수 있으므로, 교사 추론은 일회성 비용이 되고 이후 비교 실험은 학생 훈련 비용만 부담하게 된다. 전체 어휘 분포가 아니라 상위 100개 항목만 보존한다는 점에서 저장량도 제한되며, 본문의 8K 실험에서는 이 희소 캐시를 이용한 오프라인 방식과 전체 교사 분포를 사용하는 온라인 방식의 훈련 손실 곡선이 거의 정확히 겹쳤다. 다만 이 품질 일치는 본문에서 사용한 모델·데이터·설정의 실험 결과로 제시된 것이다.

4. 세 가지 KL 손실 구현의 차이

밀집형 KL은 캐시된 상위 100개 로짓을 전체 어휘 확률 격자로 복원하고, 학생의 전체 로그 확률 격자와 비교하는 기준 구현이다. 온라인 증류와 가장 유사해 정확성 비교의 기준이 되지만, 어휘×시퀀스 크기의 교사 분포와 비교 결과를 밀집 텐서로 보관해 메모리 사용량이 크다. 순전파 청크형 KL은 교사 출력을 희소한 상태로 유지하면서 시퀀스 위치를 여러 조각으로 나누어 손실을 계산하며, 벤치마크에서는 세 오프라인 구현 중 가장 빨랐다. 그러나 학생 출력층이 만든 전체 로짓은 여전히 한꺼번에 계산되고 역전파를 위해 저장되므로, 시퀀스가 길어질수록 학생 로짓이 새로운 메모리 병목이 된다. 융합 청크형 KL은 출력 투영을 손실 계산 안으로 결합해 청크별로 은닉 상태를 로짓으로 변환하고 손실에 반영한 직후 폐기하며, 역전파 시 각 청크를 다시 계산한다. 출력 투영을 순전파와 역전파에서 두 번 수행하는 계산 비용을 감수하는 대신 학생의 전체 로짓 행렬을 한 번도 만들지 않는 것이 핵심이다.

5. 단일 H200의 8K 문맥 비교

연구진은 단일 H200에서 Llama 3.1 8B Instruct를 교사로, 약 32억 매개변수의 Llama 모델을 학생으로 사용해 8K 문맥 길이에서 네 방식을 비교했다. 온라인 증류는 최대 메모리 102.8GB, 반복 시간 25.9초, 처리량 237TFLOP/s를 기록했고, 오프라인 밀집형 KL은 각각 78.3GB, 18.5초, 331TFLOP/s였다. 순전파 청크형 KL은 61.8GB, 18.4초, 335TFLOP/s로 가장 빠른 반복 시간과 가장 높은 처리량을 보였으며, 융합 청크형 KL은 58.3GB, 20.2초, 304TFLOP/s로 최대 메모리가 가장 낮았다. 네 방식의 훈련 손실은 거의 동일해, 이 구성에서는 위치별 상위 100개 로짓만 사용한 오프라인 증류가 온라인 증류와 같은 수준의 손실을 달성했다. 8K에서는 융합 방식이 역전파 중 출력 투영을 재계산하기 때문에 가장 빠르지는 않았지만, 연구진은 그 장점이 문맥 길이가 더 커질 때 뚜렷해진다고 설명한다.

6. 장문맥에서 커지는 융합 청크형의 이점

장문맥 확장 특성을 분리해 보기 위해 연구진은 트랜스포머 본체 없이 출력 투영 네트워크와 손실 커널만 사용한 벤치마크를 수행했다. 32K 토큰에서 밀집형 손실의 최대 메모리는 85.2GiB였지만 완전 청크형 방식은 5.45GiB만 사용해 15.6배 감소했으며, 밀집형 손실은 64K 토큰부터 메모리 부족으로 실행에 실패했다. 256K 토큰에서는 완전 청크형 손실이 11.6GiB를 사용한 반면 그다음으로 나은 청크형 변형은 134.2GiB를 사용했다. 같은 256K 길이에서 완전 청크형은 해당 비교 방식보다 반복당 약 3.3배 빨랐다. 이는 짧은 문맥에서 보이던 출력 투영 재계산의 부담보다 전체 로짓을 생성하고 보관하는 비용이 장문맥에서 훨씬 빠르게 커지기 때문에, 문맥이 길어질수록 융합 방식의 메모리와 실행 시간상의 이점이 확대된다는 결과다.

7. GPT-OSS 20B 증류의 실제 자원 절감

연구진은 GPT-OSS 20B 모델을 32,768토큰 문맥에서 증류하는 작업에도 융합 손실을 적용했다. 전체 학생 로짓을 보관하지 않아 확보된 메모리 덕분에 기존 GPU 노드 4대가 필요했던 구성을 GPU 노드 1대로 줄일 수 있었다. 단계당 실행 시간은 57.0초에서 12.23초로 감소해 약 5배 빨라졌고, GPU당 처리량은 74.2TFLOP/s에서 345.7TFLOP/s로 상승했다. 이 결과는 손실 커널의 메모리 최적화가 단순히 동일 장비에서 여유 공간을 확보하는 데 그치지 않고, 실제 증류 작업에 필요한 노드 수와 반복 시간을 동시에 줄일 수 있음을 보여준다. 논문이 제시한 오프라인 캐시와 융합 청크형 손실의 조합은 장문맥 복원 훈련을 더 적은 하드웨어에서 실행하고 대규모 실험을 반복할 수 있도록 비용 구조를 바꾸는 데 목적이 있다.

8. 학생 모델의 결과와 공개 범위

효율적인 오프라인 증류 설정을 바탕으로 연구진은 Llama 3.1 8B Instruct를 교사로 삼아 약 32억 매개변수의 학생 모델을 훈련했다. 이 학생은 교사보다 매개변수가 절반 미만이면서 BoolQ와 HellaSwag에서 교사의 정확도 대부분을 유지했고, MMLU에서는 교사와의 차이가 약 9점 이내였다. 따라서 본문이 제시한 결과에서는 증류 비용과 모델 크기를 낮추면서도 교사의 단문맥 평가 성능 상당 부분을 보존했다. 논문에는 손실 함수 선택과 시퀀스 패킹이 복원 품질에 미치는 영향 등 추가 절제 실험과, 융합 청크형 손실의 폐쇄형 그래디언트 및 전체 훈련 구성이 포함돼 있다. 연구진은 청크형 손실 구현을 공개 저장소에 오픈소스로 제공해 동일한 계산 방식을 다른 증류 파이프라인에서도 검토하고 적용할 수 있도록 했다.

🧾 핵심 주장 / 시사점

• 교사 로짓 캐시는 변하지 않는 교사 출력을 일회성 계산으로 분리하므로, 동일한 캐시를 여러 절제 실험에 재사용할 때 반복 연구의 비용 절감 효과가 커진다.
• 8K에서는 순전파 청크형 KL이 가장 빠르고 융합 청크형 KL이 가장 적은 메모리를 사용했으므로, 짧은 문맥에서는 속도와 메모리 중 어느 제약이 중요한지에 따라 구현 선택이 달라질 수 있다.
• 융합 청크형 방식은 역전파 재계산이라는 비용을 추가하지만, 64K 이상에서 밀집형이 실패하고 256K에서 메모리와 속도 모두 큰 차이를 보였다는 결과는 장문맥일수록 이 교환이 유리해짐을 보여준다.

✅ 액션 아이템

• 오프라인 증류에서 교사의 상위 100개 로짓 캐시를 사용해 매 훈련 단계 교사 재실행과 교사 메모리 상주 부담을 제거한다.
• 단일 H200 8K 문맥 실험 기준 온라인 대비 102.8GB에서 58.3GB로 줄인 융합 청크형 KL 손실 구성을 기본 증류 방식으로 반영한다.
• GPT-OSS 20B의 32,768토큰 증류에서 GPU 노드 4대에서 1대로 축소해 단계 시간을 57.0초에서 12.23초로 단축한 근거로 자원 구성안을 조정한다.

❓ 열린 질문

• 단일 H200 8K 문맥에서 온라인 증류와 손실 곡선이 일치한 결과가 32K 손실 커널(85.2GiB→5.45GiB)까지도 유지되는지를 어떻게 판단할 것인가?
• 교사의 상위 100개 로짓 캐시만으로 전체 어휘×시퀀스 로짓을 보관하지 않을 때 문맥이 길어질수록 메모리 억제 효과가 GPT-OSS 20B에서 일관되게 유지되는가?
• 32억 매개변수 학생 기준에서 15.6배 축소된 메모리와 12.23초 단계 시간 성능이 knowledge distillation 품질 유지와 함께 다른 distillation 설정에서도 성립 가능한가?