PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Cai, Zefan, Zhang, Yichi, Gao, Bofei, Liu, Yuliang, Li, Yucheng, Liu, Tianyu, Lu, Keming, Xiong, Wayne, Dong, Yue, Hu, Junjie, Xiao, Wen
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866913839478472704
author Cai, Zefan
Zhang, Yichi
Gao, Bofei
Liu, Yuliang
Li, Yucheng
Liu, Tianyu
Lu, Keming
Xiong, Wayne
Dong, Yue
Hu, Junjie
Xiao, Wen
author_facet Cai, Zefan
Zhang, Yichi
Gao, Bofei
Liu, Yuliang
Li, Yucheng
Liu, Tianyu
Lu, Keming
Xiong, Wayne
Dong, Yue
Hu, Junjie
Xiao, Wen
contents In this study, we investigate whether attention-based information flow inside large language models (LLMs) is aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers, progressively consolidating within specific contexts, and ultimately focusing on critical tokens (a.k.a massive activation or attention sink) in higher layers. Motivated by these insights, we developed PyramidKV, a novel and effective KV cache compression method. This approach dynamically adjusts the KV cache size across different layers, allocating more cache in lower layers and less in higher ones, diverging from traditional methods that maintain a uniform KV cache size. Our experimental evaluations, utilizing the LongBench benchmark, show that PyramidKV matches the performance of models with a full KV cache while retaining only 12% of the KV cache, thus significantly reducing memory usage. In scenarios emphasizing memory efficiency, where only 0.7% of the KV cache is maintained, PyramidKV surpasses other KV cache compression techniques, achieving up to a 20.5 absolute accuracy improvement on TREC dataset. In the Needle-in-a-Haystack experiment, PyramidKV outperforms competing methods in maintaining long-context comprehension in LLMs; notably, retaining just 128 KV cache entries enables the LLAMA-3-70B model to achieve 100.0 Acc. performance.
format Preprint
id arxiv_https___arxiv_org_abs_2406_02069
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
Cai, Zefan
Zhang, Yichi
Gao, Bofei
Liu, Yuliang
Li, Yucheng
Liu, Tianyu
Lu, Keming
Xiong, Wayne
Dong, Yue
Hu, Junjie
Xiao, Wen
Computation and Language
Artificial Intelligence
In this study, we investigate whether attention-based information flow inside large language models (LLMs) is aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers, progressively consolidating within specific contexts, and ultimately focusing on critical tokens (a.k.a massive activation or attention sink) in higher layers. Motivated by these insights, we developed PyramidKV, a novel and effective KV cache compression method. This approach dynamically adjusts the KV cache size across different layers, allocating more cache in lower layers and less in higher ones, diverging from traditional methods that maintain a uniform KV cache size. Our experimental evaluations, utilizing the LongBench benchmark, show that PyramidKV matches the performance of models with a full KV cache while retaining only 12% of the KV cache, thus significantly reducing memory usage. In scenarios emphasizing memory efficiency, where only 0.7% of the KV cache is maintained, PyramidKV surpasses other KV cache compression techniques, achieving up to a 20.5 absolute accuracy improvement on TREC dataset. In the Needle-in-a-Haystack experiment, PyramidKV outperforms competing methods in maintaining long-context comprehension in LLMs; notably, retaining just 128 KV cache entries enables the LLAMA-3-70B model to achieve 100.0 Acc. performance.
title PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
topic Computation and Language
Artificial Intelligence
url https://arxiv.org/abs/2406.02069