거대 모델의 지식을 압축하는 기술, Knowledge Distillation

Distillation이 뭔가요

Distillation은 큰 모델(Teacher)이 학습한 지식을 작은 모델(Student)에게 전달하는 기법입니다. 원래 학습 데이터로 Student를 처음부터 학습시키는 게 아니라, Teacher의 출력을 정답으로 삼아 학습하는 방식이에요.

이 아이디어는 2006년에 처음 등장했고, 2015년 Hinton 등의 논문 "Distilling the Knowledge in a Neural Network"에서 Teacher-Student 프레임워크, soft label, temperature scaling 같은 핵심 개념들이 체계화됐습니다. LLM 시대에 들어서면서 더 중요해졌는데, DistilBERT(2019)가 BERT를 40% 줄이면서 97% 성능을 유지한 게 대표적인 사례예요.

Distillation은 어떻게 이루어지는가

가장 일반적인 방식은 오프라인 distillation으로, 두 단계를 순차적으로 진행합니다.

1단계: Teacher가 데이터를 생성합니다. Teacher 모델을 먼저 실행해서 입력에 대한 출력을 저장해요. 분류 문제라면 softmax 확률 분포 전체를, LLM이라면 응답 텍스트를 디스크에 기록합니다.

2단계: Student가 해당 데이터로 학습합니다. Teacher는 이미 종료된 상태이고, 저장된 Teacher 출력을 정답처럼 사용해서 Student만 학습합니다.

두 모델을 동시에 올리면 GPU 메모리 소모가 과도해지기 때문에 이렇게 분리하는 게 일반적이에요.

LLM에서는 더 단순하게, Teacher에게 대량의 프롬프트를 보내 응답을 수집하고 그 (프롬프트, 응답) 쌍으로 Student를 fine-tuning하는 방식이 널리 쓰입니다. Alpaca(2023)가 GPT-3.5 응답으로 LLaMA 7B를 학습시킨 게 그 예예요.

Distillation이 왜 작동하는가 — Dark Knowledge

핵심은 "Dark Knowledge" 개념에 있어요.

일반 학습에서는 hard label을 씁니다. 고양이 이미지라면 정답은 단순히 "고양이 = 1, 나머지 = 0"이에요. 하지만 Teacher의 soft label에는 훨씬 풍부한 정보가 담겨 있어요:

고양이: 0.9, 호랑이: 0.05, 개: 0.03, 자동차: 0.0001 ...

이 분포 안에는 "고양이와 호랑이는 시각적으로 유사하다", "고양이와 자동차는 거의 관련 없다"는 클래스 간의 관계 정보가 내포되어 있어요. Hinton이 이걸 Dark Knowledge라고 불렀고, 이게 바로 같은 데이터로 처음부터 학습하는 것보다 distillation이 더 나은 성능을 내는 주된 이유입니다.

Temperature scaling은 이 정보를 더 잘 추출하기 위한 장치예요. Temperature를 높일수록 확률 분포가 평탄해지면서 작은 확률값들 사이의 차이가 더 뚜렷하게 드러납니다. 보통 T=2~20 사이에서 조절하고, Teacher와 Student 양쪽에 동일한 값을 적용해야 해요.

Capacity Gap 문제

직관적으로는 Teacher가 클수록 좋을 것 같지만, Teacher와 Student의 크기 차이가 너무 크면 오히려 성능이 떨어집니다. 이걸 Capacity Gap 문제라고 해요.

Student가 Teacher의 복잡한 결정 경계를 수용할 용량 자체가 부족하기 때문이에요. 이를 해결하는 대표적인 방법이 Teacher Assistant(TA) 전략입니다:

큰 Teacher → 중간 크기 TA → 작은 Student

중간 모델을 거쳐 점진적으로 지식을 전달하면 직접 전달하는 것보다 더 나은 결과를 얻을 수 있어요.

Distillation vs Pruning

비슷해 보이지만 접근이 다른 기법으로 Pruning이 있어요. Distillation은 작은 아키텍처의 새 모델을 처음부터 학습시키는 것이고, Pruning은 기존 모델에서 불필요한 가중치를 잘라내는 것입니다. 비유하자면 distillation은 선생님의 노하우를 배운 새 직원을 뽑는 거고, pruning은 기존 직원의 불필요한 업무를 덜어주는 거예요.

실무에서는 이 둘을 조합하기도 합니다. Distillation으로 작은 모델을 만든 뒤 pruning까지 적용하는 식이에요.

법적·윤리적 이슈

실무에서 distillation을 쓸 때 반드시 챙겨야 할 부분이에요.

OpenAI, Google 등 주요 API 제공자들의 이용약관은 자사 모델 출력으로 경쟁 모델을 학습시키는 것을 명시적으로 금지하고 있습니다. Alpaca도 이 때문에 "연구 목적" 단서를 달았고 상업적 사용에 제한이 있었어요.

"API 출력을 수집해서 학습 데이터로 쓰면 되지 않을까?"라는 접근은 실제 서비스에서 상당한 법적 리스크를 수반합니다. 오픈소스 모델을 Teacher로 쓰거나 자체 모델 간 distillation을 수행하는 게 안전한 선택이에요.

마치며

Distillation은 단순히 모델을 작게 만드는 기법이 아닙니다. 큰 모델이 학습한 클래스 간의 유사성, 입력 데이터의 구조, 과제 해결 전략을 효율적으로 전달하는 지식 이전 기술이에요.

LLM이 점점 커지는 지금, 거대 모델의 성능을 실제 서비스에서 쓸 수 있는 크기로 압축하는 distillation의 중요성은 앞으로도 계속 커질 겁니다. 다만 적용할 때는 capacity gap, 법적 이슈 같은 실무적 제약도 함께 고려해야 해요.