Объединение экспертов: адаптация иерархической маршрутизации к эквивалентно декомпозированному трансформеру

Union of Experts: Adapting Hierarchical Routing to Equivalently Decomposed Transformer

March 4, 2025

Авторы: Yujiao Yang, Jing Lian, Linhui Li

cs.AI

Аннотация

Модель Mixture-of-Experts (MoE) повышает производительность модели, сохраняя при этом вычислительную эффективность, что делает её хорошо подходящей для крупномасштабных приложений. Однако в существующей парадигме MoE эксперты работают как отдельные единицы, что ограничивает качество взаимодействия между ними. Более того, они не были эффективно расширены на блоки внимания, что сдерживает дальнейшее повышение эффективности. Для решения этих проблем мы предлагаем модель Union-of-Experts (UoE), которая разбивает трансформер на эквивалентную группу экспертов и реализует динамическую маршрутизацию входных данных и экспертов. Наш подход улучшает дизайн MoE за счёт трёх ключевых инноваций: (1) Мы провели эквивалентное разложение экспертов как на MLP-блоки, так и на блоки внимания на основе разделения матриц в тензорном параллелизме. (2) Мы разработали две парадигмы маршрутизации: выбор данных на уровне патчей и выбор экспертов, чтобы применять маршрутизацию на разных уровнях. (3) Мы спроектировали архитектуру модели UoE, включая Selective Multi-Head Attention (SMHA) и Union-of-MLP-Experts (UoME). (4) Мы реализовали параллельное выполнение операций маршрутизации и вычислений в UoE и оптимизировали эффективность на основе анализа аппаратной обработки. Эксперименты показывают, что модель с использованием UoE превосходит Full Attention, современные MoE и эффективные трансформеры в нескольких задачах в областях изображений и естественного языка. Исходные коды доступны по адресу https://github.com/YujiaoYang-work/UoE.

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.