[AI System] Knowledge Distillation
"Knowledge Distillation"
[Knowledge Distillation]
· 이미 학습된 큰 Model(Teacher)과, 새로 학습할 작은 Model(Student) 존재 → Training Data를 사용해서 Student Model을 학습
** 단순히 정답 Label만 보고 배우는 것이 아니라, Teacher Model의 판단 기준과 노하우를 Distill하여 Student Model에게 전송
· 학습 과정에서는 큰 Teacher Model 활용 / 추론 과정에서는 작은 Student Model만 사용
∴ 새로운 데이터가 들어왔을 때 작은 Model이 출력 값 생성 가능 → 저장 공간, 계산량, 지연 시간↓
[Different Forms of Distillation Losses]
▣ Logit Distillation
· Student Model이 Teacher Model과 비슷한 Logit 또는 Softmax Prediction을 내도록(0.731보다 ↑ / 0.269보다 ↓) 학습
→ 정답만 배우는 것이 아니라 Teacher의 확신 정도와 Class 간 상대적 관계까지 학습
· Teacher Model은 (55% vs 45%)로 애매하게 판단하고, Student Model은 (90% vs 10%)로 확신 있게 판단
→ Student가 Teacher의 불확실성, Class 간 상대적 유사도, 판단 방식 학습을 위해 Teacher에 맞추는 것이 좋음
** But, Teacher가 항상 완벽 X → Teacher의 확률 분포를 무조건 따라가면 Teacher의 실수까지 학습 가능
∴ Student를 학습할 때는 KD Loss만 사용하지 않고, 실제 정답 Label을 이용하는 Cross-Entropy Loss도 함께 사용 필요
· C : Number of Classes
· ti : One-Hot Encoded Ground Truth(Label)
· piS : Softmax Output of the Student Model
· piT : Softmax Output of the Teacher Model
- 일반적인 학습 방식 (Cross-Entropy Loss)
· Student Model이 Training Data를 입력 받고 Softmax Output 출력
· 정답 Label은 One-Hot Vector로 표현
· Student의 예측 확률이 정답 Class에 가까워지도록 학습
→ CE Loss는 Student가 실제 정답을 잘 맞히도록 만드는 기본 손실 함수
- Logit Distillation의 전체 학습 구조
· Student Model은 CE Loss로 실제 정답을 배우고, 동시에 KD Loss로 Teacher Model의 Softmax Output을 따라 학습
→ 최종 Loss L = LCE + λ · LKD (λ : Teacher Model을 얼마나 강하게 따라갈지 조절하는 값)
> Effect of Each Loss
· CE Loss : Student의 Softmax가 정답 Class는 1에 가깝고 나머지는 0에 가까워지도록 만듦
· KD Loss : Student의 Softmax가 Teacher의 Softmax 분포와 비슷해지도록 만듦
→ CE : 정답 맞히기 / KD : Teacher처럼 판단하기
<정답 = Dog>
· CE Loss : Dog 확률↑ / 다른 Class 확률↓
· KD Loss : Teacher가 가진 확률 분포를 따라감
<Teacher Model의 오답>
· Teacher의 예측이 잘못되어 있으면 KD Loss만 사용할 경우 Student도 잘못된 방향으로 학습** Teacher의 판단 방식을 학습하는 KD Loss와 실제 정답 Label을 사용하는 CE Loss 함께 사용 필요
> Temperature in Softmax
· 일반 Softmax는 Logit을 확률로 바꾸는데, 여기에 τ를 넣으면 출력 분포의 날카로움을 조절 가능
· τ↓ : 가장 큰 Logit에 확률이 더 몰림 / τ↑ : 확률 분포가 더 균등 (보통 CE Loss에서는 τ = 1 사용)
∴ KD Loss : τ > 1 적용 / CE Loss : τ = 1
→ Teacher의 출력이 부드러워져서 Student가 Teacher의 확신 정도와 Class 간 관계를 더 잘 학습 가능
** Knowledge Distillation : 학습 단계에서 Teacher를 활용하지만, 최종적으로는 Student만 사용
→ 효율적인 추론을 위한 학습 과정
→ 분류 문제에만 한정되지 않고, LLM에서도 다음 단어 확률 분포를 맞추는 방식으로 널리 사용 가능
▣ Feature Distillation
· Student가 Teacher와 비슷한 중간 Activation을 갖도록 학습
· 최종 Loss : CE Loss와 KD Loss를 함께 사용
· KD Loss : Teacher와 Student의 Normalized Activation 차이를 줄이는 방식
→ 출력뿐만 아니라 내부 표현도 비슷하게 만듦
> Feature Distillation Intuition
<Deep Learning Model>
· 낮은 Layer에서 Edge/색 같은 Low-Level Feature, 중간 Layer에서 패턴, 높은 Layer에서 더 의미 있는 High-Level Feature 학습
→ Student가 Teacher와 비슷한 이유와 표현을 바탕으로 판단하도록 중간 Layer들을 맞추는 방법
** Layer 수 불일치 : Teacher와 Student는 보통 Layer 수가 다르므로 1:1로 모든 Layer Matching 불가
→ Teacher의 일부 Layer와 Student의 일부 Layer만 선택해서 맞춤
** Dimension 불일치 : Activation 차원이 다를 수 있으므로 Projection Function(투영 함수)을 사용하여 차원을 일치시킨 후 비교
▣ Relational Distillation
· Conventional KD : 하나의 Input에 대해 Teacher와 Student의 출력 또는 Feature를 맞춤
· Relational KD : 여러 Input 사이의 관계의 구조를 Student가 학습
> Conventional KD
· 각 Sample 하나씩(xi) Teacher(fT)와 Student(fS)의 출력 차이를 최소화
> Relational KD
· Distance-Based Distillation : 2개의 Data Sample 사이의 거리를 계산
· Angle-Based Distillation : 3개의 Data Sample이 형성하는 기하학적 각도를 계산
→ Teacher와 Student의 최종 출력이나 중간 Feature를 직접 맞추는 대신, Sample들 사이의 관계를 맞춤
(이렇게 학습한 Student는 최종적으로 작은 Model만 사용하여 효율적인 추론 수행 가능)
> No Pre-Trained Teacher?
· Pre-Trained Teacher 부재 : Online Distillation / Self - Distillation
[Different Teacher Models]
▣ Model Ensemble
> Training
· 같은 구조의 Model을 서로 다른 Random Seed로 초기화해서 독립적으로 학습 시작
· 각 Model은 독립적으로 각각의 Training Data로 학습하여 서로 다른 CE Loss 계산
→ 서로 다른 Model 생성
> Inference
· 각 Model의 Output을 평균으로 결정 / 각 Model이 선택한 Class를 Majority Vote로 결정
→ 하나의 Model보다 더 안정적인 예측 가능
> Averaging
· Cat, Dog, Truck, Horse 확률의 평균을 계산하여 가장 높은 평균 값을 선택
→ 1개의 Model이 실수하더라도 전체 평균에 의해 안정적인 결정 가능
> Majority Vote
· 각 Model이 가장 높게 예측한 Class를 하나의 표로 투표 진행
→ 1개의 Model이 실수하더라도 다른 Model들의 표에 의해 안정적인 결정 가능
> Model Ensemble Feature
· 서로 다른 Random Seed로 학습한 Model들은 서로 다른 관점을 학습 → Ensemble은 이런 여러 관점을 함께 사용하여 성능↑
** 단일 Model은 특정 Feature에 의존할 수 있지만, Ensemble은 여러 Model의 판단을 모아 더 견고한 예측 가능
▣ Self Distillation
> Step 1
· Pre-Trained Teacher Model 존재 X
· 먼저 하나의 Model(W1)을 아무 지식이 없는 상태에서 CE Loss만 사용해서 학습
(Ex. Wheel만 보고 판단)
> Step 2
· W2 : CE Loss로 정답 Label 학습 + 이전에 학습한 W1의 출력을 따라가는 KD Loss도 함께 사용
→ W1이 W2의 Teacher
(Ex. Wheel + Headlight 보고 판단)
> Step 3
· W3 : CE Loss로 정답 Label 학습 + 이전에 학습한 W2의 출력을 따라가는 KD Loss도 함께 사용
→ W2가 W3의 Teacher (W2는 W1보다 성능↑)
(Ex. Wheel + Headlight + Window 보고 판단)
· 위 과정을 K번째 Model까지 반복하며 W1 → W2 → W3 → … → WK 순서로 지식 전달
→ 각 단계의 Student는 이전 단계 모델을 Teacher로 삼아 학습 (성능이 계단식으로 향상)
→ Pre-Trained Teacher가 없어도 Distillation 구조 가능
> Self Distillation 특징
① 큰 Pre-Trained Teacher가 없어도 같은 크기의 Model끼리 지식 전달
② 순차적으로 Model을 키우는 방식
③ 학습이 끝나면 보통 마지막 Model(WK)을 사용
④ W1, W2, W3, …, WK를 Ensemble하면 성능 향상 가능
→ 기존 Ensemble은 독립적으로 학습된 여러 약한 Model을 합치는 방식
→ Self Distillation 후 Ensemble은 Distillation을 거쳐 상대적으로 강해진 Model을 합치는 방식
▣ Online Distillation
· 각 Model은 CE Loss로 정답을 배우고, 동시에 상대 Model의 출력 분포를 따라가는 KD Loss를 사용
→ 2개의 작은 Model이 서로 Teacher이자 Student
L1 = LCE, 1 + λ · LKD, 1
L2 = LCE, 2 + λ · LKD, 2
· Teacher와 Student를 동시에 Online 방식으로 학습
[Summary]
· Knowledge Distillation : 큰 Model의 능력을 작은 Model에게 전수하여 Edge Device에서 효율적인 추론을 가능하게 하는 학습 Process
· Logit Distillation : 최종 예측 확률 분포를 Matching
· Feature Distillation : 중간 Layer의 Feature Map을 Matching
· Relational Distillation : Sample Data들이 형성하는 고차원적 기하학적 관계(거리, 각도)를 Matching
· Self Distillation : 하나의 Model이 세대를 거듭하며 자기 자신을 점진적으로 성장
· Online Distillation : 두 경량 Model이 실시간으로 서로를 Distillation하며 함께 성장
[Reference]
· Knowledge Distillation (AI System : Week 12) - Dongjun Han
댓글
댓글 쓰기