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) стали перспективным подходом для масштабирования моделей Transformer. Хотя первоначальные работы в основном интегрировали MoE в слои прямого распространения (feed-forward network, FFN), недавние исследования изучили возможность расширения парадигмы MoE на слои внимания для повышения производительности модели. Однако существующие MoE-слои, основанные на механизме внимания, требуют специализированных реализаций и демонстрируют менее оптимальную производительность по сравнению с их аналогами, основанными на FFN. В данной статье мы стремимся унифицировать дизайн MoE в слоях внимания и FFN, предлагая новую реформулировку механизма внимания, которая раскрывает лежащую в его основе структуру, подобную FFN. Наша предложенная архитектура, 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.