[논문 리뷰] MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts

Xi Lin, Akshat Shrivastava외 6명|arXiv (Cornell University)|2024. 07. 31.

Robotics and Automated Systems인용 수 4

한 줄 요약

MoMa는 모달리티별 전문가 그룹(텍스트 또는 이미지)으로 전문가를 그룹화함으로써 조기 융합 다중모态 언어 모델의 효율적인 사전 훈련을 위한 모달리티 인식형 믹스처 오브 응용(모드) 아키텍처를 도입한다. 이로써 각 모달리티에 특화된 파라미터 할당과 그룹 내에서의 학습된 라우팅이 가능해진다. 1조 토큰 예산 하에서 MoMa는 조밀한 기준 모델 대비 총 3.7배의 FLOPs 절감을 달성했으며, 이미지 처리 분야에선 5.2배의 절감을 기록했고, 혼합 모달 전문가를 사용하는 표준 MoE보다 뛰어난 성능을 보였다.

ABSTRACT

We introduce MoMa, a novel modality-aware mixture-of-experts (MoE) architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing expert modules into modality-specific groups. These groups exclusively process designated tokens while employing learned routing within each group to maintain semantically informed adaptivity. Our empirical results reveal substantial pre-training efficiency gains through this modality-specific parameter allocation. Under a 1-trillion-token training budget, the MoMa 1.4B model, featuring 4 text experts and 4 image experts, achieves impressive FLOPs savings: 3.7x overall, with 2.6x for text and 5.2x for image processing compared to a compute-equivalent dense baseline, measured by pre-training loss. This outperforms the standard expert-choice MoE with 8 mixed-modal experts, which achieves 3x overall FLOPs savings (3x for text, 2.8x for image). Combining MoMa with mixture-of-depths (MoD) further improves pre-training FLOPs savings to 4.2x overall (text: 3.4x, image: 5.3x), although this combination hurts performance in causal inference due to increased sensitivity to router accuracy. These results demonstrate MoMa's potential to significantly advance the efficiency of mixed-modal, early-fusion language model pre-training, paving the way for more resource-efficient and capable multimodal AI systems.

연구 동기 및 목표

• 모달리티별 희소성의 특성을 활용해 조기 융합 다중모달 언어 모델의 사전 훈련 효율성을 향상시키기 위해.
• 다중모달 기초 모델을 대규모 파라미터와 데이터로 확장할 때 조밀한 모델의 계산 비효율성을 해결하기 위해.
• 교차 모달 통합을 유지하면서도 모달리티별 특징 학습을 최적화하기 위해 모달리티 인식 희소성(MaS)을 탐색하기 위해.
• 모달리티 인식형 MoE와 믹스처 오브 디프스(MoD)를 조합했을 때 FLOPs 절감 및 추론 성능에 미치는 영향을 평가하기 위해.
• 모달리티 인식 희소성이 경험적 스케일링 법칙을 유지함으로써 향후 대규모 다중모달 사전 훈련의 확장성과 효율성을 보장할 수 있음을 입증하기 위해.

제안 방법

• 모델는 텍스트와 이미지 토큰의 통합 시퀀스를 사용하는 캐멀레온 기반 트랜스포머 아키텍처를 사용한다.
• 전문가들은 모달리티에 따라 구분된 그룹으로 나뉘며, 각각 4개의 텍스트 전문가와 4개의 이미지 전문가로 구성되며, 각각 자신의 전용 모달리티 토큰만 처리한다.
• 각 모달리티 그룹 내에서, 토큰 별 활성 전문가를 선택하는 학습된 라우팅 메커니즘이 도입되어 의미 기반 적응성을 확보하면서도 모달리티별 파라미터 할당을 유지한다.
• 라우팅은 전문가 선택 라우팅을 사용하여 부하 균형을 유지하고 고성능 훈련을 위한 정적 계산 그래프를 유지한다.
• 모달리티 인식형 MoE와 믹스처 오브 디프스(MoD)를 결합하여, 토큰이 특정 레이어를 선택적으로 건너뛸 수 있도록 허용함으로써 계산량을 줄인다.
• 모델는 1조 토큰 예산 하에서 FLOPs 제어 훈련 환경에서 훈련되며, MoMa는 조밀한 기준 모델과 혼합 모달 전문가를 사용하는 표준 MoE와 비교된다.

Figure 1 : Overview of our proposed multimodal early-fusion architecture. (a) The early-fusion architecture processes interleaved text and image data as a unified token sequence using a transformer model. (b) Detailed architecture of the mixture of modality-aware experts (MoMa) transformer block, wh

실험 결과

연구 질문

• RQ1 모달리티 인식 전문가 그룹화가 성능 저하 없이 조기 융합 다중모달 모델의 사전 훈련 효율성을 향상시킬 수 있는가?
• RQ2 혼합 모달 전문가 라우팅 대비 모달리티별 파라미터 할당이 FLOPs 절감 및 훈련 안정성 측면에서 어떻게 비교되는가?
• RQ3 모달리티 인식형 MoE와 믹스처 오브 디프스(MoD)를 조합했을 때 총 FLOPs 감소 및 추론 정확도에 어떤 영향을 미치는가?
• RQ4 모달리티 인식 희소성이 조밀한 모델에서 관찰된 경험적 스케일링 법칙을 유지하는가?
• RQ5 MoD가 강화된 모델의 성능은 인과 추론 시나리오에서 라우터 정확도에 얼마나 민감한가?

주요 결과

• 1조 토큰 훈련 예산 하에서 MoMa 1.4B 모델은 계산량이 동일한 조밀한 기준 모델 대비 총 3.7배의 FLOPs 절감을 달성했으며, 텍스트 처리에선 2.6배, 이미지 처리에선 5.2배의 절감을 기록했다.
• MoMa 모델는 혼합 모달 전문가 8개를 사용하는 표준 전문가 선택 MoE보다 뛰어나며, 총 3배의 FLOPs 절감(텍스트 3배, 이미지 2.8배)을 기록했다.
• MoMa와 믹스처 오브 디프스(MoD)를 조합하면 총 FLOPs 절감률이 4.2배로 증가했고(텍스트 3.4배, 이미지 5.3배), 더욱 높은 효율성 향상을 입증했다.
• 더 높은 FLOPs 절감에도 불구하고, MoMaD 모델는 라우터 정확도 민감도 증가로 인해 인과 추론 작업에서 성능 저하를 보였다.
• MoMa 아키텍처는 안정적인 경험적 스케일링 법칙을 유지하며, 향후 대규모 다중모달 사전 훈련에 대한 확장성을 보여주었다.
• 모달리티에 구애받지 않는 업사이클링 기법이 성능 향상에 기여함으로써 아키텍처 개선의 잠재력이 있음을 시사했다.

Figure 2 : Architecture of transformer layer consisting of MoMa combined with mixture-of-depths (MoD).