斯坦福大学研究团队如何让AI"聪明地遗忘"

这项由斯坦福大学领导的研究发表于2026年6月，论文编号为arXiv:2604.20920，有兴趣深入了解的读者可以通过该编号查询完整论文。

每当你让ChatGPT或类似的AI助手阅读一篇很长的文章并回答问题时，这个AI其实正在做一件极其耗费精力的事情——它必须把文章里每一个字都"记在脑子里"，哪怕其中大部分内容跟你的问题毫无关系。这就好比你问一个朋友"这本500页的小说里，主角最后去了哪座城市"，但你的朋友为了回答这个问题，却必须把整本书从头到尾背诵一遍，然后才能开口说话。这显然既浪费时间，又耗费巨大的精力。

现实中的AI大模型面临着完全相同的困境。随着文本越来越长，AI需要处理的信息量呈爆炸式增长，计算成本和内存占用以平方级别飞速上涨。处理一篇一万字的文章，其难度可能是处理一千字文章的一百倍，而非十倍。这个问题在业界被称为"长文本推理瓶颈"，是制约当前大语言模型能力的核心障碍之一。

为了解决这个问题，斯坦福大学的研究团队提出了一套名为"简化稀疏注意力"（Simplified Sparse Attention，简称SSA）的全新方法。这套方法的核心理念极其优雅：教会AI像聪明的读书人一样工作——先快速浏览全文做笔记，等真正需要某段内容时，再回头精读那一部分。更令人印象深刻的是，这套方法不需要对AI的基本结构做任何改动，就像给一个人传授新的阅读技巧，而不是换一个新的大脑。

一、AI为什么需要学会"做笔记再精读"

要理解SSA究竟解决了什么问题，可以先想想一个图书馆管理员的工作。假设有人来问："请帮我找出所有关于法国大革命的资料。"一个低效的管理员会把图书馆里所有书架上的每一本书都翻一遍；而一个聪明的管理员会先查询索引卡，找到相关书目，再精确地取出那几本书。前者虽然不会遗漏任何信息，但耗时耗力；后者虽然依赖索引的准确性，但效率极高。

当前AI处理长文本的方式，更接近那个低效管理员——无论你的问题多么具体，它都必须把所有内容平等地过一遍。这种方式在学术上被称为"全注意力机制"（Full Attention），其代价随文本长度急剧增加。于是，研究者们开始探索各种"稀疏注意力"方案，让AI只关注最相关的那部分内容。

然而，已有的稀疏注意力方案各有各的麻烦。有些方案只在推理时应用，完全没有经过专门训练，效果打折扣；有些方案虽然效果好，却需要在AI的基本结构上动大手术，比如增加新的神经网络模块，或者引入专用的辅助索引器，让整个系统变得复杂难以维护。斯坦福的研究团队因此提出了一个问题：有没有可能设计一种方法，既让AI真正学会这种"聪明阅读"策略，又完全不改动AI的基本架构？

答案就是SSA，而它的核心工具是一种叫做"要点标记"（Gist Tokens）的特殊符号。

二、"要点标记"究竟是什么

要理解要点标记，可以借助一个学生做笔记的场景。假设你正在读一本历史教材，每读完一章，你会在旁边的便利贴上写下几行摘要："本章核心：法国大革命起因，三个关键事件，时间1789年"。这张便利贴就是一个"要点标记"——它不包含原文的全部细节，但压缩了这一章最重要的信息。等你日后复习时，先看便利贴，就能大致判断"这章的内容和我现在的问题相关吗"，如果相关，再翻回原文精读；如果不相关，直接跳过。

SSA的做法与此如出一辙。研究团队把输入给AI的长文本切分成若干小段，每段长度固定（比如每16个词语为一段）。在每一小段的末尾，插入一个特殊的"要点标记"符号。在训练阶段，这个要点标记会通过特殊的注意力规则被"强迫"学习压缩它前面那一段内容的精华——因为训练规则规定，后续的文字不能直接回头看某段的原始内容，只能通过该段的要点标记来获取信息。这就像强制规定学生只能通过便利贴来复习，迫使他们把笔记做得精炼准确。

通过这种"被迫压缩"的训练，每个要点标记的内容就变得极具代表性。它不是简单的平均，而是经过AI自己学习后形成的语义浓缩，能够捕捉到原始段落中最关键的含义。训练所用的目标和普通语言模型完全一样——预测下一个词，使用标准的交叉熵损失函数，无需任何额外的辅助损失或新增参数。

三、推理时的"选择性展开"：聪明阅读的精髓

有了这些经过训练的要点标记，SSA在实际回答问题时的工作流程就变得既简单又高效。这个过程可以用"先看目录，再翻正文"来描述。

当用户提出一个问题时，AI首先拿着这个问题的"向量表示"（可以理解为问题的数字版本），去和文章中所有段落的要点标记逐一比对，计算相关程度的分数。这个比对过程极其快速，因为要点标记的数量远远少于原始文本的词语数量——假设原文有1000个词，分成100段，那么要点标记只有100个，而不是1000个。比对完成后，AI选出得分最高的若干段落（即Top-k选择），然后把这些被选中段落的原始内容重新"展开"，加入到AI的实际注意力范围中。对于那些没被选中的段落，AI完全不去处理它们，就像那些没被取下的书，直接忽略。

这个"展开"的过程在论文中被称为"选择性展开"（Selective Unfolding），是整个SSA方法最核心的创新。它巧妙地将"用什么来选"（要点标记，少量且精炼）和"选完后读什么"（原始内容，丰富且详细）分开处理，两全其美。

在具体的注意力计算上，被选中段落的内容会和该段的要点标记一起进入AI的视野。研究团队发现，把选中段落的要点标记和原始词语一同保留，比只保留原始词语效果更好——要点标记作为对整段内容的压缩总结，提供了一种"全局视角"，弥补了原始词语只能逐词表达的局限。而那些没被选中的段落，连它们的要点标记也会被排除在外，这样可以把有限的"注意力预算"完全集中在最相关的内容上。

值得一提的是，由于问题只需与要点标记比对，而不需要与完整的KV缓存（AI记忆内容的存储形式）比对，SSA还解决了另一个工程难题：许多现有的稀疏注意力方法虽然最终只处理一部分内容，但为了"找到"应该处理哪些内容，仍然需要把全部内容读入内存，这导致内存带宽成为瓶颈。SSA通过要点标记直接定位相关段落，彻底避免了这一问题。

四、如何让AI真正学会这套技巧：两阶段训练

SSA的训练过程分为两个阶段，可以类比为先"教会学生做笔记"，再"让学生在考试条件下练习"。

第一阶段叫做"持续预训练"，是必须经历的步骤。在这个阶段，研究团队用大量文本对AI进行继续训练，但特别设计了注意力掩码（一种控制AI"看"哪些内容的规则）：文章被分段加入要点标记，每段内的原始词语可以互相"看见"，也可以看到之前所有段的要点标记，但不能直接看到更早那些段的原始词语。这种限制迫使要点标记充当信息中转站，必须把过去段落的精华保留下来，以供未来词语参考。训练目标依然是普通的下一词预测，整个过程无需特殊的CUDA核心，完全可并行化。

完成第一阶段后，AI就已经具备了在推理时进行选择性展开的能力，不需要任何额外训练就能直接使用SSA。但为了让AI更好地适应"在考试时只有部分内容可用"的场景，研究团队还设计了第二阶段——"选择性微调"，属于可选步骤。

在微调阶段，训练数据同样被分为压缩上下文（前半部分）和生成上下文（后半部分）。压缩部分的处理方式和第一阶段完全相同；但在生成部分，每个位置的注意力掩码不再是固定的，而是动态地根据当前词语的向量，计算与各段要点标记的相关分数，然后只允许注意到得分最高的几段内容的原始词语。这让AI在训练时就"亲身体验"了选择性展开的过程，而不仅仅是在推理时临时应用一个没训练过的技巧。不过微调阶段需要位置依赖的稀疏掩码，实现上略复杂，对于不想增加实现难度的开发者，跳过这一步也能得到相当好的效果。

五、层层嵌套的"笔记的笔记"：H-SSA的诞生

SSA已经很强大了，但研究团队发现，这套"做笔记"的逻辑其实可以无限嵌套。既然可以对原始段落做笔记，为什么不能对笔记再做笔记呢？

以图书馆的比喻来延伸：每本书有章节摘要，图书馆有图书分类目录，城市图书馆系统有总目录。当你要找一条特定信息时，你先看总目录确定大类，再看分类目录锁定具体书目，最后取出那本书翻到相关章节——这正是分层检索的精髓。

H-SSA（层级简化稀疏注意力）就是这个原理的工程实现。在单层SSA的基础上，研究团队引入了"元要点标记"（Meta-Gist Tokens）——每隔若干个普通要点标记，就插入一个元要点标记，负责压缩这一组普通要点标记的内容。类比来说，如果普通要点标记是每章的便利贴摘要，那元要点标记就是每本书的封底简介。

在推理时，H-SSA采用从粗到细的两步选择：先用当前问题与数量很少的元要点标记比对，找到最相关的几个大组；再在这些大组内，用问题与普通要点标记比对，锁定最相关的具体段落；最后才展开这些段落的原始内容。这种层级路由将每步的搜索范围大幅缩小，使得总体的路由计算量从线性（正比于段落总数）降低到对数级别（正比于段落总数取对数）。

从计算复杂度来讲，这意味着什么？假设文章有一百万个词，单层SSA在路由时仍需处理约六万多个段落标记，而H-SSA可能只需处理几百个元要点标记，再处理几十个普通要点标记，总计不过几百次比较。这是质的飞跃，让超长文本的处理成为现实。理论上，H-SSA的层级结构可以无限延伸，支持任意长度的文本，只需随着文本增长添加新的层级即可。

六、针对现代GPU的专属内核设计

技术上的优雅还需要工程实现的配合。研究团队专门为SSA和H-SSA设计了高效的计算内核，让理论上的效率优势在真实硬件上得以实现。

在文本处理阶段（预填充阶段），问题在于SSA的稀疏注意力模式在结构上比较特殊：要点标记这类"全局"信息出现在序列各处，导致标准的分块稀疏计算无法直接跳过空白块。研究团队用了一个聪明的技巧——"键列置换"：在计算注意力之前，先把所有频繁出现的全局信息（要点标记和注意力锚点）挪到序列的最前面，把局部信息移到后面。由于注意力计算对键的排列顺序不敏感（softmax会做归一化），这个置换不影响计算结果，却让稀疏结构变得规整：全局信息形成一个密集的薄块，局部信息形成一个对角带状结构，空白部分可以被高效跳过。

在回答生成阶段（解码阶段），研究团队设计了一套三步流水线内核。第一步是"并行压缩"：一个内核快速扫描需要保留的键值对索引，利用线程束级别的原子操作，把需要处理的内容地址紧密排列，避免稀疏数据带来的碎片化。第二步是"分割K部分注意力计算"：另一个内核在压缩后的紧密列表上并行计算注意力，多个线程块同时处理不同区段，每个线程块内部会把同一KV组中多个查询头共享的键值块只加载一次，最大化内存复用率，这与NSA（原生稀疏注意力）的分组设计思路类似。第三步是"合并"：将各分块的部分注意力结果通过标准的对数求和指数技巧合并为最终输出。整个解码过程只读取被选中的若干段内容的键值对，加上少量元数据，完全不触碰完整的键值缓存，从根本上消除了内存带宽瓶颈。

七、实验结果：不只是"省力"，有时还更准

研究团队在两个模型家族上验证了SSA的效果：规模较大的Qwen2-7B-Instruct（一个经过指令微调的70亿参数模型）和规模较小的Llama3.2-1B（一个10亿参数的基础模型）。测试涵盖两大场景：长文本理解（使用LongBench基准测试）和检索增强生成（使用多个多文档问答数据集）。

在长文本理解场景下，SSA与其他方法在相同压缩比下进行比较。以8倍压缩为例，这意味着AI实际处理的内容只有原文的八分之一。SSA在持续预训练后平均得分46.20，而同类的ActivationBeacon方法只有42.52，UniGist方法为43.40，差距在2.8到3.7分之间。即使将压缩比提高到16倍和32倍，SSA依然保持领先优势：32倍压缩下SSA得44.07，而对手方法普遍只有38分左右。经过选择性微调后，SSA在某些具体任务上甚至超越了不做任何压缩的完整注意力模型——例如在MF-en任务（一种多字段英文问答）上，SSA得54.24，而全注意力基线只有50.33。这说明选择性展开不只是弥补压缩损失，有时还能主动过滤掉干扰信息，反而提升准确率。

与只在推理时应用稀疏策略（不做任何额外训练）的方法相比，SSA的优势更加显著。H2O方法在8倍"压缩"下得44.20，StreamingLLM得38.81，Quest只有17.32，而SSA（仅经过预训练）就已经达到46.20，且这些对比是在相同的计算预算下进行的。

检索增强生成场景则展现出SSA最令人惊讶的优势。在这个场景下，用户提供多篇文档（其中大部分是无关文档）和一个问题，AI需要找到正确答案。仅经过持续预训练的SSA在8倍压缩下就达到了33.68的平均分，不仅远超KVLink（21.58）和UniGist（22.53）超过11分，甚至超过了原始未压缩模型（27.99）和经过完整注意力持续预训练的模型（27.14）。这个结果表明，当大多数文档都是无关干扰时，强迫AI"专注于少数段落"反而有助于它忽略噪声，找到真正相关的信息。经过微调后，SSA在8倍压缩下达到53.39，比KVLink高出近12分，比UniGist高出超过8分，与全注意力微调模型（57.76）的差距也只有4.4分。

H-SSA在高压缩比下表现出更强的竞争力。在32倍压缩微调后，H-SSA得44.94，而单层SSA只有43.35。在16倍压缩的检索增强场景下，H-SSA达到50.72，优于SSA的49.29。这与直觉一致：层级结构在压缩率越高时，粗粒度到细粒度的分级筛选越能有效分配有限的注意力预算。

研究团队还专门测试了一种叫做"KV缓存复用"的场景：同一批文档被不同的问题反复查询，通过提前缓存文档的键值对来避免重复计算。SSA在此场景下同样出色，8倍压缩下达到48.07，16倍压缩下H-SSA达到46.34，均大幅领先KVLink和UniGist。

在效率测试方面，研究团队在单张NVIDIA H100 GPU上测量了SSA在16倍压缩下的端到端延迟。结果非常直观：全注意力解码（Flash-Decoding）的每词输出时间随上下文长度线性增长，从8K词时的21.9毫秒一路爬升到44K词时的76.4毫秒；而SSA的每词输出时间几乎保持平坦，始终在21到23毫秒之间，H-SSA也稳定在25毫秒附近。在44K词时，SSA比全注意力快3.37倍，H-SSA快3.05倍。在预填充阶段，SSA从约33K词开始就比全注意力FlashAttention更快，44K词时达到0.90倍延迟。单独看注意力算子（排除其他模型组件的影响），H-SSA的预填充在200K词时比稠密FlashAttention快8.3倍，解码快1.58倍。

八、细节决定成败：关于设计选择的实验对比

研究团队进行了一系列精细的消融实验，验证SSA中各个设计选择的必要性。

关于选中段落后的注意力上下文组成，研究团队比较了三种方案：只包含选中段落的原始词语（不含要点标记）、包含所有段落的要点标记加上选中段落的原始词语、只包含选中段落的要点标记加上其原始词语。结果显示第三种方案效果最好（平均53.39），第二种方案次之（53.27），纯原始词语最差（52.76）。这说明保留选中段落的要点标记有助于AI同时拥有"压缩的全局视角"和"详细的局部信息"，而把所有不相关段落的要点标记也纳入会引入噪声，稍微降低性能。

关于Top-k（固定选取前k个段落）与Top-p（累积概率超过阈值p时停止选取）的比较，研究表明自适应Top-k始终优于各种阈值的Top-p方案，且稳定性更好。Top-p方案的问题在于，不同问题的相关性分数分布差异很大，一个固定的概率阈值在某些问题上会选太多段落，在另一些问题上又选太少，导致效果波动明显。Top-k通过直接控制token预算，提供了更一致的压缩率保证。

研究团队还测试了一项极端验证——"密钥检索"任务：在超长干扰文本中隐藏一个简短密码，测试AI能否找到它。SSA和H-SSA在两种模型上均达到100%的完美准确率，覆盖从5千词到5万词（Qwen2，超出训练长度2.5倍）以及从4千词到4万词（Llama3.2，超出训练长度10倍）的全部范围，且密钥位于任意位置时均无遗漏。这表明SSA的选择性展开机制在极端检索任务下非常可靠，并且具备良好的长度泛化能力，训练时没见过这么长的文章，推理时依然能够正确处理。

说到底，SSA这项研究的价值在于它把一个看似需要大量架构改造的复杂问题，用一个优雅且简单的训练技巧给解决了。它告诉我们，让AI"学会聪明阅读"并不需要给它装新的零件，只需要用合适的方式训练它，它自己就会学会做摘要、学会判断相关性、学会在需要细节时精确回头查阅。这对未来的AI应用具有实际意义——无论是阅读超长法律文书、分析大型代码库、还是在海量文档中回答专业问题，SSA都可以在几乎不改变现有系统架构的前提下，让AI变得既快又准。当然，该研究也存在一定局限：训练需要额外的预训练步骤，对计算资源有一定要求；选择性展开的chunk大小等超参数需要根据具体任务调整；在某些需要高度连续上下文的任务上，强制分段可能会割裂重要的跨段依赖。但这些都是未来可以继续优化的方向，并不影响这项工作所展示的核心可行性。

感兴趣的读者可以通过arXiv编号2604.20920查阅完整论文，代码也已在GitHub公开，检索"simplified-sparse-attention"即可找到。

Q&A

Q1：SSA和普通的稀疏注意力方法有什么区别？

A：普通的稀疏注意力方法通常直接在推理时决定"看哪些内容"，没有经过专门训练，相当于用一个未受过训练的人来做文档筛选。SSA则通过持续预训练，让AI真正学会把每段内容压缩进"要点标记"，再用这些标记来精确筛选相关段落。更重要的是，SSA不需要修改模型架构，只需在标准训练流程中调整注意力掩码即可实现，工程实现门槛低得多。

Q2：H-SSA层级结构的层数越多越好吗？

A：不一定。层数越多，每步的路由计算成本越低，理论复杂度越接近对数级别，但每新增一层都会让压缩信息经历更多次的"二次压缩"，可能损失更多细节。论文中测试的两层H-SSA在16倍和32倍压缩场景下效果优于单层SSA，但具体层数需要在压缩率、信息保真度和计算效率之间权衡选择，并非无限叠加层数就最好。

Q3：SSA在检索增强生成场景下为何能超越不压缩的全注意力模型？

A：这是因为检索增强生成通常会提供多篇文档，但只有少数文档真正与问题相关，其余都是干扰信息。全注意力模型会平等对待所有文档，注意力被大量无关内容分散，反而难以聚焦在正确答案上。SSA的选择性展开机制通过要点标记打分，主动过滤掉低相关性的文档段落，让AI的注意力高度集中在真正相关的内容上，相当于自动去除了噪声，因此在这类场景下效果反而更好。