Dual-Flow Transformer: новая архитектура снижает затраты на инференс LLM

На arXiv опубликована работа, описывающая новую архитектуру нейросетей Dual-Flow Transformer. Её цель — сократить накопительные затраты на инференс больших языковых моделей (LLM), которые становятся всё более значимыми по сравнению с одноразовыми расходами на обучение.

В стандартных LLM выделяются две фазы: prefill — параллельная обработка промпта, обычно ограниченная вычислительной мощностью, и decode — последовательная генерация токенов, где узким местом часто является пропускная способность памяти. Традиционное масштабирование модели по ширине или глубине увеличивает нагрузку на обе фазы, так как каждый новый слой задействуется и в prefill, и в decode.

Авторы предлагают иной подход: добавить вычисления, которые направлены только на предсказание продолжения, сохраняя неизменным основной поток обработки промпта и единый персистентный KV-кэш. Архитектура Dual-Flow Transformer состоит из основного потока — полноценной каузальной модели, которая обрабатывает промпт и записывает KV-кэш, и вспомогательного потока, который активируется лишь с финальной позиции промпта. Этот дополнительный поток добавляет вычисления для генерации продолжения, не влияя на основной поток и не создавая дополнительного персистентного состояния.

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

В ходе экспериментов Dual-Flow показал более низкое значение функции потерь на валидации по сравнению с сопоставимыми по количеству токенов моделями в различных архитектурных конфигурациях. Особый интерес представляют результаты для моделей на смеси экспертов (MoE): разделение позволяет независимо управлять числом активных экспертов на этапах prefill и decode, что даёт гибкость в оптимизации затрат.

Исследователи изучили два режима: увеличение вычислительной нагрузки на decode при фиксированной экспертной нагрузке на prefill, и перераспределение фиксированного бюджета decode-экспертов между потоками. Эти эксперименты выявили компромисс между качеством предсказаний и распределением ресурсов по фазам инференса.

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