Разреженная архитектура MoE увеличивает параметры без роста нагрузки, но требует балансировки экспертов

Разработчик Авинош Сурьяраччи (Avinash Sooriyarachchi) опубликовал открытый код проекта makeMoE, демонстрирующий пошаговое создание языковой модели с архитектурой «разреженная смесь экспертов» (Sparse Mixture of Experts). Проект базируется на методологии Андрея Карпатии (Andrej Karpathy) и использует ту же задачу генерации текста в стиле Шекспира на уровне символов. Ключевое отличие от стандартных трансформеров — замена обычного полносвязного слоя на набор специализированных подсетей («экспертов»), где для обработки каждого токена активируется лишь небольшая часть из них.

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

Архитектурные особенности и механизмы маршрутизации

В основе реализации лежит стандартный механизм причинного само-внимания (causal self-attention), который обеспечивает корректную генерацию последовательностей, скрывая от модели будущие токены. Изменения коснулись только вычислительного блока внутри каждого слоя трансформера. Вместо одного универсального нейронного слоя модель использует пул из нескольких независимых сетей.

Для управления потоком данных применяется механизм маршрутизации (router), который решает, какие именно эксперты будут обрабатывать текущий токен. В проекте реализованы два подхода к этому процессу:

• Top-k Gating: Алгоритм выбирает топ-k экспертов с наивысшими оценками релевантности для конкретного токена. Остальные эксперты пропускаются, что экономит ресурсы.
• Noisy Top-k Gating: К оценкам релевантности добавляется случайный шум. Это предотвращает ситуацию, когда маршрутизатор постоянно выбирает одних и тех же «любимых» экспертов, заставляя модель исследовать возможности других подсетей и обеспечивая балансировку нагрузки.

Каждый активированный эксперт представляет собой простую многослойную перцептронную сеть (MLP). Итоговый результат формируется как взвешенная сумма выходов выбранных экспертов, где веса определяются маршрутизатором.

Практическая реализация и результаты обучения

Код написан на фреймворке PyTorch и полностью воспроизводит цикл обучения, включая предобработку данных, инициализацию весов и процесс вывода. Для инициализации весов использован метод Кайминга Хэ (Kaiming He), что обусловлено применением функции активации ReLU в экспертных сетях. Обучение велось на видеокарте NVIDIA A100 в среде Databricks с использованием системы трассировки MLflow для мониторинга метрик.

Модель содержит примерно 9 млн параметров. В ходе обучения наблюдалось снижение функции потерь (loss) как на обучающей, так и на валидационной выборке. Оптимальная точка остановки обучения была достигнута на шаге около 4500, после чего качество на валидационном наборе начало незначительно ухудшаться.

Генерация текста моделью демонстрирует узнаваемые паттерны шекспировских пьес, хотя и содержит ошибки, характерные для моделей такого масштаба. Пример вывода показывает попытки имитации диалогов и имен персонажей, что подтверждает работоспособность архитектуры.

Важный нюанс: В текущей реализации отсутствует функция потерь для балансировки нагрузки (load balancing loss). Без этого механизма существует риск, что в процессе обучения активируется лишь 10–15% доступных экспертов, тогда как остальные останутся неиспользованными, что снижает эффективность архитектуры.

Операционные последствия и технические ограничения

• Зависимость от памяти GPU: Несмотря на то, что в каждом шаге вычислений активны только несколько экспертов, все параметры модели должны быть загружены в память видеокарты. При масштабировании до сотен миллиардов параметров это становится критическим ограничением для развертывания.
• Сложность балансировки: Эффективность архитектуры напрямую зависит от качества работы маршрутизатора. Если модель не научится равномерно распределять токены, часть вычислительных ресурсов будет простаивать, сводя на нет преимущества разреженности.
• Потенциал для экспериментов: Код ориентирован на читаемость и возможность модификации («hackability»), а не на максимальную производительность. Это делает его удобным полигоном для тестирования новых стратегий инициализации, методов токенизации (переход от символов к суб-словам) и поиска оптимальной архитектуры.
• Необходимость доработки: Для промышленного применения потребуется реализовать механизм «емкости экспертов» (Expert Capacity), который ограничивает количество токенов, поступающих к одному эксперту, предотвращая перегрузку отдельных узлов.

Проект служит практическим руководством по пониманию внутренней работы современных больших языковых моделей, таких как Mixtral, и показывает, как базовые принципы трансформеров адаптируются под новые требования к эффективности.