[PR] Masked Autoencoders Are Scalable Vision Learners

본 리뷰는 아래 유투브 링크에 실린 이진원 님의 페이퍼 리뷰를 바탕으로 개인적으로 공부하기 위해 제작되었습니다.

https://www.youtube.com/watch?v=mtUa3AAxPNQ&t=1463s

 

 

레즈넷의 저자인 카이밍허가 1저자로 참여한 논문

FAIR에서 2021년 11월에 발표함.

 

논문에서 말하고자하는 것은 아래 이미지가 전부임

이미지에 마스크를 씌워놓고 이를 복원하는 과정에서 모델을 사전학습하겠다는 것.

 

 

인트로던션

딥러닝 모델의 capablility와 capacity가 커짐으로써 백만장 정도의 이미지는 쉽게 오버피팅이 되기 때문에 더 많은 데이터가 필요로 하게 됨. NLP 쪽에서는 이를 self-supervised pretraining 방법을 통해 다루었음.

그래서 GPT나 BERT같은 모델들이 나왔고, GPT는 다음 단어를 예측하는 형태로, BERT는 중간에 단어를 가려놓고 맞추는 방식으로 문제를 풀었음. 그래서 모델 사이즈를 엄청나게 키우면서 엄청나게 많은 데이터로 self-supervised learning으로 학습을 할 수 있었음.

computer vision분야에서도 그런 접근방법이 많이 있었지만, 지도학습에 비해 성능이 잘 나오지 않아서 NLP에 비해 이러한 연구가 뒤쳐져있다는 것을 시사하면서 시작함.

 

왜 그럴까? 

첫번째로 아키텍쳐가 다름.

비전 분야는 CNN을 사용하기 때문에 mask token이나 positional embedding 같은 것을 통합시키기 어려웠음. 하지만 ViT의 등장으로 이제 가능하게 됨.

둘째로, information density를 언급함.

language는 사람이 만든것이고 이는 단어 하나하나의 정보량이 이미지에 비해 굉장히 큼. 그렇기 때문에 단어를 가려놓고 맞추는 방식으로 학습해도 높은 수준의 언어 이해를 하는 것처럼 보이지만, 이미지는 픽셀 하나하나가 어떤 특별한 의미가 있는 것이 아니고, 공간적인 redundancy가 있음. 그렇기 때문에 missing patch를 복원시키는게 그렇게 어려운 일은 아님.

셋째로, 디코더의 역할이 다름.

BERT의 디코더는 간단한 MLP로도 해결할수있지만, 이미지에서의 디코더의 디자인은 중요한 역할을 담당하고 있음. latent representation으로부터 semantic level을 결정해야하기 때문.

 

 

관련 연구

앞서 말했듯이 BERT와 GPT가 대표적임. 이것의 장점은 쉽게 scale up 할 수 있다는 점, 그로 인해 성능도 올라가고. 더불어 zero, few shot에 대한 성능도 좋아짐.

 

이미지 쪽에서는 denosing autoencoder과 같은 오토인코더 계열이 있음. 지워진 부분을 inpainting하는 context encoder와 모델도 있음. NLP쪽에서 영감을 받아 iGPT나 BEiT도 있음.

BEiT는 BERT Pre-training of Image Transformers 라는 것으로, 

 

이미지를 패치로 나누고, 패치 중 일부를 가리고, 이를 에측하는 방향으로 학습이 됨. 

다른 점이 있다면 픽셀단위를 예측하는 것이 아니고, 패치를 tokenize해서 그 토큰을 예측하는 방식으로 되어 있음.

그래서 사실은 pretraining하기 전에 더 pretraining을 통해 tokenizer를 학습해야함. 이미지 생성 모델인 DALLE에서 사용하는 그 방식을 사용함. 8000개 정도의 코드북을 만들어서 거기에 토크나이즈 됨.

BEiT가 예측하는 것을 결국에 원본 이미지로부터 나온 토큰이 됨.

 

그 다음 관련 연구로는 Self-supervised learning이 있음.

기존에는 pretext task와 같이 이미지를 회전시키고 얼마나 회전시켰는지를 맞춘다거나, 직소퍼즐을 푼다든지 등의 방식을 사용했지만, 이는 실제로 다운스트림 테스크에서 우리가 사용하려는 방식과 차이가 있음

또한 대조학습 기반의 방법도 있지만 이는 augmentation 방법에 디펜던시가 있음.

 

이 논문은 오토인코더 기반의 방법을 사용하지만, 기존과 다른점은 asymmetric design을 채택함.

마스크 토큰을 사용하지 않고, 마스킹된 부분은 제외하고 남아있는 패치들만 가지고 인코딩을 하고, 디코더에서는 마스킹된 패치를 그 자리에 끼워넣어, 디코더에서만 마스크 패치를 사용함. (아래 그림 참고)

 

 

 

복원 결과

loss function은 단순히 mse를 사용하고 있고, 당연히 결과값은 blurry하지만 우리는 이걸 하려는게 아니고, 그럴듯하게 복원만 해낼 수 있다면 좋은 representation이 학습돼서 다운스크림 태스크에서 좋은 성능을 낼수 있을 것이다 하는 기대를 할 수 있음.

(사람이 봤을때 굉장히 많이 가려져있는데 복원을 생각보다 잘하는게 놀랍다)

 

 

디테일

마스킹은 오버랩되지 않고, 랜덤하게 마스킹을 함(정규 분포 활용).

굉장히 많은 부분은 마스킹(75%)하기 때문에, 단순히 extrapolation 하는 것만으로 복원하는 것은 어려움. 때문에 어런식으로 학습이 된다면 좋은 representation을 학습할 수 있을 것이라 기대함.

 

인코더는 ViT를 사용했고, 패치들은 linear projection를 통해 임베딩 되고 positional embeddings이 더해짐.

실제로 인코더에 들어가는 부분은 25%밖에 되지 않기 때문에 기존 방법보다 3배 정도 빠르거나 큰 모델로 학습할 수 있다는 점은 강점.

 

디코더에서는 mask token도 함께 넣어 들어감. mask token은 학습 가능한 토큰임. 더불어 positional embeddings도 또 사용됨.

디코더는 pretraining때만 사용되고 실제로 다운스크림 테스크에서는 사용하지 않기 때문에 인코더와 독립적으로 구성하여 유연한 디자인을 채택함. 더불어 MAE의 디코더는 인코더에 비해 10% 미만의 토큰 당 연산량을 가지고 있기 때문에 pretraining에 대한 부담이 적음.

 

loss에는 복원의 대상이 되는 패치들만 고려됨. 

 

 

실험

ImageNet-1K(IN1K)를 통해 pretraining을 함.

그 이후에 end-to-end fintuning을 하거나 linear probing으로 사전 학습 모델의 성능을 테스트함.

베이스라인은 ViT-large. 일반적으로 이 모델은 굉장히 커서 쉽게 오버피팅되고 스크래치부터 학습시키기가 까다롭다.

그래서 강한 정규화 방법이 필요하다고 함.

그래서 중간은 본 저자(FAIR)가 여러가지 실험을 했을때의 결과이고, 우측은 MAE을 사용했을때 결과.

 

 

실험을 하나씩 보면, 마스크 비율에 따른 결과에서, 마스킹 비율이 높아질수록 성능이 올라가다가 75%에서 포화되는걸 볼 수 있음. 생각보다 높은 수준에서 마스킹이 필요한것은 놀라운 일.

BERT는 15% 정도 마스킹을 한것에 비해 상당히 높은 수준. 더불어 다른 비전쪽에서 사용했던 20~50% 수준과 비교해봤을때도 높음.

 

 

디코더 디자인도 살펴보면, 디코더의 depth와 dimension에 대한 실험임.

 

왼쪽에서 ft는 finetuning이고 lin은 linear probing을 의미하는데, lin같은 경우에는 block이 많아질수록 성능이 높아짐.

이는 recontruction task와 recognition task와의 차이가 있기 때문이라고 설명하고 있음.

왜냐하면 MAE의 끝단에 있는 layer들은 recon에 특화된 애들이지 recognition에 특화된게 아니기 때문에, 더 많은 block을 필요로하는것임.

당연히 ft에서는 디코더가 필요없기 떄문에 block이 얼마나 있던 큰 차이가 없는것을 볼 수 있음.

심지어 block을 1개만 사용하였더라도 ft에서 높은 성능을 보임.

이는 pretraining에 대한 부담을 낮춰줄 수 있고, ViT-large에 비해 9% 정도의 연산량.

 

참고로 회색으로 된 부분이 MAE에서 디폴트로 선택한 옵션.

 

 

마스크 토큰에 대한 이야기를 하자면, 마스크 토큰을 사용하면 오히려 성능이 낮아짐.

왜냐? 실제로 inference 할때와의 차이가 있기 때문인것으로 보임.

실제 inference할때는 마스킹이 안된 이미지들의 패치를 뜯어서 넣을텐데, 학습할때 마스크 토큰을 껴서 하게되면, 학습과 inference 사이의 괴리가 생기게 됨.

 

 

더불어 마스크 토큰을 안 쓰게 되면, computation 역시 굉장히 줄일 수 있음.

ViT-H 같은 경우에 4배 이상의 속도가 나타남. 25%만 남겨두고 학습하니까..

 

 

Reconstruction Target

재구성을 우리가 original pixel을 예측하게 할것인지, 아니면 PCA coefficient를 예측하게 할건지, 앞서 설명한 BEiT의 사용된 token(dVAE)을 예측하게 할것인지 실험을 해봤음.

결과적으로 픽셀을 직접 예측하되 normalizaion된 것을 예측하는 것이 성능이 제일 좋았음.

dVAE token도 좋았지만, dVAE encoder도 굉장히 큰 모델이고 (ViT-L 대비 40%), 여러가지로 오버헤드가 있음.

 

 

Data augmentation

단순히 random size cropping 방식만 사용해도 성능이 굉장히 잘 나왔고, 아무런 augmentation 없이도 성능이 잘 나옴.

이유는 이미 높은 수준의 마스킹을 하는 것 자체가 augmentation의 효과가 있기 때문.

더불어 contrastive leanrning 방식은 augmentation 방법에 의존적인 것에 비하면 더 좋은 결과.

 

 

Masking sampling strategy

마스킹하는 방법도 여러가지 시도를 해봄.

랜덤하게 하는 것이 성능이 제일 좋았음. 

 

 

Training schedule

얼마나 오래 훈련을 시킬것인지? 놀라운것은 계속해서 성능이 올라간다는 것.

MoCo v3 같은 경우는 ViT-L 기준 300에폭에서 saturation이 되는데, MAE는 1600에폭까지 올라감.

더불어 MAE는 25%만 보기때문에 훈련 부담도 contrastive learning에 비해 훨씬 덜함.

 

 

다른 SSL과의 비교

ViT-B나 ViT-L는 비슷하지만, 모델이 점점 커지면 커질수록 성능 차이가 벌어지는 것을 볼수 있음.

저 ViT-H448의 87.8%는 그 당시 SOTA 성능이었음. 그때까지 512size를 사용하는 87.1%가 최고 성능이었음.

 

 

다른 SL 모델들과의 비교

 

제일 아래있는게 원래 ViT 논문 성능인데, 여기서는 ViT-B에서 ViT-L로 갈때 오히려 성능이 떨어짐. (데이터가 부족하여 오버피팅됨)

여기서 본 논문의 저자들이 여러 실험을 통해 끌어올렸고, 거기서 MAE 기반의 모델들은 모델이 커질수록 점점 더 성능이 좋아짐.

(진원님의 개인적인 아쉬움은 MAE에 IN1K 말고 더 많은 데이터를 사용해서 평가했더라면 어땠을까 하는.. 어디까지 올라갈수있는지 한번 보여줬으면 하는..)

 

 

Partial fine-truning

 

이 표를 보면 ft과 lin의 결과가 uncorrelated 되어있다. 그래서 논문에서는 linear probing에 대한 비판적인 시각을 이야기함.

 

linear probing이 popular protocol이지만 non-linear feature에 대한 것은 잘 잡지 못하는 단점이 있음. 

여기서는 그래서 마지막 레이어 몇개를 추가하여 ft하는 것도 필요하다고 함.

 

그래서 ft을 위한 block 몇개를 가지고 테스트를 해봤는데, 1개만 써도 성능이 엄청 올라감을 볼 수 있음.

 

 

Transfer learning Experiments

다른 테스크에 대해서도 성능을 확인해보니, 잘 나왔다~ 하는 이야기.

COCO 데이터셋을 사용했고, Mask R-CNN이 사용됨. ViT 백본으로 FPN 구조를 사용하기 위해 변형함.

 

Sementic segmentation은 ADE20K라는 데이터셋에 UperNet이라는 것으로 실험함. (BEiT와 같은 코드)

 

픽셀과 토큰 중에 어떤게 좋은가? 높거나 거의 같음. 그래서 굳이 dVAE 토큰을 사용할 필요가 없음.

 

 

Discussion & Conclusion

어쨋든 이렇게 간단한 알고리즘으로 굉장히 스케일 업 잘 된다!

NLP에서는 SSL이 잘 되지만, 여전히 Cv에서는 SL이 도미넌트함. 그래서 CV 에서도 SSL을 통해 더 잘 나올수있음!