利用 CItruS 解决长上下文大语言模型中的信息忽视问题

像 Llama 2 和 Mistral 这样的大语言模型 (LLM) 彻底改变了我们与文本交互的方式。然而，它们存在一个显著的局限性: 上下文窗口。虽然模型在处理更长序列方面越来越出色，但在处理整本书或海量法律文档时，计算成本依然高昂且占用大量内存。

为了解决这个问题，研究人员开发了“状态驱逐” (State Eviction) 方法。随着文本变长，这些技术通过丢弃“不重要”的信息来压缩模型的内存 (即键值缓存/Key-Value Cache) 。这听起来像是一个完美的解决方案: 在处理无限文本的同时保持低内存占用。

但这其中存在一个陷阱。目前的方法是根据什么有助于模型预测下一个单词 (流畅性) 来决定删除哪些内容。它们并没有考虑回答特定用户查询实际需要哪些信息。最近一篇论文的作者将这种现象称为信息忽视 (Information Neglect) 。

在这篇文章中，我们将深入探讨论文 “CItruS: Chunked Instruction-aware State Eviction” (CItruS: 分块指令感知状态驱逐) 。 我们将探索为何当前的压缩方法在下游任务中表现不佳，CItruS 如何通过分离“阅读”与“求解”过程来解决这一问题，以及它如何在无需任何模型再训练的情况下实现最先进的检索性能。

问题所在: 流畅性与实用性

要理解 CItruS 的核心创新，我们首先需要了解 LLM 通常如何处理内存，以及当前的效率优化方法错在哪里。

键值缓存与驱逐

当 Transformer 模型生成文本时，它会将先前 Token 的表示存储在键值 (KV) 缓存中。这可以防止模型在生成每个新词时重新计算整个历史记录。然而，随着输入长度的增加，该缓存呈线性增长，最终耗尽所有可用的 GPU 内存。

“状态驱逐”是一种防止这种内存溢出的策略。它通过监控注意力权重 (本质上是模型在多大程度上“关注”过去的特定 Token) 来工作。如果缓存中的某个 Token 没有受到太多关注，驱逐策略就会认为它不重要并将其删除。

现有的方法如 StreamingLLM、H2O 和 TOVA 都采用了这种方法。它们有效地保持了模型的流畅性 (维持低困惑度) ，即使在数百万个 Token 上也是如此。

定义信息忽视

CItruS 背后的研究人员发现了这些方法的一个关键缺陷。他们认为流畅性不等于实用性 。

当 LLM 在阅读过程 (语言建模) 中决定哪些 Token 是“重要”的时，它会优先考虑那些有助于维持语法和局部连贯性的 Token。然而，解决下游任务所需的信息——比如回答关于第 3 章细节的具体问题——在模型仅仅试图预测第 4 章的下一个单词时，可能看起来并不“重要”。

因此，标准的驱逐方法往往会丢弃用户在“干草堆”中寻找的确切“针”。作者将此称为信息忽视 (Information Neglect) 。

为了证明这一点，作者进行了一项实验，比较了模型在仅仅阅读上下文时的注意力模式与试图遵循指令时的注意力模式。

图 1: Mistral 7B Instruct 模型应用于 Qasper 数据集第 16 层的注意力分布样本。注意力分布是根据文档上下文和指令文本对键值缓存计算得出的。X 轴代表键值缓存内的不同位置，而 y 轴代表注意力权重。位置按来自上下文的注意力权重降序重新排列，为清晰起见省略了低注意力权重的位置。

如图 图 1 所示，注意力分布截然不同。紫线显示了源自文档上下文的注意力 (语言建模) ，而粉红线显示了源自指令的注意力 (任务求解) 。上下文关注一组 Token，但指令关注的是完全不同的一组。如果我们仅仅根据紫线来驱逐状态，就会丢失粉红线所代表的信息。

量化差距

研究人员通过计算 top-k 重要状态的“交集”将这一观察结果形式化。他们选取了一篇长文档，并检查了哪些隐藏状态被文档上下文选为重要，而哪些被指令文本选为重要。

图 2: 我们要应用交集计算来探索状态驱逐模型中信息忽视问题的实验示意图。

图 2 展示了实验设置: 两条并行路径处理文档。一条路径根据上下文注意力选择状态，另一条根据指令注意力进行选择。然后他们测量重叠部分。

图 3: 使用 Mistral 7B Instruct 进行实验，指令文本与文档上下文选出的 top-k 隐藏状态之间的差异，k 设置为 20。上下文-指令交集表示长文档中某段上下文的注意力分布与指令文本对键值缓存选出的 top-k 隐藏状态之间的重叠。

图 3 中的结果非常鲜明。在模型的中间层 (第 10-25 层) ，交集比率降至 0.2 以下。这意味着指令所需的超过 80% 的信息被标准的基于上下文的驱逐方法丢弃了。 这是实用性的巨大损失，也解释了为什么压缩模型经常在阅读理解任务中失败。

解决方案: CItruS

为了解决信息忽视问题，作者提出了 CItruS (Chunked Instruction-aware State Eviction，分块指令感知状态驱逐) 。 CItruS 的核心理念是解耦两个认知过程: 语言建模和任务求解 。

CItruS 并没有试图强迫一个缓存同时服务于这两个目的，而是承认“阅读”过程 (编码文档) 和“思考”过程 (回答问题) 具有不同的注意力偏好。

图 4: 我们要提出的用于特定任务长序列建模的不同子过程示意图。每个过程充当不同的角色。

如 图 4 所示，系统被分开了。语言建模轨道专注于保持流畅性和编码文档结构。任务求解轨道则使用特定的指令 (问题) 来搜索信息。

方法详解

CItruS 引入了两个主要的技术组件: 标准分块状态驱逐 (CSE) 和指令感知缓存 。

1. 标准分块状态驱逐 (CSE)

逐个处理 Token 效率低下。CItruS 以较大的块 (chunks) (例如，一次 256 或 1024 个 Token) 来处理文档。

对于每个块，模型计算当前在缓存中的状态的重要性。重要性分数源自当前块中所有 Token 的平均注意力分数。状态 \(c\) 相对于块 \(s\) 的重要性公式为:

重要性分数公式

这里，\(Q_t\) 是当前块中 Token 的查询向量，\(K_c\) 是存储状态的键向量。这个公式本质上是在问: “平均而言，当前块中的 Token 有多关注这个特定的缓存状态？”

基于这个分数，模型选择 top-\(k\) 个状态进行保留:

选择函数公式

并更新缓存 \(\hat{C}\):

缓存更新公式

这创建了一个基准的高效模型 (标准 CSE) ，其行为类似于之前的状态驱逐方法，专注于流畅性。

2. 指令感知驱逐

这是 CItruS 与现状分道扬镳的地方。为了解决信息忽视问题，CItruS 将指令文本 (用户的问题) 直接纳入驱逐过程。作者提出了两种架构设计，如 图 5 所示。

图 5: 我们提出的标准 CSE 和 CItruS 的不同缓存设计示意图。

图 5(a) 展示了上述的标准 CSE。在决定保留块 1 的哪些内容时，它只考虑上下文 (块 2) 。

图 5(b): 独立缓存 (Individual Cache) 。 在这个设计中，模型维护两个独立的缓存:

  1. 缓存 \(C\) : 由块注意力管理的标准语言建模缓存。
  2. 指令缓存 \(C^I\) : 专门用于任务的缓存。

当处理一个块时，模型使用指令文本作为查询来计算一组单独的注意力分数。然后，它选择那些专门与指令相关的 top 状态并将其存储在 \(C^I\) 中。

指令选择公式

指令缓存更新公式

这确保了即使某条信息对文档的一般流程来说看似无关紧要 (并从 \(C\) 中被丢弃) ，如果问题问到了它，它也会被保留在 \(C^I\) 中。

图 5(c): 共享缓存 (Shared Cache) 。 维护两个缓存会增加内存使用量。作者假设可能存在足够的重叠部分来合并它们。在共享缓存设计中，只有一个缓存。但是，决定保留什么的决策受到了指令的影响。Top-\(k\) 状态是基于指令文本的注意力来选择的。这些状态随后同时用于任务求解和语言建模。

令人惊讶的是，正如我们将在结果中看到的，共享缓存方法效果非常好，这表明“交集”状态 (对上下文和指令都有用) 通常足以维持流畅性。

实验与结果

作者在三个主要类别上评估了 CItruS: 长文档阅读理解、知识检索和语言建模流畅性。他们将其与 StreamingLLM、H2O 和 TOVA 等强基准进行了比较。

1. 知识检索 (大海捞针)

对长上下文模型的终极测试是“大海捞针” (Needle-in-a-Haystack) 或“密钥检索”任务。模型能否在海量文档中找到隐藏的特定随机数字？

图 6 展示了不同上下文长度 (对数刻度) 下的密钥检索准确率。

图 6: 使用 Llama 2 7B Chat、Mistral 7B Instruct 和 Llama 2 13B Chat 进行密钥检索任务的结果。

  • 蓝线 (标准) : 代表标准状态驱逐。随着文档变长 (在 x 轴向右移动) ，性能崩溃。当文档非常长时，准确率接近于零。
  • 红线 (CItruS - 独立 & 共享) : 性能保持在完美的 100% 。

无论是使用 Llama 2 还是 Mistral，CItruS 都能完美地检索信息，即使序列长度呈指数增长。这从经验上证明了指令感知机制成功防止了信息忽视。

作者还使用 ROUGE 分数 (测量文本重叠) 在“大海捞针”任务上进行了测试。

表 2: “大海捞针”任务的结果。最佳结果加粗显示。R-1、R-2 和 R-L 分别代表 ROUGE-1、ROUGE-2 和 ROUGE-L。

表 2 证实了 CItruS 的优势。无论使用 768 还是 1024 的缓存预算，CItruS (独立和共享变体) 都显著优于标准 CSE。例如，对于 Mistral 7B，ROUGE-1 分数从大约 15.17 (标准) 跃升至 63.47 (共享缓存) 。

2. 阅读理解 (下游任务)

合成检索是一回事，但真实问题呢？作者在 Qasper、HotpotQA 和 TriviaQA 等数据集上进行了测试。

表 5: 六个不同长序列任务的平均结果。结果按不同源文本长度分组呈现。最佳结果加粗显示。

表 5 显示了汇总结果。“平均排名” (Avg Rank，8 为最好) 显示 CItruS 始终名列前茅。

  • 标准 CSE 的表现与 H2O 和 StreamingLLM 等基准相当。
  • CItruS (共享缓存) 始终获得最高排名 (例如，对于 Llama 2 13B 在 4k-8k 长度上为 7.33/8) 。

这表明 CItruS 不仅仅是在寻找简单的密钥；它保留了回答问题所需的复杂语义信息。

3. 语言建模流畅性

根据“指令”而不是“上下文”来更改缓存的一个主要担忧是，模型可能会失去生成流畅英语的能力 (困惑度) 。如果我们过度优化答案，可能会破坏语法。

图 7: Llama 2 7B chat 和 Mistral 7B Instruct 模型的语言建模结果。折线图采用 4096 的窗口大小进行了平滑处理以提高清晰度。

图 7 绘制了随 Token 长度变化的困惑度 (越低越好) 。

  • 蓝色 (标准 CSE) 和 橙色 (Streaming LLM) 代表针对流畅性优化的基准。
  • 绿色 (共享缓存 CSE) 代表 CItruS。

绿线几乎与基准线完全重合。这是一个重大发现: 使用指令来选择缓存状态不会降低模型的一般语言建模能力。 事实证明，有助于回答问题的状态也足以维持文档处理的语法和流畅度。

分析: 位置偏差

“针”的位置重要吗？先前的研究表明，LLM 患有“迷失在中间” (Lost in the Middle) 综合症。

图 8: 使用 Mistral 7B Instruct 在“大海捞针”任务上使用共享缓存 CItruS (\\(l_s=1024, k=1024\\)) 的位置相关结果。X 轴代表插入“针”的位置，而 y 轴代表文档的长度。网格的颜色代表 ROUGE-1 分数。

图 8 作为性能热力图。x 轴是针的位置，y 轴是文档长度。深蓝色表示性能更好。 虽然 CItruS 表现非常好 (主要是深蓝色) ，但在中间列有轻微的褪色。这表明，虽然 CItruS 极大地改善了检索，但 LLM 倾向于关注文档开头和结尾的固有偏差在一定程度上仍然存在。然而，与在这些长度下完全变白 (零检索) 的基准相比，这是一个巨大的进步。

为什么 CItruS 很重要

CItruS 对使用 LLM 的学生和研究人员具有重大意义:

  1. 无需训练: CItruS 是一种推理时技术。你不需要微调 Llama 或 Mistral。你只需在生成过程中改变 KV 缓存的管理方式。这使得它易于获取且实施成本低廉。
  2. 内存效率: 通过证明共享缓存有效，CItruS 表明我们不需要加倍内存需求就能获得良好的检索效果。我们只需要更聪明地决定保留什么。
  3. 解决信息忽视: 该论文提供了一个清晰的理论框架 (信息忽视) 和一个实用的解决方案。它强调了神经网络中的“重要性”是相对于目标 (流畅性 vs. 任务) 而言的，而不是 Token 的绝对属性。

结论

“CItruS”论文解决了长上下文建模的关键瓶颈: 内存效率与信息保留之间的权衡。通过识别“信息忽视”——即标准模型倾向于为了局部流畅性而丢弃任务相关细节——作者提出了一种新颖的、感知指令的驱逐策略。

通过使用分块状态驱逐和指令感知缓存 , CItruS 使得标准的开源模型能够处理包含多达一百万个 Token 的文档，并以近乎完美的准确率检索特定细节。它架起了 LLM “阅读”思维与“求解”思维之间的桥梁，确保当模型阅读一本书时，它能准确记住你要求它寻找的内容。