딥러닝 모델 경량화를 위한 Knowledge Distillation 기본 개념 정리 | 지식 증류, KD

지식 증류의 기본 개념

수업 관련 생성 이미지 (by Diffusion)

Knowledge Distillation(KD, 지식 증류)는 고성능의 모델(Teacher)에서 지식을 전달 받아 상대적으로 간단한 모델(Student)을 학습시키는 방법입니다. 전달하는 지식의 종류나 Teacher 모델의 유형에 따라 분류가 가능합니다.

우선 Teacher 모델이 내부 구조, 파라미터를 공개 여부에 따라 White-box, Black-box, Gray-box로 구분됩니다. Black-box는 결과만 확인 가능한 경우며 White-box는 Llama 처럼 오픈소스로 모델의 내부 구조나 파라미터를 전부 알 수 있는 경우입니다. Gray-box는 그 중간으로 일부만 공개되어 있는 경우죠. 이러한 Teacher 모델의 특징에 의해 KD에 활용할 수 있는 정보의 종류가 달라지게 됩니다.

증류할 지식의 종류에 따른 분류는 다음과 같습니다. Response-based는 Black-box나 gray-box에 활용하고, 별도로 특징을 뽑을 수 있는 White-box는 Feature-based 접근 방식이 가능합니다.

• Response-based : Teacher 모델의 logit이나 output 같은 응답을 활용
• Feature-based : Teacher 모델의 중간 레이어의 feature, representation을 활용

Logit-based KD

Logit-based KD는 Response-based의 한 갈래로 정규화되지 않은 예측값인 logit을 활용하는 방법입니다. 오답에도 약간의 점수를 주기 때문에 모델의 강건성을 향상시키고, 연속된 형태이기이에 더 높은 엔트로피를 갖고 있으며, 하나의 클래스에서 다양한 인스턴스 차이나 클래스 간 유사도를 인식하는 등 장점이 있기 때문에 선호되는 방식입니다.

학습 방법

Teacher 모델은 Hard label(정답 클래스)로 학습을 진행하고 logit을 별도로 뽑아둡니다. 우리가 관심있는 Student 모델은 Hard label에 대한 학습(BCE loss 등)과 Teacher 모델의 logit 값을 Soft label로 해서(KL divergence) 두 개의 loss를 바탕으로 학습합니다. 두 손실을 섞는 방법은 하이퍼 파라미터로 실험을 통해 적절한 값을 찾는 것이 필요합니다.

여기서 클래스 간 logit의 차이가 극명하게 나거나 별로 나지 않을 수도 있습니다. 목적에 따라서 제대로 훈련이 되지 않을 수도 있으니 모든 클래스의 logit에 동일한 값(temperature, T)를 나누고 softmax를 적용하는 방법이 있습니다. 이는 기존 확률분포의 차이를 극명하게 또는 완만하게 될 수 있도록 변형을 가합니다.

• T < 1 : 차이를 극명하게 만들 경우
• T > 1 : 차이를 완만하게 만들 경우

Feature-based KD

Feature-based는 Teacher 모델이 특징을 Student 모델에 직접 활용하는 방식으로 모델의 구조, 파라미터 등을 알아야 사용이 가능합니다. (즉, White-box 경우)

학습 방법

Student 모델은 Teacher 모델보다 작기 때문에 차원을 맞춰주기 위한 regressor 층에 매핑합니다. 그리고 Teacher 모델과 Student 모델의 중간 층 피쳐맵에서 Mean Squared Error를 계산해서 훈련을 진행합니다. 여기서 중간 층의 피쳐맵을 사용하는 이유는 수행하려는 작업 데이터의 특징들이 층마다 다르게 분포하게 되는데, 중간 층이 양 극단의 중간 성격으로 섞여 있기 때문에 지식 증류에 효과적이기 때문입니다.

다른 방법들

Imitation Learning

앞서 언급한 방법들은 logit이나 특징들을 가져다 써야 가능한 방법이지만, ChatGPT처럼 거대하면서 좋은 성능을 가진 모델들은 Black-box인 경우가 많습니다. 따라서 이러한 거대 모델들을 지식 증류에 활용하기 위해 고안된 방법론으로 모방이 등장했습니다.

Imitation Learning은 LLM처럼 거대한 모델을 Teacher 모델로 삼고 질문에 대한 답변을 (질문-답변) 쌍을 데이터베이스로 저장합니다. 그리고 저장한 데이터 중 불필요한 내용들을 전처리하여 Student 모델의 학습에 활용하는 방식입니다. 좋은 성능을 가진 거대모델들이 등장한 현 시점에서 데이터 증강의 한 방식으로 활용하는 것이죠.

Multi-Teacher

여러 데이터 분석에서 성능 향상을 위해 흔하게 사용되는 방법으로 앙상블이 활용되곤 합니다. 여러 모델들의 결과를 종합해서 결과를 내는 방식인데, 이를 지식 증류에서도 활요이 가능합니다. 단일 모델의 한계를 극복할 수 있도록 여러 Teacher 모델을 앙상블하는 방식을 채택해 이를 학습하도록 하는 것입니다.

Cross-modal

최근 딥러닝 분야의 트렌드 중 하나는 여러 감각이 종합된 Multi-modal입니다. 이러한 트렌드에 맞춰 다른 modality(시각 - 청각, CV - NLP 등)를 가진 Teacher 모델에서 학습을 진행하는 것이 Cross-modal 입니다.

참고자료

[1] 조현수. "모델 최적화 및 경량화". boostcamp AI Tech.