오늘 리뷰할 논문은 Facebook AI Research(FAIR)에서 Kaiming He의 주도로 발표된 Masked Autoencoders Are Scaleable Vision Learners입니다. 해당 논문에서 제안한 MAE(Masked Autoencoder)는 기존의 다른 Transformer 기반의 vision encoder보다 훨씬 적은 연산으로 우수한 성능을 보이는 encoder구조를 제안하였습니다. MAE에서는 입력 값으로 들어오는 Visual token을 전부 사용하지 않고 일부만 입력 값으로 사용합니다.
아래 보이는 figure가 MAE의 구조입니다. 위에서 말한 것처럼 일부 visual token만 가지고 encoder에 태워주고 decoder에 넣어줄 때는 구멍이 뚫려있던 칸들에 masked token을 채워주면서 decoder에 태워주게 됩니다. 그리고 decoder가 masked된 부분을 원래 이미지도 복원하도록 학습을 진행하게 됩니다. 학습을 할 때에만 decoder를 사용하고 실제로 inference를 할 때는 visual encoder이기 때문에 encoder만 가지고 와서 입력 데이터의 feature를 뽑아내게 됩니다.
(https://arxiv.org/abs/2111.06377)
(해당 글은 개인적인 기록을 목적으로 잘못된 내용이 있을 수 있음을 알립니다.)
0. Abstract
Masked Autoencoder(MAE)는 이미지만 있으면 학습이 가능한 self-supervised learner입니다. 학습은 간단하게 말하자면 입력으로 들어오는 이미지의 일부 패치를 안 보이게 마스킹하고, 결과 값으로 해당 부분을 복원하는 방식으로 진행됩니다.
MAE는 크게 두 가지의 구조적 특징을 가집니다.
- 비대칭적인 Encoder-Decoder구조: Encoder는 전체 이미지의 patch 중 일부 만을 사용합니다. 또한 Deocoder는 Encoder에 비해 훨씬 파라미터 수(Transformer layer의 수라고 봐도 됨)가 훨씬 적습니다. 또한 Decoder는 Encoder에 들어갔던 위치에는 token은 latent를, encoder에 들어가지 않은 위치에 해당하는 부분에 masked patch로 채워서 feed forward해줍니다.
- MAE에서는 이미지 전체 patch 중 75%를 가리고 학습을 하여도 학습이 잘 되는 모습을 보입니다. 이런 masking 방법론은 모델을 학습할 때 굉장히 효율적으로 작용하게 됩니다.
이 두 특징 덕분에 scalability가 굉장히 개선된 모델을 만들 수 있었고, 결과적으로 학습 시간을 3배 정도 단축하는 동시에 정확도까지 개선이 되었다고 합니다.
여기서 주로 비교 대상이 되는 모델은 Vision Transformer(ViT)라는 모델입니다. 현재까지도 vision 쪽에서 transformer 구조를 사용할 때 해당 모델의 영향을 굉장히 많이 받고 있습니다. 혹시 ViT에 친숙하지 않으신 분은 아래 포스트를 함께 참고하시면 좋을 것 같습니다.
Vision Tranforer관련 포스트(추후 업로드 예정)
1. Introduction
딥러닝의 최근 연구 방향은 모델의 capacity를 키워서 성능을 개선 시키는 방향으로 진행되어 왔습니다. (모델의 크키, 파라미터의 수를 늘려왔다는 뜻) 이러한 모델을 학습 시키지 위해서는 점점 더 많은 양의 데이터가 필요합니다. 굳이 다르게 표현해보자면 모델의 capacity(용량)이 커지니까 모델이 암기할 수 있는 데이터가 많아지고 이에 따라 모델의 generalizability(다양한 데이터를 잘 처리하는 능력)이 상승하게 되는 구조입니다. 하지만 일반적으로 방대한 양의 데이터를 labelling 하는 것이 굉장히 비싼 작업이기도 하고 JFT-300M 같은 대용량 데이터셋은 대중적으로 접근이 불가능합니다. 즉 모델의 크기가 커져도 이를 뒷 받침할 수 있는 방대한 데이터 필요한데, 이를 구하는게 쉽지 않다는 것이죠.
NLP 분야에서는 이러한 문제를 self-supervised learning 방식으로 성공적으로 풀어나갔었습니다. 별도의 Labelling이 필요 없는 문서들을 적은 노력으로 모으고, 이 데이터의 일부를 가리는 식으로 학습을 진행한 바 있습니다. GPT 같은 경우는 현재 토큰 이후를 모두 가리는 autoregressive 방식, Bert는 임의의 토큰을 가리는 방식을 사용하였습니다. 그리고 여기에서 파생되어 다양한 방법을 사용해서 거대 언어 모델을 학습하는 것에 성공하였습니다.
저자들은 위에서 사용된 masked autoencoding 방식의 방식을 vision 분야에도 적용하게 위한 시도가 있었지만 NLP 분야만큼 엄청난 성과를 내진 못했다고 합니다. 그래서 NLP와 vision의 어떤 차이점이 masked autoencoding의 효율을 다르게 하는지를 먼저 답하고자 합니다.
- Vision 분야에서 지배적으로 사용되던 CNN기반의 구조는 현재 NLP의 압도적인 비중을 차지하는 Transformer 기반 구조와 비교했을 때 차이가 큽니다. CNN 계열 모델들은 masked autoencoding에 필수적인 요소들인 masked token이나 positional embedding 같은 요소들을 적용하기에 구조적인 한계가 있었습니다. 하지만 이는 Vision Transformer(ViT)에서 Transformer를 vision 분야에 성공적으로 적용하면서 어느정도 해결이 되었다고 봐도 무방합니다.
- 이미지와 텍스트 데이터의 구조적인 차이가 굉장히 큽니다. 텍스트 데이터를 이루는 단위는 단어 혹은 토큰 하나로, 각 토큰이 함축하고 있는 의미가 굉장히 큽니다. 하지만 이미지는 픽셀로 이루어져있고, 이 픽셀은 따로 두고 보았을 때 텍스트 토큰처럼 유의미한 의미를 갖지 못합니다. (Redundancy가 크다고 표현) 쉽게 말하면 텍스트 데이터는 토큰 몇 개를 빼면 문장을 이해하는데 어려움이 생기지만 이미지에서는 픽셀들이 어느정도 날아가도 어떤 이미지인지 알아볼 수 있다는 점을 생각해보시면 됩니다. 그래서 해당 페이퍼의 저자들도 이미지의 이런 높은 redundancy 때문에 굉장히 높은 비율(75%)의 이미지를 가리면서 redundancy를 최대한 낮추려는 시도를 합니다.
- BERT 같은 모델을 MLM pretrain을 할 때 masked token이 원래 어떤 token인지를 맞추기 위해서 얕은 MLP를 사용하여 token을 유추하게 됩니다. 반면에 masked image patch(token)을 복원하는 것은 이산적인 token을 맞춘다기 보다는 patch 안의 픽셀 하나하나를 복원해야합니다. (이미지 쪽이 훨씬 어렵겠죠) 그래서 vision 분야에서 autoencoder의 decoder가 모델을 학습하는데 굉장히 큰 영향을 주게 됩니다.
2. Related Work
Masked language modeling, Autoencoding, Masked image encoding, Self-supervised learning에 대한 간단한 소개를 하고 있습니다. 문서가 너무 길어질 것 같아서 해당 부분은 생략하도록 하겠습니다.
3. Approach
MAE는 간단히 말하면 데이터의 일부만 보여준 후 온전한 원본을 맞추는 방식으로 학습을 합니다. Encoder와 deocder가 대칭적인 구조를 가지는 일반적인 autoencoder와는 다르게 MAE는 두 가지 측면에서 비대칭적인 autoencoder 구조를 가집니다. 일단 decoder가 훨씬 lightweigth합니다. 또한 encoder는 이미지 patch 수 *(1-masking ratio) 만큼의 토큰을 사용하고, decoder는 이미지 patch 갯수 만큼의 토큰을 사용합니다.
가벼운 decoder는 학습 시 전체 모델 파라미터 수를 줄여주기 때문에 학습이 용이하다는 장점이 있으며, encoder에 들어가는 토큰 수가 줄어드는 것은 실제 사용하는 부분인 encoder의 계산 부하가 줄어든다는 장점을 가지게 됩니다.
3.1. Masking
일단 이미지를 패치로 만들어주어야합니다. ViT처럼 이미지를 겹치지 않는 격자로 나누어 하나의 token으로 취급합니다. ViT 페이퍼에서는 나눠진 토큰을 linear projection한다고 했던 것으로 기억합니다. 그래서 MLP를 사용하여 linear projection을 한다고 생각을 하실 수 있습니다. 하지만 실제 구현을 할 때는 계산적으로 부하가 적고, 결과에 지장이 별로 없는 1 layer CNN을 patch 간에 겹치지 않도록 하는 경우가 많은 것으로 알고 있습니다. MAE에서는 random sampling하여 임의로 patch를 가렸습니다. 또한 높은 masking ratio를 통해서 남아있는 patch들이 더 이상 redundancy가 높은(중요도가 상대적으로 낮은) 데이터가 아니게 만들어 주었습니다.
3.2. MAE encoder
Encoder는 ViT와 동일한 구조입니다. Transformer encoder를 쌓은 구조이며, 앞에서 말했던 것처럼 이미지를 겹치지 않는 격자(patch)들로 나눕니다. 여기서 patch들 중에서 25%만 사용됩니다. 모델의 파라미터 숫자가 같더라도 입력으로 들어오는 데이터의 양이 1/4가량 감소했기 때문에 계산할 때 사용되는 메모리가 더 적습니다.
위 특징은 모델이 transformer 계열 구조이기 때문에 그런 것입니다. 그리고 추후에 리뷰를 하겠지만 MAE와 비슷한 페이퍼로 BEIT라는 페이퍼가 있는데 여기서는 masked patch까지 모델에 넣어줍니다.
3.3. MAE decoder
MAE decoder는 이미지의 patch 갯수 만큼이 입력 값으로 들어가게 됩니다. Encoder에서 feature를 뽑은 patch는 그대로 들어가고, 나머지 자리에는 mask token(0으로 차있음)이 들어가게 됩니다. 앞에서 섹션에서 말씀드렸 듯이 입력으로 들어오는 patch의 수가 늘어나면 파라미터의 갯수가 같아도 연산 부하가 커지게 됩니다. 하지만 MAE deocoder는 transformer 블럭을 얕게 쌓기 때문에 계산의 부담이 적어지게 됩니다. 얕은 decoder는 적은 수의 파라미터라는 의미도 되는데, 이렇게 되면 학습해야할 파라미터 수가 상대적으로 적어지기 때문에 학습하는데 시간을 절약할 수도 있습니다. 정량적으로, decoder에서 하나의 token(patch)을 계산하는데 드는 computational cost가 encoder의 10% 정도 밖에 되지 않는다고 합니다.
Decoder는 원본 이미지를 복원하도록 학습(reconstruction task)을 합니다. 이를 통해 MAE encoder가 뽑아내는 feature가 유의미한 feature라는 것을 보여줄 수 있기도 합니다. 실제로 학습이 끝난 단계에서는 image feature extractor를 만드는게 목적이기 때문에 decoder는 떼고 encoder만 사용하게 됩니다.
MAE를 포함해서 다양한 모델에서 timm라이브러리를 사용해서 transformer block을 쉽게 불러옵니다. 모델 정의 코드 부분은 간결하게 짜여져있으니 한 번 확인해봐도 좋을 것 같습니다. (https://github.com/facebookresearch/mae/blob/main/models_mae.py)
3.4. Reconstruction target
Masked patch에 대해서만 pixel별로 mean square error(MSE) loss를 구해서 학습을 진행합니다.
3.5. Simple implementation
이미지를 패치로 나눈 뒤 리스트의 담아 섞어준 뒤 정해둔 비율만큼을 날린다는 내용입니다.
4. ImageNet Experiments
ImageNet 1K를 가지고 pretrain을 진행하였습니다. 성능은 ImageNet 1K Validation set을 기준으로 리포트 되어있습니다.
Baseline comparison (ViT-Large)
MAE는 encoder가 ViT와 동일한 구조이기 때문에 ViT와 classification accuracy를 비교하였습니다. 기존 ViT 페이퍼에서 리포트 한 성능이 scratch, original이고 MAE 저자들이 직접 학습해서 리포트 한 성능이 scratch, our impl입니다. 이 대목에서 transformer based model은 같은 모델이라도 학습을 어떻게 하느냐에 따라 성능 차이가 크게 날 수 있음을 알 수 있습니다.
그리고 마지막 baseline MAE가 해당 페이퍼에서 제시한 모델이며 성능이 가장 높았습니다.
4.1. Main properties
다양한 실험을 보여 줍니다.
Masking Ratio
Masking ratio를 바꾸어가면서 시험을 해본 결과 75%대의 masking을 했을 때 성능이 가장 좋았다고 합니다.
Decoder Depth
Decoder는 8층을 쌓을 때 성능이 가장 좋았다고 합니다. 다른 실험 테이블을 보아하니 ViT-L, ViT-H backbone 세팅에서 모두 8층짜리 Decoder를 사용한 것으로 보입니다. 정확히 기억은 안나지만 ViT-L이 transformer block 12층을 쌓았던 것으로 기억합니다.
Transformer Block Dimension
Decoder의 Transformer block 내부는 512차원일 때가 가장 좋았으며 encoder보다 dimension이 더 작습니다.
Mask Token
Mask token을 encoder에서 넣느냐 마느냐에 따른 성능 차이입니다. FLOPs는 연산 부하라고 보시면 되겠습니다.
Reconstruction Target
Decoder를 통해 원본이미지를 복원할 때 어떤 objective를 주는지에 따라 결과가 어떻게 달라지는지 실험한 부분입니다. 첫 번째와 두 번재가 위에서 말한 masked 됐던 부분의 pixel을 MSE loss 통해서 구해주는 방법입니다. 단 두 번째 실험에서는 픽셀들을 normalization을 한 후에 loss를 계산했습니다. 세 번째 실험은 patch 단위로 PCA를 진행해서 원본 patch의 주성분과 복원된 patch의 주성분이 같아지도록 학습하는 방법입니다. 네 번째 실험은 VQVAE를 사용한 실험입니다. VQVAE에 대해 소개할 기회가 된다면 소개를 해보도록 하겠습니다. 해당 방법은 VQVAE의 encoder를 따로 학습해서 모델학습 과정이 2 stage가 되기도 하고, 추가적인 연산 부하가 생긴다는 단점이 있었습니다.
그 외
그 외의 실험들입니다. 크게 언급할 만한 부분은 없는 것 같고, 바로 위 테이블만 보겠습니다. 마지막 테이블을 보면 ViT-H448 같은 경우는 MAE에서만 존재하는걸 확인할 수 있습니다. 이게 앞에서 말씀드렸던 이미지의 25%만 사용하는데에서 오는 연산적 이득 때문에 거대 모델을 학습할 수 있었기 때문에 가능했던 것 같습니다. 물론 이런 대형 연구 그룹은 연산 자원이 부족해서 못시키진 않겠지만..아무튼 본인들의 모델이 "scalable"하다는 것으로 보여주고자 굳이 ViT-H448을 하나 더 넣은 것 같습니다.
추가적인 정보가 궁금하시다면 페이퍼를 보시면 좋을 것 같습니다.
5. Transfer Learning Experiments
Object detection, object segmentation, semantic segmentation 등의 task에 pretrained model을 사용해서 학습한 결과 준수한 결과가 있었다고 합니다.