PuzzleKV сжимает KV-кэш LLM: 96% точности при 60% памяти
Рост контекстных окон в больших языковых моделях (LLM) упирается в ограничения памяти: ключевое значение для хранения промежуточных состояний имеет KV-кэш. Его объем растет линейно с длиной контекста, что затрудняет инференс на длинных последовательностях. Существующие методы сжатия, включая низкоранговые, часто используют фиксированные проекционные пространства, которые могут терять важные детали.
Группа исследователей предложила метод PuzzleKV, который разбивает каждый per-head KV-кэш на логические страницы фиксированной длины. Внутри отдельных страниц обнаруживается выраженная низкоранговая структура, что позволяет сжимать их как независимые единицы. Такой подход отличается от глобальных методов вроде SVD, которые строят общий базис для всего кэша.
PuzzleKV работает без обучения и калибровки на данных. Он раскладывает страницы внутри каждого слоя и головы внимания, вычисляет внимание непосредственно над плотными и факторизованными страницами, а также инкрементально сжимает новые страницы в процессе авторегрессионной генерации. Это делает метод пригодным для практического применения в реальном времени.
В экспериментах на моделях и бенчмарках RULER и LongBench метод показал эффективность при одинаковых бюджетах хранения. При использовании около 60% исходного объема KV-кэша PuzzleKV сохраняет более 96% производительности полного кэша на всех протестированных моделях и настройках. Это заметно лучше глобального SVD на RULER и сравнимо на LongBench.
Для более агрессивного сжатия PuzzleKV можно комбинировать с квантованием. В таком режиме удается сохранить более 93% производительности при использовании лишь 18.7% исходного объема памяти. Это открывает возможности для значительного увеличения эффективной длины контекста без существенного роста требований к GPU-памяти.
Работа опубликована на arXiv, авторы подчеркивают, что метод решает проблему детализации представлений, свойственную прежним подходам. PuzzleKV может быть полезен для развертывания LLM на устройствах с ограниченными ресурсами и для систем, работающих с длинными документами.
