STAPLE добавляет миллиарды параметров в OLMo-2 и декодирует вдвое быстрее MoE

Команда исследователей опубликовала на arXiv метод STAPLE, который позволяет нарастить число параметров уже обученной языковой модели без дорогого разреженного претрейна с нуля. На OLMo-2 1B он дает +4,27 пункта на HumanEval, +2,80 на MBPP, +2,27 на GSM8K и +1,65 на DROP относительно простого продолжения dense-обучения. При этом, по данным авторов, генерация замедляется не более чем на 1,5%, а обучение не более чем на 3%, хотя в модель поздно добавляются миллиарды параметров.

Почему не MoE

Разреженные вычисления давно считаются способом увеличить емкость модели, не раздувая активные вычисления на каждый токен. Проблема в цене: претрейн разреженных архитектур требует много данных и денег. Существующие методы апсайклинга, то есть превращения готового dense-чекпоинта в разреженную модель, в основном нацелены на Mixture-of-Experts. Там авторы указывают три узких места: однородность экспертов, низкая sample efficiency и системные накладные расходы на маршрутизацию и память.

Статическая разреженная память устроена иначе. Доступ к ней предсказуем: параметры выбираются по идентификатору токена, а не по решению роутера. Обновления тоже привязаны к конкретному токену. Но, как отмечают авторы, наивная вставка такой памяти в предобученную модель дает слабую отдачу.

Как устроен STAPLE

Метод построен на STEM, где up-проекции в FFN-блоках заменяются токен-индексированными эмбеддингами. Для обучения с нуля это работает, но для готового чекпоинта замена ломает уже накопленные вычисления. STAPLE делает по-другому: up-проекции остаются на месте, а эмбеддинги добавляются как остаточный путь. Предобученная часть продолжает работать как раньше, а новая память учит то, что модель знает о токенах хуже: дополнительную токен-специфичную информацию.

Важная инженерная деталь: реализация использует выгрузку эмбеддингов на хост. На GPU держится лишь около 5% самых частых эмбеддингов, остальное лежит в памяти хоста. Из-за распределения токенов этого хватает, чтобы пропускная способность почти не просела.

Цифры

Основные результаты из препринта:

• OLMo-2 1B, прирост к continued dense training: HumanEval +4,27, MBPP +2,80, GSM8K +2,27, DROP +1,65.
• Средний результат растет вместе со статической емкостью, улучшения сохраняются на OLMo-2 7B.
• STAPLE превосходит статические разреженные базовые варианты: STEM и модель в стиле Engram.
• Против MoE-апсайклинга: выше средний балл, декодирование в 2 раза быстрее, менее половины GPU-памяти под модель и меньше энергии на токен.
• Накладные расходы: генерация в пределах 1,5% от dense-базы, обучение в пределах 3%.

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

Что это даст на практике

Для команд, которые дообучают открытые модели, STAPLE предлагает дешевый путь вверх по емкости: берется существующий чекпоинт, добавляется статическая память, и модель продолжает обучаться с небольшой потерей скорости. Особенно заметен прирост на коде и математике, где точное знание токенов и шаблонов особенно полезно. Для инференса интересен профиль памяти: основная масса параметров живет на хосте, а на GPU нужен небольшой кэш. Это снижает требования к видеопамяти по сравнению с MoE сопоставимого качества.

Прежде чем переносить подход в продакшен, стоит проверить три вещи: поведение на моделях крупнее 7B, устойчивость к смене домена, где распределение частых токенов другое, и задержку при малых батчах. Если эти проверки пройдут, апсайклинг через статические эмбеддинги станет серьезной альтернативой MoE для расширения уже обученных моделей.