TriRoute: контроллер для совместного управления вниманием, экспертами и KV-кэшем

Исследователи представили TriRoute — единый обучаемый контроллер, который для каждого токена на каждом слое трансформера принимает скоординированное решение по трём осям: режиму внимания (пропуск, локальное или полное), набору экспертов FFN (включая «нулевого» эксперта для пропуска блока) и разрядности квантования KV-кэша. Это первый подход, объединяющий три ключевые техники условных вычислений в одну систему, что позволяет добиться лучшего баланса между качеством и вычислительными затратами. Контроллер обучается сквозным образом с использованием гетерогенной релаксации — Gumbel-Softmax с straight-through оценкой для категориальных решений и load-balanced top-k gating для экспертов — под лагранжевым бюджетным ограничением, которое плавно регулирует среднюю стоимость вычислений и памяти.

Как работает TriRoute

TriRoute основан на идее, что решения о внимании, экспертах и квантовании взаимосвязаны. Например, токен, требующий полного внимания, вероятно, также нуждается в высокоточном кэшировании, независимо от выбранного эксперта. Контроллер анализирует скрытое состояние токена и выдаёт три скоординированных решения, оптимизируя совместно все три аспекта. Обучение происходит с помощью лагранжевой релаксации, что позволяет гибко настраивать бюджет вычислительных ресурсов.

Как TriRoute сочетает три оси управления

Каждая ось имеет свои опции. Для внимания: пропуск (не вычислять внимание), локальное внимание (ограниченный контекст) или полное внимание. Для экспертов FFN: выбор одного из нескольких экспертов или пропуск блока через «нулевого» эксперта. Для KV-кэша: низкая, средняя или высокая разрядность квантования. Контроллер принимает решение для каждого токена на каждом слое, что позволяет адаптировать вычисления под конкретные потребности.

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

Существующие методы условных вычислений — Mixture-of-Experts (MoE), Mixture-of-Depths (MoD) и квантование KV-кэша — действуют независимо, хотя их решения сильно связаны. Например, токен, требующий полного внимания, скорее всего, также нуждается в высокоточном кэшировании. TriRoute впервые предлагает совместное управление всеми тремя осями, что позволяет добиться лучшего баланса между качеством и вычислительными затратами. Это особенно актуально для развёртывания больших языковых моделей на устройствах с ограниченными ресурсами.

Результаты тестирования

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

Какие задачи решает TriRoute лучше всего

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

Кого затронет TriRoute

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

Что пока неизвестно

В статье не указано, как метод ведёт себя на моделях больше 1.3B параметров, а также не приведены результаты на задачах, требующих длинного контекста (более 8K токенов). Неясно, насколько сложно внедрить TriRoute в существующие фреймворки. Дополнительные исследования должны прояснить масштабируемость и практическую применимость метода.