TEXAS: новый способ дообучения MoE-моделей побеждает в 17 из 18 тестов

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

Авторы представляют метод Task-Expert-Aware Supervision (TEXAS, «учёт задачных экспертов при обучении с учителем»), который соединяет два шага. Первый, поиск задачных экспертов с учётом правильности ответа: TEXAS сравнивает активации экспертов на примерах, которые исходная модель решает верно, и на тех, где она ошибается, и оставляет тех экспертов, которые сильнее активируются именно на успешных примерах. Второй шаг, распределение внимания на уровне токенов: при дообучении TEXAS повышает вес токенов ответа в примерах, где модель ошиблась, если эти токены активируют отобранных экспертов. При этом метод не ограничивает дообучение фиксированным набором экспертов и не навязывает модели заранее заданное распределение маршрутизации, он опирается на то, как модель уже маршрутизирует токены.

Метод проверили на трёх MoE-моделях и шести бенчмарках (конкретные названия моделей и тестов в тексте не приводятся). TEXAS показал лучший или не хуже лучшего результат в 17 из 18 комбинаций «модель × бенчмарк» и в среднем превзошёл самый сильный из сравниваемых методов на 1,3, 1,5 балла. Дополнительные эксперименты по удалению отдельных частей метода (ablation), по словам авторов, подтверждают пользу и от найденных экспертов, и от предложенной схемы распределения внимания при обучении.

Ключевые факты

TEXAS отбирает «задачных» экспертов MoE-модели, сравнивая их активацию на примерах, решённых верно, и на тех, где модель ошиблась, оставляет тех, кто сильнее активируется при успехе, а не просто чаще используется
При дообучении метод повышает вес токенов ответа в неверно решённых примерах, если эти токены активируют отобранных экспертов
Метод не ограничивает дообучение фиксированным набором экспертов и не задаёт модели целевое распределение маршрутизации, работает поверх уже существующего поведения роутера
На трёх MoE-моделях и шести бенчмарках TEXAS показал лучший или не хуже лучшего результат в 17 из 18 сочетаний и в среднем обошёл сильнейший из сравниваемых методов на 1,3, 1,5 балла
Дополнительные эксперименты по удалению частей метода подтвердили вклад и отбора экспертов, и схемы распределения внимания при обучении

Почему это важно

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

Кому это важно

Тем, кто дообучает MoE-модели под конкретные задачи, исследовательским командам и инженерам, работающим с разреженными архитектурами (Mixture-of-Experts), где важно повышать качество на целевой задаче без изменения архитектуры и без полного дообучения всех параметров.

Как это применить

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

Можно ли доверять

Текст, аннотация к статье на arXiv, то есть препринт без независимого рецензирования на момент публикации. В доступном тексте не указаны имена авторов, их организации и дата публикации, не назван конкретный код или данные для воспроизведения, а также не названы конкретные проверенные модели и бенчмарки, только их количество (три модели, шесть бенчмарков). Заявленные результаты (17 из 18, 1,3, 1,5 балла), собственные цифры авторов по итогам их же экспериментов; независимого подтверждения пока нет.

Риски и подводные камни

Результат представлен только как относительное улучшение, 1,3, 1,5 балла в среднем над сильнейшим из сравниваемых методов, без абсолютных цифр точности, поэтому по одной аннотации нельзя оценить, насколько высок базовый уровень качества и на сколько в реальности заметен выигрыш. Раз конкретные модели и бенчмарки не названы, не ясно, насколько результаты переносятся на другие MoE-архитектуры и задачи за пределами шести проверенных тестов. Отсутствие упоминания о выложенном коде или данных также усложняет независимую проверку метода.