[논문 리뷰] Masked autoencoders are effective solution to transformer data-hungry

Jiawei Mao, Honggu Zhou외 2명|arXiv (Cornell University)|2022. 12. 12.

Advanced Neural Network Applications인용 수 8

한 줄 요약

이 논문은 소규모 데이터 셋 환경에서 비전 트랜스포머(ViTs)의 성능을 향상시키기 위해 MAE 디코더를 단순화하고, 위치 예측 작업 및 클래스 토큰 인식형 대비 학습 작업을 도입한 소규모 데이터셋 마스킹 오토인코더(SDMAE)를 제안한다. SDMAE는 ViT 아키텍처를 수정하지 않고도 과적합을 줄이고 국소 불변성 및 표현 학습을 향상시켜, Tiny-ImageNet, CIFAR-100 및 의료 영상 기준치와 같은 소규모 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Vision Transformers (ViTs) outperforms convolutional neural networks (CNNs) in several vision tasks with its global modeling capabilities. However, ViT lacks the inductive bias inherent to convolution making it require a large amount of data for training. This results in ViT not performing as well as CNNs on small datasets like medicine and science. We experimentally found that masked autoencoders (MAE) can make the transformer focus more on the image itself, thus alleviating the data-hungry issue of ViT to some extent. Yet the current MAE model is too complex resulting in over-fitting problems on small datasets. This leads to a gap between MAEs trained on small datasets and advanced CNNs models still. Therefore, we investigated how to reduce the decoder complexity in MAE and found a more suitable architectural configuration for it with small datasets. Besides, we additionally designed a location prediction task and a contrastive learning task to introduce localization and invariance characteristics for MAE. Our contrastive learning task not only enables the model to learn high-level visual information but also allows the training of MAE's class token. This is something that most MAE improvement efforts do not consider. Extensive experiments have shown that our method shows state-of-the-art performance on standard small datasets as well as medical datasets with few samples compared to the current popular masked image modeling (MIM) and vision transformers for small datasets.The code and models are available at https://github.com/Talented-Q/SDMAE.

연구 동기 및 목표

• 의료 및 과학적 영상과 같은 소규모 데이터셋에서 비전 트랜스포머(ViTs)의 데이터 집약적 성향을 해결하기 위해.
• 마스킹 오토인코더(MAE)가 아키텍처 변경 없이 ViT의 데이터 의존도를 효과적으로 줄일 수 있는지 조사하기 위해.
• 소규모 데이터셋에서 과적합을 방지하기 위한 최적의 경량 디코더 구성(configuration)을 특정하기 위해.
• 자기 지도 학습 작업을 통해 국소화 및 불변성과 같은 인덕티브 바이어스를 MAE에 통합하기 위해.
• MAE 사전 훈련 중에 대비 학습 작업을 통해 ViT의 클래스 토큰을 효과적으로 사전 훈련할 수 있도록 하기 위해.

제안 방법

• 표준 ViT 모델의 소규모 데이터 훈련에 특화된 수정된 MAE 프레임워크인 소규모 데이터셋 마스킹 오토인코더(SDMAE)를 제안한다.
• 소규모 데이터셋에서 과적합을 방지하기 위한 최소한의 디코더 구성(configuration)을 특정하기 위해 광범위한 아블레이션 연구를 수행한다 (깊이=1, d_model=128).
• 사람의 주석 없이 부가적인 파rameter가 거의 없는 가벼운 위치 예측 헤드를 도입하여 마스킹된 토큰의 위치를 예측한다.
• 특징의 불변성을 향상시키고 ViT의 클래스 토큰에 대한 전용 사전 훈련을 가능하게 하는 대비 학습 작업을 설계한다.
• 마스킹된 이미지 모델링과 위치 예측, 대비 학습을 통합된 사전 훈련 목표로 결합한다.
• 모든 실험에서 일관된 사전 훈련 및 미세조정 프rotocol를 사용하여 공정한 비교를 보장한다.

실험 결과

연구 질문

• RQ1 ViT 아키텍처를 수정하지 않고 MAE를 소규모 데이터셋에 효과적으로 적용할 수 있는가?
• RQ2 소규모 데이터셋에서 훈련할 경우 과적합을 방지하기 위해 MAE에 최적의 디코더 구성(configuration)은 무엇인가?
• RQ3 소규모 데이터 설정에서 국소 인덕티브 바이어스 복원에 가벼운 위치 예측 작업이 얼마나 효과적인가?
• RQ4 대비 학습 작업이 표현 학습을 향상시키고 MAE 기반 ViT에서 클래스 토큰 사전 훈련을 가능하게 할 수 있는가?
• RQ5 마스킹된 이미지 모델링에 추가 자기 지도 학습 작업을 결합하면 소규모 및 의료 데이터셋에서 최신 기술 수준 성능을 달성할 수 있는가?

주요 결과

• SDMAE는 Tiny-ImageNet에서 상위-1 정확도 72.24%를 기록하여, 표준 MAE 및 기존의 ViT와 CNN 기반 모델을 모두 압도하며 소규모 데이터셋에서 최고 성능을 달성한다.
• 아블레이션 연구 결과, 위치 예측과 대비 학습을 모두 결합할 경우 성능이 가장 높으며, 위치 예측을 제거하면 정확도가 1.3% 감소함을 확인하였다.
• 최적의 디코더 구성(configuration) (깊이=1, d_model=128)는 과적합을 크게 줄이고 소규모 데이터셋에서 일반화 성능을 향상시킨다.
• SDMAE는 APTOS 2019와 같은 의료 영상 작업에서 최신 기술 수준 성능을 달성하여 고급 CNN 및 기타 MIM 방법을 능가한다.
• 의료 데이터에 대한 주의력 시각화 결과, SDMAE가 임상적으로 관련성이 높은 영역에 집중하는 것을 확인하여 실용적 유용성을 입증한다.
• 대비 학습 작업 덕분에 클래스 토큰의 효과적인 사전 훈련이 가능해졌으며, 이는 일반적인 MAE 접근 방식에서 자주 간과되는 점이다.