Новый метод сжатия KV-кэша для LLM: фрактальные коды ускоряют инференс

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

Что такое KV-кэш и почему его сжатие критически важно

KV-кэш (Key-Value cache) — это структура данных, используемая в трансформерах для хранения промежуточных представлений ключей и значений внимания. При генерации текста модель обращается к этому кэшу, чтобы не пересчитывать предыдущие токены. Однако с увеличением длины контекста размер кэша растет линейно, что создает серьезные ограничения по памяти. Например, для модели с 70 миллиардами параметров и контекстом в 128 тысяч токенов KV-кэш может занимать более 100 гигабайт. Это делает инференс дорогим и замедляет работу приложений, требующих длительного диалога или анализа больших документов.

Существующие методы сжатия, такие как квантование, удаление малозначимых токенов или выгрузка на диск, решают проблему лишь частично. Квантование снижает точность, удаление теряет информацию, а выгрузка увеличивает задержки. Новый метод, описанный в препринте, предлагает принципиально иной подход: фрактальные итеративные коды, которые преобразуют последовательность символов в последовательность низкоразмерных векторов. Это позволяет создать архив, поддерживающий произвольный доступ за O(1) и вставку за амортизированное O(1).

Как работает новый метод сжатия KV-кэша

В основе метода лежит векторное квантование с остатком (residual VQ), применяемое к каждому слою внимания. Сначала модель выделяет небольшое точное окно, которое включает несколько начальных токенов (так называемые attention sinks) и несколько последних токенов. Остальные токены сжимаются с помощью фрактальных кодов. Эксперименты проводились на модели GPT-2 с контекстом 1024 токена. При точном окне из 4 attention sinks и 32 последних токенов сжатие составило от 36 до 54 раз по сравнению с хранением в формате fp16. При этом увеличение перплексии (меры качества генерации) составило от 11 до 15 процентов.

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

Почему этот метод может изменить инференс LLM

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

Кроме того, произвольный доступ за O(1) означает, что модель может обратиться к любому элементу кэша без последовательного сканирования. В сочетании с амортизированной вставкой за O(1) это делает метод пригодным для потоковой обработки, когда контекст постоянно расширяется. Разработчики чат-ботов и систем с длинным контекстом смогут развертывать модели на устройствах с ограниченной памятью, таких как смартфоны или периферийные серверы.

Какие остаются вопросы и ограничения

Несмотря на многообещающие результаты, метод пока не проверен на крупных моделях и контекстах. Эксперименты проводились на GPT-2 с 1024 токенами, в то время как современные системы работают с контекстами до 128 тысяч токенов и более. Неясно, насколько хорошо фрактальные коды масштабируются на такие размеры. Также не проводилось прямое сравнение с современными методами сжатия, такими как KIVI или StreamingLLM. Без таких бенчмарков сложно оценить практическую ценность.

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

Когда ждать внедрения и кого это коснется

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

Как новый метод сжатия KV-кэша повлияет на производительность LLM

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

Однако стоит помнить, что любое сжатие связано с потерей информации. Увеличение перплексии на 11-15 процентов может быть приемлемым для многих приложений, но для задач, требующих высокой точности, может потребоваться компромисс. Гибридное распределение битов между ключами и значениями дает возможность тонкой настройки, но окончательное решение будет зависеть от конкретного сценария использования.

Заключение

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