TriRoute: единая обучаемая маршрутизация внимания, MoE-экспертов и KV-кэша

Условные вычисления (conditional computation), приём, который позволяет тратить на каждый токен ровно столько ресурсов, сколько ему нужно, а не одинаковый объём для всех, и тем самым отвязать качество языковой модели от стоимости обработки одного токена. Три известные техники решают эту задачу порознь, каждая на своей оси: Mixture-of-Experts (MoE, «смесь экспертов») разрежает полносвязный слой (FFN), направляя токен не во всю сеть, а к нескольким «экспертам»; Mixture-of-Depths (MoD, «смесь глубин») пропускает токен целиком через часть блоков трансформера, экономя вычисления; квантование KV-кэша сжимает память, в которой хранятся ключи и значения внимания для уже обработанных токенов. Авторы новой работы утверждают, что эти три решения, насколько подробно обработать токен вниманием, каким экспертом его обработать и с какой точностью хранить в кэше, тесно связаны между собой и должны приниматься совместно, а не порознь: например, редкий токен, которому нужно полное внимание, скорее всего потребует и хранения в кэше с высокой точностью, независимо от того, какой эксперт его обрабатывает.

Для этого предложен TriRoute, единый лёгкий контроллер, общий сразу для всех трёх осей управления. Для каждого токена на каждом слое трансформера он выдаёт согласованное решение по трём параметрам одновременно: (1) режим внимания, пропустить токен, обработать его локально (по ближайшему контексту) или уделить ему полное внимание; (2) разреженный набор экспертов FFN, причём среди «экспертов» есть и «нулевой», его выбор фактически воспроизводит пропуск блока, как в MoD; (3) битность хранения токена в KV-кэше.

Контроллер обучается end-to-end (единым сквозным процессом вместе с моделью) с помощью неоднородной релаксации дискретных решений: для категориальных решений вроде режима внимания используется техника Gumbel-Softmax со straight-through-оценкой, она позволяет обучать дискретный выбор градиентным спуском, как если бы он был непрерывным, а для выбора экспертов применяется top-k-гейтинг с балансировкой нагрузки, стандартный приём MoE, не позволяющий токенам скапливаться у одних и тех же экспертов. Обучение идёт под лагранжевым ограничением бюджета: средние затраты вычислений и памяти превращаются в управляемый параметр, который можно задать заранее и получить контроллер, укладывающийся в нужный бюджет.

При наивном совместном обучении всех трёх осей авторы обнаружили эффект «каскадного коллапса маршрутизации» (routing-collapse cascade): если контроллер вырождается по одной оси, например, начинает всегда выбирать один и тот же режим внимания, этот коллапс распространяется и на две другие оси. Проблему устранили нормировкой, отдельной для каждой оси, и специальной функцией потерь, учитывающей связь между осями (coupling-aware balancing loss).

Метод проверили на декодерных моделях (архитектура, аналогичная GPT) от 160 миллионов до 1,3 миллиарда параметров, обученных на вычислительно-оптимальном (compute-optimal) числе токенов. При одинаковом бюджете инференса, по числу операций (FLOPs) и по памяти, TriRoute по Парето превосходит лучшую комбинацию из трёх независимых техник (MoD + MoE + квантование KV-кэша): на любой точке допустимого бюджета TriRoute даёт не худшее, а как правило лучшее качество, чем раздельное применение тех же трёх приёмов. Кроме того, TriRoute лучше сохраняет качество на «хвостовых» сложных случаях, редких сущностях (именах, терминах), коде и арифметике, которые обычно проседают, если модель оптимизируют только под среднюю перплексию (усреднённую метрику качества предсказания текста).

Апостериорный анализ показал интерпретируемую структуру решений контроллера: полное внимание и хранение в кэше с высокой точностью он выделяет токенам в начале предложения, редким частям слов (subwords) и именованным сущностям, тогда как служебные слова, предлоги, союзы и другие часто встречающиеся токены, обрабатываются по «дешёвому» маршруту с минимальными затратами.

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

TriRoute, единый контроллер, который для каждого токена на каждом слое трансформера совместно выбирает режим внимания (пропустить/локально/полностью), набор MoE-экспертов FFN (включая «нулевой» эксперт, эквивалентный пропуску блока как в MoD) и битность хранения в KV-кэше.
Обучение end-to-end: Gumbel-Softmax со straight-through-оценкой для категориальных решений, load-balanced top-k-гейтинг для выбора экспертов, лагранжево ограничение бюджета превращает средние вычисления и память в управляемый параметр.
Авторы обнаружили и устранили «каскадный коллапс маршрутизации», эффект, при котором коллапс на одной оси управления (внимание/эксперты/кэш) распространяется на остальные, с помощью нормировки по каждой оси в отдельности и coupling-aware balancing loss (функции потерь, учитывающей связь между осями).
На декодерных моделях от 160 млн до 1,3 млрд параметров TriRoute по Парето превосходит лучшую независимую комбинацию MoD + MoE + квантование KV-кэша при равном бюджете по FLOPs и памяти.
Метод лучше сохраняет качество на редких сущностях, коде и арифметике, которые обычно проседают при оптимизации только под среднюю перплексию; апостериорный анализ показал, что контроллер выделяет полное внимание и точный кэш токенам в начале предложения, редким подсловам и именованным сущностям, а служебные слова обрабатывает дёшево.

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

Инференс больших языковых моделей стоит дорого, и последние несколько лет индустрия ищет способы тратить вычисления неравномерно, больше на «сложные» токены, меньше на «простые». До сих пор для этого использовали три независимых техники (MoE, MoD и квантование KV-кэша), каждую по своей оси, и просто складывали их вместе. TriRoute показывает, что решения по этим трём осям на самом деле сильно связаны между собой, и совместное управление ими даёт лучшее соотношение качества и стоимости инференса, чем любая комбинация раздельных методов при том же бюджете. Отдельно ценно, что метод лучше сохраняет качество именно там, где агрессивная экономия вычислений обычно вредит больше всего, на редких сущностях, коде и арифметике.

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

Исследователям и инженерам, разрабатывающим архитектуры больших языковых моделей и техники условных вычислений (MoE, MoD, квантование кэша). Командам, которые проектируют и обслуживают инференс LLM (inference-инфраструктуру) и ищут способ снизить стоимость обслуживания моделей без потери качества. Разработчикам моделей, ориентированных на код, работу с редкими именами и терминами и на арифметику, именно в этих задачах выигрыш TriRoute над раздельными методами заметнее всего, по данным авторов.

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

TriRoute, не готовый продукт или библиотека, а архитектурное решение, которое нужно встраивать в модель на этапе предобучения: контроллер обучается end-to-end вместе с самой моделью под заданный лагранжев бюджет вычислений и памяти. Из текста работы не следует, что для TriRoute опубликован код или веса, судя по всему, это исследовательский результат, который командам, разрабатывающим следующее поколение эффективных LLM-архитектур, предстоит воспроизводить самостоятельно. Эксперименты проведены на моделях до 1,3 млрд параметров, поэтому перед использованием на моделях промышленного масштаба (десятки-сотни миллиардов параметров) потребуется отдельная проверка, что эффект сохраняется при увеличении размера.

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

Работа опубликована на arXiv как препринт, независимого рецензирования по тексту источника не видно. При этом методология сравнения выглядит корректной: TriRoute сравнивают не с «сырыми» базовыми моделями, а с лучшей возможной комбинацией трёх существующих независимых техник (MoD + MoE + квантование KV-кэша) при строго одинаковом бюджете по FLOPs и памяти, то есть сопоставление по Парето, а не по одной цифре. Проверка проведена на диапазоне размеров моделей, от 160 млн до 1,3 млрд параметров, что снижает риск случайного результата на одной конкретной конфигурации. Ограничение в том, что абстракт работы не приводит точные числа (проценты, значения перплексии), поэтому оценить величину выигрыша количественно по одному только тексту нельзя.

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

Сами авторы описывают, что наивное совместное обучение трёх осей маршрутизации приводит к «каскадному коллапсу», это показывает, что подход технически хрупкий и требует специальных приёмов (нормировка по каждой оси в отдельности, coupling-aware balancing loss), без которых система ломается; насколько сложно воспроизвести эти приёмы в других условиях, из текста не ясно. Все три решения, внимание, эксперты, кэш, принимает один общий контроллер, а значит, ошибка или смещение в этом контроллере способны одновременно испортить все три оси сразу, в отличие от независимых систем, которые в среднем ломаются по отдельности. Наконец, результаты получены на моделях до 1,3 млрд параметров, заметно меньше моделей, которые реально используются в продакшен-сервисах, поэтому применимость на промышленном масштабе пока не подтверждена.

«Токен, достаточно редкий, чтобы заслуживать полного внимания, может также нуждаться в кэшировании с высокой точностью, независимо от того, какой эксперт его обрабатывает.»

— TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation (arXiv:2607.06601)