ThinK: Thinner Key Cache by Query-Driven Pruning
Large Language Models (LLMs) have revolutionized the field of natural
language processing, achieving unprecedented performance across a variety of
applications by leveraging increased model sizes and sequence lengths. However,
the associated rise in computational and memory costs poses significant
challenges, particularly in managing long sequences due to the quadratic
complexity of the transformer attention mechanism. This paper focuses on the
long-context scenario, addressing the inefficiencies in KV cache memory
consumption during inference. Unlike existing approaches that optimize the
memory based on the sequence lengths, we uncover that the channel dimension of
the KV cache exhibits significant redundancy, characterized by unbalanced
magnitude distribution and low-rank structure in attention weights. Based on
these observations, we propose ThinK, a novel query-dependent KV cache pruning
method designed to minimize attention weight loss while selectively pruning the
least significant channels. Our approach not only maintains or enhances model
accuracy but also achieves a reduction in memory costs by over 20% compared
with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and
Mistral models across various long-sequence datasets confirm the efficacy of
ThinK, setting a new precedent for efficient LLM deployment without
compromising performance. We also outline the potential of extending our method
to value cache pruning, demonstrating ThinK's versatility and broad
applicability in reducing both memory and computational overheads.
66a9b4a435c896924429516b
Yuhui Xu
claimed_verified
2024-07-31T07:44:00.316Z
6602869253a0518b2a98cafd
/avatars/c14b5953a716f42c83ad28147f8308ae.svg
Yuhui Xu
user
yuhuixu
66a9b4a435c896924429516c
Zhanming Jie
admin_assigned
2024-07-31T07:58:57.282Z
5f729da2f1e7ef6e919a1c27
/avatars/bdf48f0d3290fab54d7e9636a3d14158.svg
Zhanming (Allan) Jie
user
allanjie
66a9b4a435c896924429516d
Hanze Dong
admin_assigned
2024-07-31T07:59:10.358Z
63a3ff69f91ad3ea5703841d
/avatars/69227c4bce01d33747c1377b6f9672db.svg
Hanze Dong
user
hendrydong
66a9b4a435c896924429516e
Lei Wang
claimed_verified
2024-07-31T12:49:35.556Z
646def60df618b303b419323
/avatars/97aa761d5255abf230304cfeade87835.svg
Lei Wang
user
demolei
66a9b4a435c896924429516f
Xudong Lu
claimed_verified
2024-07-31T07:44:02.548Z
642e686bbe01b88c9446db8b
https://cdn-avatars.huggingface.co/v1/production/uploads/642e686bbe01b88c9446db8b/tb1DKe5xt50ykOeXiUuTE.jpeg
Lu Xudong
user
lucky-lance
66a9b4a435c8969244295170
Aojun Zhou
claimed_verified
2024-10-11T07:45:31.264Z
637de1520d5bb06fbe5207a9
/avatars/1090851217270c5a858b13e013356d4f.svg
AJ.Zhou
user
AJZhou
66a9b4a435c8969244295171
Amrita Saha
admin_assigned
2024-07-31T08:01:09.678Z
6461c2905dba83471db3be53
/avatars/6e36cf86201d590ac729a75d4a439cde.svg
Amrita Saha
user
amritasaha87
66a9b4a435c8969244295172
Caiming Xiong
admin_assigned
2024-07-31T08:01:15.606Z
649dbcc4e0fff1ed099dc80a
/avatars/c87c273ca628dbcddccbf1ee19b2ce33.svg
Caiming Xiong
user
cxiong
66a9b4a435c8969244295173
Doyen Sahoo
admin_assigned
2024-07-31T08:01:22.415Z
65f84fd980481173afd91233
/avatars/6ac7bd6beba24d1476c5179b88c9e3fa.svg
Doyen
user
doyensahoo
ThinK: Thinner Key Cache by Query-Driven Pruning
Large Language Models (LLMs) have revolutionized the field of natural
language processing, achieving unprecedented performance across a variety of
applications by leveraging increased model sizes and sequence lengths. However,
the associated rise in computational and memory costs poses significant
challenges, particularly in managing long sequences due to the quadratic
complexity of the transformer attention mechanism. This paper focuses on the
long-context scenario, addressing the inefficiencies in KV cache memory
consumption during inference. Unlike existing approaches that optimize the
memory based on the sequence lengths, we uncover that the channel dimension of
the KV cache exhibits significant redundancy, characterized by unbalanced
magnitude distribution and low-rank structure in attention weights. Based on
these observations, we propose ThinK, a novel query-dependent KV cache pruning
method designed to minimize attention weight loss while selectively pruning the
least significant channels. Our approach not only maintains or enhances model
accuracy but also achieves a reduction in memory costs by over 20% compared
with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and
Mistral models across various long-sequence datasets confirm the efficacy of
ThinK, setting a new precedent for efficient LLM deployment without
compromising performance. We also outline the potential of extending our method
to value cache pruning, demonstrating ThinK's versatility and broad
applicability in reducing both memory and computational overheads.
ThinK, a novel query-dependent KV cache pruning method, effectively reduces memory costs by over 20% in long-context scenarios without compromising the accuracy of Large Language Models (LLMs).
Large Language Models (LLMs) computational costs memory costs long sequences transformer attention mechanism KV cache memory channel dimension redundancy unbalanced magnitude distribution low-rank structure attention weights ThinK query-dependent KV cache pruning model accuracy LLaMA3 Mistral models long-sequence datasets value cache pruning