UMoE: Vereinheitlichung von Aufmerksamkeit und FFN durch gemeinsame Experten

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

Autoren: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

Zusammenfassung

Sparse Mixture of Experts (MoE)-Architekturen haben sich als vielversprechender Ansatz zur Skalierung von Transformer-Modellen erwiesen. Während frühe Arbeiten MoE hauptsächlich in Feed-Forward-Netzwerk (FFN)-Schichten integrierten, haben neuere Studien die Erweiterung des MoE-Paradigmas auf Attention-Schichten untersucht, um die Modellleistung zu verbessern. Allerdings erfordern bestehende Attention-basierte MoE-Schichten spezialisierte Implementierungen und zeigen im Vergleich zu ihren FFN-basierten Gegenstücken eine suboptimale Leistung. In diesem Artikel streben wir an, die MoE-Designs in Attention- und FFN-Schichten zu vereinheitlichen, indem wir eine neuartige Neuformulierung des Attention-Mechanismus einführen, die eine zugrunde liegende FFN-ähnliche Struktur innerhalb von Attention-Modulen offenlegt. Unsere vorgeschlagene Architektur, UMoE, erreicht eine überlegene Leistung durch Attention-basierte MoE-Schichten und ermöglicht gleichzeitig eine effiziente Parameterfreigabe zwischen FFN- und Attention-Komponenten.

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.