Reference

He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. B. (2021). Masked autoencoders are scalable vision learners. arXiv (Cornell University). https://doi.org/10.48550/arxiv.2111.06377


Problem Statement (Focused Issue)

 

  • Aided by the rapid gains in hardware, models today can easily overfit one million images and begin to demand hundreds of millions of - often publicly inaccessible - labeled images
    • This issue has been successfully addressed in natural language processing (NLP) by self-supervised pre-training: based on autoregressive language modeling in GPT and masked autoencoding in BERT
      • They remove a portion of the data and learn to predict the removed content
      • 이 masked autoenchoder 아이디어를 computer vision에도 적용시킨다면?
Self-supervised pre-training
Self-Supervised Learning은 레이블이 없는 데이터에서 유용한 표현을 학습하는 방법입니다. Pre-Training은 이러한 자기 지도 학습을 통해 초기 모델을 훈련시키는 단계로, 나중에 다른 특정 작업(Supervised Task)을 수행하기 위해 이 모델을 Fine-Tuning할 수 있습니다.
  • BUT! 이 아이디어를 적용시키는 데 NLP만큼 뛰어난 성과를 얻진 못함.
    • Why? What makes masked autoencoding different between vision and language?
      1. 아키텍처의 차이
        • Vision에서는 주로 합성곱 신경망(CNN)이 사용되었으며, 이 합성곱 신경망은 규칙적인 grid에서 작동하므로 mask token이나position embedding을 통합하기 어려웠다. 반면, Natural language processing에서는 BERT와 같은 모델이 사용되었다.
        • 그러나 최근에는 Vision Transformers (ViT)의 도입으로 이 아키텍처의 차이가 해결되었다.
      2. 정보 밀도의 차이
        • 언어는 인간이 생성한 신호로서 높은 의미적 정보를 포함하고 있습니다. 문장에서 몇 개의 단어만 마스킹하더라도 모델은 복잡한 언어 이해 능력을 요구한다.
        • 반면, 이미지에서는 spatial redundancy(공간적 중복)이 많아 높은 비율로 마스킹하더라도 인접한 패치로부터 쉽게 복원될 수 있다.
          • Patch간 상관관계가 존재하기 때문!
        • 이를 극복하기 위해, MAE에서는 높은 비율의 무작위 패치를 마스킹하는 전략을 사용한다.
      3. 디코더의 역할 차이
        • 언어 모델에서는 디코더가 의미적 정보를 많이 포함한 단어를 예측하지만, 시각적 모델에서는 픽셀을 복원한다.
          • 이는 시각적 모델의 디코더 출력이 언어 모델의 디코더 출력보다 의미적 수준이 낮다는 것을 의미한다.
          • BERT에서는 디코더가 매우 단순할 수 있지만, 이미지에서는 디코더의 설계가 중요한 역할을 한다.

Method (Approach)

  • Masked AutoEncoder (MAE)
    • Mask random patches of the input image and reconstruct the missing pixels
    • Based on:
      1. Asymmetric encoder-decoder architecture
        • Encoder: operates only on the visible subset of patches (without mask tokens)
          • Token: 이미지 처리에서 토큰은 이미지의 일부분이나 패치를 의미한다. 이미지를 일정한 크기의 작은 조각으로 분할하여 각각을 토큰으로 간주한다.
          • 즉, 마스킹된 패치는 인코더에 입력되지 않고, 대신 마스킹된 위치를 나타내는 mask token으로 대체됩니다. Encoder 이후 Mask tocken은 introduced 된다.
        • Decoder: reconstructs the original image in pixels  from the latent representation (encoded patches) and mask tokens 
          • Latent Representation은 입력 이미지의 추상적이고 압축된 표현, 즉, 인코더를 통해 입력 이미지가 변환된 후의 상태를 나타냅니다.
          • After pre-training, the decoder is discarded and the encoder is applied to uncorrupted images (full sets of patches) for recognition tasks
      2. Masking a high proportion of the input image (e.g., 75%) yields a nontrivial and meaningful self-supervisory task

 

MAE architecture

 

  • Mechanism of MAE
    1. During pre-training, a large random subset of image patches (75%) is masked out
      • Divide an image into regular non-overlapping patches
      • Then, sample a subset of patches and mask the remaining ones
    2. The encoder is applied to the small subset of visible patches and mask tokens are introduced after the encoder
      • Encoder는 Vision Transformer(ViT)를 기반으로하며, Masking되지 않은 패치에만 적용
        • Patch Embedding: 패치를 linear projection하고 positional embeddings을 추가
        • Transformer blocks: 여러 트랜스포머 블록을 통해 처리
      • 마스킹된 75%의 패치는 인코더에 입력되지 않으며, 대신 마스킹 토큰으로 대체된다.
        • 전체 패치 중 일부(25%)만 처리하기 때문에, 계산 및 메모리 효율성이 높다.
    3. The full set of encoded patches and mask tokens is processed by a small decoder that reconstructs the original image in pixel
      • Positional embeddings 추가: 모든 토큰에 위치 임베딩을 추가하여, 마스크 토큰이 이미지 내 위치 정보를 가질 수 있도록 한다.
        • Without this, mask tokens would have no information about their location in the image
      • Reconstruction target: MAE reconstructs the input by predicting the pixel values for each masked patch
      • Loss function은 원본 이미지와 재구성된 이미지 간의 Mean Squared Error(MSE)를 계산하며, 이는 마스킹된 패치에 대해서만 계산된다.
    4. After pre-training, the decoder is discarded and the encoder is applied to uncorrupted images (full sets of patches) for recognition tasks

Left: the masked image, Middle: MAE reconstruction, Right: the ground-truth (Images are from ImageNet validation images)

 

ViT (Vision Transformer)
Vision Transformer(ViT)는 자연어 처리(NLP)에서 성공을 거둔 Transformer 아키텍처를 Computer Vision 작업에 적용한 모델입니다. Transformer는 주로 텍스트 데이터를 처리하기 위해 설계되었지만, ViT는 이를 이미지 데이터에 적용하여 이미지 분류와 같은 작업을 수행합니다. ViT는 CNN(Convolutional Neural Network)과 달리 전통적인 합성곱 연산을 사용하지 않으며, 대신 Transformer의 Self-Attention 메커니즘을 사용합니다.
  • 왜 Masking ratio를 75%로 했을까?
    • 75% works well for both fine-tuning and linear probing


ImageNet Experiments

  • Self-supervised pre-training on the ImageNet-1K training set => Supervised training to evaluate the representations with:
    1. End-to-end fine-tuning: 사전 학습된 모델의 모든 파라미터를 다시 학습시켜 특정 작업에 맞게 최적화하는 방법
    2. Linear probing: 사전 학습된 모델의 파라미터를 고정시키고, 마지막 레이어에 새로운 선형 분류기를 추가하여 학습하는 방법

  • ViT-Large(ViT-L/16) 모델을 사용하여 MAE를 적용한 결과, 기존 학습된 모델보다 성능이 크게 향상됨.
    • ViT-L을 처음부터 학습하는 것보다 MAE로 사전 학습 후 파인 튜닝하는 것이 성능이 훨씬 높았다 (76.5% → 84.9%).

Main Properties from Ablation experiments

  • Masking ratio:
    • A high masking ratio (75%) works well for both fine-tuning and linear probing

  • Decoder Design:
    • Depth: 깊은 decoder가 linear probing 정확도를 향상시킴.
      • The last several layers in an autoencoder are more specialized for reconstruction but are less relevant for recognition
      • Linear Probing은 사전 학습된 모델의 최종 레이어에 선형 분류기를 추가하는 방법입니다. 우리의 최종 레이어는 디코더와 관련이 있기 때문에, 디코더가 깊어질수록 Linear Probing에서 더 높은 성능 개선이 나타난다.
      • How about fine-tuning?
        • The last layers of the encoder can be tuned to adapt to the recognition task
        • 이는 파인 튜닝 과정에서 전체 모델이 재학습되기 때문에, 디코더의 깊이가 성능에 큰 영향을 주지 않는다.
  • Mask token:
    • 인코더에서 마스크 토큰을 건너뛰고 디코더에서 이를 적용. 이는 훈련 시 계산량을 크게 줄이고 정확도를 향상시킴.
  • Reconstruction Target:
    • Pixel-Based Reconstruction: 이미지 패치를 픽셀 단위로 재구성
      • Normalization: 픽셀 값을 정규화하면 정확도가 향상됩니다. 이는 패치 간의 대비를 높여주기 때문이다.
        • Normalization은 각 패치 내의 픽셀 값들을 일정 범위(예: 0에서 1 사이)로 스케일링하면, 패치 내부의 대비가 증가함에 따라, 동일한 패치 내에서 어두운 부분은 더 어두워지고 밝은 부분은 더 밝아진다.
      • PCA (주성분 분석): 주성분 분석(PCA)을 통해 가장 큰 PCA 계수를 재구성 목표로 사용했지만, 이는 정확도를 저하시켰습니다. 이는 High-frequency component이 재구성에 중요한 역할을 하기 때문이다.
        • PCA: 고차원 데이터의 차원을 축소하여 데이터의 주요 정보를 보존하면서 데이터의 구조를 단순화하는 기법
        • High-frequency component: 밝기나 색상이 급격하게 변하는 부분. 혹은 이미지의 세밀한 부분이나 복잡한 패턴이 포함되는 부분이다.
    • Token-Based Reconstruction: BEiT와 같이 토큰화를 사용한 재구성 방식
      • MAE 디코더가 cross-entropy loss을 사용하여 토큰 인덱스를 예측합니다. 이는 fine-tuning 정확도를 0.4% 향상시키지만, Normlaized pixels과 비교했을 때는 별다른 장점이 없다.
      • 토큰화는 추가적인 사전 학습 단계를 필요로 하며, 큰 계산 비용(예: ViT-L의 40% FLOPs)을 추가합니다. 따라서, 토큰화는 전이 학습(transfer learning)에서 특별한 이점을 제공하지 않았다.
  • Data Augmentation:
    • MAE는 데이터 증강이 거의 필요하지 않으며, 무작위 마스킹이 데이터 증강의 역할을 대체함.
      • Contrastive learning과는 대조적!
        • Contrastive Learning은 인스턴스 간의 유사성과 차이를 학습하여 데이터 표현을 개선하는 자기 지도 학습(self-supervised learning) 기법입니다. 이 접근법은 주로 데이터 포인트 쌍을 사용하여 유사한 쌍을 가까이, 다른 쌍을 멀리 위치시키는 방식으로 작동한다.
        • Contrastive Learning에서 첫번째 단계가 Augmentation인만큼 Augmentation에 의존적.
  • Mask Sampling Strategy:
    • Random Sampling: MAE에서 가장 좋은 성능을 보이며, 높은 마스킹 비율을 사용할 수 있게 함.
    • Block and Grid Sampling: 블록 샘플링은 50% 비율에서만 잘 작동하며, 그리드 샘플링은 더 쉬운 작업을 제공하지만 성능이 낮음.
      • Block-wise masking: reasonably well at a ratio of 50% but degrades at a ratio of 75%
        • Higher training loss is observed
        • The reconstruction is also blurrier
      • Grid-wise sampling: regularly keeps one of every four patches
        • Easier task and has lower training loss
        • The reconstruction is sharper
        • However, the representation quality is lower

  • Training schedules:
    • A longer training schedule gives a noticeable improvement

 

Comparisons with Previous Results

Comparisons with previous results on ImageNet-1K

  • 비교 Model: ViT-B, ViT-L, ViT-H
  • IN1K (ImageNet-1K) 데이터로 사전 학습 후, end-to-end fine-tuning 진행
ViT-B 모든 방법이 비슷한 성능을 보임
ViT-B 방법들 간의 차이가 더 큼, 큰 모델의 경우 과적합을 줄이는 것이 도전 과제
ViT-H MAE가 성능 향상을 보여줌
MAE ViT-H 모델을 사용하여 IN1K 데이터로만 학습했을 때 86.9%의 정확도 달성, 448 크기로 미세 조정했을 때 87.8% 정확도 달성 (이전 최고 기록인 87.1%를 초과)

 

MAE pre-training vs. supervised pre-training

  • MAE 사전 학습 했을 경우, IN1K 데이터로만 학습했을 때 일반화 성능 향상, 더 높은 용량의 모델에서 더 큰 이득을 보임

 

Partial Fine-tuning

  • 사전 학습된 모델의 특정 부분만을 미세 조정하는 방법
    • 사전 학습된 모델의 특정 레이어(예: 초기 레이어 또는 특정 블록)를 고정한 후, 고정되지 않은 나머지 레이어만을 학습하여, 전체 모델을 다시 학습하지 않고도 모델의 성능을 최적화한다.
    • 모델의 특정 블록만을 미세 조정하여, 전체 모델을 다시 학습하는 것보다 효율적으로 모델을 최적화할 수 있음을 확인하였다.
  • MoCo v3와 MAE를 비교하여, MoCo v3가 Linear Probing (Tuning 0 blocks is linear probing)에서 더 높은 정확도를 보이지만, 부분 미세 조정에서는 MAE가 더 나은 성능을 발휘함을 보여준다.
    • MAE의 표현은 Linear Separability이 낮습니다. 이는 Linear Probing에서 MoCo v3보다 낮은 성능을 보이는 이유이다.
      • Linear Separability은 데이터를 선형 경계를 사용하여 두 개 이상의 클래스나 그룹으로 명확하게 구분할 수 있는 능력을 의미한다.
    • MAE는 비선형 특징이 강하여, 비선형 헤드를 튜닝할 때 좋은 성능을 발휘한다.

 

Transfer Learning Experiments

  • MAE는 COCO 객체 탐지 및 분할, ADE20K의 semantic segmentation, 그리고 iNaturalist 및 Places 데이터셋에서 분류 작업에서 강력한 transfer learning 성능을 보이며, 픽셀 기반 MAE는 간단하면서도 빠르고 우수한 성능을 보여줌을 알 수 있다.

Discussion & Conclusion

  • We observe on ImageNet and in transfer learning that an autoencoder - a simple self-supervised method similar to techniques in NLP - provides scalable benefits
  • We note that images and languages are signals of a different nature and this difference must be addressed carefully
    • Image: merely recorded light without a semantic decomposition into the visual analogue of words
      • Remove patches that most likely do not form a semantic segment

Overall Reflection

 

최근에 읽은 논문에서 NLP에서 사용되던 아이디어를 Computer Vision 분야에 적용한 점이 참신하게 느껴졌음. 이를 통해 특정 분야만 공부하는 것이 아니라, 다른 분야도 끊임없이 탐구해야 한다는 중요성을 깨달음. 특히, 논문에서 언급한 것처럼 기존 방법을 응용할 때, 인풋 데이터가 다르다면 그 데이터의 성질을 제대로 파악하고 활용하는 것이 매우 중요하다는 점을 실감했음.

 

여러 논문을 읽으면서 각 분야(Computer Vision, Reinforcement Learning, NLP 등)마다 발전된 정도와 사용하는 방법이 다를 수 있지만, 서로 다른 분야의 아이디어를 융합하는 연구가 활발하게 진행되고 있다는 것을 느꼈음. 이를 통해 다양한 분야의 지식을 통합하여 새로운 연구 방향을 모색하는 것이 중요하다는 생각이 들었음.

 

 

+ Recent posts