본문 바로가기

LLM

LLM 지식 증류(Knowledge Distillation) 완전 정리 — 큰 모델의 지식을 작은 모델에 이식하는 법

반응형

GPT-4급 성능을 원하지만 H100 서버 비용은 감당 못해요. 이 모순을 해결하는 게 지식 증류예요.

"큰 모델(선생)이 배운 걸 작은 모델(학생)한테 가르친다."

DeepSeek-R1-Distill-Llama-8B가 이 방식으로 만들어졌어요. DeepSeek R1(671B)의 추론 능력을 8B 모델에 이식한 거예요.


핵심 개념 — 소프트 레이블이 왜 강력한가

일반 학습은 정답만 알려줘요.

하드 레이블: "답은 A야" (0 또는 1)

소프트 레이블: "A가 70%, B가 20%, C가 10%"
→ A가 맞지만 B와도 가까움을 알 수 있음
→ 클래스 간 관계 정보가 담겨 있음

큰 모델의 소프트 레이블(로짓 분포)에는 단순 정답보다 훨씬 풍부한 정보가 있어요. 학생 모델은 이걸 배워서 실제보다 훨씬 강해져요.


증류의 3가지 방식

1. 블랙박스 증류 (Black-Box KD)

선생 모델의 출력 텍스트만 사용해요. API 접근만 있어도 가능해요.

흐름:
선생 모델(GPT-4) → 질문에 대한 답변 생성
                 → 답변 데이터셋 구축
                 → 학생 모델 파인튜닝 (SFT)

가장 쉽지만 품질은 세 방식 중 가장 낮아요.

주의: OpenAI, Google, Anthropic 등 대부분의 상용 API는 경쟁 모델 학습에 출력 사용 금지예요. 오픈소스 모델(DeepSeek, Llama, Qwen)만 써야 해요.

# 핵심 코드 — 선생 모델로 데이터셋 생성
from transformers import pipeline

teacher = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B")

questions = ["피보나치 수열이란?", "TCP/IP 동작 원리는?"]
dataset = []

for q in questions:
    response = teacher(f"<|user|>{q}<|assistant|>", max_new_tokens=512)
    dataset.append({"input": q, "output": response[0]["generated_text"]})

# 이후 학생 모델 SFT 파인튜닝 (LoRA/QLoRA 방식 동일)

2. 화이트박스 증류 (White-Box KD)

선생 모델의 로짓(확률 분포)까지 활용해요. 오픈소스 모델에서만 가능하고 품질이 가장 높아요.

손실 함수:
L = α × CE(학생 출력, 정답) + (1-α) × KL(학생 로짓, 선생 로짓)

α: 두 손실 간 균형 (보통 0.3~0.7)
KL: KL 다이버전스 (두 분포의 차이)
import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.5):
    # 소프트 레이블 손실 (선생 지식)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    kd_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (temperature ** 2)

    # 하드 레이블 손실 (정답)
    ce_loss = F.cross_entropy(student_logits, labels)

    return alpha * kd_loss + (1 - alpha) * ce_loss

Temperature 파라미터 — 값이 클수록 소프트 레이블이 더 부드러워져요(낮은 확률들도 살아남음). 보통 2~6 사이 사용.


3. Chain-of-Thought 증류 (추론 능력 이식)

2026년 현재 가장 주목받는 방식이에요. 선생 모델의 **사고 과정(CoT)**까지 학생한테 가르쳐요.

일반 증류:
Q: "1+1은?" → A: "2"

CoT 증류:
Q: "1+1은?" → 
<think>
1에 1을 더하면 2가 됩니다.
</think>
A: "2"

→ 학생이 사고 과정 자체를 학습

DeepSeek-R1이 이 방식으로 작은 모델들에게 추론 능력을 이식했어요.

# CoT 데이터셋 형식
{
    "input": "정렬 알고리즘 시간복잡도 비교해줘",
    "chain_of_thought": "버블정렬은 O(n²)... 퀵정렬은 평균 O(n log n)...",
    "output": "퀵정렬이 평균적으로 가장 효율적입니다."
}

실전 증류 파이프라인

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer
import torch

# 선생 모델 (32B) — 데이터 생성용
teacher_model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# 학생 모델 (7B) — 학습 대상
student_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-7B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# SFT로 학생 파인튜닝 (블랙박스 증류)
trainer = SFTTrainer(
    model=student_model,
    train_dataset=distillation_dataset,  # 선생이 생성한 데이터
    args=TrainingArguments(
        output_dir="./student-distilled",
        num_train_epochs=3,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=8,
        learning_rate=2e-5,  # SFT는 파인튜닝보다 낮게
        bf16=True,
    )
)
trainer.train()

증류 vs 파인튜닝 vs 양자화

세 방식은 목적이 달라요.

양자화:   같은 모델, 정밀도 줄임 → 속도/메모리 개선
파인튜닝: 같은 모델, 행동 변경   → 특정 도메인 적응
증류:     작은 모델, 능력 이식   → 모델 자체를 경량화

보통 조합해서 써요.

권장 순서: 증류 → 파인튜닝 → 양자화
(P-KD-Q: Pruning → KD → Quantization)

실제 증류 모델 사례들

학생 모델 선생 모델 방식 결과

DistilBERT (66M) BERT (110M) 화이트박스 97% 성능, 40% 작음
DeepSeek-R1-Distill-8B DeepSeek-R1 (671B) CoT 증류 추론 능력 대부분 이식
Microsoft Phi-3 GPT-4 류 블랙박스 3.8B로 GPT-3.5 수준
Llama-3.2-1B Llama-3.1-8B+ 화이트박스 1B로 실용적 성능

언제 증류를 써야 하나

✅ 엣지/온디바이스 배포 (스마트폰, 임베디드)
✅ 추론 비용을 대폭 낮춰야 할 때
✅ 특정 도메인 전용 소형 모델 필요할 때
✅ 선생 모델보다 빠른 응답이 필요할 때

❌ 선생 모델 자체가 성능이 낮을 때
❌ 증류 학습 데이터가 충분하지 않을 때
❌ 단순히 비용만 줄이려면 → 양자화가 더 쉬움

마무리

지식 증류의 핵심은 이거예요.

"작은 모델이 혼자 배우면 멍청하지만, 큰 모델한테 배우면 영리해진다."

2026년에 가장 주목받는 건 CoT 증류예요. DeepSeek R1 같은 강력한 추론 모델의 사고 과정을 소형 모델에 이식해서, 폰에서 돌아가는 추론 모델을 만드는 게 현재 핫한 연구 방향이에요. 😄


 

반응형