NVIDIA выложила в открытый доступ метод KVzap. Он сжимает кэш ключей и значений в трансформерах. Сжатие почти без потерь достигает 2-4 раз.
Кэш для модели Llama1-65B занимает 335 гигабайт. Это на 128 тысячах токенов. Память ограничивает размер батча и замедляет ответ.
Другие методы сжимали кэш по другим осям. Они не трогали последовательность токенов. Проблема оставалась.
Команда NVIDIA собрала два десятка методов в проекте KVpress. Они выложили таблицу лидеров на Hugging Face. KVzip был самым сильным базовым методом.
KVzip оценивал важность каждой записи в кэше. Он просил модель повторить исходный контекст. Метод был точным, но медленным.
KVzap заменяет сложную оценку маленькой моделью-суррогатом. Она работает напрямую со скрытыми состояниями. Архитектура — линейный слой или небольшая MLP.
Модель обучали на 27 тысячах промптов. Длина промптов — от 750 до 1250 токенов. Для каждого ключа и значения получили 1,2 миллиона пар для обучения.
Суррогат учился предсказывать оценку KVzip+. Корреляция предсказаний с оригинальными оценками — от 0,63 до 0,77. MLP-вариант стабильно обгонял линейный.
Новый метод быстрый. Он не требует удвоения контекста. Его можно использовать в реальном времени.
