Union of Experts: Anpassung des hierarchischen Routings an äquivalent zerlegte Transformer

Union of Experts: Adapting Hierarchical Routing to Equivalently Decomposed Transformer

March 4, 2025

Autoren: Yujiao Yang, Jing Lian, Linhui Li

cs.AI

Zusammenfassung

Mixture-of-Experts (MoE) verbessert die Modellleistung bei gleichzeitiger Beibehaltung der Recheneffizienz, was es besonders für groß angelegte Anwendungen geeignet macht. Allerdings arbeiten Experten im bestehenden MoE-Paradigma als Einzelpersonen, wodurch hochwertige Interaktionen zwischen den Experten fehlen. Darüber hinaus wurden sie nicht effektiv auf Attention-Blöcke erweitert, was weitere Effizienzsteigerungen einschränkt. Um diese Probleme zu lösen, schlagen wir Union-of-Experts (UoE) vor, das den Transformer in eine gleichwertige Gruppe von Experten zerlegt und dann dynamisches Routing auf Eingabedaten und Experten implementiert. Unser Ansatz erweitert das MoE-Design mit drei wesentlichen Innovationen: (1) Wir haben eine gleichwertige Expertenzerlegung sowohl auf MLP-Blöcke als auch auf Attention-Blöcke basierend auf Matrixpartitionierung in Tensor-Parallelität durchgeführt. (2) Wir haben zwei Routing-Paradigmen entwickelt: patchweise Datenauswahl und Expertenauswahl, um Routing auf verschiedenen Ebenen anzuwenden. (3) Wir entwerfen die Architektur des UoE-Modells, einschließlich Selective Multi-Head Attention (SMHA) und Union-of-MLP-Experts (UoME). (4) Wir entwickeln eine parallele Implementierung des Routings und der Berechnungsoperationen von UoE und optimieren die Effizienz basierend auf der Hardware-Verarbeitungsanalyse. Die Experimente zeigen, dass das mit UoE ausgestattete Modell Full Attention, state-of-the-art MoEs und effiziente Transformer in mehreren Aufgaben sowohl im Bild- als auch im natürlichen Sprachbereich übertrifft. Die Quellcodes sind unter https://github.com/YujiaoYang-work/UoE verfügbar.

English

Mixture-of-Experts (MoE) enhances model performance while maintaining computational efficiency, making it well-suited for large-scale applications. However, expert in exist MoE paradigm works as an individual, thereby lacking high-quality expert interactions. Moreover, they have not been effectively extended to attention block, which constrains further efficiency improvements. To tackle these issues, we propose Union-of-Experts (UoE), which decomposes transformer into an equitant group of experts, and then implement dynamic routing on input data and experts. Our approach advances MoE design with three key innovations: (1) We conducted equitant expert decomposition on both MLP blocks and attention blocks based on matrix partition in tensor parallelism. (2) We developed two routing paradigms: patch wise data selection and expert selection, to apply routing across different levels. (3) We design the architecture of UoE model, including Selective Multi-Head Attention (SMHA) and Union-of-MLP-Experts (UoME). (4) We develop parallel implementation of UoE's routing and computation operation, and optimize efficiency based on the hardware processing analysis. The experiments demonstrate that the model employed with UoE surpass Full Attention, state-of-art MoEs and efficient transformers in several tasks across image and natural language domains. The source codes are available at https://github.com/YujiaoYang-work/UoE.