[논문 리뷰] MIM-Refiner: A Contrastive Learning Boost from Intermediate Pre-Trained Representations

Benedikt Alkin, Lukas Miklautz외 2명|arXiv (Cornell University)|2024. 02. 15.

Machine Learning in BioinformaticsBiochemistry, Genetics and Molecular Biology인용 수 3

한 줄 요약

MIM-Refiner는 고품질 인코더 블록에 부착된 다중 대비 학습 헤드를 통해 사전 훈련된 마스킹 이미지 모델(MIM)의 중간 표현을 개선함으로써 성능을 햖थन한다. 이는 새로운 최근접 이웃 정렬(NNA) 목적함수를 사용하는 것으로, ImageNet-1K에서 선형 탐색 정확도 84.7%와 1-shot 분류 정확도 64.2%를 기록하며 최신 기술 수준(SOTA)에 도달한다. 이는 훨씬 더 많은 데이터로 훈련된 더 큰 모델들을 능가하는 성능이다.

ABSTRACT

We introduce MIM (Masked Image Modeling)-Refiner, a contrastive learning boost for pre-trained MIM models. MIM-Refiner is motivated by the insight that strong representations within MIM models generally reside in intermediate layers. Accordingly, MIM-Refiner leverages multiple contrastive heads that are connected to different intermediate layers. In each head, a modified nearest neighbor objective constructs semantic clusters that capture semantic information which improves performance on downstream tasks, including off-the-shelf and fine-tuning settings. The refinement process is short and simple - yet highly effective. Within a few epochs, we refine the features of MIM models from subpar to state-of-the-art, off-the-shelf features. Refining a ViT-H, pre-trained with data2vec 2.0 on ImageNet-1K, sets a new state-of-the-art in linear probing (84.7%) and low-shot classification among models that are pre-trained on ImageNet-1K. MIM-Refiner efficiently combines the advantages of MIM and ID objectives and compares favorably against previous state-of-the-art SSL models on a variety of benchmarks such as low-shot classification, long-tailed classification, clustering and semantic segmentation.

연구 동기 및 목표

• 사전 훈련은 강력한데도 불구하고 저데이터 환경에서의 다운스트림 성능이 열악한 MIM 모델의 문제를 해결하기 위해.
• 디코더의 제한으로 인해 MIM 모델의 표현 품질이 최종 레이어가 아니라 중간 인코더 블록에서 최고에 도달한다는 것을 규명하기 위해.
• 모델을 다시 훈련하지 않고도 기존 MIM 특징을 의미론적 클러스터로 개선할 수 있는 경량이고 효율적인 방법을 개발하기 위해.
• 선형 탐색, 클러스터링, 소수의 샘플 분류와 같은 다운스트림 작업에 대한 일반화 능력과 이식성 향상을 위해.
• ImageNet-1K에서 data2vec 2.0로 사전 훈련된 ViT-H 모델을 개선함으로써, 수천 배 더 많은 데이터로 훈련된 더 큰 모델들보다도 극도로 낮은 샘플 수 환경에서 승승을 거두기 위해.

제안 방법

• MIM-Refiner는 MIM 인코더의 중간 블록에 다중 대비 학습 헤드를 부착한다. 특히 표현 품질가 최고에 이르는 블록들에 집중한다.
• 각 대비 헤드는 각 샘플을 배치 내에서 가장 가까운 이웃과 정렬하고, 나머지 모든 샘플은 밀어내는 방식으로 수정된 최근접 이웃 정렬(NNA) 목적함수를 적용한다.
• NNA 목적함수는 음성 쌍을 밀어내는 대신, 가장 가까운 이웃을 선택함으로써 긍정 쌍에 집중함으로써 신호 품질을 향상시킨다.
• 개선 과정은 순차적이고 짧으며, 사전 훈련에 사용된 동일한 ImageNet-1K 데이터로 몇 에포크만 훈련하면 된다.
• 배치 정규화 레이어는 대비 헤드에서 필수적이다. 이는 훈련을 안정화시키고 성능 저하를 방지하기 때문이다.
• 이 방법은 원본 MIM 모델 가중치를 유지하며, 오직 대비 헤드와 정규화 레이어만 피니팅한다.

Figure 1 : Linear probing state-of-the-art on ImageNet-1K over the last four years.

실험 결과

연구 질문

• RQ1 왜 MIM 모델은 강력한 사전 훈련을 거친 후에도 특히 소수의 샘플이 있는 상황에서 다운스트림 작업에서 성능이 열악한가?
• RQ2 MIM 인코더의 어디에서 의미적으로 더 유의미한 표현이 생성되는가? 그리고 이를 어떻게 활용할 수 있는가?
• RQ3 모델 전체를 다시 훈련하지 않고도 경량이고 짧은 개선 과정이 MIM 모델의 특징을 크게 향상시킬 수 있는가?
• RQ4 최종 레이어에서의 단일 헤드 대비 학습보다 중간 레이어에 다중 대비 헤드를 적용하는 것이 성능을 더 좋게 하는가?
• RQ5 개선된 MIM 모델은 훨씬 더 많은 데이터로 훈련된 더 큰 모델들을 초월할 수 있는가? 특히 극도로 낮은 샘플 수 환경에서.

주요 결과

• MIM-Refiner는 data2vec 2.0로 사전 훈련된 ViT-H 모델을 사용해 ImageNet-1K에서 선형 탐색 정확도 84.7%의 새로운 최신 기술 수준(SOTA)을 달성했다.
• 이 방법은 1-shot ImageNet-1K 분류에서 64.2%의 정확도로 새로운 SOTA를 수립했으며, DINOv2-g, OpenCLIP-G, MAWS-6.5B 등 최대 2000배 더 많은 데이터로 훈련된 더 큰 모델들을 능가했다.
• 개선된 모델은 k-NN 정확도 82.3%와 k-means 클러스터링 정확도 67.3%를 기록하여 잠재 공간 내에서 강력한 의미론적 클러스터링을 보였다.
• 개선 과정은 매우 효율적이며, 몇 에포크만으로도 열악한 품질의 특징을 최신 기술 수준의 품질로 향상시킬 수 있었다.
• 100%의 레이블로 피니팅했을 때도 개선된 모델은 원본 모델을 略로 뛰어넘었으며, 이는 고데이터 환경에서의 성능 저하가 없음을 확인했다.
• ADE20K 및 VTAB와 같은 다른 데이터셋으로의 전이 성능도 우수했으며, MAE-Refined는 벤치마크 전반에서 높은 mIoU와 평균 순위를 기록했다.