Ada-KV
Ada-KV: 通过自适应预算分配优化KV缓存淘汰以提高高效LLM推理
摘要
大型语言模型在各个领域表现出色，但由于长序列推理所需的键值（KV）缓存不断扩大，面临内存和时间效率的挑战。最近的努力试图在运行时通过淘汰大量非关键缓存元素，将KV缓存大小减少到给定的内存预算，同时保留生成质量。我们重新审视当前的淘汰方法，发现它们基本上最小化了多头自注意力机制预淘汰和淘汰后输出之间的L1淘汰损失的上界。此外，我们的分析表明，跨注意力头均匀分配预算的常见做法损害了它们淘汰后生成质量。鉴于这些发现，我们提出了一个简单而有效的自适应预算分配算法。该算法不仅优化了理论损失上界，而且通过与不同头部的不同特性对齐，实际上减少了L1淘汰损失。通过将这个算法集成到两种最先进的方法中，我们展示了使用自适应预算分配来优化KV缓存淘汰的有效性。在16个数据集和“大海捞针”测试上的广泛评估证实了各种任务的显著性能提升。我们的代码可在 https://github.com/FFY0/AdaKV 获取。
1 引言
自回归大型语言模型（LLMs）取得了显著的成功，并在多种自然语言处理应用中得到广泛应用，包括对话系统（Yi et al. 2024）、文档摘要（Laban et al. 2023）和代码生成（Gu 2023）。LLMs的广泛部署推动了它们处理扩展序列的能力的发展。例如，GPT-4支持长达128K的序列（Achiam et al. 2023），Claude3支持长达200K（Anthropic 2024），Gemini-Pro-1.5（Reid et al. 2024）支持长达1M个标记。现代LLMs通常构建有多个transformer块层，每个层都包含一个多头自注意力层。这种自注意力机制的计算成本随着输入序列的标记数量呈二次方增加。因此，LLMs在处理长序列输入时面临重大的效率挑战。
对于每个多头自注意力层，推理过程包括两个阶段：预填充和解码。在预填充阶段，LLMs计算并存储来自提示的所有输入标记的键值（KV）缓存元素。在随后的解码中，LLMs自回归地使用最后一个标记从过去的KV缓存元素中检索信息以生成每个输出标记，直到达到最大长度或主动停止该过程。然而，在处理长序列输入时，生成的缓存大小可能会显著增加，甚至远远超过模型的参数大小（Sun et al. 2024）。这种增长导致了在预填充和解码阶段的重大内存挑战。此外，解码过程中大量的KV缓存I/O会产生显著的延迟。这种延迟因此成为解码的瓶颈，甚至超过了计算时间（Tang et al. 2024）。
为了解决大KV缓存大小带来的挑战，已经开发了各种缓存淘汰方法，正如最近文献中所强调的（Ge et al. 2023; Zhang et al. 2024b; Yang et al. 2024; Zhang et al. 2024a; Li et al. 2024）。如图1所示，这些方法仅在每个头部保留预算大小的缓存元素，同时淘汰其他元素。因此，它们减少了内存负担并提高了解码速度，从而促进了高效的长序列推理。这些淘汰方法是设计为即插即用的，允许轻松集成到任何LLM中，无需微调。它们通常采用基于注意力权重的Topk选择策略，以有效地区分关键和非关键缓存元素，决定保留哪些和淘汰哪些。尽管取得了这些进展，但在使用这些淘汰方法的同时最小化质量损失的挑战仍未解决。
我们的研究首先重新审视当前的Top-k淘汰方法，从理论角度揭示它们的基本原理。我们揭示了它们相当于最小化自注意力机制预淘汰和淘汰后输出之间的L1距离的上界。此外，我们发现跨不同注意力头均匀分配预算的常见做法导致了资源的错配，限制了现有方法的有效性。基于这些发现，我们提出了第一个自适应预算分配算法。如图2所示的示例，它通过根据它们不同的集中度自适应地在不同注意力头之间分配整体预算，有效地提高了预算利用率，从而提高了淘汰后的生成质量。我们进一步从理论和实证上展示了自适应预算分配的优势。从理论上讲，与均匀分配相比，它降低了L1淘汰损失的上界。从实证上讲，它与头部间不同的集中度对齐，有效地减少了实际中的L1淘汰损失。
通过将这个自适应分配算法集成到两种最先进的（SOTA）方法中，我们开发了两种自适应淘汰方法，Ada-SnapKV和Ada-Pyramid。在16个标准数据集上的广泛评估，涵盖了LongBench（Bai et al. 2023）中的多种任务，证明了Ada-SnapKV和AdaPyramid在生成质量上的显著改进。此外，“大海捞针”基准测试进一步证实了自适应分配在上下文检索能力上的改进。主要贡献总结如下：
- 通过将淘汰损失定义为预淘汰和淘汰后输出之间的L1距离，我们揭示了当前淘汰方法在分配的预算内等效地最小化其理论上限。我们进一步确定，缓存淘汰中普遍存在的均匀预算分配阻碍了这种最小化。
- 我们引入了第一个自适应预算分配算法，用于缓存淘汰，它既优化了理论上限，也优化了实际的淘汰损失。通过将这个算法集成到两种SOTA方法中，我们开发了Ada-SnapKV和Ada-Pyramid，通过全面的评估证实了显著的增强。
- 我们提供了一个基于淘汰损失上限的理论框架，用于分析缓存淘汰，并通过高效的CUDA实现展示了实际进展。这个框架和实现为通过自适应分配策略优化缓存淘汰铺平了道路。
2 相关工作
在长序列推理中，KV缓存元素的庞大规模导致了内存受限的情况，引起了显著的内存负担和I/O延迟（Wang和Chen 2023）。许多研究试图通过减少缓存大小来缓解这个问题，特别是通过淘汰非关键的缓存元素。这些缓存淘汰方法主要分为两类：滑动窗口淘汰和Top-k淘汰方法。滑动窗口淘汰方法（2020；2024；2023），以StreamingLLM（Xiao等人，2023）为例，简单地保留一些初始缓存元素和滑动窗口内的元素，同时淘汰其他元素。然而，这种不加选择的滑动淘汰缓存元素导致生成质量大幅下降。相比之下，Top-k淘汰方法（2023；2024b；2024；2024b；2024；2024a；2024）基于注意力权重识别和保留一组关键的缓存元素，以提高淘汰后的生成质量。本文提出的自适应预算分配，针对Top-k淘汰方法，在同一总体预算内提高了淘汰后的生成质量。
在Top-k淘汰方法领域，早期的工作如FastGen（Ge等人，2023）搜索并结合了多种策略，例如维护特殊元素、标点元素、近期元素和Top-k选定元素的缓存，基于注意力头的特征。H2O作为作品的代表（Zhang等人，2024b；Liu等人，2024b；Ren和Zhu，2024），开发了一种基于Top-k的淘汰方案，利用所有标记的查询状态来识别关键缓存元素。然而，由于单向注意力机制，这些方法中的最新标记积累的观察次数较少，导致了错误的淘汰。最近的工作，如SnapKV（Li等人，2024）和Pyramid（Yang等人，2024；Zhang等人，2024a），通过使用观察窗口内的查询状态来解决这个问题，从而实现了SOTA性能。然而，据我们所知，现有的Top-k淘汰方法通常在不同头部之间均匀分配总体预算，导致错配。相比之下，我们的自适应分配算法在理论和实证上都展示了优越的结果，为优化现有方法提供了一种新的方法。
3 方法
在本节中，我们首先提供多头自注意力层的正式描述（第3.1节）。基于此，我们通过引入 L 1 L_1 L1淘汰损失度量，从理论上重新审视现有Top-k淘汰方法的基本原理（第3.2节）。受到理论发现的启发，我们提出了一个简单而有效的自适应预算分配算法，该算法在理论和实践中都被证明优于传统的均匀预算分配（第3.3节）。我们进一步通过将算法集成到两个最先进的工作中，验证了其与现有Top-k淘汰方法的兼容性，从而提高了淘汰后的生成质量（第3.4节）。
3.1 预备知识
LLMs的特征是自回归生成模式，其中每一步都涉及使用最后一个标记来预测下一个标记。定义 X ∈ R n × d X ∈ R^{n×d} X∈Rn×d为序列中所有标记的嵌入矩阵， x ∈ R 1 × d x ∈ R^{1×d} x∈R1×d作为当前时间步用作输入的最后一个标记。为了澄清后续的理论阐述，我们采用（Liu等人，2023）中的符号系统，在假设h个注意力头的情况下，对一个多头自注意力层进行正式描述。每个头部 i ∈ [ 1 , h ] i ∈ [1, h] i∈[1,h]的变换矩阵 W i Q , W i K , W i V ∈ R d × d h W^Q_i, W^K_i, W^V_i ∈ R^{d×d_h} WiQ,WiK,WiV∈Rd×dh将标记嵌入映射到它们各自的查询、键和值，最终输出矩阵 W i O ∈ R d h × d W^O_i ∈ R^{d_h×d} WiO∈Rdh×d将中间结果转换为输出隐藏状态。在每个时间步，头部i上的先前KV缓存元素已初始化为：
K
i
=
X
W
i
K
,
V
i
=
X
W
i
V
K_i = XW^K_i, V_i = XW^V_i
Ki=XWiK,Vi=XWiV
然后，输入标记x被映射到每个头部的相应{查询，键，值}，先前的KV缓存相应更新：
q i = x W i Q , k i = x W i K , v i = x W i V q_i = xW^Q_i, k_i = xW^K_i, v_i = xW^V_i qi=xWiQ,ki=xWiK,vi=xWiV
K i = C a t [ K i : k i ] , V i = C a t [ V i : v i ] K_i = Cat[K_i : k_i], V_i = Cat[V_i : v_i] Ki=Cat[Ki:ki],Vi=Cat[Vi:vi]
最后，输出 y ∈ R 1 × d y ∈ R^{1×d} y∈R1×d使用注意力权重 A i ∈ R 1 × n A_i ∈ R^{1×n} Ai∈R1×n计算如下(为了简化起见，省略了缩放因子√dh)：
y = ∑ i ∈ [ 1 , h ] ( A i V i W i O ) , 其中 A i = s o f t m a x ( q i K i T ) y = ∑_{i∈[1,h]} (A_i V_i W^O_i) ,\text{其中 } A_i = softmax(q_i K_i^T) y=i∈[1,h]∑(AiViWiO),其中 Ai=softmax(qiKiT)
3.2 重新审视Top-k缓存淘汰
一组指示变量 { I i ∈ R 1 × n } \{\mathcal{I}_i ∈ R^{1×n}\} {Ii∈R1×n}表示所有头部 i ∈ [ 1 , h ] {i ∈ [1, h]} i∈[1,h]的分配预算 { B i } \{B_i\} {Bi}的淘汰决策：
I i j = { 1 如果 K i j 和 V i j 被保留 0 否则，驱逐 K i j 和 V i j \mathcal{I}_i^j=\begin{cases}1&\text{如果 }K_i^j\text{ 和 }V_i^j\text{ 被保留}\\0&\text{否则，驱逐 }K_i^j\text{ 和 }V_i^j&\end{cases} Iij={10如果 Kij 和 Vij 被保留否则，驱逐 Kij 和 Vij
其中每个元素 I i j \mathcal{I}_i^j Iij指示第 j j j个 ∈ [ 1 , n ] ∈ [1, n] ∈[1,n]的KV缓存元素在 K i , V i ∈ R n × d h K_i, V_i ∈ R^{n×d_h} Ki,Vi∈Rn×dh中是否被淘汰。因此，只有预算大小 B i B_i Bi的缓存元素被保留在头部 i ： ∑ j ∈ [ 1 , n ] I i j = B i i：∑_{j∈[1,n]} \mathcal{I}_i^j = B_i i：∑j∈[1,n]Iij=Bi，并且一个注意力层的总体预算是 B = ∑ i ∈ [ 1 , h ] B i B = ∑_{i∈[1,h]} B_i B=∑i∈[1,h]Bi。然后，我们可以获得多头自注意力机制的淘汰后输出 y ^ \hat y y^：
y ^ = ∑ i ∈ [ 1 , h ] ( A ^ i V i W i O ) \hat y = ∑_{i∈[1,h]} (\hat A_i V_i W^O_i) y^=i∈[1,h]∑(A^iViWiO)
A i ^ = s o f t m a x ( − ∞ ⊙ ( 1 − I i ) + q i K i T ) \hat {A_i}= softmax(−∞ ⊙ (1 − \mathcal{I}_i) + q_i K_i^T) Ai^=softmax(−∞⊙(1−Ii)+qiKiT)
其中⊙表示逐元素乘法。定理1通过用 A i A_i Ai表示 A ^ i \hat A_i A^i进一步简化了淘汰后的输出 y ^ \hat y y^。详细证明见附录A.3。
定理1 给定分配的预算
{
B
i
}
\{B_i\}
{Bi}，淘汰后的输出
y
^
\hat y
y^可以重写为：
y
^
=
∑
i
∈
[
1
,
h
]
A
i
⊙
I
i
∣
∣
A
i
⊙
I
i
∣
∣
1
V
i
W
i
O
\hat y = ∑_{i∈[1,h]} \frac{A_i ⊙ I_i }{||A_i ⊙ I_i||_1 }V_i W^O_i
y^=i∈[1,h]∑∣∣Ai⊙Ii∣∣1Ai⊙IiViWiO
生成质量在缓存淘汰后的下降源于注意力输出的变化。因此，我们量化淘汰损失为自注意力机制预淘汰和淘汰后输出之间的
L
1
L_1
L1距离：
L
1
Eviction Loss
=
∣
∣
y
−
y
^
∣
∣
1
L_1 \text{ Eviction Loss} = ||y − \hat y||_1
L1 Eviction Loss=∣∣y−y^∣∣1
利用矩阵的行范数，我们在定理2中推导出
L
1
L_1
L1淘汰损失的上界
ϵ
ϵ
ϵ。详细证明见附录A.4。
定理2 给定分配的预算
{
B
i
}
\{B_i\}
{Bi}，由缓存淘汰引起的
L
1
L_1
L1淘汰损失可以被限制为
ϵ
ϵ
ϵ：
L
1
Eviction Loss
≤
ϵ
=
2
h
C
−
2
C
∑
i
∈
[
1
,
h
]
∑
j
∈
[
1
,
n
]
I
i
j
A
i
j
L_1 \text{ Eviction Loss}≤ ϵ = 2hC − 2C ∑_{i∈[1,h]} ∑_{j∈[1,n]} \mathcal{I}_i^j A_i^j
L1 Eviction Loss≤ϵ=2hC−2Ci∈[1,h]∑j∈[1,n]∑IijAij
其中
C
=
M
a
x
∣
∣
V
i
W
i
O
∣
∣
∞
C = Max ||V_i W^O_i||_∞
C=Max∣∣ViWiO∣∣∞是常数，代表所有矩阵中的最大行范数。Top-k淘汰方法通常假设关键缓存元素在将来的生成过程中保持稳定，以便于缓存淘汰（2024b；2024；2024；2024a）。SOTA方法（Li等人，2024；Yang等人，2024；Zhang等人，2024a）利用观察窗口中的几个标记的查询状态来计算与过去
K
V
KV
KV缓存元素的观察注意力权重，这与Top-k选择相结合，近似识别在后续生成中关键的缓存元素。为了简化，我们假设窗口大小为1，这意味着淘汰程序仅依赖于与最后一个标记
x
x
x相关联的单个查询状态
q
i
q_i
qi进行关键缓存检测。如算法1所示，Top-k选择只保留每个头部
i
i
i中对应于
B
i
B_i
Bi个最高观察权重
A
i
j
∈
Top-k
(
A
i
,
B
i
)
A^j_i ∈ \text {Top-k}(A_i, B_i)
Aij∈Top-k(Ai,Bi)的缓存元素，而淘汰其他元素。显然，给定任何分配的预算
B
i
{B_i}
Bi，Top-k选择的淘汰决策最大化了以下方程：
{
I
i
T
o
p
}
=
arg
max
{
I
i
}
∑
i
∈
[
1
,
h
]
∑
j
∈
[
1
,
n
]
I
i
j
A
i
j
.
\begin{Bmatrix}\overset{Top}{\mathcal{I}_i}\\\end{Bmatrix}=\arg\max_{\left\{\mathcal{I}_i\right\}}\sum^{i\in[1,h]}\sum^{j\in[1,n]}\mathcal{I}_i^jA_i^j.
{IiTop}=arg{Ii}max∑i∈[1,h]∑j∈[1,n]IijAij.
因此，我们建立了定理3，它表明Top-k淘汰的原则是最小化
L
1
L_1
L1淘汰损失的上界
ϵ
ϵ
ϵ。
定理3 给定预算分配结果
{
B
i
}
\{B_i\}
{Bi}，Top-k缓存淘汰方法基本上最小化
L
1
L_1
L1淘汰损失的上界
ϵ
ϵ
ϵ：
{
I
i
T
o
p
}
=
arg
min
ϵ
{
I
i
}
.
\begin{Bmatrix}\overset{Top}{\mathcal{I}_i}\\\end{Bmatrix}=\underset{{\{\mathcal{I}_i\}}}{\operatorname*{\arg\min\epsilon}}.
{IiTop}={Ii}argminϵ.
3.3 使用自适应预算分配进行优化
定理2和3表明最小化目标 ϵ ϵ ϵ取决于淘汰决策 { I i T o p } \begin{Bmatrix}\overset{Top}{\mathcal{I}_i}\\\end{Bmatrix} {IiTop}，因此间接受到预算分配 { B i } \{B_i\} {Bi}的影响。这表明适当的分配可以进一步优化 ϵ ϵ ϵ，与简单的均匀分配相比，后者忽视了自注意力机制中每个头部的不同特征——这一差异在其他领域已被注意到（Voita等人，2019；Michel, Levy和Neubig，2019；Clark等人，2019）。为了解决这个问题，我们提出了第一个自适应预算分配算法，用于头部之间的Top-k淘汰。如算法2所述，它最初从一层中的所有头部选择B个最大的观察注意力权重。每个头部的选择频率随后通知自适应预算分配 { B i = B i ∗ } \{B_i=B^*_i\} {Bi=Bi∗}，从而优化淘汰过程中的Top-k选择结果$\begin{Bmatrix}\overset{Top}{\mathcal{I}_i}\\end{Bmatrix} $。
理论视角：自适应与均匀分配
分配结果 B i B_i Bi重塑了算法1的选择结果 { I i T o p } \begin{Bmatrix}\overset{Top}{\mathcal{I}_i}\\\end{Bmatrix} {IiTop}，影响了理论上的淘汰损失上界ϵ。在均匀预算分配 B i = B / h {B_i = B/h} Bi=B/h下，修订后的上界ϵ’为：
ϵ
′
=
2
h
C
−
2
C
∑
i
∈
[
1
,
h
]
∑
j
∈
[
1
,
n
]
I
i
′
T
o
p
A
i
j
ϵ' = 2hC − 2C \sum^{i\in[1,h]}\sum^{j\in[1,n]} \overset{Top}{\mathcal{I}_i^{'}} A_i^j
ϵ′=2hC−2C∑i∈[1,h]∑j∈[1,n]Ii′TopAij
相反，在自适应预算分配
{
B
i
=
B
∗
∗
i
}
\{B_i = B^**i\}
{Bi=B∗∗i}下，其上界
{
I
i
∗
T
o
p
}
\begin{Bmatrix}\overset{Top}{\mathcal{I}_i^*}\\\end{Bmatrix}
{Ii∗Top}给出为：
ϵ
∗
=
2
h
C
−
2
C
∑
i
∈
[
1
,
h
]
∑
j
∈
[
1
,
n
]
I
i
∗
T
o
p
A
i
j
ϵ^* = 2hC − 2C \sum^{i\in[1,h]}\sum^{j\in[1,n]} \overset{Top}{\mathcal{I}_i^{*}} A_i^j
ϵ∗=2hC−2C∑i∈[1,h]∑j∈[1,n]Ii∗TopAij
算法2：自适应预算分配
输入：总预算B，观察到的注意力权重 { A i } \{A_i\} {Ai}；
输出：分配的预算 { B i ∗ } \{B^*_i\} {Bi∗}。
-
跨头部连接 A = C a t ( { A i } , d i m = 1 ) A = Cat(\{A_i\}, dim=1) A=Cat({Ai},dim=1)。
-
创建头部指示器 I = [ 1 … 1 : … : h … h ] I = [1…1 : …: h…h] I=[1…1:…:h…h]，每个索引 i {i} i重复 n n n次。
-
识别顶部索引 T = T o p − k ( A , B ) . i n d i c e s T = Top-k(A, B).indices T=Top−k(A,B).indices。
-
选择相应的头部指示器 I ∗ = I [ T ] I^* = I[T] I∗=I[T]。
-
统计 I ∗ I^* I∗中每个i的频率以确定 { B i ∗ } \{B^*_i\} {Bi∗}。
-
返回分配的预算 { B i ∗ } \{B^*_i\} {Bi∗}。
图3：不同头部的注意力集中度变化。（在单个文档QA数据集Qasper的第一个样本上使用Mistral-instruct-v0.2。）我们聚合不同比例的顶部注意力权重，以分析不同头部i的注意力集中度。大多数头部集中度较低，只需要较小的缓存比例，例如5%，就能聚合接近1的权重，而其他分散的头部则需要更大的比例，如50%。
定理4 自适应预算分配确保其
ϵ
∗
ϵ^*
ϵ∗始终等于或低于均匀预算分配的ϵ’。
ϵ
∗
≤
ϵ
′
ϵ^* ≤ ϵ'
ϵ∗≤ϵ′
定理4支持自适应预算分配的理论优势，实现了优化的淘汰损失上界。证明见附录A.5。
实证视角：自适应与均匀分配
实证证据进一步证实，自适应预算分配利用头部间不同的注意力集中度。利用这种变异性对于优化预算效率和最小化实际场景中的淘汰损失至关重要。图3表明，集中度较低的头部需要显著较小的预算比例，就能实现接近最优的顶部权重聚合 ∑ j ∈ [ 1 , n ] I i T o p A i j \sum^{j\in[1,n]} \overset{Top}{\mathcal{I}_i }A_i^j ∑j∈[1,n]IiTopAij，与更分散的头部相比。在这种情况下，之前的均匀预算分配面临困境：它要么在集中度较低的头部上浪费过多不必要的预算，要么在分散的头部上承受大量的淘汰损失。这显著削弱了总体预算与淘汰后生成质量之间的权衡性能。相比之下，自适应预算分配算法为分散的头部分配大预算，同时控制其他稀疏头部的预算大小，有效地保持总体预算大小，并减轻生成质量的下降。图4进一步说明了自适应预算分配在所有样本中一致地降低了实际淘汰损失，强调了其在优化预算分配中的有效性。
3.4 无缝集成的实现
我们通过将自适应预算分配算法集成到两种最先进的方法，SnapKV和Pyramid中，展示了其与现有Top-k淘汰框架的无缝兼容性，创建了增强版本：Ada-SnapKV和Ada-Pyramid。SnapKV和Pyramid都使用来自最近观察窗口（通常大小为32）的标记 X w i n ∈ R w i n s i z e × d X^{win} ∈ R^{winsize×d} Xwin∈Rwinsize×d来识别和淘汰过去KV缓存中的较不重要元素。SnapKV在较大的预算场景下表现出色，而Pyramid针对预算受限条件进行了优化。这是因为Pyramid采用预设的超参数，通过金字塔形式的预算分配，偏好较浅的层，而SnapKV则均匀分配预算。因此，对于具有总体预算B的特定层，它们的淘汰算法是相同的。然而，这两种方法传统上都在一层内的所有头部之间均匀分配预算。将我们的自适应分配集成进来，如算法3所述，我们修改了这些方法，以更好地管理头部级别的预算分配。这种集成发生在每层的淘汰过程之前，我们的算法根据头部间的观察注意力权重自适应调整预算分配，如图6所示。总的来说，它们首先使用观察窗口内的查询状态计算过去缓存元素的观察注意力权重 A ˉ i \bar{A}_i Aˉi。一个最大池化层处理这些权重以保留关键信息（Li等人，2024），然后对观察窗口外的过去缓存元素进行Top-k选择。这些选定的元素，以及观察窗口内的其他元素，被保留，其余的被淘汰以减少缓存大小。此外，我们引入了一个保护性的超参数 α α α（默认为0.5），以防止对高度稀疏的头部分配过小的预算，从而增强对关键元素稳定性的假设的容错性（Li等人，2024；Zhang等人，2024b）。为了有效处理自适应分配的可变数量缓存元素，我们使用自定义CUDA内核实现了扁平化存储架构。这种方法，结合Flash Attention技术（Kwon等人，2023），确保了Ada-SnapKV和Ada-Pyramid的计算效率与传统方法一致。有关计算效率的更多细节，请参见附录A.2。
4 实验
4.1 设置
数据集
首先，我们使用LongBench（Bai等人，2023）中的16个数据集进行全面评估，这是一个长序列基准测试，涵盖了单文档问答（Koˇcisk`y等人，2018；Dasigi等人，2021）、多文档问答（Yang等人，2018；Ho等人，2020；Trivedi等人，2022）、摘要（Huang等人，2021；Zhong等人，2021；Fabbri等人，2019）、少样本学习（Joshi等人，2017；Gliwa等人，2019；Li和Roth，2002）、合成任务（Bai等人，2023）和代码生成（Guo等人，2023；Liu, Xu, 和 McAuley，2023）等多任务领域的数据集。这些数据集的平均输入长度从1,235到18,409个标记不等，需要大量的KV缓存用于推理，因此适合评估不同内存预算下的KV缓存淘汰方法。每个数据集都使用LongBench推荐的指标进行评估，每个质量分数最高为100分。附录A.6提供了数据集的详细信息。此外，我们还使用广泛使用的“大海捞针”测试来特别检查所提出的自适应预算分配对模型基本长上下文检索能力的影响。
基线
我们选择SnapKV（Li等人，2024）和Pyramid（Yang等人，2024；Zhang等人，2024a）作为主要的基线，因为它们是领先的方法，也是我们Ada-SnapKV和Ada-Pyramid方法的基础。此外，StreamingLLM（Xiao等人，2023）代表滑动窗口淘汰方法，H2O（Zhang等人，2024b）则是基于所有查询状态的早期Top-K淘汰方法的代表。
基础模型
在实验中，我们使用了两个开源的基础模型：Mistral-7B-instruct-v0.2（Jiang等人，2023）和LWM-Text-Chat-1M（Liu等人，2024a）。Mistral 7B模型具有32K的上下文长度，并已在相关研究中被采用为主要模型（Li等人，2024；Zhang等人，2024a），因为它的参数规模适中，且在长序列任务中表现出色。同时，LWM 7B模型以其1M的上下文长度成为最先进的模型，有助于在“大海捞针”测试中评估极端上下文长度。
参数
考虑到所有数据集的最小平均长度为1235，我们在四种不同的层预算大小B ∈ {128 × h, 256 × h, 512 × h, 1024 × h}下评估所有方法进行全面评估。我们遵循先前研究（Li等人，2024；Yang等人，2024；Zhang等人，2024a）的常见做法，在每层的预填充阶段后进行缓存淘汰方法的比较。在所有实验中，自适应预算分配中的超参数α设置为0.5。Ada-SnapKV和Ada-Pyramid以及SnapKV和Pyramid都使用与（Li等人，2024）中描述相同的配置设置，确保与观察窗口大小为32和最大池化核大小为7的比较性。StreamingLLM和H2O的参数符合（Zhang等人，2024b；Xiao等人，2023）中的默认设置。所有实验都在单个A100-80G GPU上进行。更多细节，请参阅补充材料中的代码。
4.2 在16个数据集上的评估
基于Mistral模型的每个数据集的详细结果见表1，基于LWM模型的其他结果由于空间限制放在附录A.7中。我们以预算B = 128h为例，展示表1中的改进。通过集成自适应预算分配，Ada-SnapKV在16个数据集中的15个提高了质量分数，与原始SnapKV相比，平均分数从35.09提高到36.71。同样，Ada-Pyramid在16个数据集中的14个超过了原始Pyramid，将平均分数从35.58提高到36.81。图5总结了基于Mistral和LWM的所有方法在16个数据集上的平均分数。值得注意的是，作为滑动窗口淘汰方法代表的StreamingLLM，由于在识别重要缓存元素方面的效率低下，产生了显著较低的质量输出。采用Top-k选择和观察窗口的SnapKV和Pyramid表现出紧密匹配的性能，超过了先前的H2O。此外，我们的Ada-SnapKV和Ada-Pyramid方法在不同预算下提高了生成质量。这两种Ada增强方法交替领先并超过了基础版本，特别是在小预算下。这种一致的改进强调了自适应预算分配的必要性和有效性，得到了理论和实证证据的支持。
4.3 在“大海捞针”测试中的评估
如图6所示，我们使用“大海捞针”测试来展示自适应预算分配通过增强长上下文检索能力。与之前的实验一致，所有配置都保持观察窗口大小为32和池化核大小为7，最大推理长度限制为在A100-80G上全缓存情况下的37K。在B = 128h的缓存预算下，Ada-SnapKV和SnapKV将最大长度扩展到429K，而Ada-Pyramid和Pyramid扩展到365K。显著的是，Ada-SnapKV和Ada-Pyramid与之前的SnapKV和Pyramid相比，提高了长文本检索能力。特别是，Ada-SnapKV和Ada-Pyramid在原始的37K长度内实现了近乎无损的检索，这是标准SnapKV和Pyramid无法复制的。在平均分数方面，Ada-SnapKV从94.84提高到95.99，而Ada-Pyramid从96.02提高到96.69。在“大海捞针”测试中，不同长度的内存和时间计算效率的额外评估可在附录A.2中找到，证明了自适应分配保持了与原始版本一致的计算效率。
5 结论
在本研究中，我们重新审视了用于高效长序列推理的主流缓存淘汰方法，揭示了它们主要最小化预淘汰和淘汰后输出之间的L1距离的上界。基于这一洞见，我们提出了第一个自适应预算分配算法，用于优化KV缓存淘汰，理论上降低了与以往方法相比的上界。我们的实证研究还表明，这种自适应算法利用了多头自注意力机制中注意力集中度的不同程度。开发的两种新型自适应淘汰方法，Ada-SnapKV和Ada-Pyramid，在综合评估中表现出显著的改进。我们的工作突出了通过理论框架和自适应预算分配实施，特别是针对LLMs中不同注意力头的独特特性，推进缓存淘汰方法的实质性潜力。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
