Учёные нашли способ сократить кэш трансформеров до 97% без потери качества

Классическая архитектура трансформеров использует три отдельные проекции для Query, Key и Value (QKV). Однако, как показало новое исследование, опубликованное на arXiv, от двух из них можно отказаться без серьёзной потери качества. Учёные протестировали три варианта объединения проекций: Q-K=V (общий Key-Value), Q=K-V (общий Query-Key) и Q=K=V (единая проекция).

Наибольший практический интерес представляет вариант Q-K=V. При таком подходе Key и Value используют одну матрицу, что вдвое сокращает объём KV-кэша — узкого места при инференсе длинных последовательностей. На языковых моделях объёмом 300 млн и 1,2 млрд параметров (обученных на 10 млрд токенов) перплексия ухудшилась всего на 3,1%.

Более того, объединение проекций совместимо с техниками группового внимания (GQA/MQA), где несколько голов внимания разделяют Key и Value. При комбинации Q-K=V с GQA-4 снижение кэша достигает 87,5%, а с MQA — 96,9%. Такая экономия открывает путь к запуску крупных моделей на мобильных устройствах и другом периферийном оборудовании.

Симметричные варианты Q=K-V и Q=K=V, напротив, нарушают направленность внимания и дают худшие результаты. Для компенсации авторы пробовали двумерные позиционные кодировки, но качество всё равно оставалось ниже.

Исследователи также выяснили, почему Q-K=V работает хорошо. Оказалось, что в обученных моделях Key и Value занимают схожие представления, а внимание работает в низкоранговом пространстве. Таким образом, разделение проекций — это неотъемлемое свойство архитектуры, от которого можно обоснованно отказаться.

Код экспериментов доступен на GitHub. Работа показывает, что упрощение трансформеров может быть не только теоретически интересным, но и практически значимым для развёртывания ИИ на периферии.