Исследователи из Allen Institute for AI и UC Berkeley представили EMO — mixture-of-experts-модель, которая формирует модульную структуру уже на этапе предварительного обучения. Авторы утверждают, что модель можно сократить до небольшой части экспертов почти без заметной потери качества.
Архитектуры mixture-of-experts, или MoE, уже стали стандартом для языковых моделей вроде DeepSeek-V4 и Qwen3.5. В них для каждого токена активируется лишь несколько экспертов, что позволяет наращивать число параметров без резкого роста вычислительных затрат. Однако весь набор экспертов всё равно нужно держать в памяти, потому что разные токены внутри одной задачи обращаются к разным частям модели.
По словам авторов работы, причина в том, что в обычных MoE эксперты часто реагируют на поверхностные языковые признаки — предлоги, знаки препинания или артикли, а не на темы вроде математики или программирования. Из-за этого трудно выделить действительно полезное подмножество экспертов для конкретной задачи.
EMO решает эту проблему за счёт более простого разбиения данных. Вместо заранее заданных тематических категорий, как в BTX или FlexOlmo от Ai2, исследователи используют границы документов. Предполагается, что токены внутри одного документа обычно относятся к одной теме.
Модель заставляет все токены документа выбирать активных экспертов из общего пула. Для этого она усредняет предпочтения роутера по всем токенам документа и оставляет экспертов, которые выбираются чаще всего.
Чтобы обучение оставалось стабильным, авторам пришлось внести два изменения. Во-первых, балансировку нагрузки между экспертами они считают не локально по каждому батчу, а глобально по множеству документов. Иначе две цели обучения мешали бы друг другу: одна объединяет токены документа, другая старается распределить их по максимально широкому набору экспертов.
Во-вторых, исследователи случайно меняют размер пула экспертов для документа во время обучения, а не фиксируют его. Это позволяет модели работать с подгруппами экспертов разного размера на этапе вывода.
Команда обучила MoE-модель с 1 млрд активных и 14 млрд общих параметров, 128 экспертами и активацией восьми экспертов на токен. Для обучения использовали 1 трлн токенов из корпуса предварительного обучения OLMoE. В полном виде EMO показывает тот же результат, что и стандартная MoE, обученная в тех же условиях. Авторы также пишут, что модель превосходит OLMoE, хотя для обучения использовалось в пять раз больше данных.
Затем исследователи начали убирать экспертов, чтобы проверить предел сокращения. При сохранении 25% экспертов — 32 из 128 — EMO теряет около одного процентного пункта по нескольким бенчмаркам. При 12,5%, то есть 16 экспертах, снижение составляет примерно три пункта.
В такой же конфигурации стандартная MoE резко теряет качество — на 10–15 процентных пунктов. В ряде случаев она опускается ниже уровня плотной модели с тем же числом активных параметров. На математическом бенчмарке GSM8K подмножества EMO с 12,5% экспертов после дообучения снова достигают уровня полной модели.
Авторы также отмечают, что для выбора подходящих экспертов нужно немного данных. По их словам, одного few-shot-примера достаточно, чтобы подобрать подгруппу, работающую сопоставимо с вариантом, выбранным по полному валидационному набору. EMO поддерживает как простой выбор через роутер, так и более специализированный метод Easy-EP.
Чтобы понять, чему именно научилась модель, исследователи проанализировали внутреннее распределение токенов по экспертам. Для каждого токена они записывали вероятности, с которыми роутер отправляет его к каждому эксперту, а затем группировали токены с похожими профилями активации.
Разница оказалась заметной. В стандартной MoE кластеры экспертов соответствуют простым языковым категориям, таким как предлоги, имена собственные и артикли. В EMO кластеры связаны с темами: здравоохранение и медицина, политика США, кино и музыкальные рецензии. Токены из одного документа в EMO сходятся в одном кластере, тогда как в обычной MoE они распределяются по многим. Для этих кластеров опубликована интерактивная визуализация.
На выборке из 20 млн документов из датасета WebOrganizer с 24 размеченными человеком доменными метками авторы проверили, активируют ли близкие темы схожих экспертов. В EMO разделение оказалось заметно чище, особенно в глубоких слоях модели. В стандартной MoE пересечения были сильнее.
Один из главных сценариев применения — запуск моделей в условиях ограниченной памяти, когда загружаются только эксперты, связанные с нужной темой. В прямом сравнении подгруппы экспертов EMO показывают результат на уровне или выше, чем стандартная MoE с 32 экспертами и плотная модель с восемью активными параметрами, если обе обучать с нуля.
Исследователи также рассматривают дообучение отдельных частей модели во время работы. В качестве примера они приводят детское приложение, где можно отключить кластеры, реагирующие на спам, азартные игры или контент для взрослых. В начальном тесте авторы дообучили подгруппу из 32 экспертов и затем встроили её обратно в 128-экспертную модель. Это улучшило полную модель, но не дотянуло до качества отдельной подгруппы.
По мнению авторов, EMO может быть полезна и для мониторинга, поскольку по экспертам видно, какие части модели использует конкретный входной запрос. Ai2 публикует саму модель EMO, сопоставимую по условиям обучения базовую стандартную MoE, а также код обучения на Hugging Face и GitHub. Исследователи также выпустили интерактивную демонстрацию активаций токенов.
Авторы отмечают, что остаются открытые вопросы: как лучше выбирать и комбинировать подгруппы экспертов, как переобучать отдельные модули под конкретные задачи и как использовать модульную структуру для лучшей интерпретации моделей.
Источник: The Decoder, работа исследователей Allen Institute for AI и UC Berkeley.
