Новый метод TriRoute объединяет три техники оптимизации LLM, сокращая затраты без потери качества

Команда исследователей опубликовала на arXiv препринт статьи, в которой представила новый метод оптимизации работы больших языковых моделей (LLM) под названием TriRoute. Разработка направлена на снижение вычислительных затрат и объёма используемой памяти без ущерба для качества генерации.

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

Контроллер обучается сквозным образом с помощью гетерогенной релаксации (Gumbel-Softmax и нагрузочно-сбалансированного top-k gating) под бюджетным ограничением на вычислительные затраты и память. Это позволяет настраивать средние расходы ресурсов как регулируемый параметр. В ходе обучения учёные выявили проблему «каскадного разрушения маршрутизации» при наивном совместном обучении, когда коллапс в одной оси распространяется на другие. Для её решения были введены поосевая нормализация и функция потерь, учитывающая балансировку нагрузки между осями.

Эксперименты проводились на моделях только с декодером от 160 млн до 1,3 млрд параметров при оптимальном количестве токенов. TriRoute показал превосходство по Парето-критерию над лучшей независимой комбинацией MoD, MoE и квантования кэша KV при одинаковых вычислительных затратах и объёме памяти. Важно, что метод лучше сохраняет качество на редких сущностях, коде и арифметических задачах, которые часто страдают при оптимизации, нацеленной лишь на снижение перплексии.

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