[논문공부] (자세한 리뷰) Masked Autoencoders are Scalable Vision Learners

들어가며:

Masked Autoencoders Are Scalable Vision Learners(링크)는 Facebook AI Research(아직은 Facebook으로 되어있는데 meta로 바뀌겠죠?)에서 나온 논문으로, 현재 CVPR 2022에 submit되어 review중인 듯합니다.

 

간단하게 요약하면, Masked autoencoder는

Masked patch를 복원하는 task로 학습한 것을 pretrained weight로 삼아서
downstream task로 finetuning했을 때 예측 성능이 더 좋다.

는 내용입니다.

Language model에서 Transformer 이후 BERT가 등장한 것과 마찬가지로 Vision Transformer(이하 ViT)의 BERT격이 이 MAE라고 생각합니다.

 

0. introduction: ViT

최근 vision 분야의 연구는 Transformer 구조로 학습하는 것이 트렌드입니다.

 

기존의 CNN 기반의 ResNet 등의 모델과는 다르게 self-attention 모듈을 여러 block으로 쌓아서 표현력을 높이는 구조입니다.

CNN기반(ex.ResNet)도 좋지만, ViT는 더 좋다!

다만, CNN 기반 model과는 입력을 사용하는 방법에서 조금 차이가 있습니다.

Language에서 sequence의 sample 단위는 word token인데, 이것을 일정 크기의 patch(16x16)으로 나눠서 입력에 사용합니다.

아무래도, 픽셀 하나하나를 sequence sample로 사용하기에는 너무 micro하여 attention weight를 고려하기엔 정보가 적기 때문인 듯합니다.

 

 

Left: Vision Transformer / Right: Transformer(Attention is all you need)

왼쪽이 ViT 논문(링크)에 있는 구조입니다.

위 사진에서 오른쪽 그림은 익숙하실 텐데, Transformer의 original 논문의 구조 그림입니다.

오른쪽의 그림의 Transformer에서 오른쪽 절반을 날리고 encoder에서 decoder로 들어가는 layer입력을 MLP head로 빼서 classification을 수행한 것입니다.

(normalization layer가 먼저 수행되는 사소한 차이점이 있긴 합니다.)

 

추가로, self-attention은 위치에 대해 invariant하기 때문에 positional encoding을 입력에 더해주어 위치정보를 부여한 다음 encoder 입력에 사용하게 됩니다.

만약 positional encoding을 더하지 않으면 각 patch가 random하게 뒤섞인 이미지를 넣더라도 self-attention 결과는 원래 이미지를 사용한 것과 동일하게 됩니다.

 

여기까지가 ViT에 대한 간략한 설명입니다.

 

1. Masked autoencoder 구조 및 학습 방법

전체적인 구조는 ViT의 형태를 따릅니다.

1. 다만, ViT와 다른 점은 입력으로 들어갈 때 입력 patch 중에서 일부를 날립니다.
2. 그 이후에 positional encoding을 입력에 더해주고, encoder로 들어가게 됩니다. 이 encoder는 self-attention 기반의 Transformer encoder 구조입니다.
3. 그 후 Encoder output에서 아까 mask를 씌운 위치에 대응되는 masked token을 다시 붙입니다.
4. 그리고 decoder에 들어가기 전에 positional encoding을 한번 더 먹인 다음 decoder를 거쳐 나온 output을 reshape(unflatten)합니다.
5. 이렇게 구한 decoder output인 predicted target과 원래 이미지인 target 사이의 Mean square error loss를 계산하여 backpropagation함으로써 pretraining 학습이 진행됩니다.
6. Decoder부분은 떼고 linear layer(혹은 MLP)를 달아서 원하는 task로 finetuning training합니다

Masked Autoencoder 학습 방법

(위 그림에서 mask token을 붙인 것을 사용하는지 아니면 그 전단계인 encoder output 직후의 layer를 사용하는지는 언급되지 않았습니다만, 제 생각엔 그림과 다르게 mask token 붙이기 전 단계를 사용할 것이라 생각합니다).

 

 

모델 구조 및 학습 방법의 detail은 논문에 설명되어있습니다만, 제 기준 좀 중요해보이는 detail만 추려보면

• masking을 random하게/덩어리로/규칙적으로 실험하는 세가지 세팅으로 했는데 random한 게 가장 성능이 좋다고 합니다.
• 2.의 실험 결과에서 설명하겠지만, masking ratio를 75%로 한 것이 성능이 제일 좋았다고 합니다. BERT에서 통상 15%정도를 masking한 것과 비교하면 dramatic하게 차이가 있습니다. 
• (이것도 2.에서 언급)pretrain(reconstruction) 때 사용하는 decoder의 design에 따라 finetuning task 성능이 다르다고 합니다(decoder는 finetuning에서 사용하지 않음에도). 

 

2. Experiment

Language 분야에서 Transformer 이후 발표된 BERT에서의 main contribution은

Masking된 부분을 reconstruction하는 pretraining을 했더니 다른 task로 finetuning한 게 성능이 좋았다.

는 것입니다.

MAE도 마찬가지입니다.

 

2.1) Base 실험

일단 제일 기본 실험 결과를 살펴보면

실험 세팅은 ImageNet 1K dataset(이하 IN1K)으로 실험하고, 다른 Large dataset으로 pretraining한 weight는 사용하지 않고,

오직 IN1K만 사용해서  classification 했을 때 얼마나 성능이 좋느냐입니다.

위 표에서 설명하듯, ViT 논문에서의 성능은 76.5%이고 이번 논문을 위한 baseline으로 실험한 결과는 82.5%입니다

(실험 configuration이 둘이 좀 다르긴 합니다).

ViT의 성능은 원래 논문에서보다 더 좋다는 이야기를 잠깐 하며, 그것보다 MAE가 더 좋다는 말을 합니다(상대를 드높여주는 척 하다가 사실 내가 더 개쩖을 보이는 전략).

 

결국, 이 논문의 핵심 주장 중 하나인

다른 외부 dataset 없이 실험했을 때, ViT에 비해 classification 성능이 좋다.

는 것을 보여주고 있습니다.

 

 

2.2) Masking ratio

Masking ratio를 변경해가며 실험한 결과

Masking ratio를 75%까지 높인 것이 finetuning 했을 때의 classification 성능이 좋다고 합니다.

위와 아래는 fine-tuning, linear probing의 차이인데, classification으로 downstream task 학습 때 encoder를 trainable하게 두느냐, freeze하여 고정하느냐의 차이입니다.

BERT에서는 통상 15%정도의 masking ratio면 적절하다고 하는데, 그에 비하면 굉장히 높은 수치입니다.

이 실험 결과의 해석으로 저자는

Language는 인간이 만든 signal로 문장을 이루는 word token 하나하나가 의미론적으로 정보가 굉장히 dense하다.
그렇기 때문에 단어 한 두개만 빼도 이것을 예측하는 것이 쉽지 않고 더 high-level의 추론능력으로 학습하는 것을 필요로 한다.
하지만, Image의 자연적인 signal로 정보가 비교적 sparse하다. 그래서 patch 조금 빼는 것은 주변 patch로부터 interpolation해도 적당히 될만큼 쉬운 task이다.
그렇기 때문에, high-level의 추론능력을 갖게 하기 위해서는 language에서보다 더 높은 masking ratio를 필요로 한다.

라고 언급하였습니다.

 

2.3) Decoder depth

Decoder depth에 따른 실험결과

decoder의 depth가 finetuning 실험에서는 영향을 별로 안끼치지만, linear probing에 대해서는 크게 끼친 것을 확인할 수 있습니다. 특히 충분한 depth가 없으면 성능이 유의미하게 감소합니다.

그 이유로는 reconstruction와 classification의 task간 거리에 의한 것이라고 논문에서는 언급하였습니다.

위 그림은 일반적인 딥러닝 모델의 학습과정을 나타낸 것입니다.

입력에 가까운 layer는 data 자체의 feature에 specific한 정보를 가집니다.

반대로, model의 최종 output에 가까운 layer는 task에 specific한 정보를 가집니다.

(정보 변화의 정도가 linear하다는 뜻은 아니고, 편의를 위해 저렇게 나타낸 것일 뿐입니다)

따라서, decoder depth가 1이면 decoder를 뗀 모델의 마지막 layer가 reconstruction task에 specific한 정보가 많이 남아있을 것입니다.

결국, classification으로 linear probing시에는 reconstruction task 정보가 많은 layer를 갖고 학습을 해야합니다.

다만, encoder도 trainable한 finetuning 상황이라면 데이터로부터 다시 classification task에 specific하게 변화할 수 있기에 성능에 큰 변화가 없습니다.

결과적으로, 불가피하게 linear probing을 해야하는 상황이라면, pretrain을 할 때 decoder depth를 충분히 확보해야 한다는 메시지를 주고 있습니다.

 

2.4) Data augmentation

Data augmentation별 결과

나름 직관적입니다. Data augmentation을 하는게 아무 것도 안한 것보다 조금이지만 성능이 더 좋았다는 이야기를 하고 있습니다.

이 부분에서 논문에서 주장하는 MAE의 좋은 점은

Contrastive learning 기반 실험 기법들은 성능이 data augmentation 설계에 심하게 좌우되는데, MAE는 data augmentation 유무에는 성능이 robust하다.

입니다.

 

추가로,

"random masking이 어느 정도 data augmentation의 역할을 간접적으로 해주기 때문이다"

라고도 설명하였습니다.

 

 

3. Discussion and conclusion

MAE는 masked patch를 복원하는 self-supervised learning을 통한 pretrain을 수행하고, 이를 pretrained weight로 initialize하여 원하는 target task 성능을 올리는 실험 세팅으로 유의미한 결과를 보여주었습니다.

Language domain에서 BERT의 발상을 차용하여 만든 것이 Vision domain의 MAE인데, '서로 전혀 다른 nature를 지니는 두 domain에서 같은 접근법(Transformer-to-BERT)으로 유의미한 결과를 도출한 것에 의의가 있다'고 개인적으로 생각합니다. 그 이유는 다른 domain에서도 이같은 접근법을 사용할 수 있다는 힌트를 주는 논문이 되기 때문입니다.

 

 

+ 실험 결과에서 Transfer learning 실험 결과는 언급하지 않았습니다(저의 게으름으로..)

+ Transfer learning에 대한 결과를 더 담아야할 것 같다는 comments가 있었습니다(arxiv에 있는데 이게 CVPR로부터 받은건지 어디서 받은건지는 불분명합니다).

혹시 포스팅 내용에 에러가 있다면 알려주시면 감사하겠습니다 :)