会议背景与演讲主题引入

本次伦敦机器学习研讨会由 Evolution AI 赞助,特邀 Meta AI 的研究科学家兼经理 Yuan Dong Tan 进行主题演讲。演讲的核心议题聚焦于大语言模型(LLM)的高效推理,特别是如何在支持长上下文(Long Context)的同时提升模型性能。演讲者指出,当前大模型已广泛应用于对话 AI、内容生成、智能体交互、推理规划及摘要生成等多个领域,其背后的核心驱动力是 Transformer 架构中的注意力机制(Attention Mechanism)。该机制通过查询(Query)与键(Key)的交互,决定哪些之前的键最重要,并整合其贡献以形成下一个 token 的表示。然而,这种机制也带来了显著的计算瓶颈,因此如何在不重新训练模型的前提下,通过优化推理过程来降低内存成本、减少延迟并支持更长的上下文窗口,成为本次演讲的主要探讨方向。

长上下文与推理效率的双重挑战

当前大模型面临两大核心挑战:自回归生成的低效性与长上下文支持的局限性。首先,自回归生成是一个逐 token 生成的过程,这不仅速度慢,而且随着生成内容的增加,KV Cache(键值缓存)会迅速膨胀,占用大量内存,限制了模型在计算资源受限设备上的应用。其次,关于长上下文问题,大多数模型在预训练时具有固定的上下文窗口大小(例如 Llama v2 为 4k)。如果生成的内容超出这个窗口,由于模型无法从现有预训练窗口中进行外推,性能会出现严重下降。演讲者提出,理想的解决方案是在推理时动态调整模型行为,使其既能支持长上下文,又能保持高效推理,从而兼顾两者的优势,而无需在训练阶段进行额外的修改或重新训练。

利用稀疏性优化推理性能

为了解上述效率问题,演讲者引入了稀疏性(Sparsity)这一关键特性作为优化切入点。通过观察 Transformer 中的前馈循环,研究发现注意力头(Attention Heads)和多层感知机层(MLP Layers)中存在大量的“沉默”状态。具体而言,在许多情况下,大量的注意力头在特定输入下几乎没有被激活;同样,在 MLP 层中,大多数隐藏神经元也是沉默的。对于使用 ReLU 激活函数的 OPT 模型,这些隐藏神经元通常为零。然而,当遇到特定的输入时,特定的神经元或注意力头会被激活。这种输入特定的激活模式表明,并非所有计算单元在所有时刻都参与计算。演讲者指出,正是基于这种稀疏性特征,我们可以设计算法来跳过不必要的计算,从而显著提升推理速度。这一思路源自发表于 ICLR 2023 的论文,旨在通过识别和利用这种结构化的激活模式,实现更高效的模型推理。

预测上下文稀疏性以动态加载参数

在大型语言模型(LLM)的推理过程中,并非所有的注意力头(Attention Heads)或神经元在所有时刻都是活跃的。例如,在特定的层中,Head 43 可能不会提供任何显著的信息,而 Head 42 和 44 则会在新的 token 进入时展现出强烈的模式,分别关注如 "in" 或 "shipping" 等特定词汇,而 Head 23 则处于未激活状态,被视为非激活头。针对这种非激活头,研究者提出了一种预测上下文稀疏性(Contextual Sparsity)的思路,即在输入实际通过计算之前,预测哪些注意力头将被激活。这样做的好处是可以只加载所需的参数子集,从而节省大量的缓存内存并减少计算量。例如,如果在执行第 T4 步之前已知第 K 层的 MLP 层中只有两个神经元会被激活,系统可以预先加载这两行和两列的参数,仅计算这两个神经元的输出,而将剩余的激活头视为常数或均匀分布。这种动态加载机制能够显著降低延迟,尤其是在批量大小(batch size)为 1 时,因为此时激活模式最为稀疏;而在较大的批量中,由于几乎所有神经元都会被使用,这种方法的收益会减弱。

"So the good thing about that is that you can actually load the desired set of parameters and save a lot of cache memories."

为了实现这种预测,研究团队使用评估数据集通过整个管道,获取所有激活头的真实标签(Ground Truth),即哪些 MLP 神经元被激活。随后,他们训练最近邻搜索模型或预测模型,利用前几层提取的特征来预测后续层的激活模式。这种方法利用了嵌入向量在层间缓慢变化的特性,通过前几层的特征来预测当前的激活模式,从而为动态加载权重争取时间。实验表明,这种方法(称为 DJU)在多个下游任务中保持了与原始性能相当的准确性,同时显著降低了延迟。此外,该技术与量化等其他优化技术正交,可以组合使用以进一步提升服务速度。

长文本推理中的注意力分数缓存优化

除了参数稀疏性,长文本推理还面临线性复杂度带来的性能瓶颈。在 Transformer 中,预测下一个 token 需要计算与所有先前 token 的注意力,当序列很长时,这会导致推理速度极慢。通过分析注意力图结构,研究者发现并非所有先前 token 都同等重要。通过对大量数据集的注意力分数进行累积分析,发现共现时间(Co-occurrence time)与累积注意力分数高度一致:对于给定的查询 token,大多数 token 从未出现,只有少数 token 频繁出现,呈现出明显的峰值分布。这种分布特性表明,可以通过缓存这些高注意力的 token 来优化推理过程。

"The answer is no. All these previous tokens equally important. The answer is no."

基于这一发现,H2O 论文从不同角度解决了长文本推理问题。通过识别并缓存那些具有高注意力分数的关键 token,模型可以避免重复计算所有先前 token 的注意力,从而大幅减少计算开销。这种优化方法特别适用于生成大量内容的场景,因为它直接针对了长序列推理中的主要性能瓶颈。通过将注意力分数与 token 共现频率关联,研究者证明了缓存高注意力 token 的有效性,为长上下文支持的高效推理提供了新的解决方案。

H2O算法:基于高频Token的KV缓存优化

在Transformer的推理过程中,高频出现的Token(Heavy Hitters)对下游任务的准确性至关重要。实验表明,如果移除这些高频Token,模型在各项任务上的准确率会大幅下降,这证明了它们在注意力计算中的核心地位。基于这一发现,研究者提出了H2O(Heavy Hitter Oracle)技术,旨在通过局部贪婪算法筛选出重要的Token,从而优化自注意力机制。具体而言,算法会追踪之前Token的注意力得分:若某Token在过去获得了极高的注意力分数,则将其保留为“Heavy Hitter”;反之,若其累积注意力得分极低,说明其很少被关注,则将其丢弃。这种方法将注意力机制的复杂度从原来的 $O(N^2)$ 降低为常数级缓存大小,即序列长度乘以缓存大小,从而显著提升了推理速度并减少了内存占用。

“If you do that then we basically keep like a constant cache of the previous tokens and previous KV caches then this attention uh mechanism is not becoming um O square anymore right so become like constant”

性能表现与副作用:速度提升与多样性增加

H2O技术在保持模型性能的同时,带来了显著的效率提升和意外的生成效果改善。当KV缓存预算保持在20%以上时,模型性能几乎没有下降,这意味着可以节省大量内存和时间。在不同GPU上的基准测试显示,由于减少了注意力计算量,H2O确实加快了推理速度。此外,该技术还与量化技术兼容,进一步加速了模型运行。更有趣的是,H2O产生了减少重复并增加创造性的副作用。以Llama 7B为例,使用标准方法生成时,容易出现“begin to cry”这样的重复句子;而使用H2O后,由于丢弃了部分不相关的历史Token,限制了未来生成的约束,使得输出更加多样化。这种多样性提升可能是因为丢弃了部分低相关性Token,从而为后续生成提供了更广阔的空间。

注意力同步(Attention Syncs)现象解析

为了深入理解H2O中保留的“Heavy Hitters”究竟是谁,研究者分析了Transformer内部的注意力模式,发现了一个有趣的现象:注意力同步(Attention Syncs)。在Transformer的前两层中,注意力分布较为均匀,但从第二层开始直到高层,初始Token(如第一个Token)会吸引大量的注意力,尽管这些Token在语义上可能并不重要。例如,第一个Token可能是“ivory”,与当前生成的内容无关,却依然获得了高注意力分数。研究者假设,由于注意力分数总和必须为1,当模型不想关注所有内容时,无法直接消除注意力,因此倾向于将不必要的注意力分数集中分配给始终存在且相关性最低的初始Token。这种机制导致了“注意力同步”现象,即Token无论相关性如何,都会不恰当地吸引注意力。这一发现揭示了Transformer在处理长上下文时的一种内在偏差,也为进一步优化长上下文推理提供了理论依据。

“we actually see a lot of attention actually uh concentrated uh into the first tokens... we actually call these tokens as attention sync right this means that the token that disappropriately attract attention irrespective of their relevance”

注意力同步机制对模型容量的决定性影响

研究团队通过对比实验发现,绝对位置信息比具体的Token内容对模型容量(Plasticity)的影响更为关键。当保留前四个注意力同步(Attention Sync)Token但将其内容替换为任意字符(如全部替换为换行符)时,模型的容量并未发生显著变化。这表明,只要保留绝对位置结构,具体的语义内容并非维持模型稳定性的必要条件。然而,如果彻底移除这些注意力同步Token,模型的容量会瞬间爆炸式增长,导致性能崩溃。这是因为移除同步机制破坏了原有的注意力结构,使得Softmax无法在正常状态下运行,分母中缺失了原本用于平衡的大数值,从而引发数值不稳定。这一现象证实了注意力同步机制在隐藏注意力质量、维持模型推理稳定性方面的核心作用。

"if we change this first four attention sync uh with arbitrary content... the propacity does not really change... However, if you get rid of the attention sync... all sudden the capacity blows up"

Streaming ARMS:线性复杂度与无限长度生成的平衡

基于上述发现,团队提出了Streaming ARMS方法,旨在实现无限长度序列生成的同时保持线性计算复杂度。传统方法如Dense Attention具有$O(t^2)$的复杂度,而滑动窗口注意力若采用重计算机制,虽然能保持容量稳定,但计算成本极高。Streaming ARMS通过保留初始的注意力同步Token,并采用滚动KV缓存(Rolling KV Cache)机制,动态淘汰中间Token并添加新Token,从而在两者之间取得平衡。其核心设计在于位置滚动(Positional Rolling):在计算下一个Token的自注意力分数时,重新计算所有位置编码和RoPE(旋转位置编码)。关键在于,位置编码是基于缓存内的相对位置而非原始绝对位置计算的。只要缓存大小小于预训练上下文窗口,就能避免位置编码超出预训练范围导致的容量爆炸问题,从而实现稳定的无限长度生成。

性能基准测试与消融实验数据

为了验证Streaming ARMS的有效性,团队进行了多项基准测试,包括与Dense Attention、滑动窗口注意力(带重计算)的对比,以及关于“专用同步Token”的训练实验。实验结果显示,Streaming ARMS在VRAM使用上保持恒定,且在生成400万个Token的过程中速度稳定,未出现容量爆炸。与需要重计算所有先前注意力图的滑动窗口方法相比,Streaming ARMS由于无需重复计算,速度提升了高达22倍。此外,消融实验表明,通过引入可学习的专用同步Token(Learnable Sync Tokens),即使允许其嵌入随时间变化,模型也能保持合理的容量;而固定嵌入的零同步Token(Zero Sync)同样有效,证明了注意力同步机制可以独立于语义内容存在。

对比维度 Dense Attention Sliding Window (w/ Recomputation) Streaming ARMS (Proposed)
计算复杂度 $O(t^2)$ 高(需重计算注意力图) 线性 $O(t)$
VRAM使用 随长度增长而爆炸 恒定(受窗口大小限制) 恒定
容量稳定性 超出预训练长度后爆炸 稳定 稳定
生成速度 慢(随长度平方级下降) 较慢(重计算开销大) 最快(快22倍于重计算窗口)
最大生成长度 受限于预训练窗口 受限于窗口大小 无限(实验达400万Token)

专用注意力同步Token的训练可行性

进一步的研究探讨了是否可以通过预训练专门优化注意力同步机制。实验引入了专用同步Token(Dedicated Sync Tokens),分为两类:一类是嵌入可随时间变化的可学习同步Token,另一类是嵌入在训练期间保持不变的零同步Token。结果显示,只要保留注意力同步机制,模型就能在长序列生成中保持合理的容量和稳定性。这证实了注意力同步Token确实能够存储那些对于不同生成任务而言非必要的注意力质量,从而解耦了位置信息与语义内容的影响,为长上下文推理提供了新的优化方向。

现有长上下文方案的局限与“中间丢失”现象

尽管Meta等机构提出了包含Registers的Transformer变体,试图通过累积注意力质量来优化长文本处理,但长上下文问题尚未得到根本解决。核心痛点在于“中间丢失”(Loss in Middle)现象,即模型在处理超长文本时,会遗忘位于上下文中间部分的关键信息。为了验证这一结论,研究团队进行了专门的实验,通过输入大量合成数据,设定查询(Query)与答案(Answer)之间的行距(Line Distance)作为变量。实验结果显示,随着行距的增加,模型的检索性能逐渐下降,当行距超出特定范围时,性能骤降至零。这表明,虽然当前的流式生成技术(如Streaming LLM)解决了显存爆炸问题,实现了无限流式生成,但模型依然无法有效记忆被驱逐出滑动窗口的中间数据。因此,当前技术仅适用于主要关注近期Token的聊天场景,而非真正的无限上下文理解。

H2O与流式注意力机制的融合优化

针对上述局限,研究者提出将H2O(Heavy Hitter Oracle)与流式注意力机制结合,以进一步提升性能。H2O的核心思想是动态选择重要的Token(Heavy Hitters),而流式注意力机制(Streaming ARMs)则固定关注前四个Token作为Heavy Hitters。从逻辑上看,流式注意力机制可以被视为H2O的一个子集。通过引入滚动位置编码(Rolling Positional Encoding)技术,并将H2O的动态选择机制应用于流式生成中,实验表明该方法在问答和推理任务上取得了比单一机制更优的结果。这种改进不仅保留了无限流式生成的能力,避免了因扩展预训练上下文窗口而导致的显存溢出或生成垃圾内容的问题,还证明了动态选择重要Token结合流式处理是可行的优化路径。

预训练上下文窗口的外推困境

要彻底解决长上下文问题,必须扩展预训练模型的上下文窗口,但这面临巨大的技术障碍。核心问题在于注意力分数函数的外推(Extrapolation)失效。在RoPE等位置编码中,注意力分数是通过嵌入向量与正弦/余弦函数计算得出的。在预训练阶段,注意力跨度(即两个Token位置索引之差)从未超过固定常数(例如Llama 2为4k,Llama 1为2k)。这意味着模型仅在有限窗口内学习了注意力函数的行为。实验数据显示,在2k窗口内,注意力函数表现良好且平滑;然而,一旦位置差超出预训练窗口范围,注意力分数函数会突然变得“疯狂”且不可预测。这是因为函数近似器擅长插值(Interpolation)但不擅长外推(Extrapolation)。这种外推失败导致模型在生成超出预训练窗口的Token时,注意力机制失效,从而无法正确生成内容。因此,简单的扩展上下文而不重新预训练或改进位置编码机制,无法解决根本问题。

位置插值解决外推崩溃问题

在长上下文大语言模型(LLM)的推理中,当序列长度超出预训练窗口时,模型性能往往会急剧下降。传统的线性外推(Extrapolation)方法在处理这种分布外(Out-of-Distribution)数据时表现极不稳定,容易“崩溃”。相比之下,位置插值(Position Interpolation, PI)虽然看似简单,但在理论上具有更严格的界限。论文中的定理表明,插值在实际界限上比外推准确度高出约600倍。这一理论依据构成了该研究的核心思路:将原本困难的外推问题转化为相对稳定的插值问题。

实现这一转化的具体操作非常直观:将所有嵌入(embeddings)的位置编码(positional encodings)按一定常数比例压缩,使得扩展后的位置编码仍然落在预训练窗口的范围内。这种“挤压”操作保留了相对位置关系的结构,从而让模型能够适应更长的序列。尽管方法简单,但其效果显著,为后续的微调奠定了良好的归纳偏置(Inductive Bias)基础。

微调效率与性能对比数据

为了验证位置插值的有效性,研究团队在Llama-1等模型上进行了多项实验。实验对比了三种情况:直接使用原始模型、仅进行微调、以及结合位置插值进行微调。结果显示,如果不采用任何特殊方法,当上下文长度超过2k窗口时,模型复杂度会突然爆发,导致性能崩溃。然而,一旦引入位置插值并进行微调,模型能够显著扩展上下文窗口。

关键的效率优势体现在微调步数上。研究指出,在许多情况下,仅需不到1000步的微调,甚至只需100步,即可将上下文窗口扩展至原始长度的8倍。相比之下,如果不使用位置插值而直接进行微调,模型需要消耗大量的微调步数,但窗口大小的提升却非常有限。这表明位置插值极大地提高了样本效率(Sample Efficiency),使模型更容易被调整以适应长窗口。

以下是不同方法在上下文扩展能力上的对比数据:

方法 微调步数 上下文窗口扩展倍数 性能稳定性
原始模型 (无微调) 0 1x (受限) 超过2k后崩溃
仅微调 (无PI) 大量 (未具体量化) 小幅提升 扩展效果不佳
位置插值 + 微调 < 1000步 (通常~100步) 最高 8x 稳定

此外,研究还验证了下游任务的性能保持情况。尽管进行了长上下文适配,在大多数下游任务中,模型的短文本处理能力保持相同或近似相同,证明了该方法不会以牺牲基础能力为代价。

后续进展与H2O缓存机制

位置插值工作发布后,社区迅速跟进并提出了更优的改进方案。在ArXiv发布仅两天后,就有研究者提出了一种混合策略:在低频部分进行插值,而在高频部分进行外推。具体而言,他们不再使用全局缩放因子 $\alpha$,而是将其放入余弦函数的内部参数中,仅增加 $c$ 的值。这种方法有时甚至无需微调,模型即可自动外推到非常大的上下文窗口,显示了该领域的快速进步。

在问答环节,针对H2O(Heavy Hitters Oracle)缓存机制的细节,演讲者澄清了位置编码的处理方式。在原始的H2O中,所有位置嵌入都被保留,这限制了其无法超过预训练窗口。而在新的流式H2O(Streaming H2O)中,位置编码是基于缓存中剩余位置重新计算的,从而实现了无限长度的扩展。此外,用于预测“热点”注意力位置的神经网络极其轻量,仅由一层MLP组成,进行二分类预测。预测机制主要依赖推理时的统计信息,即计算注意力图中各列的总和,若数值较小则剔除对应Token,无需预先训练复杂的预测网络。

“插值在实际界限上比外推准确度高出约600倍。”

“在许多情况下,仅需不到1000步的微调...即可将上下文窗口扩展至原始长度的8倍。”

“在流式H2O中,位置编码是基于缓存中剩余位置重新计算的,从而实现了无限长度的扩展。”

启发式Token选择与位置插值原理

在讨论输入文本类型(如数学、编程或文学)对“Heavy Heater”Token(即高注意力权重的Token)的影响时,演讲者指出目前尚未进行更细致的分类测试。现有的方法主要基于经验发现,简单的启发式规则在实际应用中表现良好,尽管未来仍有通过训练模型来预测哪些Token最可能成为热点的改进空间。针对如何扩展LLM上下文窗口的问题,位置插值(Positional Interpolation)的核心逻辑是通过缩放位置索引来实现。例如,若将2k的训练窗口扩展至4k,只需将所有位置索引除以2,使重缩放后的索引仍落在原始上下文窗口范围内,从而兼容预训练时的位置分布。

“We just empirically find that this simple heuristics actually works quite well.”

位置编码机制对比:Alibi与RoPE

在位置编码的具体实现上,Alibi与RoPE(Rotary Positional Embeddings)均属于相对位置编码,但衰减机制不同。Alibi在Token距离较远时呈现指数衰减,而RoPE则是指数衰减结合正弦函数。从直觉上看,正弦函数能为距离极近的Token提供更高的分辨率,从而在局部细节上可能更具优势。演讲者提到,有证据表明用Alibi替换RoPE后模型仍能良好工作,这为位置编码的多样性提供了实证支持。此外,关于注意力同步(Attention Sync)现象,研究发现注意力同步主要发生在前几个Token上,且这与Token的具体内容无关,而主要取决于位置编码。即使将前几个Token替换为固定的占位符(如前四个Enter键),只要保留其位置信息,模型仍能获得足够的存储容量来承载过度的注意力质量,证明位置而非内容是触发同步的关键因素。

长依赖问题与评估指标的局限

解决长上下文中的长依赖问题仍是社区面临的开放性问题。当前的注意力机制假设所有Token间存在连接,导致计算复杂度呈平方级增长。一种潜在的优化思路是利用验证数据集计算注意力分数,在计算完整矩阵前预测Query可能关注的Token,从而仅保留相关Token。然而,这种方法在深层网络中可能失效,因为高层Token携带的上下文信息可能使原本不相关的Token变得相关。目前的研究主要利用相似层之间的相似性来预测下一个Token。在评估方面,尽管困惑度(Perplexity)和F1分数(F1 score)被广泛使用,但它们无法充分反映长上下文中的问答能力。对于长上下文检索,社区常采用Pass@k、Needle-in-a-Haystack等特定指标来衡量一致性,但目前仍缺乏统一的评估标准,研究者需根据具体场景自定义指标。

“The reason why we choose to use attention syncs for the tokens is that it's actually compatible with the existing architectures.”

缓存依赖性与困惑度爆炸的临界点

在关于长上下文大语言模型(LLM)高效推理的讨论中,一个核心观察点在于输入长度与模型困惑度(Perplexity, PPL)之间的关系。当横轴表示输入长度,纵轴表示对数困惑度时,不同的流式处理方法展现出截然不同的性能轨迹。特别是窗口注意力机制(Window Attention)与流式LLM(Streaming LLM)以及本团队提出的方法之间,存在显著的性能分化。关键在于KV缓存(Key-Value Cache)的保留策略。对于窗口注意力而言,其性能高度依赖于对早期Token的缓存保留。一旦输入序列超过缓存窗口范围,且未进行重新计算(Recompute),模型的困惑度会急剧上升,出现“断崖式下跌”的现象。相比之下,流式LLM通过保留初始部分的Token,避免了这种爆炸性的性能退化。这种差异揭示了在长上下文场景中,缓存管理的完整性直接决定了模型能否维持稳定的推理质量。

窗口注意力的脆弱性与重新计算的代价

窗口注意力机制的致命弱点在于其对前几个Token的绝对依赖。如果这些初始Token被移出缓存且未重新计算,模型将彻底失效。正如讨论中指出的:“如果你失去了前四个Token,基本上你就完了(you are dead basically)。”这是一种极其脆弱的状态。只要输入序列在缓存窗口内滑动,模型尚能维持运行;但一旦滑动窗口移出初始关键区域,性能便会“掉下悬崖”。为了恢复性能,系统必须执行重新计算(Recompute)操作,这将引入额外的计算开销。重新计算虽然能重置初始Token的状态,使其重新处于安全窗口内,但这一过程带来了显著的额外计算负担。因此,窗口注意力在长上下文处理中,往往需要在缓存命中率与计算成本之间做出艰难权衡,而流式方法则通过持续保留关键状态,规避了这一两难困境。

流式策略的鲁棒性与最终结论

流式LLM的核心优势在于其持续保留初始Token的机制。与窗口注意力的“滑动即死亡”不同,流式方法确保关键信息始终处于可用状态,从而避免了困惑度的爆炸。这种策略使得模型在长序列处理中保持鲁棒性,无需频繁触发高成本的重新计算。在讨论的最后,演讲者Yangd Tian与主持人Dej以及观众进行了总结。Yangd感谢了Dej的主持和观众的参与,并指出演讲已录制,将在伦敦机器学习研讨会(London Machine Learning Meetup)的YouTube频道发布。这一交流不仅澄清了技术细节,也强调了高效推理架构在长上下文场景中的重要性。流式方法通过牺牲少量的内存以换取计算稳定性和低延迟,为LLM的长上下文应用提供了一条可行的技术路径。