Mixture of Experts

Mixture of Experts (MoE, рус. «смесь экспертов») — это архитектурный подход в машинном обучении, при котором модель состоит из множества специализированных подсетей («экспертов») и управляющего механизма («гейта» или «маршрутизатора»), распределяющего входные данные между этими подсетями. В отличие от плотных (dense) моделей, где для обработки каждого примера активируются все параметры, MoE активирует лишь подмножество экспертов, что позволяет значительно увеличить общую ёмкость модели (количество параметров) без пропорционального роста вычислительных затрат во время инференса.

История

Идея комитетов или ансамблей моделей, где разные алгоритмы отвечают за разные области входного пространства, возникла ещё в 1980-х годах. В 1991 году группа учёных, включая Роберта Джейкобса, Майкла Джордана и Джеффри Хинтона, формализовала концепцию адаптивного смешивания экспертов (Adaptive Mixture of Local Experts). Предлагалось, что для разных участков данных (например, разных типов изображений или разных фонем в речи) должны обучаться отдельные нейронные сети, а «управляющая сеть» (gating network) должна решать, какой эксперт или какая комбинация экспертов лучше всего подходит для данного входа.

Долгое время MoE оставалась нишевой техникой, применяемой в основном в задачах регрессии и классификации с небольшим числом экспертов. Прорыв произошёл в 2010-х годах с ростом масштаба нейросетей и появлением более эффективных методов обучения. В 2017 году исследователи Google представили Sparse MoE — вариант, где для каждого примера активируется лишь небольшое количество экспертов (например, 2 из сотен). Это позволило обучать модели с триллионами параметров, такие как Switch Transformer (2021) и GLaM (2021), которые показали впечатляющие результаты в обработке естественного языка.

Архитектура и принцип работы

Базовая архитектура MoE состоит из двух ключевых компонентов:

1. Набор экспертов — это, как правило, несколько нейронных сетей одной архитектуры (например, полносвязных слоёв или блоков трансформера). Каждый эксперт специализируется на определённом подмножестве данных или типе вычислений. Количество экспертов может варьироваться от нескольких штук до тысяч.
2. Маршрутизатор (гейт) — это обучаемая функция, которая для каждого входного вектора вычисляет веса (оценки) для всех экспертов. Чаще всего это линейный слой с функцией softmax, выдающий вероятности выбора каждого эксперта.

В плотном (dense) MoE выход модели является взвешенной суммой выходов всех экспертов, где веса берутся от маршрутизатора. В разреженном (sparse) MoE маршрутизатор выбирает только k экспертов с наибольшими весами (обычно k=1 или k=2), а остальные игнорируются. Это и есть ключевое преимущество: при огромном общем числе параметров (например, 1 триллион) на обработку одного примера тратится вычислительная мощность, соответствующая лишь небольшой части этих параметров (например, 10 миллиардов).

Балансировка загрузки экспертов

Основная проблема разреженного MoE — дисбаланс загрузки: маршрутизатор может начать отправлять большинство примеров одним и тем же экспертам, оставляя других без обучения. Для борьбы с этим используются специальные техники:

• Вспомогательные функции потерь (auxiliary loss): к основной функции потерь добавляется штраф за неравномерное распределение нагрузки. Например, поощряется, чтобы каждый эксперт получал примерно одинаковое количество примеров.
• Шум (noise): добавление случайного шума к логитам маршрутизатора во время обучения, чтобы стимулировать исследование разных экспертов.
• Дополнительные эксперты (capacity factor): каждому эксперту выделяется фиксированная «ёмкость» (максимальное количество примеров, которое он может обработать за один шаг). Если эксперт переполнен, лишние примеры отбрасываются, что заставляет маршрутизатор распределять нагрузку более равномерно.

Классификация

MoE-архитектуры можно классифицировать по нескольким признакам:

• По типу активации экспертов:
• Плотный (Dense): активируются все эксперты, их выходы комбинируются с весами.
• Разреженный (Sparse): активируется только k экспертов (k << общего числа экспертов). Наиболее распространён в современных больших моделях.
• По способу маршрутизации:
• Токен-уровень (Token-level): каждый токен (или элемент последовательности) маршрутизируется независимо. Используется в моделях-трансформерах.
• Эксперт-уровень (Expert-level): целые подзадачи или типы данных направляются конкретным экспертам (например, один эксперт для текста, другой для изображений).
• По расположению экспертов:
• Глубокий (Deep MoE): эксперты заменяют один или несколько слоёв в глубокой сети (например, слой Feed-Forward в трансформере).
• Широкий (Shallow MoE): эксперты располагаются на одном уровне, обрабатывая выход предыдущего слоя.

Применение

MoE нашёл широкое применение в областях, где требуется высокая ёмкость модели при ограниченных вычислительных ресурсах:

• Обработка естественного языка (NLP): большинство современных больших языковых моделей (LLM), таких как Mixtral 8x7B от Mistral AI, Gemini от Google, GPT-4 (по некоторым данным), используют разреженный MoE. Это позволяет им иметь сотни миллиардов или триллионы параметров, сохраняя при этом скорость инференса, сравнимую с моделью в 10-20 раз меньшего размера.
• Компьютерное зрение: MoE применяется для обработки изображений, где разные эксперты могут специализироваться на распознавании разных объектов, текстур или сцен. Например, модель V-MoE (Vision MoE) от Google.
• Мультимодальные модели: MoE позволяет объединять обработку данных разных модальностей (текст, изображения, аудио), назначая каждому типу данных своего эксперта или группу экспертов.
• Рекомендательные системы: MoE используется для персонализации, где разные эксперты моделируют предпочтения разных групп пользователей.

Преимущества и недостатки

Преимущества

• Масштабируемость: позволяет создавать модели с огромным числом параметров без пропорционального роста вычислительных затрат.
• Специализация: каждый эксперт может выучить уникальные закономерности, что повышает общую точность модели.
• Эффективность: во время инференса используется лишь малая часть параметров, что снижает задержки и энергопотребление.

Недостатки

• Сложность обучения: требуется тщательная настройка балансировки нагрузки и вспомогательных функций потерь. Нестабильное обучение — распространённая проблема.
• Высокие требования к памяти: хотя вычисления разрежены, все параметры экспертов должны храниться в памяти (например, на GPU или в распределённой системе), что увеличивает требования к объёму RAM/VRAM.
• Проблемы с батчами: из-за того, что разные примеры в батче могут попадать к разным экспертам, эффективная упаковка данных для параллельных вычислений может быть затруднена.
• Трудности с выводом (inference): на этапе вывода разреженный MoE может быть менее эффективен на обычных GPU, чем плотная модель, если не использовать специальные оптимизации (например, Expert Parallelism).

Интересные факты

• Одна из самых известных MoE-моделей, Switch Transformer (2021), использовала 1.6 триллиона параметров, но для каждого токена активировала только один эксперт.
• В модели Mixtral 8x7B (2023) используется 8 экспертов, каждый из которых имеет 7 миллиардов параметров. При этом во время инференса активируется только 2 эксперта, что даёт вычислительную стоимость, эквивалентную модели с ~12 миллиардами параметров, но при этом качество сопоставимо с моделями в 30-70 миллиардов.
• Термин «эксперт» в данном контексте не означает, что модель обладает человеческой экспертизой. Это просто специализированная подсеть, обученная на определённом подмножестве данных.

Источники

1. Jacobs, R. A., Jordan, M. I., Nowlan, S. J., & Hinton, G. E. (1991). Adaptive mixtures of local experts. Neural computation, 3(1), 79-87.
2. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538.
3. Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120), 1-39.
4. Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., ... & Dean, J. (2022). Glam: Efficient scaling of language models with mixture-of-experts. International Conference on Machine Learning, 5547-5569.
5. Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., ... & Sayed, W. E. (2024). Mixtral of experts. arXiv preprint arXiv:2401.04088.