지식 증류(Knowledge Distillation)

오늘부터는 지식 증류(Knowledge Distillation)에 대해서 다뤄보도록 하겠습니다. 지금까지 파인튜닝(Fine-tuning)이나 강화학습(Reinforcement Learning)에 대해 몇 번 다뤄봤지만, 지식 증류에 대해서는 처음 다뤄보게 되겠습니다.

최근 대규모 인공지능 모델(Large Language Models, LLM)은 자연어 처리, 이미지 인식 등 다양한 분야에서 높은 성능을 보이고 있습니다. 그러나 이러한 모델들은 수십억에서 수천억 개의 매개변수를 가지며, 이는 막대한 컴퓨팅 자원을 요구합니다. 학습 과정뿐만 아니라 실제 서비스 환경에서의 추론(Inference) 단계에서도 높은 연산 능력과 메모리가 필요합니다. 이로 인해 모델의 배포 및 유지보수 비용이 증가하고, 실시간 응답이 필요한 서비스나 제한된 자원을 가진 기기(예: 스마트폰, 엣지 디바이스)에 적용하는 데 제약이 큽니다.

이러한 문제를 해결하기 위한 여러 최적화 기술 중 하나가 지식 증류(Knowledge Distillation)입니다. 지식 증류는 크고 복잡한 모델의 지식을 작고 효율적인 모델로 효과적으로 전달하는 방법 중 하나입니다.

지식 증류의 개념

Knowledge Distillation는 대규모 인공지능 모델이 학습한 Knowledge를 더 작은 모델에 Distillation하는 과정입니다. 여기서 Knowledge는 모델이 데이터를 통해 습득한 패턴 인식 능력과 예측 결과를 의미합니다.

교사 모델과 학생 모델의 역할

교사(Teacher) 모델: 이 모델은 일반적으로 이미 학습이 완료되어 높은 성능을 발휘하는 대규모 모델입니다. 교사 모델은 방대한 데이터로부터 복잡한 특징과 패턴을 학습하며, 이를 통해 입력에 대한 상세하고 정교한 예측을 생성합니다. Llama 3.2 1B와 같은 모델이 교사 모델의 예시가 될 수 있습니다.
학생(Student) 모델: 교사 모델에 비해 매개변수 수가 적고 구조가 단순하여 연산 자원을 덜 소모하는 소형 모델입니다. 학생 모델의 목표는 교사 모델이 생성하는 예측 결과를 모방함으로써, 작은 크기에도 불구하고 교사 모델에 준하는 성능을 달성하는 것입니다.

지식 증류의 필요성

지식 증류는 다음과 같은 실용적인 이점으로 인해 모델 최적화의 중요한 방법론으로 사용됩니다.

효율적인 배포: 모델 크기가 작아지므로 스마트폰, 임베디드 시스템, 엣지 컴퓨팅 장치와 같이 하드웨어 자원이 제한적인 환경에서도 인공지능 모델을 구동할 수 있습니다.
추론 속도 향상: 소형 모델은 대형 모델에 비해 연산량이 적어 동일한 작업을 더 빠르게 처리합니다. 이는 실시간 응답이 요구되는 애플리케이션의 성능을 개선합니다.
운영 비용 절감: 모델 구동에 필요한 컴퓨팅 자원 및 에너지 소모가 줄어들어, 클라우드 컴퓨팅 비용이나 하드웨어 투자 비용을 절감할 수 있습니다.
데이터 프라이버시 및 보안: 온디바이스 AI 구현이 가능해지면, 민감한 사용자 데이터를 외부 서버로 전송할 필요 없이 기기 내부에서 처리할 수 있어 개인정보 보호에 기여합니다.

지식 증류의 작동 원리

지식 증류는 학생 모델이 교사 모델의 소프트 타겟을 모방하도록 학습시키는 방식으로 진행됩니다. 이 과정에서 특정 하이퍼파라미터와 손실 함수가 중요한 역할을 합니다.

온도(Temperature) 하이퍼파라미터

소프트 타겟의 확률 분포는 온도(Temperature)라는 하이퍼파라미터를 통해 조절될 수 있습니다. 이 온도는 모델의 로직(logits, 최종 소프트맥스 함수를 거치기 전의 원시 출력 값)에 적용되어 소프트맥스(softmax) 함수의 출력 분포를 변화시킵니다.

낮은 온도: 로직에 낮은 온도를 적용하면 확률 분포가 ‘뾰족(sharper)’해집니다. 즉, 가장 높은 로직 값을 가진 클래스에 확률이 집중되고, 나머지 클래스에는 매우 낮은 확률이 할당됩니다. 이는 정답에 가까운 ‘하드 타겟’과 유사한 효과를 냅니다.
높은 온도: 로직에 높은 온도를 적용하면 확률 분포가 ‘평평(smoother)’해집니다. 즉, 모든 클래스에 걸쳐 확률이 더 고르게 분산됩니다. 이 경우 교사 모델이 예측한 미세한 가능성까지 학생 모델에게 전달되어, 학생 모델이 교사 모델의 ‘불확실성’이나 ‘대안적 판단’에 대한 정보를 더 풍부하게 학습할 수 있도록 돕습니다.

손실 함수

학생 모델의 학습 목표는 교사 모델의 예측을 최대한 유사하게 모방하는 것입니다. 이를 위해 다음과 같은 손실 함수를 결합하여 사용합니다.

증류 손실(Distillation Loss): 이 손실은 학생 모델이 예측한 로직(또는 확률 분포)과 교사 모델의 소프트 타겟 간의 차이를 측정합니다. 주로 KL-다이버전스(Kullback-Leibler Divergence)가 사용됩니다.
학생 손실(Student Loss): 만약 학생 모델을 학습시킬 때 원본 데이터의 실제 정답 레이블(Hard Label)이 있다면, 이 레이블에 대한 학생 모델의 예측 오차를 계산하는 손실입니다. 예를 들어, 분류 문제에서는 크로스 엔트로피(Cross-Entropy) 손실이 사용됩니다.

효과적인 증류를 위한 고려사항

지식 증류의 성능은 여러 요인에 의해 영향을 받으며, 최적의 결과를 얻기 위해서는 다음과 같은 사항들을 고려해야 합니다.

교사 모델의 성능: 교사 모델의 예측 정확도와 학습된 지식의 품질은 증류된 학생 모델의 최종 성능을 결정하는 중요한 요소입니다. 교사 모델의 성능이 낮다면 학생 모델의 성능도 제한될 수 있습니다.
학생 모델의 용량: 학생 모델의 구조적 복잡도와 매개변수 수는 교사 모델의 지식을 얼마나 효과적으로 흡수할 수 있는지를 결정합니다. 너무 작은 학생 모델은 교사 모델의 복잡한 지식을 충분히 학습하지 못할 수 있습니다.
데이터셋의 특성: 교사 모델이 소프트 타겟을 생성하고 학생 모델이 학습하는 데 사용되는 데이터셋의 규모, 다양성, 그리고 실제 배포 환경과의 유사성은 증류 학습의 성과에 직접적인 영향을 미칩니다.
하이퍼파라미터 튜닝: 온도(Temperature), 증류 손실 가중치(alpha), 학습률(Learning Rate) 등 지식 증류 관련 하이퍼파라미터들은 학생 모델의 학습 안정성과 최종 성능에 큰 영향을 미칩니다. 이들을 체계적으로 탐색하고 최적화하는 과정이 필요합니다.
증류 방법론의 선택: 소프트 타겟 기반의 증류 외에도, 교사 모델의 중간 계층 특징(feature)을 모방하거나, 어텐션 메커니즘을 모방하는 등 다양한 지식 증류 기법이 존재합니다. 특정 문제와 모델 아키텍처에 가장 적합한 증류 방법론을 선택하는 것이 중요합니다.

결과적으로,

위와 같은 이론적인 배경을 염두하고, 실제 실험을 해보기 위해서 교사 모델과 학생 모델을 선별해야 하겠습니다. 저사양 PC라는 환경적인 제약조건을 고려해 본다면 가능한 교사 모델의 범위가 그리 넓지 못하죠. 우선 모델부터 선정하고 나서 실제 지식 증류를 실험해 보도록 해야 겠습니다.

지식 증류(Knowledge Distillation) – AI 모델 경량화를 위한 방법 #1