[논문 리뷰] Task-customized Masked AutoEncoder via Mixture of Cluster-conditional Experts

Zhili Liu, Kai Chen외 5명|arXiv (Cornell University)|2024. 02. 08.

Advanced Clustering Algorithms Research인용 수 6

한 줄 요약

MoCE를 제안하는 분할-가이드 혼합 전문가 확장(MoCE), 작업 맞춤형 자기지도 학습 사전학습을 위한 MAE의 클러스터링 기반 확장으로 11개의 다운스트림 작업에서 전이 성능 향상 및 탐지/분할에서의 최첨단 결과를 제공합니다.

ABSTRACT

Masked Autoencoder~(MAE) is a prevailing self-supervised learning method that achieves promising results in model pre-training. However, when the various downstream tasks have data distributions different from the pre-training data, the semantically irrelevant pre-training information might result in negative transfer, impeding MAE's scalability. To address this issue, we propose a novel MAE-based pre-training paradigm, Mixture of Cluster-conditional Experts (MoCE), which can be trained once but provides customized pre-training models for diverse downstream tasks. Different from the mixture of experts (MoE), our MoCE trains each expert only with semantically relevant images by using cluster-conditional gates. Thus, each downstream task can be allocated to its customized model pre-trained with data most similar to the downstream data. Experiments on a collection of 11 downstream tasks show that MoCE outperforms the vanilla MAE by 2.45\% on average. It also obtains new state-of-the-art self-supervised learning results on detection and segmentation.

연구 동기 및 목표

• 표준 MAE를 semantically 다른 다운스트림 작업으로 전이할 때 부정 전이가 발생하는지 입증합니다.
• Self-supervised 설정에서 MAE에 naive TokenMoE를 적용하는 한계점을 식별합니다.
• 의미적으로 유사한 이미지 클러스터에서 전문가를 학습시키기 위해 MoCE를 도입합니다.
• MoCE가 여러 다운스트림 작업 및 배포에서 우수한 전이를 달성하는지 보여줍니다.
• 효율적 추론을 위한 작업 맞춤형 MoCE 모델 배포에 대한 지침을 제공합니다.

제안 방법

• 사전 학습 데이터를 사전 학습된 MAE의 특징을 사용해 의미적으로 일관된 그룹으로 클러스터링합니다.
• 선택적 MLP 계층을 MoCE 게이트로 대체해 이미지 토큰을 클러스터 조건의 전문가로 라우팅합니다.
• 각 전문가는 할당된 클러스터의 데이터에서 학습하도록 하는 증류 손실을 사용해 학습을 안정화합니다.
• 배포 시 클러스터링 모듈을 사용해 다운스트림 작업과 가장 의미론적으로 관련된 전문가를 선택합니다.
• 자신감 있는 게이트 라우팅을 장려하기 위해 불균형성/정규화 손실을 도입합니다.

실험 결과

연구 질문

• RQ1 MAE가 의미적으로 다양한 다운스트림 작업에서 부정 전이를 겪을 수 있는가?
• RQ2 MAE에 Tokens MoE를 나열적으로 적용하면 전이 성능이 개선되는가 아니면 저하되는가?
• RQ3 클러스터 조건 라우팅으로 MoCE가 일반 MAE 및 TokenMoE보다 더 나은 작업 특화 사전학습을 가능하게 하는가?
• RQ4 MoCE가 탐지 및 분할을 포함한 다양한 다운스트림 작업에서 어떻게 성능을 발휘하는가?

주요 결과

Aircraft	Caltech	Cars	C10	C100	DTD	Flowers	Food	Pets	SUN	VOC	Avg.
72.71	91.24	84.47	96.15	77.33	75.05	96.25	80.49	92.78	62.46	85.02	83.09
70.51	89.70	81.40	95.18	76.44	73.67	95.09	77.45	90.71	61.12	80.15	81.04
75.77	89.73	86.65	95.31	83.60	73.62	95.53	84.77	91.25	64.64	83.51	84.03
76.57	90.04	86.95	96.92	81.42	73.09	96.14	82.90	92.65	64.40	85.37	84.22
78.73	90.61	88.56	97.79	84.68	74.04	96.94	86.24	93.07	65.05	85.26	85.54

• MoCE는 11개의 다운스트림 작업에서 MAE*에 비해 Top-1 정확도에서 평균 최대 2.45% 향상을 달성합니다.
• MoCE는 탐지 및 분할 벤치마크에서 최첨단의 자기지도 학습 결과를 제공합니다.
• TokenMoE는 의미론 비적합한 라우팅으로 인해 MAE를 개선하지 못하거나 오히려 저하될 수 있습니다.
• MoCE의 클러스터 조건 전문가들은 의미적으로 응집력 있는 특화 기능(예: 의상, 새, 개)을 학습합니다.
• 하나의 활성 MoCE 전문가로의 배포는 테스트 시 파라미터 수를 절반으로 줄이고 효율을 향상시킬 수 있습니다.