2020년 Vision Transformer가 등장하면서 Vision task에서도 Transformer가 큰 격변을 가져오기 시작했습니다. 하지만, 이러한 Vision task에서의 Transformer들은 1. 거대 데이터셋으로 학습 필요 2. 매우 무거움 등의 문제로 인해 활발히 활용되기 어려웠습니다. 이번에는 이러한 Transformer를 실질적으로 활용하기 위해서 Efficiency에 집중한 Efficient Transformer들의 시초 격인 DeiT : Training data-efficient image transformers& distillation through attention를 살펴보도록 하겠습니다.
1. Abstract
최근에는 NLP뿐 아니라 Audio, Computer Vision 등 다양한 분야에서 attention 기반의 neural network가 주목을 받고 있습니다. Vision 분야에서는 NLP에서 혁신적인 결과를 가져왔던 ‘Transformer is all you need’의 Transfoermer architecture를 그대로 가져와 이미지를 patch단위로 input으로 넣어 처리하는 Vision Transformer(이하 VIT로 칭함)로 이어졌습니다.
하지만, 이러한 High Performance Transformer 모델들은 상당히 많은 양의 데이터로 학습되어야하고 상당히 많은 양의 파라미터수를 가지고 있기에 무겁습니다. 따라서 큰 infrastructure( + hard computing resource )가 뒷받침되어야 했습니다. 실제로 VIT논문에서는 거대한 데이터셋(JFT-300M)으로 학습했을 때 유의미한 결과를 얻을 수 있었고, ( Inductive Bias - VIT 게시글 확인 ) “do not generalize well when trained on insufficient amounts of data” 와 같이 결론을 내렸습니다. 그렇기에 이렇게 거대 데이터셋으로 pretrain된 네트워크를 fine tune해서 사용하는 방식이 만연했습니다. 이러한 SOTA를 달성하고 있는 Transformer의 학습에 거대한 데이터가 필요하고, 모델이 무겁다는 점은 좋은 성능에도 불구하고 사용을 제한적이게 만들었습니다.
본 논문에서는 적은 수의 파라미터와 간단한 컴퓨팅 자원( Cost Efficiency ) 그리고 효율적인 수의 데이터( Data Efficiency)만으로도 기존의 무거운 high performance transformer에 가까운 혹은 이상의 성능을 낼 수 있도록 하는 것을 목표로 삼았습니다.
적은 수의 파라미터와 간단한 컴퓨팅 자원을 위해서 저자들은 Knowledge Distillation의 아이디어를 가져옵니다. 하지만, 기존의 distillation방식 대신에 Transformer에 적합한 token-based Distillation방식을 통해 이러한 목표를 달성하고자 했습니다. 이를 통해 현실적인 모델의 사용과 기존의 transformer간의 gap을 줄일 수 있는 efficient transformer를 제시해 downstream task 및 다양한 문제 해결에 사용할 수 있도록 했다는 점에서 의의가 있습니다. ( 물론 제시된 Knowledge Distillation 아이디어에서 Convnet을 teacher network로 사용한다는 점은 data-efficient training에도 도움을 줬습니다. )
효율적인 수의 데이터를 위해서 저자들은 strong data augmentation과 improved training기법들(optimizer regularization 등)을 통해서 VIT가 data-efficient하게 학습되기 위해서 어떻게 configuration되어야 하는지를 제시하고 있습니다. 중요한 점은 저자들은 Data Efficiency를 위해서 novel architecture나 method를 제안하는 것이 아니라 기존의 data problem을 해결하기 위해 많이 사용해오던 data augmentation으로 충분히 해결될 수 있는 문제라는 것을 입증하고 이를 위해 더 기존의 기법들에 대한 ablation study를 통해 개선된 학습을 하도록 합니다.
2. Related Works
2.1. Knowledge Distillation
(*우선 기본적인 Knowledge Distillation의 컨셉에 대한 해당 글을 참고해주세요.)
다시 한 번 Knowledge Distillation에 대해서 정리하자면 결국 기존의 무거운 모델을 사용하는 것보다 Knowledge Distillation을 통해서 동일한 혹은 그 이상의 성능을 내는 비슷한 가벼운 모델을 구성함으로써 적은 수의 파라미터와 간단한 컴퓨팅 자원으로 목표를 달성하고자 했습니다.
본 논문에서는 SOTA를 달성하고 있는 무거운 Transformer모델들을 실질적으로 사용하기 위한 transformer에 적합한 novel distillation strategy를 제시합니다. 또한 기존의 Knowledge Distillation 방식(soft label)과는 다른 의견인 hard label의 사용에 힘을 싣습니다.
2.2. Learned with Efficient Data
기존의 Vision Transformer는 거대 외부 데이터셋을 통해 학습이 선행되어야 합니다. 다양한 해석이 존재할 수 있지만, 지배적인 입장은 transformer가 inductvie bias가 적기 때문에 그렇다는 것입니다. 결국 Vision Transformer의 경우 잘 학습되기 위해서는 거대 외부 데이터셋으로 학습되어야 하고 이걸 실제로 사용하기 위해서는 이 VIT를 Imagenet등으로 finetune시켜서 사용하게 되는데, 이렇게 무거운 VIT를 거대한 데이터셋으로 학습시키기 위해서는 많은 학습시간과 Computing resource가 소요됩니다. 실제로 Finetune을 하더라도 현실적으로 쉽지 않은 요구사항이 충족되어야 합니다.
그렇기에 실질적인 사용에 있어서 Vision Transformer에는 Efficiency가 요구되었습니다. 본 논문에서는 이러한 Efficient Vision Transformer의 시초로 large trainig dataset없이 high performance 달성을 위해서 오로지 imagenet만으로 학습해 cost가 적지만 비슷한 성능을 내는 efficient transformer를 제시합니다.
3. Distillation Method - cost efficiency
본 논문에서는 크게 2축으로 1.hard distillation vs soft distilaltoin 2.classical distillation vs distillation token 를 설명하며 Distillation에 대한 인사이트를 제공하고 무거운 Transformer모델들을 실질적으로 사용하기 위한 transformer에 적합한 novel distillation strategy를 제시합니다. ( Cost Efficieny의 관점 ) 각 축에 대해서 살펴보고 본 논문에서 결과적으로 어떤 distillation strategy를 사용했고 왜 사용했는지를 이해해봅시다.
3.1. Hard Distillation vs Soft Distilaltion
사실 Knowledge Distillation에서 어떤 label을 사용할 것인지는 이전부터 논의되어 오던 주제입니다.
- soft label의 사용을 주장하는 측에서는 label smoothing의 효과를 제공해주기에 soft label을 사용하는 것이 좋다고 주장합니다.
- hard label의 사용을 주장하는 측에서는 soft label은 data augmentatoin에 따라서 label과 image 사이의 mis alginement를 발생 시킬 수 있기에 hard label을 사용하는 것이 안전하다고 주장합니다.
예를 들어, teacher model 학습 과정에서 crop augmentation을 사용했는데 특정 이미지(cat label되어 있고 lanscape에 cat이 이미지에 구석에 존재하는 ) 만약 soft label을 사용한다면 crop된 이미지에서 만약 cat이 속해있지 않게 augmentation이 된다면 implicit하게 label을 바꾸게 될 것입니다. 즉, 해당 이미지의 label은 cat인데도 teacher model학습과정에서 이러한 specific augmentation때문에 soft label의 기준에서는 다른 label들의 값이 높아져 오히려 왜곡된다는 것입니다.
이때 hard label을 사용하면 이렇게 specific augmentation에 depending되어 implicit하게 label을 바꾸는 문제를 해소해줍니다. 본 논문에서는 data-efficient transformer를 위해서 strong data augmentation을 적용하고 있기에 저자들은 hard label을 선택함으로써 이러한 문제를 해소합니다. ( 추가적으로 저자들은 hard label을 사용하고 label smoothing을 다음과 같이 해줌으로써 soft label과 같은 label smoothing의 효과를 가져올 수 있어 앞서 언급된 문제도 완화하고 label smoothing의 효과도 가져오기 위해서 true label에는 1-p 나머지 p는 남은 class들에 나눠주는 방식을 추가적으로 원한다면 사용할 수 있다고 제시했습니다. )
또한 저자들은 Knowledge Distillation에서 hard label을 사용하는 것이 더 좋다는 것을 다음과 같이 실증적으로 soft distillation(usual distillation) 방식보다 대부분의 상황에서 hard distillation이 높은 성능을 보인다는 것을 통해 주장하고 있습니다.( distillation token 아이디어 없이 기존 VIT처럼 그냥 class token만 사용) 따라서 본 논문에서 추후 서술되는 모든 dsitillation 방식들은 hard label을 활용하게 됩니다.
3.2. classical distillation vs distillation token
본 논문에서는 무거운 Transformer모델들을 실질적으로 사용하기 위한 transformer에 적합한 novel distillation strategy로써 Distillation Token을 사용하는 방식을 제시합니다. 방식 자체는 아주 간단합니다. 기존의 VIT 구조에 Distillation Token을 하나 추가함으로써 Distillation Loss에 해당하는 부분을 담당하도록 하는 것입니다. 즉, 기존의 Class Token고 Patch Embedding들 뿐 아니라 Distillation Token까지 추가해 이들 간의 관계를 self-attention을 통해서 학습하고, distillation head를 통해서 distillation embedding을 Softmax function을 취한 값과 Pretrained Teacher Model의 prediciton결과를 loss로 measure해 Tearcher model의 knowledge를 학습하는 Knowledge distillation을 하는 것입니다. 물론 당연하게도 class head에서는 그대로 class embedding으로 prediction한 결과와 true label간의 loss를 measure합니다.
저자들은 Distillation Token의 효과를 확인하고, Class Head와 얼마나 정보적으로 차별성이 있고 성능 향상에 영향을 주는지 확인하기 위해서 각 token의 convnergence vector간의 유사도를 비교하고 test과정에서 각 head의 정보로 ablation study를 진행합니다.
먼저 convnerge vector간의 유사도를 비교해보면, 실제로 이 class token과 distillation token은 서로 다른 벡터로 converge합니다. 이때 두 벡터의 유사도를 계산해보면 0.93으로 1보다 작은 값이 나오게 됩니다. 즉, aim인 true label과 유사하지만, producing target이기 때문에 identical하지는 않은 결과를 만들게 되고 이는 오히려 성능적으로는 긍정적인 영향을 가져온다는 것입니다! 반면에 만약 distillation token 대신 그냥 class token을 하나 추가해본다고 하면 두 벡터 간의 유사도는 0.9999로 성능에 유의미한 영향을 미치지 않는 거의 identical한 결과를 만들어내게 됩니다.
다음으로 Test time에서 class head의 결과만을 이용해서 prediction을 할 것인지 distillation head의 결과만을 이용해서 prediction을 할 것인지, 두 head를 모두 결과 도출에 이용할 것인지를 ablation study를 진행하게 됩니다. 결과로 inductive bias 때문에 distillation head만 사용하는 것이 오히려 class head만을 사용했을 때 보다 좋은 결과를 내고, 두 개를 같이 사용하는 late fusion방식도 나쁘지 않은 결과를 가져왔다는 것 정도를 확인할 수 있습니다.
3.3. Conv Teacher
추가적으로 Vision Transformer가 거대 데이터셋으로 학습되어야 하는 이유를 inductive bias가 적어서라고 언급했었습니다. DeiT에서도 뒤에서 data-efficient한 학습을 시킨다고는 하지만 적은 데이터셋으로도 efficient한 distilled transformer를 얻기 위해서는 inductive bias를 가지고 있는 것이 좋습니다.
저자들은 이러한 inductive bias가 teacher model로부터 상속받는 것이 성능에 유의미한 영향을 미친다는 것을 다양한 실험을 통해 확인했습니다. 그 결과 CNN을 사용했을 때 transformer를 teacher로 썼을 때보다(비슷한 성능) 더 좋은 성능 향상을 얻을 수 있었습니다. 이를 저자들은 CNN을 teacher model로 사용했을 때 inductive bias를 상속받아 더 많은 정보를 습득할 수 있기 때문이라고 해석했습니다. 따라서 저자들은 SOTA를 달성하고 있는 Transformer를 효율적으로 활용하기 위해서 student를 transformer로 사용하고 더 많은 inductive bias를 전달기 위해서 Convnet을 teacher model로 사용하는 것을 적극 권장합니다.
실제로 비슷한 성능을 내고 있는 DeiT-B(Transformer)와 RegNetY-8GF(Convnet)를 각각 Teacher model로 사용했을 때 distilled 모델의 성능을 비교해보면 RegNetY-8GF을 teacher model로 사용했을 때가 더 높은 성능을 내고 있는 것을 볼 수 있습니다.
이러한 점에서 Conv teacher의 Knowledge Distillation 방식은 large amount of dataset 없이도 학습 가능한 즉 data-efficient transformer에 기여하기도 합니다.
3.4. Overview of DeiT Distillation
따라서 결과적으로 본 논문에서 제시하는 Distillation Strategy는 다음과 같은 구조를 가집니다. 아래 이미지를 통해 한 번에 이해해보면 좋을 것 같습니다.
4. Data Efiiciency & Improved Training
본 논문의 제목처럼 결국 저자들은 Data-efficient Transformer를 제안합니다. 하지만, 지금까지 언급된 내용 중 Knowledge Distillation에서의 Conv Teacher 사용을 제외하고는 data-efficient에는 초점이 맞춰져있지 않았습니다. 그 이유는 사실 저자들은 Data-Efficient를 위해 특별한 아이디어 설계가 들어갔다기보다는 기존의 연구들의 방법과 효과를 분석하면서 ablation study를 통해 data-efficient한 학습을 하도록 했기 때문입니다.
저자들은 Strong Data Augmentation을 통해서 data-efficient training이라는 목표를 도달할 수 있다고 주장합니다. 새로운 설계보다도, Transformer가 data-efficient한 학습을 위해서 필요한 것은 기존의 방법론들만으로도 충분하다는 것입니다. 물론 이를 돕기 위해서 improved training strategy가 필요하고 이에 대해서는 실험을 통해 확인한 인사이트들을 제시하고 있습니다.
결국 VIT를 data-efficient하게 학습시키기 위한 DeiT 학습전략은 Novel Method가 있기보다는 다양한 실험들을 통해 실증적이고 간단히 확인할 수 있는 Configuration에 대한 인사이트에 초점이 맞춰져 있습니다.
4.1. Data augmentation
앞서도 이야기했듯 Transformer는 more prior가 결합된 즉 inductive bias가 있는 mdoel들 (ex. convnet)과 비교해서 더 많은 데이터를 필요로 합니다. 따라서 data-efficient한 학습이 필요하고, 다른 모델들과 비교해서 같은 사이즈의 데이터셋으로 학습시키기 위해서 즉, 이 정보의 gap을 매우기 위해서는 data augmentation을 잘 활용하는 것으로 충분히 해결될 수 있다는 것입니다. 결국 저자들은 특별한 방법보다는 기존의 strong data augmentation을 통해 충분히 기존의 데이터양으로 거대 데이터셋 pretrain의 성능을 넘을 수 있다고 주장하며 다음과 같은 실험 결과를 제시합니다.
다양한 ablation study가 존재하지만, Strong data augmentation 중 하나인 Rand-Augment를 사용했을 때와 안 했을 때를 비교해보면 확실한 성능적 향상이 존재한다는 것입니다. 물론 Rand-Augment대신에 다른 Strong data augmentation인 Auto-Augment를 사용해도 마찬가지입니다. 이러한 것이 가능한 것은 앞선 Knowledge Distillation에서 hard label을 사용했기에 specific data augmentation으로 인해 발생하는 문제를 최소화했기 때문이기도 합니다.
결국 data-efficieent하게 vision transformer를 학습하기 위해서 간단히 strong data augmentation을 사용하면 된다는 것입니다. 이는 Efficient Transformer에 핵심 및 시작을 제시한 것뿐으로 추후 transformer가 adaptive augmentation을 통해 더 개선될 여지 또한 남겨두고 있습니다.
4.2. Imroved Training
저자들은 이러한 data-efficient한 학습을 돕기 위해서 개선된 학습 방식이 필요하다고 생각했고, Transformer에서 중요한 몇 가지들을 확인 및 실험을 통해 인사이트를 제공합니다. 예를 들자면 Transformer 학습 자체는 Initalization이 중요하기에 기존의 연구 방식을 차용해 truncated normal intialization으로 고정하기도 하고 VIT와 비교하면서 Regularization 기법들, Optimizer들, 하이퍼파라미터들로 ablation study를 진행해 각 choice의 영향을 분석해가면서 개선된 Vision task에서의 Transformer 학습을 위해서 참고해보면 좋을 만한 인사이트를 제시하고 있습니다.
- Initalization and hyper parameters
- Regularizaton & Optimizers
- Progressive Growing Training
- 이전 연구결과(VIT)를 활용해 Fine Tunning 시 low resolutoin으로 학습 후 higher resolution으로 fine-tunning 했을 때 성능이 더 좋아졌고 이를 활용합니다. 물론 이 과정에서 positional encoding을 어떻게 처리할지 등의 문제가 있는데 이는 VIT에서 하던 positoinal encoding interpolation 방법을 그대로 사용한다고 합니다.
학습과정에서 다음과 같은 choice들에서 각 선택의 영향을 분석하고 인사이트를 제시하고 있으니 seection 6 같은 경우 논문을 한 번 참고해보셔도 좋을 것 같습니다.
5. Conclusion
그래서 본 논문에서 제시하는 DeiT는 구조적으로는 결국 VIT의 구조를 거의 그대로 가져오고 추가적으로 Knowledge Distillation을 위해서 Distillation token, head만 추가하는 방식입니다. Knowledge Distillation과 Data augmentation 및 Improved Training을 통해서 기존처럼 거대한 양의 데이터로 학습될 필요가 없고 상대적으로 가벼운 Image Transformer를 제시하고, 앞으로 Efficiency의 관점에서 Transformer 모델들이 풀어나가야 할 부분들을 제시했다는 데에 의의가 있는 것 같습니다.
References
