본문 바로가기

IT

지식 증류(knowledge distillation)란 무엇인가?

728x90
반응형
728x170

AI 를 사용하다보면 정확도는 마음에 드는데… 모델이 너무 무거워서 서버가 감당하기 어려운 경우가 있습니다.


여기서 등장하는 핵심 용어가 바로 지식 증류(knowledge distillation)입니다. 한 줄 요약을 하자면, 거대한 모델(선생, teacher)이 알고 있는 요령과 뉘앙스를 더 작은 모델(학생, student)에게 전수해서, 가볍고 빠르게 만들면서도 성능을 최대한 유지하는 모델 경량화/모델 압축 방법입니다.

 

이 글에서는 지식 증류의 원리를 일반인도 이해할 수 있게 예시로 풀어 설명하고, 개발 관점에서 “대체 어떻게 학습시키는지”까지 정확하게 정리해 보겠습니다.


왜 지식 증류가 필요할까

대형 딥러닝 모델은 보통 성능(정확도, F1, BLEU 등)을 끌어올리기 위해 파라미터 수와 연산량을 크게 늘립니다. 문제는 그 다음입니다.

  • 추론(inference) 비용 증가: 응답 시간이 늘고, 서버 비용이 커집니다.
  • 온디바이스(on-device) 적용 난이도: 모바일, IoT, 브라우저 같은 환경에 넣기 어렵습니다.
  • 전력/발열 문제: 특히 엣지 디바이스에서 치명적입니다.
  • 배포 운영 부담: 모델 크기가 커질수록 롤백/캐시/메모리 전략도 복잡해집니다.

지식 증류는 이럴 때 “큰 모델의 성능 감각은 최대한 유지하면서, 더 작은 모델을 만드는” 실전적인 선택지입니다.


지식 증류의 핵심 아이디어는 정답만 주지 말고, 풀이 느낌을 같이 준다는 겁니다.

일반적인 지도학습은 정답 라벨이 딱 하나입니다.

  • 입력 이미지: 어떤 동물 사진
  • 정답 라벨: “고양이” (1), 나머지(0)

하지만 선생 모델은 더 많은 정보를 알고 있습니다. 예를 들어 선생 모델이 내놓는 확률 분포가 다음과 같다고 해봅시다.

  • 고양이 0.55
  • 여우 0.25
  • 개 0.15
  • 늑대 0.05

정답만 보면 “고양이”로 끝인데, 이 분포에는 중요한 힌트가 숨어 있습니다.
“이 고양이는 여우랑 좀 헷갈릴 만해. 귀 모양이나 털 색이 그쪽 계열이야.” 같은 미세한 지식이죠.

 

지식 증류는 학생 모델이 이 분포(부드러운 라벨, soft targets)를 따라 배우게 해서, 단순 암기(하드 라벨)보다 더 풍부한 학습 신호를 받도록 합니다. 이 숨은 뉘앙스를 종종 dark knowledge라고 부르기도 합니다.


“온도(Temperature)”가 왜 나오지? 뜨거우면 더 솔직해진다

선생 모델의 출력은 보통 소프트맥스(softmax)로 확률로 바뀝니다. 그런데 큰 모델은 종종 너무 확신에 차서(확률이 0.999 같은 식으로) 분포가 뾰족해집니다. 그러면 “미세한 뉘앙스”가 잘 안 보입니다.

 

그래서 지식 증류에서는 소프트맥스에 온도(Temperature) T를 넣어 분포를 더 부드럽게 만듭니다.

  • T가 커질수록: 확률 분포가 평평해져서, 2등/3등 후보가 더 잘 드러납니다.
  • T가 1이면: 일반적인 소프트맥스와 동일합니다.

이 과정이 지식 증류에서 “증류” 느낌을 가장 잘 살리는 부분입니다. 쓴맛(노이즈)을 줄이고 향(관계 정보)을 남기는 느낌이라고 보면 됩니다.


실제 학습은 어떻게 하나: 손실 함수(loss)를 두 개 섞는다

대부분의 지식 증류는 다음 두 가지를 함께 최소화합니다.

  1. 정답 라벨을 맞히는 손실(일반적인 Cross Entropy)
  2. 선생의 확률 분포를 따라가는 손실(보통 KL Divergence)

개념적으로는 이렇게 섞습니다.

  • 전체 손실 = (1 - α) * CE(하드 라벨, 학생) + α * KL(선생 소프트 분포, 학생 소프트 분포)

여기서 α는 “선생 말 얼마나 믿을래?” 비율입니다.


또한 온도 T를 쓰는 경우, 관례적으로 안정적인 학습을 위해 KL 항에 T^2를 곱해 스케일을 맞추는 구현이 흔합니다(논문/레퍼런스에서 자주 등장하는 패턴).

파이토치 스타일 지식 증류 예시 코드

아래는 분류 문제에서 가장 흔한 형태의 지식 증류 손실 구현 예시입니다.

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    # 1) 하드 라벨 손실
    ce = F.cross_entropy(student_logits, labels)

    # 2) 소프트 타깃(선생 분포) 따라가기: KL divergence
    # teacher/student 모두 temperature 적용
    teacher_probs = F.softmax(teacher_logits / T, dim=1)
    student_log_probs = F.log_softmax(student_logits / T, dim=1)

    kl = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")

    # 관례적으로 T^2를 곱해 gradient scale을 보정
    kd = kl * (T * T)

    return (1 - alpha) * ce + alpha * kd

 

이 코드에서 중요한 포인트는 세 가지입니다.

  • 선생은 보통 학습 고정(freeze) 상태로 둡니다.
  • KL은 “학생이 선생 분포를 얼마나 잘 따라가나”를 측정합니다.
  • α, T는 성능과 안정성에 영향을 크게 줘서 실험이 필요합니다.

지식 증류에도 종류가 있다: 무엇을 증류할 건가?

지식 증류라고 해서 항상 “최종 확률(logits)”만 따라 하는 건 아닙니다. 중간 표현까지 증류하면 더 좋아지는 경우도 많습니다.

분류 무엇을 따라 배우나 장점 단점
로짓(출력) 증류 최종 logits/확률 분포 구현이 단순, 적용 범위 넓음 중간 표현 지식은 놓칠 수 있음
특징(feature) 증류 중간 레이어 feature 성능 개선 여지가 큼 레이어 매칭/차원 정렬 필요
관계(relation) 증류 샘플 간/토큰 간 관계 구조적 지식 전수에 강점 설계가 복잡해질 수 있음
자기(self) 증류 같은 모델/이전 스냅샷 큰 선생 없이도 개선 가능 개선 폭이 제한적일 수 있음

 

실무에서는 “logits 증류 + 필요하면 feature 증류” 조합이 가장 자주 보입니다.


다른 모델 경량화 방법과 비교: 증류는 어디쯤에 있나?

모델을 가볍게 만드는 방법은 여러 가지가 있고, 지식 증류는 그중에서도 “학습 전략”에 가깝습니다.

방법 핵심 아이디어 장점 단점
지식 증류 큰 모델의 출력/표현을 작은 모델이 모방 성능 유지에 유리, 구조 제약 적음 선생 학습 비용, 하이퍼파라미터 튜닝 필요
프루닝(pruning) 중요하지 않은 가중치/채널 제거 기존 모델을 직접 줄일 수 있음 하드웨어 최적화/재학습이 까다로울 수 있음
양자화(quantization) FP32 대신 INT8 등으로 압축 추론 속도/메모리 효율에 매우 강함 정확도 손실 가능, 지원 런타임 영향
경량 아키텍처 MobileNet 등 가벼운 구조 사용 배포 친화적, 예측 가능 설계 선택이 성능 상한을 제한할 수 있음

 

현실적인 전략은 보통 “지식 증류로 학생 모델 성능을 끌어올리고, 필요하면 양자화까지 얹어서” 배포 최적화를 마무리하는 흐름입니다.


예시로 한 번 더: 왜 “소프트 라벨”이 도움이 되나

이해를 위해 아주 단순한 상황을 상상해봅시다.

  • 학생은 “고양이 vs 개”만 배우는 중
  • 그런데 데이터엔 늑대, 여우, 코요테처럼 애매한 애들이 섞여 있습니다

정답 라벨만 보면 고양이/개로 강제 분류해야 하니 학생은 “애매함”을 학습할 기회가 없습니다. 반면 선생의 소프트 분포는 이런 정보를 줍니다.

  • “이건 개 쪽이긴 한데 늑대랑도 닮았어”
  • “이건 고양이긴 한데 여우 느낌이 있어”

학생은 이 신호를 통해 결정 경계를 더 매끈하게 만들고, 데이터가 조금 흔들리는 상황에서도 일반화 성능이 좋아질 수 있습니다. 지식 증류가 단순 압축이 아니라 “학습을 더 똑똑하게 하는 방법”으로도 자주 언급되는 이유입니다.


지식 증류, 해보고 망하지 않으려면

  1. 선생 모델의 품질이 먼저다
    선생이 애매하게 틀리면 학생은 그 애매함까지 훌륭하게 복제합니다. 선생 성능이 일정 수준 이상이어야 효과가 큽니다.
  2. 학생 용량(capacity)을 현실적으로 잡아라
    학생이 너무 작으면 따라 하고 싶어도 못 따라 합니다. “가능한 최소 크기”를 한 번에 찾기보다 단계적으로 줄이는 접근이 안전합니다.
  3. T(온도)와 α(혼합 비율)는 반드시 튜닝 대상
    정답 라벨(CE)과 선생 신호(KL)의 균형이 무너지면 학습이 불안정해지거나 성능이 떨어집니다.
  4. 데이터가 부족하면 무라벨 데이터가 오히려 도움이 된다
    지식 증류는 선생이 라벨러 역할을 할 수 있어서, 무라벨 데이터에도 적용하기 좋습니다(단, 도메인 차이가 크면 주의).
  5. 평가 지표는 “정확도 단독”이 아니라 “정확도-지연시간-비용” 3종 세트
    지식 증류의 목적은 대개 배포 최적화입니다. 정확도 0.3% 올리자고 지연시간이 2배면 목적을 잊은 셈입니다.

지식 증류는 AI 시대의 “현실적 마법”이다

지식 증류(knowledge distillation)는 큰 모델을 작게 만드는 모델 압축/모델 경량화의 대표 주자입니다. 핵심은 단순히 정답을 베끼는 게 아니라, 선생 모델의 확률 분포와 뉘앙스를 학생에게 전달해 성능을 최대한 유지하며 가볍게 만드는 것입니다.

  • 대형 모델의 배포 비용이 부담될 때
  • 모바일/엣지/온디바이스 AI를 만들고 싶을 때
  • 서버 응답 시간을 줄이면서 품질을 유지하고 싶을 때

지식 증류는 “모델 다이어트” 분야에서 가장 실전적인 선택지 중 하나가 됩니다. 다만 다이어트가 그렇듯, 식단(T, α) 조절과 운동(학생 구조 선택, 데이터, 평가)이 함께 가야 결과가 예쁩니다.

728x90
반응형
그리드형