본문 바로가기

논문 리뷰

ComKD-CLIP: Comprehensive Knowledge Distillation for Contrastive Language-Image Pre-traning Model 리뷰

출처: ComKD-CLIP: Comprehensive Knowledge Distillation for ContrastiveLanguage-Image Pre-traning Mode

 

(본 논문에는 pre-traning(=pre-training), pipline(=pipeline) 등 오타로 보이는 부분들이 존재하지만,

본 리뷰에서는 고증을 위해 이를 그대로 표기합니다. 혼란스러워하지 마시고 일반적인 의미로 읽어주시면 감사하겠습니다.)

Abstract

Figure 1: (b) Comparison of the Harmonic Mean (HM) for generalization from base to novel categories. ComKD-CLIP is proud to achieve the best performance in 8 out of 11 diverse recognition datasets.

Approach (Preliminaries, Pipline)

- CLIP: 서로 독립적인 이미지 인코더와 텍스트 인코더 분기로 구성

- 이미지 인코딩 분기:

  • input: 레이블이 있는 시각 인식 데이터셋 D (이미지 x와 레이블 y가 존재하고, 이에 대응하는 N개의 클래스 이름 c가 존재) $$D={x_j, y_j}_{j=1}^M$$
  • 각 이미지 x는 이미지 인코더 f_I를 통해 처리되며 정규화된 이미지 특징 u을 얻음: 벡터의 방향은 유지한 채 크기(길이)를 1로 만드는 L2 정규화 사용 $$u = \frac{f_I(x)}{||f_I(x)||_2} \in \mathbb{R}^d$$

- 텍스트 인코더 분기:

  • input: 템플릿 "a photo of a {c_i}"를 이용하여 생성한 텍스트 설명 t_i 사용
  • t_i는 텍스트 인코더를 f_T를 통해 인코딩된 후, 다음의 정규화된 텍스트 특징 w_i를 생성 $$w_i =\frac{f_T(t_i)}{||f_T(t_i)||_2} \in \mathbb{R}^d$$
  • 모든 텍스트 특징의 집합 W는 분류를 위한 가중치 행렬(classification weight matrix)로 사용 $$W = [w_1, w_2, \dots, w_N] \in \mathbb{R}^{N\times d}$$

Figure 1: (a) Schematic of proposed ComKD-CLIP framework. All methods utilize the ViT-B/16 image encoder from the pretrained CLIP model.

Q. 텍스트를 a photo of class 라는 템플릿, 라벨링 정보로만 사용한다면, 기존 cv의 supervised 학습과 다른 점이 무엇일까?
A. 겉보기에는 Cat Image → Cat Label / a photo of a cat 이 비슷해보이지만 학습 방식은 완전히 다르다.
기존 supervised learning에서는 클래스가 단순한 정수(label index)이지만, CLIP에서는 클래스를 자연어 문장으로 변환한 뒤 Text Encoder를 통해 의미 정보를 가진 임베딩으로 표현한다.
   (Ex. 기존 CNN 분류기를 생각하면 라벨로 이용되는 Cat, Dog를 단순 클래스 번호로 사용하여 의미적 관계가 전혀 없다.
   반면, CLIP에서는 클래스가 언어 임베딩으로 표현되어 cat, dog, wolf가 임베딩 공간에서 서로 가까워질 수 있고 cat, airplane는 멀어진다.)
즉, 클래스 간 의미 관계까지 함께 학습할 수 있다

Q. 단순 라벨 대신 텍스트 행렬(classification weight matrix)을 사용함으로써의 이점은?
A. 기존 CNN는 학습 과정에서 FC Weight이 새로 만들어지며 새로운 데이터셋이 올 경우 FC Layer도 다시 학습해야 한다.
반면, CLIP은 FC Weight 대신 Text Embedding을 사용하여 새로운 클래스가 생기면 FC Layer를 다시 만들 필요 없이 해당 클래스 정보만 Text Encoder에 넣으면 된다. (Zero-shot Classification)

텍스트 행렬을 사용하면 클래스 간의 의미적 관계를 활용할 수 있으며, 새로운 클래스를 추가할 때 FC layer를 다시 학습하지 않아도 텍스트만 인코딩하여 분류할 수 있다. 따라서 Zero-shot 일반화멀티모달 활용이 가능해진다는 장점이 있다.

추가로, 해당 논문처럼 레이블이 있는 데이터셋에서 템플릿만 사용하는 경우에는, 일반 CLIP의 Zero-shot 장점은 크게 활용되지 않는다. 
이 논문에서 텍스트를 사용하는 가장 큰 이유는 새로운 클래스를 추가하기 위해서가 아니라, 교사의 텍스트 임베딩을 학생에게 증류(distill)하는 데 활용하기 위해서이다. 즉, 텍스트가 단순한 라벨을 넘어 의미 공간의 기준 역할을 한다.

- 분류 출력 확률 계산: 1번 식

논문의 8번 식 아래에 where q^t and q^s represent the logits predicted by the teacher and student models, respectively, which is calculated by the corresponding image features and text features using Eq. 1. 라는 문장이 나오는데. 이는 잘못된 참조로 보입니다.

  • softmax 분류기와 동일한 형태
  • uw^T: output logit. 이미지와 해당 클래스 텍스트 사이의 유사도 점수 (q)
  • tau: temperature parameter. 작을수록 softmax 결과가 더 sharp, 클수록 더 smooth (모델이 얼마나 확신을 갖고 예측할지를 조절하는 역할)

Knowledge Distillation (KD)

2번 식

원래 Hinton et al.에 의해 제안되었으며, 사전 학습된 대규모 교사 모델의 지식을 더 작고 경량화된 학생 모델로 전달하는 방법이다

  • 효과: 학생 모델이 교사 모델의 지식을 흡수하여 효율적으로 배포될 수 있도록 도움
  • 두 모델의 특징 분포를 정렬하기 위해 KL divergence loss를 사용(식 2) $$L_{kd}(q^t, q^s, \tau) = \tau ^2KL(\sigma (q^t / \tau), \sigma(q^s/\tau))$$
    • q^t , q^s : 각각 교사 모델과 학생 모델이 예측한 logits
    • \sigma(\dot) : softmax함수 (이를 전개하여 표현한 식이 1번 식)
    • \tau^2 : softmax 내부에서 logit을 \tau로 나누는데, 이는 gradient의 크기에도 영향을 준다. 그래서 Hinton은 이 영향으로 작아지는 gradient를 보정하기 위해 loss 앞에 해당 값을 곱해줌 
    • (KL Divergence는 두 확률분포가 얼마나 다른지 측정하는 함수)

2 principal stages: figure 2 - (a), (b)

Figure 2: Overview of our proposed ComKD-CLIP framework. (a) Utilization of the cue learning method with a well-trained large teacher CLIP model; (b) A smaller student CLIP model, which is trained with learnable cues and reuses the text features from the teacher model, requiring training only for the image encoder branch; (c) The schematic of IFAlign module; (d) The inference process within the trained student model, where the text encoder branch reuses the text features of the teacher model; (e) The schematic of EduAttention module.

 

대규모 CLIP teacher 모델의 사전 학습 단계

  • 레이블이 있는 도메인 데이터셋 $$D_{labeled} = {x_i, y_i}_{i=1}^M$$
  • concatenation 전략을 통해 학습 가능한 프롬프트 (learnable prompts)를 교사 모델의 이미지 인코더와 텍스트 인코더 분기 모두에 추가   PromptSRC (Khattak et al. 2023b) 및 PromptKD (Li et al. 2024b)
  • 각 input은 각 인코더를 통과하여 각 특징을 생성한다 $$u_p^t \in \mathbb{R}^d , w _p^t \in \mathbb{R}^d$$
  • 최종 출력 로짓은 q^t
  • 학습: 예측된 확률 분포와 실제 레이블 사이의 Cross-Entropy Loss 

Figure 2 - (a)

  • Prompt + Text/Image→ Concat Encoder
  • Text Feature Save
  • Logits → KD
  • Image Feature IFAlign

소규모 CLIP student 모델의 학습 단계

  • 교사 모델에서 사전학습된 text feature를 직접 활용하여 텍스트 인코더 분기의 학습 비용을 크게 줄임
  • 경량화된 CLIP 이미지 인코더 분기 설계
    • 레이블이 없는 도메인 데이터셋 D_unlabeled
    • IFAlign 모듈: 학생 모델의 이미지 특징과 교사 모델의 이미지 특징을 정렬하는 역할 수행
  • EduAttention 모듈: 학생 모델이 추출한 이미지 특징과 교사 모델이 제공하는 텍스트 특징 사이의 cross-relationship을 탐색
  • KL divergence: 교사 모델과 학생 모델이 생성한 로짓 사이의 차이를 최소화

Figure 2 - (b)

  • Text Encoder 점선 표시 : 새로운 Text Encoder를 학습하는 것이 아니라 Teacher가 만든 Text Feature를 그대로 사용
  • Image Feature  IFAlign
  • Text Feature, Image Feature  EduAttention
  • Logits → KD

Figure 2 - (d)

  • Inference
  • Image → Student Image Encoder → Project Layer → Image Feature → Prediction
Q. 프롬프트를 컨캣하는 의미? 효과?
A. Prompt를 Concat하는 이유는 입력 자체를 현재 도메인에 맞게 조정하기 위해서이다.
Transformer는 입력을 토큰 시퀀스로 처리하므로 Prompt Token을 앞에 붙이는 방식으로 쉽게 적용할 수 있으며,
Encoder 전체를 재학습하지 않고도 도메인 적응 성능을 향상시킬 수 있다.

Q. 학생 모델의 lightweight image encoder 추가 설명

A. Lightweight Image Encoder는 새로운 구조를 의미하는 것이 아니라, Teacher(Ex. ViT-L)보다 파라미터 수가 적은 Image Encoder(Ex. ViT-B)를 사용하는 것이다.
이 논문에서는 Teacher의 Text Feature를 그대로 활용하므로 Student는 Image Encoder만 학습하면 되고, 이를 경량화하여 계산량과 메모리 사용량을 줄이면서 Teacher의 성능을 최대한 유지하는 것이 목적이다.

Q. Project Layer?
A. Student와 Teacher의 Feature는 차원과 표현 공간이 다를 수 있으므로, 학습 가능한 Projection Layer(일반적으로 Linear Layer)를 통해 Student Feature를 Teacher Feature 공간으로 변환한 뒤 IFAlign을 수행한다.

Q. KD에 (u,W) 하드 라벨 대신 로짓을 사용하는 이유?
A. Hard Label은 정답 클래스만 알려주지만, Logit은 각 클래스에 대한 상대적인 유사도까지 포함한다.
따라서 학생 모델은 정답뿐 아니라 Teacher가 어떤 클래스를 얼마나 비슷하게 판단했는지까지 학습할 수 있으며,
특히 CLIP에서는 이러한 유사도가 텍스트 임베딩의 의미 관계를 반영하기 때문에 Hard Label보다 훨씬 풍부한 정보를 전달할 수 있다.


Q. 학생 모델에 대한 설명으로 '텍스트 인코더 분기의 학습 비용을 크게 줄임' 이라고 서술한 부분과 figure 2의 (b)를 보면 학생 모델도 텍스트 인코더에 대한 학습을 진행하는 것처럼 보인다. 그런데 이후 EduAttention 수식을 포함하여 해당 논문에서는 student model이 text encoder를 업데이트한다는 내용은 나오지 않는다.
그렇다면 student에는 text encoder가 아예 필요 없는 것 아닌가?
A. 논문에 Text Encoder를 freeze한다거나 Student Text Encoder를 제거한다는 문장이 명시적으로 쓰여있지 않아, 정확한 답은 알 수 없지만, 개인적으로 이런 의문을 해소하고자 몇 가지 정리를 하고자 한다.
- (가설) 아마 저자들이 figure 2에 student 모델 쪽에도 text encoder를 그린 이유는 student 모델도 CLIP 구조라는 것을 보여주기 위해 남긴 것 아닐까?
- 논문에서 학습 과정에 대해 서술한 부분을 보면 student 쪽의 텍스트 학습은 <save된 teacher text feature - eduattention - logits> 로 이어지며, student text encoder는 forward조차 하지 않는 것으로 보인다.
- 따라서 student 모델은 teacher처럼 text encoder를 새롭게 학습하지 않았으며 teacher가 미리 생성한 text feature를 그대로 사용하여 text encoder branch의 연산과 학습 비용을 줄인다. fig2에는 구조를 보여주기 위해 text encoder가 남아 있지만 이후의 과정과 식을 보면 student가 사용하는 것은 teacher의 text feature이며, text encoder 자체를 학습하는 것은 아니다. 
(해당 답변은 개인적인 의견이므로 참고용으로 보시면 좋을 것 같습니다.)

 

IFAlign

학생 모델이 추출한 이미지 특징이 교사 모델이 추출한 이미지 특징과 최대한 일치하도록 만드는 것이 목표

(더 정확히는, 학생 모델이 교사 모델이 중요한(salient) 이미지 특징을 추출하는 방법에 대한 지식을 학습할 수 있도록 한다 )

  • 추출된 특징의 평균과 분산 통계량을 정렬

계산 과정: 식 3,4

3,4번 식

  • u^p: 학생 모델과 교사 모델의 프롬프트가 적용된 이미지 특징
  • P(\dot): 특징의 차원을 효율적이고 비용 효율적인 방식으로 조정하여, 정확한 정렬이 이루어지도록 설계된 학생 모델의 이미지 인코더 분기 내에 존재하는 학습 가능한 프로젝터
  • L1 Loss 사용

구체적인 정렬 과정: 식 5

5번 식

  • alignment loss = (평균값 차이) + (분산값 차이)

Figure 2: (c) The schematic of IFAlign module.

 

Q. 왜 L1 Loss를 사용할까?
A. 논문에서는 L1 Loss를 선택한 이유를 명시적으로 설명하지는 않는다.
다만 평균과 분산이라는 통계량을 정렬하는 목적에서는 L1 Loss가 두 값의 절대적인 차이를 직접 최소화하기 때문에 직관적이다. 또한 L2 Loss보다 큰 오차(outlier)에 덜 민감하여, 일부 Feature 차이가 매우 크더라도 전체 통계량 정렬이 지나치게 왜곡되는 것을 방지할 수 있다.

Q. 왜 IFAlign에서는 KL Divergence를 사용하지 않을까?
A. KL Divergence는 확률분포를 비교하는 Loss이고, IFAlign에서 다루는 것은 Feature 자체이므로 직접 적용하기 어렵다.
또한 IFAlign의 목적은 Feature의 평균과 분산이라는 통계량을 맞추는 것이므로, 통계량의 차이를 직접 계산하는 L1 Loss가 더 적합하다.

Q. 왜 feature 자체를 맞추지 않을까?
A. 우선, Teacher와 Student는 모델 크기가 다르기 때문에 표현 능력(representation capacity)이 다르다.
논문에서 평균과 분산만 맞추는 이유를 명시적으로 설명하지는 않지만, 수식을 통해 의도를 유추하면,
feature를 원소 단위로 모두 동일하게 만들면
student의 표현 능력이 지나치게 제한될 수 있고, teacher와 student의 구조가 다를 때 학습이 어려워질 수 있다.
반면, 평균과 분산만 맞추면 전체 feature 분포는 teacher와 비슷하게 유지하면서도 student가 자신의 구조에 맞는 표현을 학습할 자유를 갖을 수 있다. 또한 평균과 분산은 특징 분포를 대표하는 중요한 통계량이므로, 비교적 적은 비용으로 teacher의 표현 특성을 전달할 수 있다.

Q. figure2 - (c)를 보면 학생 이미지 피쳐 쪽에 또 학습가능한 프롬프트가 그려져 있는데 이건 무엇인가?
A. 논문에서는 이에 대한 직접적인 설명이 없다. 다만, IFAlign은 Student Feature가 Teacher Feature를 모방하도록 학습하는 과정이고, 그림상 Student Prompt는 학습 가능한 요소로 표시되어 있다.
따라서 IFAlign Loss의 gradient가 Prompt까지 전달되어 Prompt 역시 Teacher의 Feature를 잘 생성하도록 함께 최적화된다는 의도로 표현했을 가능성이 있다.
이는 그림을 기반으로 한 개인적인 해석이며, 본문에서는 이를 명시적으로 설명하지는 않습니다.

 

EduAttention

Attention 메커니즘을 이용하여 학생 모델이 추출한 이미지 특징과 교사 모델이 제공하는 텍스트 특징 사이의 교차 관계를 탐색

(이를 통해, 학생 모델은 교사 모델이 텍스트-이미지 특징을 통합하기 위해 사용하는 unanced 전략을 학습할 수 있다)

Figure 2: (d) The inference process within the trained student model, where the text encoder branch reuses the text features of the teacher model.

계산 과정: 식 6

6번 식

  • Q는 교사 모델이 추출한 텍스트 특징
  • K,V는 학생 모델이 추출한 이미지 특징
  • f_att: 두 특징 사이의 cross-relation ship
  • C: 하이퍼파라미터
  • FC(\dot): fully connected layer

계산 과정: 식 7

7번 식

  • IFAlign과 EduAttention을 통해 교사 모델로부터 학습한 지식을 통합하기 위한 과정
  • \alpha: 학습 가능한 파라미터 (0으로 초기화되며, 학습이 진행됨에 따라 점진적으로 더 큰 가중치를 부여하도록 학습됨)
  • f_e: 최종 이미지 특징 (추출된 이미지 특징에 attention 결과를 더해줌)
  • residual connection
Q. 왜 Query가 Text Feature일까?
A. Teacher의 Text를 기준으로 Student Image Feature 중 어떤 정보가 중요한지를 찾기 위해서이다.
즉, Teacher가 텍스트와 이미지를 연결하는 방식을 Student가 학습하도록 하기 위한 Cross-Attention 구조이다.

Q. C는 무엇인가?
A. C는 Attention 계산에서 사용하는 Feature 차원에 해당하는 하이퍼파라미터이다.
Transformer의 \sqrt{d_k} 와 동일한 역할을 하며, Dot Product의 크기가 지나치게 커져 Softmax가 포화되는 것을 방지하기 위해 사용된다.

Q. 왜 \alpha를 0으로 초기화할까?
0으로 초기화하면 학습 초기에 EduAttention이 기존 Image Feature를 불안정하게 만드는 것을 방지할 수 있다.
초기에는 Student가 원래의 Image Feature만 사용하고, 학습이 진행되면서 Attention이 유용한 정보를 제공할 경우에만 해당 값이 증가하여 점진적으로 Attention 정보를 반영하도록 설계하였다.
즉, 처음부터 Attention 결과를 강제로 사용하는 것이 아니라, 학습 과정에서 모델이 "Attention 정보를 얼마나 사용할지"를 스스로 결정하도록 만든다는 것이다.

 

Distilled Knowledge Refinement

교사 모델의 feature fusion 결과를 기반으로 학생 모델이 학습한 지식을 더욱 정제하고자 함

 

KL divergence: 식 8

8번 식

  • 교사 모델과 학생 모델이 생성한 특징 분포 사이의 차이를 최소화
  • (KL divergence의 자세한 과정은 2번 식에서 볼 수 있다)

최종 손실 함수: 식 9

9번 식

  • L_align: 학생 모델의 정렬 손실
    • Mean Loss + Variance Loss
  • L_stu: 학생 모델의 증류 손실
    • KD Loss 
Q. 앞에서 이미 feature도 맞추고 cross-attention도 했는데 왜 다시 logit을 맞출까?
A. IFAlign은 이미지 표현을, EduAttention은 텍스트-이미지 관계를 학습하는 과정이다.
하지만 두 과정을 거쳤다고 해서 최종 예측 분포가 Teacher와 동일해지는 것은 아니다.
따라서 마지막으로 Logit까지 정렬하여 Teacher의 최종 의사결정(decision behavior)까지 모방하도록 한다.

Q. 여기서는 왜 logit을 맞추는데 IFAlign과 달리 KL divergence를 사용할까?
A. IFAlign은 Feature의 통계량을 비교하는 단계이므로 L1 Loss를 사용한다.
반면 Logit은 Softmax를 거쳐 확률분포 (모든 원소의 합이 1인 분포) 가 되므로, Teacher와 Student의 예측 분포 차이를 비교하는 KL Divergence가 적합하다.

Q. EduAttention Loss는?
A. EduAttention에는 별도의 Loss가 정의되어 있지 않다.
EduAttention은 Teacher의 Text Feature와 Student의 Image Feature를 결합하여 새로운 Feature를 생성하는 모듈이며, 이 Feature를 통해 계산된 Logit에 대해 KL Divergence를 적용한다.
따라서 EduAttention은 KL Loss의 역전파를 통해 간접적으로 학습된다.