SharedParam MoE: когда 4 эксперта лучше 12

Вот уже несколько месяцев работаю над оптимизацией смеси экспертов для LLM. Задача проста на словах: найти архитектуру, которая даст лучшую точность при меньшем количестве параметров. На деле всё оказалось намного интереснее.

Стартовали с классического подхода — baseline из Phase 7a с 12 независимыми экспертами и 4.5M параметров показывал точность 70.45%. Это был мой ориентир. Но такой размер модели дорог в инференсе. Нужно было поискать.

В эксперименте 10a протестировал три подхода сразу. Condition A — просто выключить MoE, использовать одну сеть без маршрутизации. Условно «No-MoE», 2.84M параметров. Результат: 69.80%. Маршрутизация между несколькими путями дала всего лишь 1.15pp — почти ничего.

Condition B — вот здесь началось интересное. SharedParam MoE: четыре эксперта, но ключевая идея в том, чтобы они делили общие слои параметров. Гейтинг работает только на последних слоях, а основная вычислительная масса — одна на всех. Плюс Loss-Free Balancing: нет штрафов на балансировку, просто следим за утилизацией и регулируем bias. На 130-й эпохе вижу 70.71%, а по финалу получилось 70.95% при 2.91M параметров. Выше baseline! Все четыре эксперта были живы ВСЁ обучение.

Condition C — Wide Shared, более агрессивный шаринг параметров. На финале 69.96%, отстал немного. Но главное: 2.86M параметров, инфернс на 25.3ms против 29.2ms у B.

Пока ждал результатов 10a, запустил 10b с MixtureGrowth — идеей вырастить сеть из маленького seed’а 182K параметров путём добавления новых слоёв и экспертов. Классный подход для прогрессивного расширения. Seed стартовал с 53.23%, потом во время freeze-фазы за 10 эпох скакнул на 58.97%. Смотрю — рост работает! На Stage2 с 2.84M параметров получилось 69.65%, всего на 0.80pp ниже original baseline.

Что здесь самое странное? Выращенная из крошечного seed модель на 5.57pp превосходит ту же архитектуру, обученную с нуля! Scratch-baseline на тех же 2.84M параметрах показал только 64.08%. Обучение в длину оказалось эффективнее, чем в глубину с нуля.

На 10c изменил расписание learning rate с cosine scheduler’ом — может быть, это даст ещё лучше. Seed уже на эпохе 50 показывает 52.44% против 48.78% в 10b без cosine. Пока расписание работает.

Вердикт текущий: SharedParam MoE — наш путь вперёд. Не просто потому что точнее, а потому что эффективнее: на 35% меньше параметров, на 50pp точнее baseline, все эксперты живы, Loss-Free Balancing не создаёт артефактов. Маршрутизация имеет смысл только если эксперты действительно специализируются. Шеринг параметров сокращает их эго.

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