Новый механизм внимания Step-снижение KV-кэша

Введение

Растущее использование больших языковых моделей (LLM) и появление новых парадигм в выводе поставили на первый план задачу эффективного крупномасштабного вывода. Значительным узким местом является кэш Key-Value (KV) в традиционных механизмах внимания, который линейно расширяется с размером пакета и длиной последовательности, становясь "пожирателем памяти", что препятствует масштабированию и расширению LLM.

В то время как появились варианты, такие как MQA, GQA и MLA, для решения этой проблемы, они часто испытывают трудности с поддержанием производительности в условиях строгих ограничений памяти или вводят сложности, которые создают инженерные проблемы и проблемы совместимости.

Многоматричное факторизационное внимание (MFA)

Недавняя статья исследователей из Stepes, Университета Цинхуа и других учреждений представляет новую архитектуру механизма внимания: Multi-matrix Factorization Attention (MFA) и ее вариант MFA-Key-Reuse (MFA-KR). Этот механизм значительно снижает стоимость вывода языковой модели, одновременно улучшая производительность.

MFA и MFA-KR не только превосходят MLA по производительности, но и соответствуют традиционной производительности MHA, снижая использование KV-кэша до 93,7%.

MFA разработан для простоты, легкого воспроизведения, низкой чувствительности к гиперпараметрам и совместимости с различными методами Pos-встраивания.

Подход и анализ MFA

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

Они ввели концепцию Generalized Multi-Head Attention (GMHA) в качестве единой основы для понимания различных вариантов MHA.

Команда также изучила вычисление и хранение ключевых значений с точки зрения вывода и изучила возможности модели с точки зрения декомпозиции.

Fully Parameterized Bilinear Attention (FPBA) был установлен в качестве теоретического верхнего предела производительности. Они обнаружили, что MHA и его варианты являются низкоранговыми разложениями FPBA.

Сравнение с MQA и MLA

Анализ был сосредоточен на двух представительных схемах улучшения: Multi-Query Attention (MQA) и Multi-Head Latent Attention (MLA).

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

Ключевые инновации MFA

Разработка MFA была обусловлена целью создать механизм внимания, который минимизирует потребление ресурсов, приближаясь к теоретическим пределам производительности. Конструкция MFA включает в себя три ключевые инновации:

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

Измерение и сравнение возможностей

Для дальнейшего анализа MFA и других механизмов внимания команда ввела две ключевые метрики:

Total Effective Rank (TER): произведение количества голов внимания и ранга факторизации на голову (FRH).
Shared Latent Subspace Dimension (SLSD): размерность скрытого пространства, общего для всех голов внимания.

MFA достигает более высокого SLSD и TER по сравнению с MQA.

По сравнению с MLA, MFA достигает меньшего размера KV-кэша и более высокого TER при аналогичных бюджетах параметров, сохраняя при этом сопоставимый SLSD.

По сравнению с традиционным MHA, MFA имеет более высокий TER, даже если его SLSD меньше.

Экспериментальные результаты

Были проведены обширные эксперименты для оценки производительности новой архитектуры в больших масштабах, тестирование моделей от 1B до 7B параметров и обучающих данных от 10B до 1T.

MFA продемонстрировал возможности масштабирования, сопоставимые с традиционными MHA, сохраняя отличную производительность даже в больших масштабах.

В то время как MFA-KR показал немного более низкую производительность, его тенденция масштабирования соответствовала MHA. Преимущества MFA и MFA-KR в экономии памяти продолжали расширяться с размером модели, при этом MFA достиг 87,5% экономии памяти, а MFA-KR снизил использование памяти до 6,25% в самом большом масштабе.

Исследования абляции

Исследования абляции подтвердили эффективность MFA и MFA-KR. Их преимущества в производительности также были подтверждены для различных основных методов позиционного кодирования.

Перспективы

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