UMoE: 공유 전문가를 통한 어텐션과 FFN의 통합

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

저자: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

초록

희소 전문가 혼합(Sparse Mixture of Experts, MoE) 아키텍처는 트랜스포머 모델의 확장을 위한 유망한 접근법으로 부상했습니다. 초기 연구들은 주로 피드포워드 네트워크(FFN) 계층에 MoE를 통합했지만, 최근 연구들은 모델 성능을 향상시키기 위해 MoE 패러다임을 어텐션 계층으로 확장하는 방안을 탐구하고 있습니다. 그러나 기존의 어텐션 기반 MoE 계층은 특수한 구현이 필요하며, FFN 기반 대비 최적의 성능을 보이지 못하고 있습니다. 본 논문에서는 어텐션 메커니즘의 새로운 재구성을 통해 어텐션 모듈 내에 내재된 FFN과 유사한 구조를 밝혀냄으로써, 어텐션과 FFN 계층에서의 MoE 설계를 통합하고자 합니다. 우리가 제안한 UMoE 아키텍처는 어텐션 기반 MoE 계층을 통해 우수한 성능을 달성함과 동시에 FFN과 어텐션 구성 요소 간의 효율적인 파라미터 공유를 가능하게 합니다.

English

Sparse Mixture of Experts (MoE) architectures have emerged as a promising approach for scaling Transformer models. While initial works primarily incorporated MoE into feed-forward network (FFN) layers, recent studies have explored extending the MoE paradigm to attention layers to enhance model performance. However, existing attention-based MoE layers require specialized implementations and demonstrate suboptimal performance compared to their FFN-based counterparts. In this paper, we aim to unify the MoE designs in attention and FFN layers by introducing a novel reformulation of the attention mechanism, revealing an underlying FFN-like structure within attention modules. Our proposed architecture, UMoE, achieves superior performance through attention-based MoE layers while enabling efficient parameter sharing between FFN and attention components.