Branch-Train-MiX — метод создания MoE-моделей

Сегодня рассмотрим статью, в которой предложен метод получения MoE-модели (Mixture-of-Experts) из обычной, dense-модели. Авторы назвали его Branch-Train-MiX (BTX).

Для начала напомним, что такое MoE. Это архитектурный подход, который предполагает использование в трансформенных слоях нескольких FNN-блоков — экспертов. У каждого из них предположительно есть собственная узкая специализация, поэтому над решением поставленной модели задачи они работают не одновременно. Роутер выбирает, какие эксперты лучше подходят для обработки токенов на каждом конкретном этапе.

Во время инференса модели, построенные на архитектуре MoE, задействуют меньше параметров, чем их dense-аналоги — и без сильной потери качества. Кроме того, они прекрасно масштабируются. Поэтому понятно стремление превратить обычную модель в MoE.

В основе идеи авторов лежит метод BTM (Branch-Train-Merge). Суть его заключается в том, чтобы взять несколько одинаковых LLM, параллельно обучить их на разных датасетах, а затем агрегировать предсказания каждой модели во время инференса.

Рассматриваемая статья предлагает после обучения слить все модели в MoE, чтобы FNN каждой базовой модели стал экспертом в соответствующем слое, и добавить роутер. Веса аттеншена при этом следует усреднить и дообучить.

Авторы взяли три копии Llama-2 7B и дообучили каждую на своём домене: математика (на тех же данных, что и Llemma), программирование (на тех же данных, что CodeLlama) и общие знания (на «Википедии»). К финальной модели добавили оригинальную Llama-2 7B, не дообученную на чём-то конкретном. Получившуюся MoE-модель файнтюнили на источниках, которые применялись для обучения всех четырёх экспертов. В результате модель не сильно уступает тем же Llemma и CodeLlama в вопросах математики и программирования.

Интересно и то, что BTX-модели обретают интерпретируемость. Авторы показывают, что в эксперта, обученного на определённом домене, попадают токены из сэмпла этого же домена.

Рассказывайте в комментариях, что думаете про BTX!

Разбор подготовил ❣ Александр Пацация

Душный NLP