[논문 리뷰] Pre-training of Lightweight Vision Transformers on Small Datasets with Minimally Scaled Images

Jen Hong Tan|arXiv (Cornell University)|2024. 02. 06.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 경량화된 비전 트랜스포머(ViTs)가 이미지 스케일링을 최소화하고 마스킹된 자동에코더(MAE) 사전학습을 통해 CIFAR-10 및 CIFAR-100과 같은 소규모 데이터셋에서 ResNet을 능가할 수 있음을 보여준다. 학습 가능한 위치 임베딩과 75% 마스킹 비율을 사용한 MAE 사전학습을 통해, 파rameter 수가 370만 개 미만이고 계산 비용도 최소화된 모델로 96.41%의 정확도를 달성하였다. CIFAR-10에서 96.41%, CIFAR-100에서 96.01%의 최신 기준 정확도를 달성하였다.

ABSTRACT

Can a lightweight Vision Transformer (ViT) match or exceed the performance of Convolutional Neural Networks (CNNs) like ResNet on small datasets with small image resolutions? This report demonstrates that a pure ViT can indeed achieve superior performance through pre-training, using a masked auto-encoder technique with minimal image scaling. Our experiments on the CIFAR-10 and CIFAR-100 datasets involved ViT models with fewer than 3.65 million parameters and a multiply-accumulate (MAC) count below 0.27G, qualifying them as 'lightweight' models. Unlike previous approaches, our method attains state-of-the-art performance among similar lightweight transformer-based architectures without significantly scaling up images from CIFAR-10 and CIFAR-100. This achievement underscores the efficiency of our model, not only in handling small datasets but also in effectively processing images close to their original scale.

연구 동기 및 목표

• 경량화된 비전 트랜스포머(ViTs)가 이미지 스케일링을 최소화한 상태에서 CIFAR-10 및 CIFAR-100과 같은 소규모 데이터셋에서 ResNet과 같은 CNN과 동등하거나 슈퍼어리어블 성능을 낼 수 있는지 조사한다.
• 자기지도 학습 사전학습을 활용하여, 저데이터 환경에서 ViTs와 CNN 간의 성능 격차를 해소한다.
• 컨volutional 인덕티브 바이어스나 이미지 업스케일링과 같은 아키텍처 수정 없이도 ViTs가 최신 기준 성능을 달성할 수 있음을 입증한다.
• 최소한의 증강 기법과 소규모 이미지(32×32)를 사용한 마스킹된 자동에코더(MAE) 사전학습의 효과를 평가한다.
• 소규모 데이터셋에서 더 큰 또는 더 복잡한 아키텍처와 경쟁 가능한 효율적인 소규모 파라미터 ViT 모델의 벤치마크를 설정한다.

제안 방법

• 75%의 패치 마스킹 비율을 가진 마스킹된 자동에코더(MAE) 프레임워크를 채택하여, 인코더에 입력되는 패치는 총 25%에 불과하다.
• 정현적 임베딩 대신 학습 가능한 위치 임베딩을 인코더와 디코더 양쪽에 사용하여 공간 인식 능력과 적응성을 향상시킨다.
• 인코더와 독립적인 유연한 아키텍처를 가진 별도의 디코더를 설계하여 재구성 품질을 향상시킨다.
• 마스킹된 패치 재구성(MSE 사용)과 비마스킹된 패치에 대한 할인된 손실을 포함하는 병합 손실 함수를 계산하여 학습 안정성을 향상시킨다.
• 표준 교차엔트로피 손실과 표준 최적화 프로토콜을 사용하여, 다운스트림 분류 작업에서 사전학습된 ViT 모델을 피니튜닝한다.
• 이미지 스케일링을 최소화하여 32×32에서 36×36으로 해상도를 뿐만 아니라 원본 데이터 무결성을 유지하고 계산 오버헤드를 줄인다.

Figure 1: Modified MAE: Implementation using separate learnable positional embeddings for Encoder and Decoder. The Encoder and the Decoder in this figure only consist of transformer layers, a layer of layerNorm and a linear projection layer.

실험 결과

연구 질문

• RQ1 이미지 업스케일링 없이도 경량 비전 트랜스포머가 CIFAR-10 및 CIFAR-100과 같은 소규모 데이터셋에서 ResNet을 능가할 수 있는가?
• RQ2 최소한의 데이터 증강과 소규모 이미지 해상도에서 MAE 사전학습이 저데이터 환경에서 인덕티브 바이어스가 부족한 ViTs의 성능을 향상시킬 수 있는가?
• RQ3 인코더와 디코더 양쪽에 학습 가능한 위치 임베딩을 사용할 경우 재구성 품질과 다운스트림 분류 성능에 어떤 영향을 미치는가?
• RQ4 한 데이터셋에서 사전학습한 ViT 모델이 다른 데이터셋에 대해 피니튜닝할 경우 효과적으로 일반화될 수 있는가, 아니면 도메인 내 사전학습이 필수적인가?
• RQ5 병합된 손실 함수(마스킹된 패치 + 비마스킹된 패치 손실)는 학습 안정성과 모델 일반화에 어떤 영향을 미치는가?

주요 결과

• 제안된 MAE 사전학습된 ViT 모델인 Mae-ViT-C10과 Mae-ViT-C100은 각각 CIFAR-10에서 96.41%, CIFAR-100에서 96.01%의 상위-1 정확도를 달성하여 최신 기준 경량 ViT 모델을 초월하였다.
• 유사한 파라미터 수(372만 개)와 MACs(0.26G)를 가진 ViT-Lite-7/4 모델은 각각 93.57%와 73.94%의 정확도를 기록하여, MAE 사전학습 전략의 우수성을 입증하였다.
• 추가적인 컨volutional 인덕티브 바이어스 없이도 CVT-7/4(94.01% on CIFAR-100) 및 CCT-7/3×2(95.04% on CIFAR-10) 등의 다른 ViT 변종보다도 성능이 뛰어나, 성능 향상의 유의미함을 입증하였다.
• 한 데이터셋에서 사전학습하고 다른 데이터셋에 피니튜닝하는 방식은 수렴 속도를 빠르게 했지만, 최종 성능는 다소 낮아, 도메인 내 사전학습이 더 효과적임을 시사하였다.
• 학습 과정에서 재구성 품질이 점차 향상되었으며, 3020번째 에포크에 이미지 세부 정보가 뚜렷하게 개선되는 것을 관찰하여 효과적인 특징 학습이 이루어졌음을 확인하였다.
• 비마스킹된 패치에 대한 할인된 손실 사용이 학습 안정성을 향상시키고, 특히 초기 학습 단계에서 일반화 성능 향상에 기여하였다.

Figure 2: The architecture of our MAE. ’B’ stands for batch size, ’N’ for number of embeddings