자기 증류(Self-distillation)는 단일 신경망이 교사(teacher)와 학생(student)의 역할을 동시에 수행하는 지식 증류의 한 형태이다. 더 크거나 복잡한 모델이 더 작은 모델을 안내하는 기존의 증류 방식과 달리, 자기 증류는 두 역할 모두에 동일한 아키텍처를 사용한다. 모델은 훈련 데이터에 대해 소프트 예측(확률 분포)을 생성하고, 이러한 예측은 실제 레이블과 함께 추가 훈련 목표로 사용된다. 이 과정은 반복적으로 또는 다세대 방식으로 적용될 수 있으며, 각 세대의 모델은 이전 세대의 출력에서 학습한다. 이 기법은 외부 교사 모델이나 추가 레이블 데이터 없이도 모델의 정확도, 보정(calibration), 견고성을 향상시키는 것으로 나타났다.
이 개념은 2010년대 중반에 대중화된 지식 증류라는 더 넓은 분야에서 등장했다. 표준 증류에서는 고용량 교사 네트워크가 소프트 레이블을 생성하고, 더 작은 학생 네트워크가 이를 모방하여 클래스 유사성과 불확실성에 대한 지식을 전달한다. 자기 증류는 모델이 자신의 지식을 증류함으로써 별도의 교사가 필요 없게 한다. 이는 더 큰 교사를 훈련하는 계산 비용을 피하고 교사-학생 설정에 흔히 수반되는 아키텍처 제약을 제거하기 때문에 특히 매력적이다. 연구자들은 단일 모델이라도 자신의 소프트 목표로 훈련할 때 성능 향상을 얻을 수 있다는 것을 발견했으며, 이 현상은 상당한 이론적·경험적 연구를 촉발했다.
메커니즘 및 훈련 절차
일반적인 자기 증류 설정에서 훈련 과정은 단계적으로 진행된다. 먼저 신경망은 하드 레이블(one-hot 인코딩된 실제 값)과 교차 엔트로피와 같은 손실 함수를 사용하여 표준 분류 작업으로 훈련된다. 초기 훈련 후, 각 훈련 예제에 대한 모델의 소프트맥스 출력이 기록된다. 이러한 소프트 출력은 확률 분포를 날카롭게 하거나 평평하게 만들기 위해 온도 스케일링되는 경우가 많으며, 교사의 지식 역할을 한다. 다음 단계에서는 동일한 모델 아키텍처가 재초기화되거나(또는 기존 가중치가 미세 조정되어) 이전 단계의 원래 하드 레이블과 소프트 목표를 모두 포함하는 결합 손실로 훈련된다. 온도 매개변수는 소프트 목표의 매끄러움을 제어하며, 더 높은 온도는 클래스 간 관계를 드러내는 더 부드러운 분포를 생성한다.
이 절차는 여러 세대에 걸쳐 반복될 수 있다. 예를 들어, 세대 0에서 훈련된 모델은 세대 1을 위한 소프트 레이블을 생성하고, 세대 1은 세대 2를 위한 소프트 레이블을 생성하는 식이다. 흥미롭게도 성능은 처음 몇 세대 동안 종종 향상되다가 정체되거나 약간 저하된다. 이러한 이득은 소프트 목표가 하드 레이블이 제공하지 않는 클래스 유사성에 대한 정보를 인코딩하기 때문에 모델이 더 정규화된 결정 경계를 학습하는 데 기인한다. 훈련 목표는 일반적으로 교차 엔트로피 손실과 증류 손실을 결합하며, 교사의 소프트 목표 영향력을 균형 있게 조정하는 하이퍼파라미터로 가중치가 부여된다.
이론적 설명
자기 증류가 작동하는 이유를 설명하기 위해 여러 이론이 제안되었다. 한 가지 주요 설명은 소프트 목표의 암묵적 정규화 효과와 관련이 있다. 하드 레이블은 모델이 모든 확률 질량을 올바른 클래스에 할당하도록 강제하여 과신적인 예측과 과적합을 초래할 수 있다. 반면 소프트 목표는 모델이 유사한 클래스 간에 확률을 분배하도록 장려하며, 이는 레이블 스무딩의 한 형태로 작용한다. 이는 노이즈에 대한 모델의 민감도를 줄이고 일반화를 개선한다.
또 다른 연구 방향은 자기 증류를 "암흑 지식(dark knowledge)" 개념과 연결한다. 이는 잘못된 클래스 간의 상대적 확률이 유용한 정보를 담고 있다는 아이디어이다. 모델이 자신의 소프트 목표로 훈련될 때, 이 암흑 지식을 효과적으로 증폭하여 더 견고한 특징 표현을 이끌어낸다. 또한 일부 연구자는 자기 증류를 엔트로피 정규화의 한 형태로 보며, 모델이 너무 확신하는 것에 대해 패널티를 부여하여 더 부드러운 결정 경계를 촉진한다고 본다.
최적화 관점에서 자기 증류는 부트스트래핑의 한 형태로 볼 수 있다. 모델은 자신의 예측을 반복적으로 개선하고, 각 세대는 다음 세대를 위한 약간 더 나은 목표를 제공한다. 이 과정은 기대-최대화(expectation-maximization)와 유사하며, 모델은 목표 생성과 목표에 맞추는 것을 번갈아 수행한다. 이론적 분석은 특정 조건에서 자기 증류가 잘 보정된 모델과 개선된 정확도에 해당하는 고정점으로 수렴함을 보여주었다.
딥러닝에서의 응용
자기 증류는 Deep learning의 다양한 영역에서 응용되고 있다. 컴퓨터 비전에서는 특히 레이블 데이터가 제한된 시나리오에서 이미지 분류 모델을 개선하는 데 사용되었다. 예를 들어, CIFAR-10 및 ImageNet과 같은 데이터셋에서 훈련된 모델은 자기 증류 시 일관된 정확도 향상을 보여주었다. 이 기법은 모델이 레이블되지 않은 데이터에 대한 소프트 레이블을 생성하여 효과적인 훈련 세트를 확장하는 준지도 학습에서도 효과적이다.
Natural language processing에서는 자기 증류가 Transformer (architecture) 기반 모델, 특히 Large language model에 적용되었다. 예를 들어, 언어 모델은 자신이 생성한 응답으로 훈련되어 일관성과 사실적 정확성을 개선할 수 있다. 이 접근 방식은 모델이 레이블되지 않은 말뭉치에 대해 자신의 예측을 반복적으로 레이블링하는 자기 훈련과 관련이 있다. 자기 증류는 또한 기계 번역과 같은 Sequence-to-Sequence (Seq2Seq) 작업에서도 사용되며, 모델 자신의 번역이 추가 훈련 데이터 역할을 한다.
분류 및 생성 외에도 자기 증류는 Model Pruning 및 압축에 적용되었다. 모델을 축소된 너비나 깊이로 자신에게 증류함으로써 연구자들은 원래 성능의 대부분을 유지하는 컴팩트한 모델을 달성할 수 있다. 이는 Qualcomm이나 Arm Holdings과 같은 제한된 계산 리소스를 가진 엣지 디바이스에 모델을 배포하는 데 특히 유용하다.
다른 기법과의 관계
자기 증류는 여러 다른 훈련 방법과 개념적 유사성을 공유한다. 훈련 예제를 의미 있는 순서로 제시하는 Curriculum Learning과 밀접한 관련이 있다. 자기 증류에서 소프트 목표는 모델의 이해를 점진적으로 개선하므로 일종의 커리큘럼으로 볼 수 있다. 또한 소프트 목표가 효과적으로 새로운 훈련 신호를 생성하므로 Data Augmentation과도 겹친다. 일부 연구자는 정규화를 더욱 강화하기 위해 자기 증류를 Dropout 및 Batch Normalization과 결합했다.
이 기법은 모델이 자신 또는 다른 모델의 피드백을 사용하여 개선하는 Reinforcement Learning from AI Feedback (RLAIF)(AI 피드백으로부터의 강화 학습)와도 연결된다. 자기 증류에서 피드백은 소프트 보상 신호 역할을 하는 모델 자신의 확률 분포이다. 또한 자기 증류는 교사와 학생이 동일한 아키텍처를 공유하는 지식 증류의 특수한 경우로 볼 수 있으며, 일반적인 이종 설정과는 대조적이다.
경험적 발견 및 벤치마크
경험적 연구는 자기 증류가 여러 아키텍처와 데이터셋에서 테스트 정확도를 일관되게 향상시킨다는 것을 입증했다. 예를 들어, CIFAR-100에서 자기 증류로 훈련된 Residual Network (ResNet)는 하드 레이블만으로 훈련된 기준선보다 1-2%의 정확도 향상을 달성한 것으로 보고되었다. ImageNet에서는 이득이 더 작지만 여전히 유의미하며, 종종 0.5-1% 정도이다. 모델이 과매개변수화되거나 데이터셋이 작을 때 개선이 더 두드러지며, 이는 자기 증류가 정규화기 역할을 함을 시사한다.
보정은 자기 증류가 뛰어난 또 다른 영역이다. 자기 증류로 훈련된 모델은 일반적으로 더 낮은 기대 보정 오류를 가지며, 이는 예측 확률이 실제 정확도를 더 잘 반영함을 의미한다. 이는 의료 진단이나 자율 주행과 같이 결정 신뢰도가 중요한 응용 분야에서 중요하다. Tesla이나 Waymo 시스템에서 잘 보정된 모델은 과신 오류의 위험을 줄인다.
자기 증류는 또한 적대적 공격에 대한 견고성을 향상시키는 것으로 나타났다. 소프트 목표로 유도된 부드러운 결정 경계는 작은 섭동이 예측을 뒤집기 어렵게 만든다. 이는 비전 및 언어 모델 모두에서 관찰되었지만, 전용 적대적 훈련 방법에 비해 효과는 완만하다.
변형 및 확장
특정 문제를 해결하기 위해 여러 자기 증류 변형이 제안되었다. "본 어게인 네트워크(born-again networks)"라는 한 변형은 교사와 아키텍처가 동일하지만 무작위 초기화가 다른 학생 모델을 훈련하는 것을 포함한다. 학생은 교사의 소프트 목표로 훈련되며, 이 과정은 여러 세대에 걸쳐 반복된다. 이 접근 방식은 일관된 개선을 보여주었으며, 이후 세대가 이전 세대보다 더 나은 성능을 내는 경우도 있다.
또 다른 확장은 "온라인 자기 증류"로, 별도의 단계가 아닌 동일한 훈련 실행 중에 모델이 자신의 예측을 증류한다. 이는 모델 출력의 이동 평균을 유지하거나 공유 분류기를 사용하여 달성된다. 온라인 방법은 여러 훈련 패스를 피하므로 더 효율적이지만 덜 안정적일 수 있다.
자기 증류는 또한 트랜스포머의 Multi-Head Attention 메커니즘과 결합되었다. 예를 들어, 모델은 자신의 어텐션 헤드에서 단일 헤드로 지식을 증류하여 해석 가능성과 성능을 개선할 수 있다. 이는 기계 번역 및 요약에 사용되는 Encoder-Decoder Architecture 아키텍처에서 특히 관련이 있다.
문제 및 한계
이점에도 불구하고 자기 증류는 보편적으로 효과적이지 않다. 일부 경우, 특히 매우 작은 모델이나 노이즈가 많은 데이터셋에서는 이득이 무시할 수 있거나 심지어 부정적일 수 있다. 이 기법은 또한 온도와 증류 가중치와 같은 추가 하이퍼파라미터를 도입하며, 이는 튜닝이 필요하다. 부적절한 설정은 과소적합이나 과도한 스무딩으로 이어져 모델이 너무 보수적이 되어 판별력을 잃을 수 있다.
또 다른 한계는 계산 오버헤드이다. 자기 증류는 별도의 교사 훈련을 피하지만, 여전히 여러 훈련 패스 또는 소프트 목표 저장이 필요하며, 이는 대규모 데이터셋에서 메모리 집약적일 수 있다. 수십억 개의 매개변수를 가진 Large language model의 경우 모든 훈련 예제에 대한 소프트 목표를 생성하고 저장하는 것은 비현실적이다. 연구자들은 데이터 하위 집합 사용이나 소프트 목표 압축과 같은 근사치를 제안했지만, 이는 추가 복잡성을 도입한다.
자기 증류에 대한 이론적 이해는 여전히 불완전하다. 경험적 결과는 유망하지만, 관찰된 모든 현상을 설명하는 통합 프레임워크는 없다. 일부 연구는 자기 증류가 초기 모델에 존재하는 편향을 증폭하여 최적이 아닌 솔루션으로 이어질 수 있음을 보여주었다. 이는 공정성이 중요한 응용 분야에서 특히 우려된다.
향후 방향
진행 중인 연구는 자기 증류를 더 효율적이고 견고하게 만드는 방법을 탐구하고 있다. 한 방향은 훈련 중 모델의 신뢰도에 따라 조정되는 적응형 온도 일정의 개발이다. 또 다른 방향은 초점 손실이나 레이블 스무딩 변형과 같은 보정을 위해 특별히 설계된 Loss Functions와 자기 증류의 통합이다.
Generative AI의 맥락에서 자기 증류는 생성된 텍스트의 사실적 일관성을 개선하는 방법으로 조사되고 있다. 모델이 자신의 응답을 증류함으로써 환각을 피하는 법을 배울 수 있다. 이는 고위험 환경에 배포되는 OpenAI 및 Anthropic 모델과 특히 관련이 있다.
자기 증류는 또한 더 다양한 훈련 신호를 생성하기 위해 Data Augmentation 기법과 결합되고 있다. 예를 들어, 모델은 동일한 입력의 증강된 버전에 대한 소프트 목표를 생성하여 일관성을 강화하고 불변성을 개선할 수 있다. 이 접근 방식은 준지도 및 자기 지도 학습 패러다임에서 유망한 결과를 보여주었다.
마지막으로 자기 증류의 이론적 기초를 이해하려는 관심이 증가하고 있다. 연구자들은 정보 이론과 최적화 도구를 사용하여 자기 증류가 언제, 왜 도움이 되는지 특성화하고 있다. 이는 기법을 적용할 시기와 하이퍼파라미터를 설정하는 방법에 대한 원칙적인 지침으로 이어질 수 있다.
결론
자기 증류는 모델이 스스로 학습할 수 있다는 단순하면서도 강력한 아이디어를 나타낸다. 자신의 예측을 소프트 목표로 사용함으로써 외부 감독 없이 더 나은 일반화, 보정, 견고성을 달성한다. 이 기법은 작은 컨볼루션 네트워크에서 대형 트랜스포머에 이르기까지 다양한 작업과 아키텍처에서 검증되었다. 계산 비용과 이론적 이해 측면에서 문제가 남아 있지만, 자기 증류는 머신러닝 도구 상자에서 필수 요소로 남을 가능성이 높다. 모델이 계속해서 크기와 복잡성이 증가함에 따라, 자기 증류는 추가 데이터나 아키텍처 변경 없이 성능을 개선할 수 있는 확장 가능한 방법을 제공한다.