$ \cdot $ Authors: Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, Armand Joulin
$ \cdot $ Published: ICCV 2021.
1.Introduction
Transformer는 최근 visual recognition에 대해 기존의 CNN 기반의 모델들을 대신할 수 있는 모델로 떠오르고 있다.
특히, NLP domain에서 사용하는 학습전략(대규모 데이터 셋에서 학습 후 목표 데이터 셋에 미세 조정)을 적용하여 더 좋은 성과를 내고 있다.
이러한 모델 중 하나인 ViT는 기존의 CNN 모델들에 비견되거나 더 뛰어난 성능을 보이지만, 몇가지 한계점이 있다.
- Computationally more demanding.
- More training data.
- Don`t exhibit unique properties.
저자들은 ViT가 보여준 좋은 성능과 한계점이 "pre-training"에 있음에 의문을 던지며 논문을 시작한다.
NLP에서 Transformer가 성공할 수 있었던 요인은 중 하나는 바로 "BERT"나 "GPT의 language modeling"에 사용된 "Self supervised pretraining"이다.
일반적으로 Supervised objectives는 하나의 문장을 하나의 label로 축약해서 학습한다.
하지만, BERT처럼 일부 단어를 masking 한 뒤 예측한 것이나 GPT처럼 다음 단어를 예측하는 Self supervised pretraining objectives는, 하나의 문장에서 pretext task를 생성하기 때문에 전자보다 훨씬 풍부한 학습 시그널을 얻을 수 있다.
비슷하게 image domain에서도 image-level supervision은 미리 정의된 수천 개의 클래스 중 하나의 라벨만 예측하게 하여, 이미지가 가진 복잡한 시각 정보를 단일 개념으로 축소시킨다.
비록 NLP에서 사용된 self-supervised pretext task들은 언어 구조에 의존하는 text-specific한 과제들이지만, 이 개념에서 영감을 받아 이미지에서도 ConvNet 기반의 self-supervised 학습법(MoCo, BYOL, SwAV 등)이 등장하여 높은 성능을 보였다.
저자들은 ViT 기반 self-supervised 학습을 진행하는 과정에서 다음과 같은 현상을 관찰했다.
- Self-attention map이 객체의 경계를 자연스럽게 포착 (Emergence of segmentation masks)
- Fine-tuning이나 data augmentation 없이도 K-NN 기반으로 높은 성능
- Multi-crop augmentation과 momentum encoder가 성능 향상에 중요
- 작은 patch size가 더 풍부한 표현 학습에 기여
특히, Self-attention map이 객체간의 경계를 포착하는 특성은 self-supervised ViT가 별도의 supervision 없이도 이미지 내의 객체 구조를 인식하고 구분할 수 있는 representation을 학습함을 보여준다.
저자들은 이러한 발견을 기반으로, self-supervised approach를 일종의 knowledge distilaiton with no label의 형태로 간소화할 수 있었고 DINO라는 framework를 제안하게 되었다.
저자들이 제안하는 DINO 주요 특징은 다음과 같다.
- Student model이 teacherr model의 output을 직접 모방(no lables)
- Momentum encoder를 기반으로 구성
- Standard cross-entropy loss 사용
- Centering 및 Sharpening을 통해 collapse 방지
- Predictor, Advanced Normalization, Contrastive loss와 같은 다른 방법들은 필수적이지 않음.
= 안정성이나 모델 성능에 미미한 이득을 줌. - ConvNet과 ViT 모두에 별도의 구조 수정 없이 적용 가능하며, BatchNorm 등 내부 정규화가 불필요.
2. Related Works
Self-supervised learning
Self-supervised learning의 초기 접근은 주로 각 이미지를 고유한 instance로 간주하여 구분는 "instance discrimination"에 초점을 두었다.
그러나 이러한 방식은 단순히 데이터 양을 늘리는 것에 비례해 더 일반화된 표현을 얻을 수 있는 것은 아니며, 모델은 여전히 instance 간의 관계를 학습하는 데 한계를 보인다.
Wu et al. (2018)의 연구는 instance 간 feature similarity를 직접 비교하는 contrastive 학습 방식을 제시하였다.
이후 연구들은 이미지 쌍 간의 명시적 비교 없이도 유용한 표현을 학습할 수 있는 방법들(BYOL, SimSiam 등)을 제안하며 발전했다.
이중, BYOL은 momentum encoder를 사용하여 teacher representation을 생성하고, student가 이를 예측하도록 학습한다.
이 과정에서 negative pair를 사용하지 않고도 안정적인 feature alignment를 달성했다는 점이 특징이다.
DINO 역시 BYOL에서 영감을 받았지만, similarity matching loss의 형태가 다르며, Student와 Teacher는 동일한 구조를 가지되, Teacher의 파라미터는 Student의 EMA(Exponential Moving Average)로 업데이트된다는 점에서 차이가 있다.
Self-training & Knowledge distilation
Self-training은 소량의 labeled data로 학습된 모델을 기반으로, 대규모 unlabeled data에 pseudo-label을 생성하여 점진적으로 feature의 품질을 향상시키는 방법이다.
이 과정에서 pseudo-label을 hard하게 지정하는 대신, teacher의 soft probability 분포를 사용하는 방식을 "soft assignment" 혹은 "knowledge distillation"이라고 부른다.
이전 연구들에서도 self-training과 knowledge distillation을 결합하려는 시도가 있었으나, 대부분의 방법은 고정된 teacher 모델을 사용했다. 반면 DINO는 학습 과정 중 student의 EMA로 업데이트되는 dynamic teacher를 도입하여, 일종의 co-distillation과 유사하지만 teacher가 독립적으로 학습되지 않는다는 점에서 구별된다.
3. Approach
3.1 SSL with Knowledge Distillation
DINO는 최근 연구된 Self-supervised 모델들과 전반적으로 비슷한 구조를 지녔으나, Knowlege distillation과도 비슷한 사항이 많다.
Knowledge distillation은 Student network $g_{\theta_{s}}$의 출력을 teacher network $g_{\theta_{t}}$의 출력과 동일하게 나오도록 student network의 파라미터 $\theta_{s}$를 $\theta_{t}$로 parameterization 하는 것이다.
어떤 입력 $x$를 주었을 때, 두 모델의 출력을 각각 K 차원의 확률값 $P_{s}$와 $P_{t}$로 나타낼 수 있다. 이 확률값은 아래의 수식으로 표현할 수 있다.
$$ P_{s}(x)^{(i)} =\frac{\exp{(g_{\theta_{s}}(x)^{(i)}/\tau_{s})}}{\sum_{k=1}^{K}\exp{(g_{\theta_{s}}(x)^{(k)}/\tau_{s})}} $$
Knowledge distillation은 고정된 $g_{\theta_{t}}$에 대해 $g_{\theta_{s}}$를 매칭하는 것으로, Cross entropy를 사용하여 objective function을 만들 수 있다.
$$\min_{\theta_{s}}-P_{t}(x)\log{P_{s}(x)} $$
저자들은 이 objective function에 "Self-supervised learning"을 적용 했다.
먼저, 이미지 데이터 셋에 Multi-crop strategy를 사용해서 서로 다른 왜곡된 시점을 구성하거나 다른 크기로 잘라냈다.
이렇게 생성된 이미지들로 구성된 set $V$ 를 정의했다.
이때, $V$ 안에는 2개의 global image view $(x_{1}^{g}, x_{2}^{g})$와 이들보다 더 작은 크기의 local view들로 나뉘어진다.
학습이 진행될 때, $V$의 모든 image들은 $g_{\theta_{s}}$를 지나가지만, $g_{\theta_{t}}$ 에는 global view에 해당하는 데이터들만이 지나갈 수 있다.
이를 통해, teacher model이 global view로부터 안정적인 표현을 학습하는 동안, student model은 local view를 통해 동일한 이미지의 의미적 정보를 포착할 수 있다. 이를 통해, global-local consistency를 얻을 수 있다.
또한, 저자들은 $g_{\theta_{t}}$와 $g_{\theta_{s}}$를 동일한 구조로 하되, 내부 파라미터 $\theta_{t}$, $\theta_{s}$는 서로 다르게 구성했고 SGD를 사용해서 학습을 진행했다.
Teacher Network
논문에서 사용된 $g_{\theta_{t}}$는 일반적인 Knowledge distillation에서 사용하는 것처럼 well pretrained model이 아니다. 대신, $g_{\theta_{s}}$ 학습하면서 생긴 이전 iteration의 가중치의 평균을 사용해서 업데이트한다.
저자들은 Teacher network를 업데이트하기 위한 방법으로 한 epoch가 진행되는 동안 freeze 시켰다가 업데이트, Unfreeze 상태 후 student의 weight를 그대로 copy 하는 방식을 테스트해봤다.
결과적으로 전자의 경우가 훨씬 뛰어난 성능을 보여, 논문에서도 전자의 방법을 채택했다. 이때, Momentum encoder를 사용해서 student weight에 EMA를 적용하였다.
Momentum encoder는 contrastive learning에서 negative queue를 대체하기 위해 도입되었지만, DINO에서는 student의 EMA(Exponential Moving Average)를 계산하여 teacher를 업데이트하는 용도로 사용된다.
이 과정은 통계학에서 사용되는 Polak–Ruppert averaging (exponential decay 평균화)과 유사하며, 이는 일종의 online model ensembling 기법으로, 안정적 수렴과 성능 향상을 동시에 유도한다.
이처럼 Teacher가 점차 더 나은 표현을 포착함에 따라, student의 학습 신호 또한 향상되고 결과적으로, 성능 향상으로 이어진다. 이러한 모델간의 상호 보완적인 관계(mutal improvment loop)가 DINO의 성능 향상 핵심 요인으로 작용한다.
Network architecture
DINO의 architecture는 다음과 같다
- Network $g$
- Backbone: ViT or ResNet-50
- Projection head: 3 MLP Layer with L2 normalization of output projection, K-dimension FC layer (similar to SwAV)
- No predictor
- etc: ViT backbone 사용 시 batchnorm layer 사용 안함(ViT가 Batchnorm layer를 안써서)
Avoiding collapse
다른 SSL 기법들이 contrasrive loss나 clustering, predictor 등을 사용하는데 비해, DINO는 Centering과 Sharpening을 사용했다.
- Teacher 모델의 momentum output을 centering & sharpening 하는 역할
-
- Centering:
각 batch의 평균 출력을 techer logits에서 빼서 정규화하는 과정. 특정 차원이 지배적으로 작용하는 것을 방지하지만, 지나치면 uniform 분포 유도할 수 있다. 이를 위해 teacher는 bias term $c$를 유지하며, EMA를 통해 batch 통계값을 업데이트한다. - Sharpening:
teacher의 softmax 출력 분포를 temperature $\tau < 1$로 조절하여, 더욱 뾰족한 형태(평균 지점으로 밀집된 형태)로 만든다. 이 과정에서 model collapse를 방지하고 학습 안정성을 높인다.
- Centering:
- Centering & Sharpening:
결과적으로, centering과 sharpening을 함께 적용하여 다양한 batch size에서도 안정적으로 teacher의 출력을 적절하게 regularize 할 수 있다. - Trade off: Stability and independence over the batch
- Centering은 batch 평균만 사용하므로 batch size에 영향을 적게 받는다. (less dependence)
- 하지만, Batch normalization처럼 전체 분포(평균+분산)을 쓰는게 아니므로, 통계적으로는 덜 안정적이다. (stability 감소)
- 따라서, batch에 대한 의존성을 줄이는 대신 안정성이 다소 감소했다.
- Centering update: Exponentail Moving Average(EMA)
- Teacher model의 bias $c$는 각 batch마다 새로 계산되는 것이 아니라, 이전 batch의 평균값을 EMA 형태로 누적해서 갱신함.
- 즉, 이전 batch의 평균과 현재 batch의 평균을 혼합해서 업데이트.
- 따라서, batch가 작거나 구성에 따라 평균이 조금 달라져도 teacher의 출력이 크게 변동되지 않고 안정적이므로 다양한 batch size를 사용해도 안정적인 학습 가능
-
3.2 Implementation and evaluation protocols
Vision Transformer
DINO는 `DeIT` 에서 사용한 Implementation을 참고하였다. 그리고 선행 연구와 마찬가지로 별도의 학습 가능한 token인 "[CLS]" (class token)을 사용했다.
Implementation details
- Pretrainig on ImageNet datasets.
- Batch size: 1024 distributed 16 GPUs (with ViT-S/16).
- Learning rate:
- Linearly ramped up during first 10 epochs. lr = 5e-04 $\times$ batch size / 256.
- Cosine scheduling: 0.04 to 04.
- Temperature $\tau_{s}$ and $\tau_{t}$:
- $\tau_{s}$: Warm up from 0.1.
- $\tau_{t}$: Linearly warm up (0.04 to 0.07) during first 30 epochs.
- Argumentations:
- BYOL and multi-crop with bicubic interpolation to adopt positional embeddings.
4. Experiments
논문에서 주요 실험은 다음과 같다.
- Comparing with SSL frameworks on ImageNet
- Propertis of ViT trained with SSL
4-1. Comparing with SSL frameworks on ImageNet
이 실험에서는 논문에서 제안한 DINO를 ImageNet 데이터셋에 대해 학습했을 때, "Same architecture"일 때와 "Across architecture"로 setting을 나눠서 성능을 검증했다.
Same architecture
같은 architecture 일때는 ResNet-50과 ViT-S를 사용했다. ViT-S를 고른 이유는 여러 측면(파라미터 수, throughput 등)에서 ResNet과 비슷해서 골랐다고 한다. 결과적으로, 두 모델을 사용했을 때 결과에서 다른 baseline들과 비슷하거나 그 이상의 성능을 얻었다.
흥미로운 점은 DINO와 ViT의 조합을 사용했을 때만, header로 Linear classifier를 사용한 경우와 k-nn 을 classifier로 사용한 경우의 성능이 동일하게 나왔다는 것이다.
Across architecture
이 setting은 방법론간의 직접적인 성능을 비교하기 보다, 더 큰 architecture로 변경했을 때, DINO로 학습된 ViT의 성능의 한계를 보기 위함이다. 결과적으로 모델 크기를 키웠을 때 성능이 향상되는 것은 맞지만, patch size의 영향을 더 크게 받는다는 것을 발견했다.
Patch size를 8로 줄였을때 훨씬 더 성능에 큰 변화를 관찰할 수 있었고, 이러한 size 감소는 실행 시간을 크게 줄이지만 더 많은 메모리를 사용했다. 하지만, 성능적으로 더 큰 이득을 얻을 수 있었다고 한다.
4-2. Properties of ViT trained with SSL
이 실험에서는 DINO의 특성을 "nearest neighbor search", "retaining object location" and "trasferability to downstream task" 측면에서 분석했다.
4.2.1 Nearset Neighbor retrieval with DINO ViT
- Image Retrieval
- ViT를 DINO로 pretrain했을 때 feature가 retrieval task에서 어떤 성능을 갖는지 Oxford and Paris image retrieval datasets을 통해 검증
- 평가 방식: feature를 freeze하고 k-nn search. mAP를 분석
- 결과:
- DINO ViT feature가 ImageNet supervised training으로 얻은 feature보다 retrieval mAP가 높음
- Annotation 없이 어떠한 데이터로도 학습할 수 있다는 SSL의 특성을 더 분석해보기 위해 GLDv2(1.2M landmark dataset)으로 DINO를 학습한 모델도 실험.
- GLDv2로 학습한 DINO ViT는 landmark retrieval SOTA를 기록
- Copy detection
- INRIA Copydays(“strong” subset)에서 copy detection 성능(mAP) 평가.
- blur, insertion, print-scan 등 다양한 변형이 적용된 이미지에 대해 변형된 원본을 찾아내는 task.
- DINO ViT는 copy detection 에서도 SOTA에 근접하는 좋은 성능을 보임.
4.2.2 Discovering the Semantic Layout of Scenes
- Video instance segmentation
- Fig.1에서 보인 것처럼 DINO의 self-attention에는 segmentation 정보가 포함 됨.
- 이에 대한 정량적 분석을 위해 DAVIS-2017 video instance segmentation benchmark에서 평가.
- Jarbir et al. protocols 사용: 연속된 프레임 간의 patch token nearest neighbor로 segmentation
= 추가 학습 or fine-tuning 없음 - 결과:
- ViT는 dense task로 설계되지 않았음에도 좋은 성능을 보임
- 특히, patch size가 8일때 크게 성능 향상
- Probing the Self-Attention Map
- Fig.3에서 각 attention head가 작은 물체(flag), 가려진 영역(bushes) 등 다양한 semantic region에 주목하는 모습을 시각화.
- 실험에서 480p 이미지 사용 => ViT-S/8의 토큰 수는 3601개.
- Fig.4에서 supervised ViT와 비교:
- supervised ViT는 clutter가 있는 장면에서 object를 잘 attention하지 못하지만, DINO self-attention은 훨씬 더 정확하게 포착.
- 정량적 비교:
- self attention map을 60% mass threshold로 잘라 segmentation mask 생성
- Ground Truth와 Jaccard similarity 비교했을 때, DINO가 supervised ViT보다 더 좋은 성능을 보임.
4.2.3 Transfer Learning on Downstream Tasks
- Table.6 에서 DINO pretrained feature를 다양한 downstream task에 fine tuning하여 성능 비교
- baseline: ImagNet supervised preatreind ViT with same architecture
- 결과:
- ViT 계열에서는 DINO가 supervised보다 transfer 성능이 더 좋음
- 이러한 결과는 ConvNet 계열 모델에서도 비슷하게 나온 연구들이 있었음.
- ImageNet 자체 acc도 fine tuning시 1~2% 가량 향상
5. Conclusion
DINO 논문에서는 Self-supervised learning(SSL)에 대한 방법론을 제안하여 이를 ViT에 적용했을 때의 성능 변화를 연구하였고, 제안된 DINO 방식을 통해 다른 Conv 모델을 포함한 baseline 모델 대비 높은 성능을 보임을 증명했다.
특히 ViT와 DINO를 결합하였을때 발견된 중요한 특성은 다음과 같다.
1. k-nn classification 성능이 뛰어나 retrieval task로 확장 가능.
2. scene layout 정보를 feature가 내재적으로 포함 (Weakly supervised segmentation에 활용 가능).
3. SSL이 ViT 기반의 BERT-like 모델을 만드는 핵심이 될 수 있는 근거를 제시.
My opinion
DINO 논문을 읽기 전에는, DINO가 특정 모델을 지칭하는 것인줄 알았는데, SSL을 computer vision 영역에 적용하는 방법론인 것을 알고 놀랐다. 그래서 SSL의 개념을 ViT에 적용하여 DINO라는 방법론을 만드는 과정이 신기했고 특히, Supervised ViT보다 DINO ViT가 segmentation task에서 object를 더 잘 포착하는 것이 흥미로웠다.
논문을 통해 이해한 것을 정리해 보자면,
- DINO는 annotation이나 label이 필요 없다.
- 대부분의 경우 supervised learning 한것보다 더 좋은 결과를 얻을 수 있다.
- DINO와 ViT의 결합은 특히 k-nn classifier와 leanear classifier의 성능이 거의 비슷하다.
- 단순 classification task 뿐만 아니라, ViT 사용 시 attention map이 물체 간의 경계를 잘 학습해서 segmentation task에서도 잘 동작한다.
DINO 방법론을 적용해서 경량화 transformer 모델들에 적용했을 때도 같은 결과가 나올까라는 궁금증이 생겼다.
그리고, teacher와 student를 같이 학습하는 만큼 학습에 필요한 메모리가 일반적인 supervised 방법보다 더 필요하지 않을까라고 생각이 든다.