ОБЪЕКТ ИССЛЕДОВАНИЯ: Оптимизация сжатия кэша ключ-значение в генерации длинного контекста

SCOPE: Optimizing Key-Value Cache Compression in Long-context Generation

December 18, 2024

Авторы: Jialong Wu, Zhenglin Wang, Linhai Zhang, Yilong Lai, Yulan He, Deyu Zhou

cs.AI

Аннотация

Кэш ключ-значение (KV) стал узким местом для LLM при генерации длинного контекста. Несмотря на многочисленные усилия в этой области, оптимизация для фазы декодирования обычно игнорируется. Однако мы считаем, что такая оптимизация крайне важна, особенно для задач генерации длинного вывода на основе следующих двух наблюдений: (i) Чрезмерное сжатие во время фазы предварительной загрузки, которое требует конкретного полного контекста, ухудшает понимание задачи рассуждения; (ii) Отклонение от основных элементов происходит в задачах рассуждения с длинными выводами. Поэтому была представлена SCOPE, простая, но эффективная структура, которая отдельно выполняет оптимизацию кэша ключ-значение во время фаз предварительной загрузки и декодирования. Конкретно, кэш ключ-значение во время фазы предварительной загрузки сохраняется для поддержания основной информации, в то время как предложена новая стратегия на основе скольжения для выбора основных элементов для фазы декодирования. Используются адаптивные и дискретные стратегии для дальнейшей оптимизации использования памяти и передачи памяти. Обширные эксперименты на LongGenBench показывают эффективность и обобщаемость SCOPE, а также его совместимость в качестве плагина для других методов сжатия KV только для предварительной загрузки.

English

Key-Value (KV) cache has become a bottleneck of LLMs for long-context generation. Despite the numerous efforts in this area, the optimization for the decoding phase is generally ignored. However, we believe such optimization is crucial, especially for long-output generation tasks based on the following two observations: (i) Excessive compression during the prefill phase, which requires specific full context impairs the comprehension of the reasoning task; (ii) Deviation of heavy hitters occurs in the reasoning tasks with long outputs. Therefore, SCOPE, a simple yet efficient framework that separately performs KV cache optimization during the prefill and decoding phases, is introduced. Specifically, the KV cache during the prefill phase is preserved to maintain the essential information, while a novel strategy based on sliding is proposed to select essential heavy hitters for the decoding phase. Memory usage and memory transfer are further optimized using adaptive and discontinuous strategies. Extensive experiments on LongGenBench show the effectiveness and generalization of SCOPE and its compatibility as a plug-in to other prefill-only KV compression methods.