[논문리뷰] Masked Autoencoders Are Scalable Vision Learners(MAE)

 

많은 DL 모델들이 `supervised learning` (지도학습)의 방법으로 학습되면서 좋은 성능을 보여주었다.

`Label`이 있는 방대한 양의 데이터는 `LLM`과 같은 `fondation model`의 등장을 이끌어 냈고 `semi-supervised learning`과 합쳐져서 더 많은 양의 데이터로 학습이 가능해졌다.

 

그러나 실제적으로 라벨링을 수행하는 것은 시간과 돈이 많이 소요되는 작업이다.

또한 `Vision` 분야에서 LLM과 같은 거대 모델을 만들어 내기 위해서는 수백만개의 라벨링된 이미지 데이터를 넘어서 수억 또는 수십억개의 라벨링된 이미지 데이터가 요구 된다.

 

오늘 살펴볼 논문 `Masked Autoencoders Are Scalable Vision Learners`는 이러한 부분에서 효과적인 방법을 제안한다.

 

바로 `self-supervised` (자기 지도학습)을 통한 학습이다.

해당 방법을 통해서 라벨링하는 작업을 제거하 수 있으며, 어노테이션 되어 있지 않은 이미지 데이터셋에 대해서도 모델을 학습할 수 있다.

 

본글은 Masked Autoencoders Are Scalable Vision Learners 논문의 내용을 요약한 글이다.

조금더 자세한 내용을 알고 싶다면 아래의 논문을 참고하기를 바란다

 

✅Paper : https://arxiv.org/pdf/2111.06377

✅Code : https://github.com/facebookresearch/mae

GitHub - facebookresearch/mae: PyTorch implementation of MAE https//arxiv.org/abs/2111.06377

---

Intro

해당 논문에 대한 아이디어는 NLP 분야의 `GPT`와 `BERT`로 부터 시작되었다.

 

`GPT`와 `BERT`는 언어의 representation (표현)을 학습하기 위해서 데이터의 일부를 제거하고 제거된 일부를 예측하는 방법으로 모델을 학습한다.

 

예를 들어서

 

`손흥민은 월드클래스 공격수 이며 그는 대한민국 국가대표 주장이다`

 

라는 데이터가 있다고 생각을 해보자, 해당 데이터를 위의 방법으로 학습을 한다면

 

`손흥민 [] [] [] [] 그는 [] [] 주장이다` 

라는 Masked 되어진 문장을 만들어 내고 해당 Mask에 들어갈 단어를 예측하게 모델을 학습하는 것이다.

 

해당 논문에서는 이러한 부분을 Vision 분야로 가져와 적용을 한다.  바로 이미지에서 일부 픽셀을 masked 하고 이미지 모델이 masked 되어진 부분을 예측하는 방법으로 이미지의 representation을 학습하는 것이다.

 

그럼이제 논문을 통해서 해당 내용을 자세히 살펴보자

---

0. Abstract

해당 논문은 `masked autoencoders (MAE)`에 대해서 소개를한다. 이방법은 확장 가능한 self-supervised 학습을 vision 분야에 적용한 방법이다.

 

해당 MAE 방법은 input으로 주어지는 이미지에 대해서 Random 한 Patch에 Mask를 적용하고 Mask로 사라진 pixel을 재구성하는 방법으로 모델을 학습한다.

 

해당 방법은 2가지의 해심 디자인으로 구성되어있다.

 

1. 비대칭 encoder - decoder 구조

`Encoder` : 마스크없는 패치에 대해서만 연산을 수행한다. (즉, 사라지지 않은 patch에 대해서만 연산을수행)

`Decoder` : Encoder로 만들어진 latent representation과 masked 되어진 패치를 통해서 원본 이미지를 재구성한다.

 

2. 원본 이미지를 Masking 하는 비율 (75 % 정도)

 

위의 2가지 핵심 디자인을 통해서 기존 이미지 모델보다 학습속도가 3배이상 빨라졌으며 동시에 모델의 성능도 기존 모델들 보다 좋아졌다.

 

그리고 이미지의 label을 직접하는 시간도 줄일 수 있다.

 

해당 방법이 가능한 이유는 간단하게 Encoder는 마스킹되어지지 않은 패치에 대해서만 연산을 수행하기 때문에 높은 비율로 마스킹을 했을때 연산량이 대폭 감소한다. 따라서 이를 통해서 학습의 효율성을 증가시킨다.

 

---

1. Introduction

딥러닝 분야에서 모델의 용량은 지속적으로 증가하고 있다.

하드웨어의 발전에 힘입어 오늘날 100만개의 이미지를 모델에 적합시킬 수 있으며, 공개적으로 액세스할 수 없는 수억개의 라벨링된 이미지가 학습데이터로 요구되기 시작했다.

 

이러한 데이터의 욕구는 자연어 처리 분야에서 `self-supervised (자기 지도학습)`을 통해서 성공적으로 해결되었다.

`GPT` ,`BERT` 같은 모델들은 데이터의 일부를 제거하고 제거된 일부를 예측하는 방법으로 모델을 학습하였다.

이를 통해서 해당 모델들은 1000억개 이상의 파라미터를 포함하는 NLP 모델을 학습시킬 수 있었다.

 

`Masked autoencoder`는 `Denoising autoencoder`를 일반화한 방법으로 컴퓨터 비전 분야에서 자연스럽게 적용이 가능하다. 실제로 비전 분야에서 BERT 이전에 해당 방법이 연구가 되었다. 하지만 BERT의 성공 이후 해당 아이디어에 상당한 관심에도 불구하고 자연어 분야에 비해서 많이 뒤쳐져 있다. 

 

본 눈문에서는 이러한 부분을 극복하기 위해서 자연어 분야와 비전분야에서 masked autoencdoer가 어떠한 차이점을 갖는지에 대한 의문으로 시작을 한다.

 

1. `자연어 분야와 비전 분야의 모델 구조가 다르다`

최근 까지 비전 분야에서는 convolution이 주류를 차지해왔다. 해당 방법은 Grid를 통해서 연산을 하기 때문에 마스크 토큰이나 위치 임베딩 같은 지표를 conv와 통합하는 것이 쉽지 않다. 하지만 `ViT` 모델의 등장으로 해당 장애물이 제거되었다.

 

2. `자연어 분야와 비전분야의 정보 밀집이 다르다`

자연어의 경우 사람이 만들어낸 정보이기 때문에 semantic하고 정보가 매우 밀집적이다. 문장 몇개의 누락된 단어만 예측하도록 모델을 훈련할 경우 해당 작업은 정교한 언어 이해를 유도한다.

반면에 이미지의 경우 자연 신호이기 떄문에 공간적 중복성이 크다. 따라서 part, objects, scenes에 대한 이해 없이도 인접한 누락된 패치를 복구할 수 있다.

따라서 해당 논문에서는 자연어와 다르게 높은 비율의 masking을 수행한다. 이를 통해서 중복성을 크게 줄이고 낮은 수준의 이미지 이해를 넘어서 이미지의 전체적인 이해가 필요한 self-supervised 방법을 만들어 낸다. 

 

3. `autoencoder의 decoder의 역할이 이미지와 텍스트에서 다르게작동한다.`

이미지 분야에서 디코더는픽셀을 재구성하기 때문에 일반적인 인식 작업보다 출력의 의미 수준이 낮다. 이는 풍부한 의미 정보를 포함하는 누락된 단어를 예측하는 언어 분야와 대조적이다.

BERT는 디코더가 사소할 수 있지만 이미지의 경우 디코더 설계가 학습된 잠재표현의 의미 수준을 결정하는데 핵십적인 역학을 할수 있다. (즉, 디코더의 정교함이 전체 학습에 영향을 줄 수 있다)

 

해당 논문은 이러한 분석에 기반하여서 간단하고 효율적이고 확장가능한 형태의 masked autoencoder를구성하였다. 

 

방법은 아래와 같다

우선 랜덤 패치에 대해서 마스킹을 수행하고 -> 인코더는 마스킹 되어지지 않은 패치에 대해서 latent representation을 만드는 연산을 수행한다. -> 이후 마스킹된 패치와, latent representation을 활용하여 사라진 픽셀의 공간을 재구성한다. 

 

인코더의 경우 마스킹 되어지지않은 패치에 대해서만 연산을 수행하기 때문에 가중치가 가볍다. 

또한 75%의 높은 마스킹 비율을 사용하기 때문에 연산에 사용할 패치수 또한 적다. 

---

2. Related Work

자연어 분야 BERT와 GPT가 사용한 Masked 방법에 대해서 간단히 설명해주고

Autoencoding을 기존의 방법, Self - supervised learning에 대해서 설명해주고 있다.

---

3. Approach

해당 부분에서는 MAE 방법을 나누어서 구조적으로 설명해준다. 

 

`Masking`

MAE는 ViT의 관점을 그대로 가져와서 이미지를 겹치지  않은 패치로 나눈다. 그리고 일부를 mask하고 일부는 그대로 둔다. 해당 논문에서 사용하는 샘플링 전략은 랜덤 패치 샘플링이다. (uniform 분배를 따름)

 

이때 Masking은 높은 비율을 가지는데 이는 중복성을 제거하기 위해서이다. 따라서 이웃하는 패치를 통해서 이미지를 쉽게 복원할 수 없게 만든다. uniform 분배를 따르기 때문에 center bias를 방지해준다. 

 

따라서 매우 희박한 입력을 만들어내기 때문에 인코더는 효율적이게 되는 것이다.

 

`MAE encoder`

해당 인코더는 ViT 모델의 인코더 구조를 활용한다. 그러나 마스킹 되어지지 않은 패치에만 연산이 적용된다. 즉, 작은 subset에 대해서만 수행이된다. 

 

`MAE decoder`

해당 디코더에는  모든 토큰이 주어지는데 이는 임베딩된 보이는 패치와 마시킹된 패치가 들어간다. 또한 모든 토큰에 위치 임베딩을 추가하여서 어떤 위치에 해당하는지의 정보를 제공한다. 

디코더 구조는 또 다른 트랜스포머 블록이다.

 

MAE 디코더는 이미지 재구성 작업을 수행하기 위한 사전 학습에만 사용한다. 

따라서 학습이후 downstream에 맞추어 디코더를 제거한다.  

또한 디코더의 구조를 매우 경량화 하여서 인코더 보다 얕게 만들어 주었다. 이를 통해서 연산량을 크게 줄였다.

 

이러한 구조이기 때문에 비대칭 설계라고 한다.

 

`Reconstruction target`

MAE는 마스킹 되어진 패치로부터 pixel 값을 예측하여 input을 재구성한다.

디코더 출력의 각 요소는패치를 나타내는 픽셀 값의 벡터이다. 디코더의 마지막 레이어는출력 채널 수가 패치의 픽셀 값 수와 같은 선형 투영이다.

 

디코더 출력은 재구성된 이미지를 형성하도록 재구성한다. 손실함수는 픽셀 공간에서 재구성된 이미지와 원본 이미지 간으 평균 제곱 오차 (MSE)를 계산한다. 

 

BERT의 방법과 유사하게 마스킹 되어진 패치에 대해서만 손실을 계산한다.

 

또한 각 마스킹된 패치의 정규화된 픽셀 값을 재구성 대상으로 하는 변형도 연구하였다. 구체적으로 패치 내 모든 픽셀의 평균과 표준 푠차를 계산하여 이 패 치를 정규화 하는데 사용한다. 정규화된 픽셀을 재구성 대상으로 사용하면 실험에서 표현의 품질은 향상되었다.

 

---

MAE 방법은 실습하기 매우쉽다

특별한 연산이 필요하지 않기 때문이다. 

1. 모든 인풋 패치에 대해서 token을 만들어낸다. (linear projection & 위치 임베딩 추가)

2. 랜덤 셔플을 토큰 리스트에하고 마스킹 비율에 따라 목록의 마지막 부분을 제거한다.

해당 프로세스는 인코더를 위한 토큰의 작은 하위 집합을 생성하며 교체없이 패치를 샘플링하는 것과 동일하다.

3.인코딩 이후 마스킹되어진 토큰을 인코딩 되어진 패치 리스트에 추가하고 unshuffle 연산을 통해서 ((inverting the random shuffle operation)) 모든 토큰을 정렬한다.

 

이후 디코더가 모든 패치에 대해서 연산을 하여 이미지를 재구성한다.

---

정리

이후에 논문에서는 해당 방법을 통해서 이미지넷에 대한 실험을 수행한다. 

 

마스킹 비율의 영향을 보았을때 높은 비율의 마스킹이 매우 좋은 성능을 보였다.

이는 15% 비율을 사용하는 BERT와 대조적이다.

 

또한 linear probing 과 fine tuning가 다른 경향을 보인다.

(linear probing은 가중치 freeze 하고 마지막에 분류기만 학습)

linear probing의 경우 스위트 포인트까지 마스킹의 비율에 따라 꾸준히 증가하여

정확도가 격차가 최대 ~20% 까지 벌어진다

 

fine-tune의 경우 결과는 비율에 덜 민감하여 광범위한 마스킹 비율이 잘 작동한다.

모든 fine tune결과는 처음 부터 훈련하는 것보다 좋다

 

---

또한 디코더 디자인에 따른 실험도 수행을 하였다.

 linear probing의 방법을 활용할때는 깊은 디코더의 길이가 좋은 성능을 보인다. 

하지만  fine - tune으로 넘어갔을때 오히려 얇은 디코더가 더 좋은 성능을 보였다.

 

---

또한 MAE는 Data augmentation을 추가로 수행하지 않아도 잘 작동한다.

즉, 증강에 들어가는 시간 또한 줄일 수 있다.