DeiT (Data-efficient Image Transformers), 2021 — 데이터 효율을 높인 비전 트랜스포머 (딥러닝 이미지 분류 시리즈 3부 ②)

 

Transformer, 데이터의 장벽에 부딪히다

 

ViT (Vision Transformer)의 등장은 딥러닝 비전 분야에 혁명을 가져왔지만, 동시에 극복해야 할 명확한 한계를 드러냈습니다.

ViT가 ResNet과 같은 CNN 모델을 능가하는 성능을 발휘하려면 JFT-300M과 같은 수억 장 규모의 초대규모 데이터셋으로 먼저 학습해야 했습니다. ImageNet이라는 표준 데이터셋만으로는 CNN보다도 성능이 떨어졌습니다.

“순수한 Transformer는 태생적으로 지역적 귀납 편향(Inductive Bias)이 부족하다. 이를 메꾸려면 엄청난 양의 데이터가 필요하다.”

 

이 데이터 의존성은 ViT가 일반 연구자나 중소기업에서 활용되기 어려운, 접근성의 장벽을 만들었습니다. 2021년, 페이스북 AI 리서치(Facebook AI Research, FAIR)는 이 문제를 정면으로 돌파하기 위해 DeiT (Data-efficient Image Transformers)를 발표합니다.

---

해결책: CNN 스승을 둔 'Transformer 학생'

 

DeiT의 핵심 전략은 '데이터가 부족하다면, 학습 방법으로 이 문제를 해결하자'는 것이었습니다. 그 방법은 바로 지식 증류(Knowledge Distillation) 기법을 ViT 학습에 접목하는 것이었습니다.

 

ViT 구조의 학생(Student) 모델을 훈련시키기 위해, ImageNet에서 최고 성능을 내던 강력한 CNN 모델 (Teacher, 스승)을 사용합니다.

 

이 스승-학생(Teacher-Student) 구조는 다음과 같이 작동합니다.

1. 스승 모델 (Teacher): ImageNet으로 미리 학습된 강력한 CNN (예: EfficientNet)이 정답 레이블($y$) 외에 '소프트 레이블(Soft Label)'을 출력합니다.
2. 학생 모델 (Student): ViT 모델은 정답 레이블($y$)과 스승 모델의 소프트 레이블 두 가지를 동시에 학습합니다.

이를 통해 ViT 학생은 방대한 데이터에서 오는 '일반적인 지식'을 스승 CNN이 이미 학습한 '정제된 경험과 구조적 지식'으로 빠르게 대체하여 흡수할 수 있습니다.

---

DeiT의 핵심 발명: 증류 토큰 (Distillation Token)

 

단순히 소프트 레이블을 이용하는 증류 기법은 이미 존재했지만, DeiT는 Transformer 구조에 딱 맞춘 혁신적인 토큰을 추가합니다.

 

기존 ViT는 분류를 위해 이미지 패치 외에 [class] 토큰을 추가하여 사용했습니다. DeiT는 여기에 [dist] (Distillation) 토큰을 하나 더 추가합니다.

 

DeiT의 작동 원리:

1. [class] 토큰: 이 토큰은 일반적인 ViT처럼 정답 레이블(Hard Label)을 맞추도록 학습됩니다. (일반적인 분류 손실)
2. [dist] 토큰: 이 토큰은 스승 모델이 내놓은 소프트 레이블(Soft Label)을 맞추도록 학습됩니다. (증류 손실)

이 두 토큰이 Transformer 인코더 내부에서 이미지 패치 정보와 함께 Self-Attention을 주고받으며 학습됩니다. 그 결과, [dist] 토큰은 스승 모델의 지식을 흡수하는 데 특화되고, 모델 전체의 데이터 효율이 비약적으로 상승하게 됩니다.

---

DeiT의 구체적인 학습 전략

 

DeiT는 ImageNet 데이터셋의 두 가지 종류의 정답 정보를 동시에 학습합니다.

 

정답 정보	출처	학습 목표	DeiT의 역할
하드 레이블 (Hard Label)	ImageNet의 공식 정답 레이블 (예: "고양이", "개")	이미지를 정확한 범주로 분류하는 것	[class] 토큰이 이 레이블을 예측하도록 학습
소프트 레이블 (Soft Label)	ImageNet 데이터를 미리 학습된 CNN 스승 모델에 넣어 얻은 추론 확률 분포	스승 모델의 구조적 지식과 판단 경향을 모방하는 것	[dist] 토큰이 이 분포를 예측하도록 학습

 

1. 스승 모델 (Teacher Model)의 역할: 소프트 레이블 제공

DeiT 연구에서는 주로 ImageNet에서 매우 높은 정확도를 보이던 CNN 모델 (예: EfficientNet)을 스승 모델로 사용했습니다.

일반적인 분류 모델은 최종적으로 하나의 정답(Hard Label)을 출력하지만, 스승 모델은 각 클래스에 대한 확률 분포(Soft Label)를 출력합니다.

 

예를 들어, "치와와" 이미지에 대해 스승 모델은 다음과 같은 확률 분포를 내놓을 수 있습니다.

• 치와와 (정답): 90%
• 테리어: 5%
• 요크셔테리어: 3%
• 고양이: 0.1%

이 소프트 레이블에는 정답 외에도 스승 모델이 "이 이미지가 테리어와 요크셔테리어와도 약간 비슷하다"고 판단한 미묘한 구조적 지식이 담겨 있습니다. DeiT는 이 정보를 학습하여, ViT가 부족했던 지역적 패턴에 대한 CNN의 귀납 편향을 간접적으로 전수받게 됩니다.

 

2. DeiT 학생 모델의 학습 과정: 두 가지 손실 함수

DeiT(학생 모델)는 두 개의 손실 함수(Loss Function)를 최적화하며 학습합니다.

1. 분류 손실 (Classification Loss) - [class] 토큰:
   • 목표: DeiT의 [class] 토큰 출력과 하드 레이블(정답) 사이의 차이를 최소화합니다.
   • 역할: 일반적인 이미지 분류 능력을 확보합니다.
2. 증류 손실 (Distillation Loss) - [dist] 토큰:
   • 목표: DeiT의 [dist] 토큰 출력과 소프트 레이블(스승의 확률) 사이의 차이를 최소화합니다. (주로 KL-Divergence 사용)
   • 역할: 스승 모델의 섬세한 판단 경향과 구조적 지식을 전달받아, 데이터 효율성을 높입니다.

DeiT는 이 두 손실을 적절한 가중치로 합산한 총 손실(Total Loss)을 최소화하는 방식으로 파라미터를 업데이트합니다.

이러한 전략 덕분에, ViT가 수억 장의 데이터로 학습해야 겨우 얻을 수 있던 성능을, DeiT는 ImageNet 데이터만 사용하고도 효과적으로 달성할 수 있었습니다.

---

DeiT의 성과와 비전 모델의 민주화

 

DeiT의 성공은 비전 분야에 중요한 메시지를 던졌습니다. 더 이상 ViT를 쓰기 위해 수억 장의 데이터가 필요하지 않다는 것입니다.

• ViT (Base): JFT-300M 데이터로 사전 학습해야 ImageNet SOTA를 달성.
• DeiT (Base): 오직 ImageNet 1.3M 데이터만 사용하고도, 당대 최고의 CNN 모델인 EfficientNet을 능가하는 성능을 달성.

DeiT는 ViT의 장벽을 낮추고, Transformer를 누구나 ImageNet 데이터만으로 훈련할 수 있도록 비전 모델을 '민주화'시키는 데 결정적인 역할을 했습니다. 이는 이후 연구자들이 ViT 구조에 집중하고 다양한 변형(Swin, MAE 등)을 시도하는 기폭제가 되었습니다.

 

“데이터 규모가 혁신을 가로막을 순 없다. 우리는 더 똑똑하게 훈련할 수 있다.”

 

ViT가 Transformer의 잠재력을 보여주었다면, DeiT는 그 잠재력을 실용화의 영역으로 끌어내린 모델이었습니다.

---

참고 문헌

Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., & Jégou, H. (2020). Training data-efficient image transformers & distillation through attention (DeiT). arXiv:2012.12877. (ICML 2021)

---

 

 

 

이전 글: ViT (Vision Transformer), 2021 — Transformer, 이미지를 읽다 (딥러닝 이미지 분류 시리즈 3부 ①)

ViT (Vision Transformer), 2021 — Transformer, 이미지를 읽다 (딥러닝 이미지 분류 시리즈 3부 ①)

 

다음 글: Swin Transformer (Shifted Window Transformer), 2021 — 윈도우 속의 계층 구조 (딥러닝 이미지 분류 시리즈 3부 ③)

Swin Transformer (Shifted Window Transformer), 2021 — 윈도우 속의 계층 구조 (딥러닝 이미지 분류 시리즈 3