이전에 CL과 MIM의 비교 논문을 읽고 MIM에 흥미가 생겨 읽어본 논문이다. 다양한 MIM 논문 중 이를 고른 이유는 저자가 Kaiming he였다는 점이 한몫했던 것 같다.
Abstract
초록에서는 이 논문의 방법론, 실험, contribution을 짧게 설명했다. 눈에 띄는 부분만 요약하면 아래와 같다.
1. Input image를 random하게 mask하고 복원하는 masked auto encoders 소개
2. Encoder와 decoder 구조를 사용하며 encoder는 mask token을 사용하지 않는 비대칭 구조
3. 훈련시간을 3배 이상 줄이며 기존 방법론과 유사하거나 더 좋은 성능을 보여줌.
Introduction
하드웨어 발전으로 모델의 크기는 계속 커지나 그에 걸맞은 labeled images에 대한 수요가 생겼으나 구하기가 어렵다.
NLP에서 이러한 사례를 해결한 대표적인 모델들로 GPT, BERT를 예로 들었고 특히 BERT의 mask방식이 어째서 CV에서는 적용되지 않는지 의문을 제시한다.
NLP의 발전과 함께 CV에서도 NLP의 방법론을 차용하기 위한 연구가 많은데 이 논문은 그 중 BERT의 masked autoencoding방식을 중점적으로 보고 있다고 보면 될것 같다.
3가지 속성의 차이로 도메인 간의 차이가 생긴다고 제시했다.
1. Architecture적 차이
2. Data의 information density 차이
3. Decoder의 역할 차이
Architecture와 Data에서 각각 ViT 구조와 mask 방법론의 사용이유를 그리고 Decoder의 차이는 CV에서 decoder의 중요성을 강조한다.
이 중 2번에서 언어는 인간이 만든 highly semantic, information-dense한 데이터이기에 빠진 단어를 예측하기 위해 언어에 대한 높은 이해가 필요하다고 한다. 하지만 image는 반대의 속성을 가지고 빠진 patch는 주변 patch를 통해 예측가능하다. 즉 상대적으로 높은 이해가 필요 없는 task임을 의미한다.
이와 같은 이유로 image에서 매우 높은 비율을 masking하면 모델이 전체적인 이해를 필요로 하게 될 것이다. 이후에도 반복해서 나오며 이 컨셉이 MAE의 바탕이 된다고 보면 된다.
Abstract
Masking
높은 비율의 masking ratio를 사용해 패치를 랜덤 샘플링한다. 이를 통해 image를 일부 제거하면 주변 이미지를 통해 맞추는게 불가능해 더 높은 이해도를 가진 모델 학습이 가능할 것이다.
MAE encoder
방법은 ViT와 크게 다르지 않다. 하지만 encoder 학습 시 masked된 patch는 입력으로 사용하지 않으므로 자원효율성이 높다.
MAE decoder
Decoder는 pretrain시에만 사용하므로 encoder디자인과는 아무런 관련 없이 유연하게 디자인할 수 있다. 이 논문에서는 encoder와 비교해 10%이하의 연산량을 가지는 decoder를 사용했고 학습시간을 대폭 감소시켰다.
Reconstruction target
Decoder의 마지막 층은 image와 똑 같은 사이즈의 liner projection layer이다. loss계산시 기존의 patch 중 masking된 부분만을 사용해 MSE를 계산한다.
Simple implementation
MAE가 sparse연산이 없어 효율적으로 구현 가능하다는 점을 제시한다. Sparse operation이라는 단어부터 생소한데 인코딩 시에 token을 random shuffle하고 list의 끝을 자르고 디코딩 시에는 이 역연산 덕분에 가능하다고 한다. “연산구조를 효율적으로 구성했다”정도로 받아들이면 될 것 같다.
ImageNet Experiments
성능 실험을 위해 모델은 ImageNet-1K(IN1K) training set에 pretrain시켰다. 그리고 evaluation set에 fine-tuning, linear probing하여 성능을 비교한다.
진행된 모든 실험이 아래 표에 제시되어 있다.
Masking ratio
MAE는 기본 masking ratio로 75%를 사용하는데 15%를 masking하는 BERT보다 매우 높은 값이다. 이 역시 더 많은 patch를 가림으로서 모델을 조금 더 추론적으로 만들기 위함이라고 볼 수 있다.
Decoder design
Decoder가 깊을 때 linear probing성능이 더 좋음을 알 수 있는데 이는 decoder가 복원에 더 전문화되고 encoder를 더 잘 압축된 representations를 생성하는데 집중하게 할 수 있기 때문이다.
Mask token
놀랍게도 mask patch를 encoder에 input으로 주지 않을 시의 성능이 더 높다. 당연히 학습시간 역시 빠른데 즉 mask patch를 뺌으로서 성능과 학습시간 둘 모두에서 성능향상이 있었다.
Reconstruction target
PCA로 차원을 줄이거나 BEiT에서 사용한 dVAE token을 활용해도 성능 향상이 눈에 띄지는 않았다.
Data augmentation
기존의 contrastive learning 모델들과는 달리 center-crop과 rand size만을 활용한 기본적인 증강만으로 좋은 성능을 보였다. 저자는 masking이 증강의 역할을 일부 대신하기 때문일 것이라 주장했다.
Mask sampling strategy
왼쪽 사진과 같은 masking방법으로 실험했을 때 random방법의 성능이 제일 좋았다. 개인적으로 이전에 말했던 증강의 역할을 대신함 + 높은 수준의 이해가 있기 때문이라 생각한다.
Training schedule
저자는 특히 CL과의 차이를 비교하는데 CL은 multi image를 사용해도 300epoch만에 saturate하지만 MIM은 더 빠르고 오래 성능이 향상된다.
Comparisons with Previous Results
Comparisons with self-supervised methods
당시 유명한 CL, MIM방법론과 비교했을 때에도 좋은 성능을 보여준다.
Partial Fine-tuning
MAE의 linear-probing성능이 fine-tuning보다는 떨어지지만 몇 블록만 사용해도 fine-tuning과 유사한 성능을 보여준다.
Transfer Learning Experiments
Transfer Learning에 사용된 모델들은 IM1K로 사전학습하였다.
Detection과 Segmentation시 기존 모델보다 좋은 성능을 보여주고
Classfication시에도 기존 모델보다 좋은 성능을 보여준다. 추가적으로 pixel 예측과 dVAE토큰의 예측이 성능에 큰 영향을 주지 않는 점을 다시 한번 확인한다.
Discussion and Conclusion
Random한 masking 방식을 통해 모델의 전반적인 이해도를 높일 수 있었는데 이는 MAE가 rich hidden representation을 가지도록 학습되기 때문이다.
처음 그림을 볼 때는 Cutout증강이 생각났는데 차이라면 매우 높은 masking ratio에 있을 것 같다. 이런 증강을 사용할 때는 훈련시간이 느려질 수밖에 없다고 생각했는데 이런 단점을 masking한 부분을 제거하여 오히려 장점으로 만들었다.
특히 pretrain 시간을 비교하고 실험을 통해 보여준 점에서 극한의 fine tuning과 general representation을 찾는 추세 속에 오랜만에 모델의 성능과 효율성을 동시에 잡은 아키텍쳐적 연구인것 같다.
결론과 글 곳곳에서 저자는 NLP와 같은 궤도를 갈 수 있다고 언급했는데 CV의 이후 연구가능성을 보여주는 동시에 NLP에 비해 AI의 낮은 활용도에 대해 다시 한번 실감했다.
