이미지에 포함된 일부의 패치만을 가지고 원본 이미지를 복구하는 encoder-decoder 구조를 제안하고, 이러한 방식이 성능에 미치는 영향을 분석한 논문이다.
Motivation
NLP 분야의 BERT와 같은 모델들은 문장의 일부 토큰을 마스킹하고 원래 문장을 예측하도록 학습하는 masked autoencoding 방식으로 사전학습을 수행함으로써 좀더 일반화가 잘 된 모델을 얻을 수 있다.
컴퓨터 비전 분야에서도 이와 유사하게 원본 이미지의 일부가 제거된 이미지를 복원시키는 문제를 풀어내도록 학습하는 방법이 제안되었지만, NLP 분야만큼의 성공적인 개선은 이루어지지 않았다.
저자들은NLP와 컴퓨터 비전에서 masked autoencoding의 효과가 다른 이유에 대해 의문을 가지고 다음과 같은 관점들에 대한 해답을 얻기 위한 연구를 진행했다.
- 컴퓨터 비전에서 Convolution Network를 주로 사용할 때에는 구조적인 차이로 인한 것으로 분석할 수도 있었지만, 유사한 구조를 가지는 Vision Transformer에서도 동일한 문제가 나타나는 것은 이상하다.
- 컴퓨터 비전과 언어는 정보의 밀도가 다르다. 문장은 토큰 하나하나가 큰 의미(high semantic) 를 담고 있지만 컴퓨터 비전은 배경처럼 의미가 덜 중요한 부분(low semantic)과 물체처럼 중요한 부분이 분산되어있다.
- 텍스트를 복원하는 디코더와 이미지를 복원하는 디코더는 서로 다른 역할을 한다고 볼 수 있다. 텍스트를 복원하는 디코더는 복원해야하는 대상 하나하나가 중요한 의미를 담고 있는 word이다. 하지만 컴퓨터 비전에서는 픽셀을 복원하며, 복원해야하는 픽셀중에는 배경처럼 정보가 적은 것도 포함되어 있다.
Contribution
- visual representation 학습을 위한 단순하면서도 효과적이고, 확장가능한 masked autoencoder 구조를 제안
Masked Autoencoder(MAE)
전통적인 오토인코더와는 다르게 저자들이 제안한 구조는 asymmetric한 구조를 가진다. 인코더는 이미지의 일부분만 가지고 latent representation을 추출하며, 디코더는 인코딩한 패치와 마스킹된 패치를 포함한 전체 이미지에 대해 복원을 하게 된다.
Masking
ViT 의 방식과 마찬가지로 이미지를 겹치지 않는 패치들로 나눈 후 랜덤하게 선택하여 마스킹을 수행한다. 이 때 높은 비율의 패치를 선택하게 되는데, 이 과정은 redundancy가 크게 제거되는 효과를 만든다고 한다.
많은 비율이 제거된 패치들로부터 생성된 이미지는 정답과 차이가 있기는 하지만 이미지로부터 얻을 수 있는 시각적인 정보 자체는 남아있는 것을 확인할 수 있다.
마스크를 너무 적게 사용하는 경우에는 주변 패치들의 정보를 보고 복원이 가능할 수 있기 때문에 문제 자체가 쉬워질 수 있기 때문에 좋은 feature를 뽑아낼 수 없게 된다고 한다.
MAE Encoder
인코더의 구조는 ViT와 동일하지만 마스킹 되지 않은 패치들에 대해서만 적용된다는 차이를 가진다. 패치들은 positional embedding이 더해진 상태로 인코더에 입력된다. 이러한 구조는 적은 수의 패치만을 입력으로 받을 수 있기 때문에 메모리 사용량과 연샨랑을 크게 줄여주며, 큰 모델을 활용할 수 있도록 만들어주기도 한다.
MAE Decoder
디코더는 인코더의 출력으로 나온 패치들과 마스크 토큰들을 입력으로 받는다. 마스크 토큰은 사라진 패치의 존재를 나타내는 학습 대상인 공유되는 벡터이다. 인코더와 마찬가지로 위치에 대한 정보를 담아내기 위해 각 패치는 positional embedding이 적용된다.
디코더는 pre-training 에서만 사용하기 때문에 구조는 자유롭게 가져갈 수 있다. 인코더와 다른 구조를 가지거나 더 가벼운 모델을 사용하는 것도 가능하다. 실험에 의하면 가벼운 모델로도 충분히 좋은 성능을 얻을 수 있다고 한다.
Reconstruction Target
MAE는 마스킹이 적용된 패치의 픽셀을 예측한다. 디코더는 최종적으로 패치의 픽셀을 출력하며, 원본 이미지의 픽셀과의 차이가 적어지는 방향으로 학습이 된다. 손실 함수는 Mean Squared Error(MSE) 를 사용하였고, BERT와 비슷하게 마스킹된 패치들에 대해서만 손실을 계산한다.
학습할 때 정규화된 픽셀 값을 사용하는 경우 더 좋은 성능을 얻을 수 있었다고 한다.
Experiment
ImageNet
Performance
Masking Ratio
- 모델이 원본과는 다른 출력을 내기는 하지만 얼추 비슷하게 복원하는 것을 확인할 수 있다. 단순히 선이나 텍스쳐를 확장해서는 완성할 수 없는 사물과 장면을 이해한다고 볼 수 있다.
- 최적의 성능은 높은 비율에서 나타났고, 75% 비율을 마스킹 했을 때 두 테스크 모두 성능이 좋음
- e2e fine-tuning의 경우 모든 비율에서 스크래치보다 좋은 성능을 보여주는 것을 확인
- e2e fine-tuning의 경우 마스킹 비율에 덜 민감하고, linear probing은 마스킹 비율에 따라 성능이 민감하게 변하는 것을 확인
Decoder Design
- linear probing에서는 충분한 디코더의 깊이와 너비가 성능에 큰 영향을 미치는 반면, e2e fine-tuning에서는 크게 영향을 받지 않는다.
⇒ 오토인코더의 레이어는 복원이라는 문제를 해결하는 것에 특화되어 있다. 따라서 인코더에서 출력된 feature는 인식과는 관련이 적기 때문에 linear probing의 성능이 안좋게 나온다.
- fine-tuning의 성능을 보면 인코더의 레이어가 인식 문제에 튜닝이 가능하다는 것을 알 수 있다.
- 결과적으로 MAE를 사용하면 가벼운 모델으로도 높은 성능을 얻을 수 있다.
Mask Token
- 인코더에 마스크 토큰을 사용하지 않을 때 연산랑이 줄어들 뿐만이 아니라, 더 좋은 성능을 보여주는 것을 확인할 수 있다.
⇒ 인코더가 항상 실제 패치만을 보고 학습을 했기 때문에 더 성능이 좋아졌다고 분석한다. 실제 인퍼런스와 학습이 처리하는 입력의 시나리오가 다르면 결과적으로 모델이 안좋은 영향을 받는다고 볼수 있다.
- 마스킹 토큰을 제외하면 학습 시의 메모리 사용량이 크게 감소하고 속도는 4배 이상 빨라진다.
Reconstruction Target
- 정규화된 픽셀의 값을 사용할때 가장 좋은 성능을 보여주는 것을 확인
- dVAE(image token in DALL-E)를 사용하는 경우 픽셀값을 사용하는 것보다 더 좋은 성능을 보여주지만, dVAE를 얻기 위해서는 이 토큰을 계산하는 네트워크를 먼저 학습해야 한다. 정규화된 픽셀값을 사용하는 것이 더 성능이 좋기 때문에 무의미한 과정이다.
Data Augmentation
- 특별한 augmenation을 적용하지 않아도 좋은 성능을 보여주며, 랜덤한 크기의 crop 이미지를 사용했을 때 가장 좋은 성능을 얻을 수 있었다.
⇒ 기존 다른 SSL에서 많이 사용하는 contrastive learning에서는 augmentation에 따라 성능이 민감하게 변하는 것을 본다면 엄청난 이점이라고 볼 수 있다.
- 저자는 마스킹 자체가 augmentation의 역할을 수행하게 되면서 pretext task를 어렵게 만드는 요소로 작용하고 있기 때문이라고 분석한다.
Mask Sampling Strategy
- 큰 영역을 지우거나 그리드 형태로 지우는 것보다 랜덤으로 마스킹을 했을 때 가장 좋은 성능을 얻을 수 있었다.
- 그리드의 경우 주변 패치들로부터 정보를 얻을 수 있기 때문에 더 쉬운 문제가 되면서 학습의 손실이 적고 복원 자체는 잘하게 되지만 latent representation의 질이 안좋아졌다고 볼 수 있다.
vs Other SSL
- ImageNet 1k 데이터셋만 사용하고도 다른 방법들보다 더 좋은 성능을 보여준다. 성능의 차이는 모델의 크기가 클수록 더 크게 나타난다.
- 동일한 시간, 많은 epoch를 학습하더라도 모델이 saturation 되지 않고 계속해서 학습이 되는 것을 확인할 수 있다. (MAE 1600 epoch 는 31시간, MoCo v3 300 epoch 는 36 시간이 소요되며 후자의 경우 성능이 더이상 오르지 않는 경향을 보였다.)
Transfer Learning
- Object Detection, Segmentation 문제에 적용했을 때에도 성능을 개선시키는 것을 확인할 수 있다.
Share article