[논문 리뷰] Residual Mixture of Experts

Lemeng Wu, Mengchen Liu외 4명|arXiv (Cornell University)|2022. 04. 20.

Domain Adaptation and Few-Shot Learning인용 수 8

한 줄 요약

이 논문은 입력에 의존하지 않는 코어와 입력에 의존하는 잔차로 MoE 가중치를 분해함으로써 비용 효율적으로 비전 트랜스포머를 확장하는 훈련 파이프라인인 Residual Mixture of Experts (RMoE)를 제안한다. 입력에 의존하는 잔차만 미세조정함으로써 최소한의 계산으로 전체 MoE 훈련과 유사한 성능를 달성하면서도 훈련 비용이 3% 미만으로 증가한다. 이는 ADE20K에서 Swin-T, CvT-13, Swin-L과 같은 비-MoE 모델보다 1.1–1.0 mIoU 향상되고, MS-COCO에서 0.6–1.6 AP 향상된다.

ABSTRACT

Mixture of Experts (MoE) is able to scale up vision transformers effectively. However, it requires prohibiting computation resources to train a large MoE transformer. In this paper, we propose Residual Mixture of Experts (RMoE), an efficient training pipeline for MoE vision transformers on downstream tasks, such as segmentation and detection. RMoE achieves comparable results with the upper-bound MoE training, while only introducing minor additional training cost than the lower-bound non-MoE training pipelines. The efficiency is supported by our key observation: the weights of an MoE transformer can be factored into an input-independent core and an input-dependent residual. Compared with the weight core, the weight residual can be efficiently trained with much less computation resource, e.g., finetuning on the downstream data. We show that, compared with the current MoE training pipeline, we get comparable results while saving over 30% training cost. When compared with state-of-the-art non- MoE transformers, such as Swin-T / CvT-13 / Swin-L, we get +1.1 / 0.9 / 1.0 mIoU gain on ADE20K segmentation and +1.4 / 1.6 / 0.6 AP gain on MS-COCO object detection task with less than 3% additional training cost.

연구 동기 및 목표

• 대규모 비전 트랜스포머의 Mixture of Experts (MoE) 훈련에 따른 높은 계산 비용 문제를 해결하기 위해.
• 대규모 업스트ream 데이터셋에서의 전체 재훈련이 필요 없이도 비전 트랜스포머의 효율적 확장을 가능하게 하기 위해.
• MoE 모델의 훈련 비용을 줄이면서도 상한선 MoE 훈련 수준의 성능를 유지하기 위해.
• 공개된 비-MoE 체크포인트를 대규모 MoE 적응에 대한 초기화로 효과적으로 활용하기 위해.
• 전문가 가중치가 코어와 훈련 가능한 잔차로 분해될 수 있음을 활용하여 효율적인 미세조정을 가능하게 하기 위해.

제안 방법

• MoE 가중치를 입력에 의존하지 않는 코어와 입력에 의존하는 잔차로 분해하며, 코어는 사전 훈련된 비-MoE 모델에서 고정된다.
• 잔차 성분은 업스트림 데이터에서의 전체 MoE 재훈련 없이도 업스트림 데이터에서 저비용 업데이트를 통해 미세조정된다.
• 성능 저하를 방지하기 위해 중간 단계의 미세조정 동안 정지 기울기 기법을 사용한다.
• 전문가 특화를 장려하기 위해 복사된 가중치에 노이즈 초기화 스케일을 적용한다.
• MoE 레이어 선택 전략으로 기울기 점수 기반 방법을 사용하여 Last-2 및 Every-2 방법보다 우수한 성능를 보였다.
• 중간 단계 및 업스트림 미세조정을 모두 지원하며, 추론 FLOPs와 지연에 미미한 영향을 미친다.

실험 결과

연구 질문

• RQ1 훈련된 MoE 트랜스포머의 가중치를 코어와 잔차 성분으로 분해하여 효율적인 미세조정을 가능하게 할 수 있는가?
• RQ2 MoE 전문가의 잔차 성분만 미세조정하는 것이 전체 MoE 훈련과 유사한 성능를 달성하면서도 비용을 크게 줄일 수 있는가?
• RQ3 기존의 비-MoE 비전 트랜스포머 체크포인트가 대규모 MoE 모델의 초기화로 효과적으로 활용될 수 있는가?
• RQ4 MoE 레이어 선택 전략의 선택(예: Every-Last, 기울기 점수)이 업스트림 성능에 어떤 영향을 미치는가?
• RQ5 최적의 성능를 달성하면서도 추가 비용을 최소화하는 하이퍼파라미터(예: top-k, 전문가 수, 노이즈 스케일)는 무엇인가?

주요 결과

• RMoE는 ADE20K에서 Swin-T 대비 +1.1 mIoU, MS-COCO에서 +1.4 AP 향상되었으며, 훈련 비용은 3% 미만으로 증가했다.
• 비-MoE 모델 대비 RMoE는 ADE20K에서 CvT-13 기준 +0.9 mIoU, Swin-L 기준 +1.0 mIoU 향상되었으며, 비용 증가도 미미했다.
• 기울기 점수 기반 MoE 레이어 선택 전략은 Last-2 및 Every-2를 능가했으며, Every-Last 성능를 재현하면서도 더 효율적이었다.
• top-1 게이트와 8개의 전문가를 사용할 경우 최적의 성능를 달성했으며, 16개로 늘리면 로드 밸런싱 손실 간섭으로 인해 성능 향상이 없었다.
• 중간 단계의 미세조정 시 정지 기울기 기법을 적용함으로써 성능 저하를 방지했으며, ADE20K에서 노이즈 스케일 0.01이 가장 우수한 결과를 냈다.
• 추론 시간과 FLOPs는 약간만 증가하여, RMoE가 MoE 모델의 효율성 우월성을 유지하고 있음을 확인했다.