지식 증류(knowledge distillation)는 모델 증류라고도 불리며, Machine learning에서 크고 복잡한 모델(교사)에서 더 작고 단순한 모델(학생)로 지식을 전이하는 기술입니다. 목표는 더 큰 모델의 예측 정확도를 유지하면서도 평가 시 계산 효율성이 더 높은 컴팩트한 모델을 만드는 것입니다. 이는 특히 모바일 기기와 같은 자원이 제한된 하드웨어에 모델을 배포할 때 유용하며, 큰 모델을 실행하는 계산 비용이 prohibitive한 경우에 중요합니다. 이 과정은 큰 모델의 출력, 특히 소프트 확률이 최종 클래스 레이블보다 더 풍부한 정보를 포함하며, 이 정보를 사용해 작은 모델을 효과적으로 훈련할 수 있다는 아이디어를 활용합니다.
지식 증류는 모델 압축과는 구별되며, 모델 압축은 새 모델을 훈련하지 않고 기존 모델 자체의 크기를 줄이는 방법(예: 파라미터당 비트 수 감소)을 의미합니다. 모델 압축은 일반적으로 아키텍처와 파라미터 수를 유지하는 반면, 증류는 새롭고 더 작은 모델을 훈련하는 것을 포함합니다. 이 기술은 객체 감지, 음향 모델, 자연어 처리, 그리고 최근에는 비격자 데이터를 위한 그래프 신경망 등 다양한 분야에서 성공적으로 적용되었습니다.
역사적 배경
지식 증류의 개념은 Deep learning 및 Neural network 연구의 더 넓은 분야에 뿌리를 두고 있습니다. 모델 간 지식 전이 아이디어는 다양한 형태로 탐구되었지만, 현대적 공식화는 종종 2015년 Geoffrey Hinton, Oriol Vinyals, Jeff Dean의 논문 "Distilling the Knowledge in a Neural Network"로 귀속됩니다. 이 연구는 소프트맥스 함수에서 온도를 사용하여 교사의 출력을 부드럽게 함으로써 더 효과적인 지식 전이를 가능하게 했습니다. 이후 지식 증류는 머신러닝 도구 상자의 표준 기술이 되었으며, 수많은 변형과 응용이 개발되었습니다.
핵심 원리
지식 증류의 기본 원리는 매우 깊은 Neural network 또는 모델 앙상블과 같은 큰 모델이 작은 모델보다 더 높은 지식 용량을 가진다는 것입니다. 그러나 이 용량이 완전히 활용되지 않을 수 있으며, 큰 모델을 평가하는 것은 용량 사용 정도와 관계없이 계산 비용이 높습니다. 증류는 큰 모델에 인코딩된 지식을 유효성 손실 없이 작은 모델로 전이하는 것을 목표로 합니다. 작은 모델은 평가 비용이 낮기 때문에 덜 강력한 하드웨어에 배포할 수 있습니다.
핵심 통찰은 훈련된 모델의 소프트 출력이 클래스에 확률을 할당하며 모델의 내부 표현에 대한 정보를 포함한다는 것입니다. 예를 들어, 모델이 고양이 이미지를 올바르게 분류할 때 '고양이'에 높은 확률을 할당하지만 '개'나 '여우'에도 작은 비제로 확률을 할당할 수 있습니다. 이러한 작은 확률은 클래스 간의 미묘한 유사성과 관계를 인코딩하며, 하드 레이블(가장 가능성 있는 클래스)만 사용하면 손실됩니다. 지식 증류는 이 소프트 정보를 활용하여 학생 모델에 더 풍부한 표현을 가르칩니다.
수학적 공식화
일반적인 분류 작업에서 신경망의 최종 레이어는 소프트맥스 함수를 사용하여 로짓(원시 점수)을 확률로 변환합니다. 표준 소프트맥스는 다음과 같습니다:
\( y_i(x) = \frac{e^{z_i(x)}}{\sum_j e^{z_j(x)}} \)
여기서 \( z_i(x) \)는 입력 \( x \)에 대한 클래스 \( i \)의 로짓입니다. 지식 증류에서는 온도 파라미터 \( t \)가 도입되어 소프트맥스를 다음과 같이 수정합니다:
\( y_i(x|t) = \frac{e^{z_i(x)/t}}{\sum_j e^{z_j(x)/t}} \)
더 높은 온도 \( t \)는 더 부드러운 확률 분포를 생성하며, 이는 확률이 클래스 간에 더 고르게 분포됨을 의미합니다. 이 부드러운 분포는 클래스 간 관계에 대한 더 많은 정보를 드러내며, 보조 클래스에 대한 모델의 신뢰도가 더 두드러집니다.
증류 중 학생 모델은 원래 훈련 데이터 또는 새롭고 잠재적으로 레이블이 없는 데이터일 수 있는 전이 세트에서 훈련됩니다. 손실 함수는 일반적으로 높은 온도에서 계산된 학생 출력과 교사 출력 간의 교차 엔트로피입니다. 단일 입력 \( x \)에 대한 손실은 다음과 같습니다:
\( E(x|t) = -\sum_i \hat{y}_i(x|t) \log y_i(x|t) \)
여기서 \( \hat{y}_i(x|t) \)는 교사 출력이고 \( y_i(x|t) \)는 학생 출력입니다. 높은 온도를 사용하면 출력의 엔트로피가 증가하여 학생이 학습할 더 많은 정보를 제공하고 서로 다른 레코드 간 그래디언트의 분산을 줄여 더 높은 학습률을 허용합니다.
훈련 과정
지식 증류의 훈련 과정은 일반적으로 여러 단계를 포함합니다. 먼저, 큰 교사 모델이 표준 기술을 사용하여 원본 데이터 세트에서 훈련됩니다. 이 모델은 단일 큰 네트워크 또는 모델 앙상블일 수 있습니다. 교사가 훈련된 후, 소프트 출력이 종종 높은 온도에서 전이 세트에 대해 생성됩니다. 그런 다음 학생 모델은 위에서 설명한 교차 엔트로피 손실을 사용하여 이러한 소프트 출력을 모방하도록 훈련됩니다.
일부 경우, 손실은 사용 가능한 경우 학생 출력과 실제 레이블 간의 교차 엔트로피로 보강됩니다. 이 조합은 학생이 교사의 지식과 실제 대상 레이블을 모두 학습하도록 돕습니다. 온도는 일반적으로 훈련 중에 높게 시작하여 점차 1로 감소시키며, 소프트 타깃 학습에서 하드 레이블로 정제하는 것으로 전환합니다.
응용 분야
지식 증류는 Artificial intelligence의 여러 분야에서 성공적으로 적용되었습니다. 컴퓨터 비전에서는 객체 감지 및 이미지 분류에 사용되어 모델이 엣지 기기에서 실행될 수 있게 합니다. 음성 인식에서는 음향 모델이 증류를 통해 지연 시간과 메모리 사용량을 줄입니다. Natural language processing에서는 증류를 사용하여 Transformer (architecture) 기반 모델과 같은 큰 모델을 Large language model을 포함한 더 작은 버전으로 압축하여 모바일 기기 또는 실시간 응용 프로그램에 배포할 수 있습니다. 예를 들어, Generative AI 시스템과 같은 큰 모델은 특정 작업을 위해 더 작은 모델로 증류되어 원래 성능의 많은 부분을 유지할 수 있습니다.
최근에는 지식 증류가 사회 네트워크나 분자 구조와 같은 비격자 데이터를 처리하는 그래프 신경망으로 확장되었습니다. 이 확장은 다양한 데이터 유형과 모델 아키텍처에 걸친 기술의 다재다능함을 보여줍니다.
변형 및 확장
특정 문제를 해결하기 위해 여러 지식 증류 변형이 개발되었습니다. 주목할 만한 변형 중 하나는 역방향 지식 증류로, 작은 모델에서 큰 모델로 지식을 전이합니다. 이는 덜 일반적이지만 특정 데이터에 대해 작은 모델이 훈련되었고 큰 모델의 해당 데이터 성능을 향상시키는 것이 목표인 시나리오에서 유용할 수 있습니다.
다른 확장에는 주의 기반 증류(학생이 교사의 주의 맵을 모방)와 특징 기반 증류(중간 표현을 타깃으로 사용)가 포함됩니다. 이러한 방법은 교사로부터 더 상세한 정보를 포착하여 지식 전이의 충실도를 향상시키는 것을 목표로 합니다.
장점과 한계
지식 증류의 주요 장점은 정확도 손실 없이 자원이 제한된 기기에 고성능 모델을 배포할 수 있다는 것입니다. 이는 모바일 앱, 임베디드 시스템, 실시간 추론과 같은 응용 프로그램에 중요합니다. 증류는 또한 더 빠른 추론 시간과 낮은 에너지 소비를 초래할 수 있습니다.
그러나 한계도 있습니다. 훈련 과정은 잘 훈련된 교사가 필요하며, 이를 얻는 데 계산 비용이 높을 수 있습니다. 또한 학생 모델은 특히 용량 격차가 너무 큰 경우 교사와 동일한 성능 수준을 항상 달성하지 못할 수 있습니다. 온도와 전이 세트의 선택도 신중한 튜닝이 필요합니다.
다른 기술과의 관계
지식 증류는 종종 다른 모델 최적화 기술과 함께 사용됩니다. 예를 들어, Model Pruning은 증류 후 학생 모델에 적용되어 크기를 더 줄일 수 있습니다. Data Augmentation은 전이 세트를 확장하여 학생의 일반화를 개선하는 데 사용될 수 있습니다. 증류는 또한 Curriculum Learning과 관련이 있으며, 소프트 타깃은 점진적 학습의 한 형태로 볼 수 있습니다.
Deep learning의 맥락에서 증류는 한 모델에서 다른 모델로 지식을 전이하는 Transfer learning의 한 형태입니다. 또한 Reinforcement Learning from AI Feedback (RLAIF)(AI 피드백에서의 강화 학습)와 밀접한 관련이 있지만, 후자는 모델을 인간 선호도에 맞추는 데 초점을 둡니다.
미래 방향
모델이 계속 커짐에 따라, 특히 Large language model 분야에서 지식 증류는 더욱 중요해질 가능성이 높습니다. 연구자들은 OpenAI, Anthropic, Google DeepMind와 같은 기관에서 개발된 대규모 모델에서 더 작고 효율적인 모델로 지식을 증류하는 방법을 탐구하고 있습니다. 이는 고급 AI 기능에 대한 접근을 민주화하는 데 중요하며, 작은 조직과 개인 개발자가 광범위한 계산 자원 없이 강력한 모델을 사용할 수 있게 합니다.
또 다른 연구 영역은 온라인 증류로, 교사와 학생이 동시에 훈련되며, 협력 증류는 여러 모델이 서로 학습합니다. 이러한 접근 방식은 증류 과정의 효율성과 효과를 개선하는 것을 목표로 합니다.
요약하면, 지식 증류는 모델 압축 및 지식 전이를 위한 강력한 기술입니다. 큰 모델의 소프트 출력을 활용하여 정확하고 효율적인 컴팩트 모델을 생성할 수 있게 하며, 현대 머신러닝 배포 전략의 초석이 됩니다.