大模型推理窗口-从有限到无限大

【本文已在同名 微信公众号 / 知乎 / 个人博客linsight.cn 上线】


步入2024年Q2,大模型在RAG、文档对话、大模型Agent能力等方向的发展持续升温。在平时的日常生活和工作中,大模型工具提供的文档总结、文本润色、代码生成等能力已经是提高效率的必备帮手,甚至在一些复杂或者不熟悉的场景上,大模型也已经能提供一些比较专业的帮助。

在这些方向上,大模型(超)长上下文的能力都是基础。目前不少模型已经在128k+的长度上展示出比较强的能力,这些模型有的收集了超长训练数据,有的优化位置编码,有的则暴力训练,以提升模型的长上下文能力。

另外一些工作则另辟蹊径,直接让模型具备无限大窗口的能力。本篇将梳理几个理论上支持无限大上下文长度的工作。

StreamingLLM

论文:Efficient Streaming Language Models with Attention Sinks

时间:2023年9月

长度:测试4M,理论∞

阶段:无需训练,直接应用在推理阶段

这是一篇由MIT,CMU,NVIDIA,META联合发布的论文。论文提出了一个在流式输出场景下支持无限大上下文长度的方法StreamingLLM,并且不需要进行任何形式的训练或微调 -- 在有限长度上训练,在无限长度上推理。论文用LLAMA2, MPT, Falcon和Pythia作为基线模型进行验证,在4M的上下文长度,StreamingLLM仍然能够保持较低的PPL。

需要注意的是,StreamingLLM主要关注在超长上下文的情况下,让模型能够生成低PPL的内容,但是在内容的准确性上并没有保证。

背景

大模型流式应用,比如现在大部分对话机器人,在部署推理的时候,在长上下文的情况下通常会面临两个问题:
- 缓存不足:流式应用下,为了提升响应速度,KV缓存目前是不可缺少的。而随着对话交互的进行,KV缓存所需的空间也以线性速度在增长。鉴于目前模型的规模动辄百亿千亿的参数量,即使较强的GPU H100、A100拥有40G/80G的显存,推理时也捉襟见肘。在显存不足的情况下,我们就不得不抛弃部分旧的KV缓存值,这也导致了模型效果的下降。
- 训练长度:主流模型大部分使用旋转位置编码RoPE或其变体,而RoPE的外推能力一般,因此当推理输入的上下文长度远超训练的长度时,模型效果也会迅速下降。

现有一些工作针对其中的部分问题进行优化:
- 类似Longformer那样的window attention通过只保存最近的KV值来缓解缓存不足的问题,同时保持一定的窗口扩展能力。但是论文观察到,一旦上下文长度超过缓存大小,window attention的机制会把最初的部分token移出缓存,这时模型的效果依然会迅速下滑。
- sliding window with recomputation(https://github.com/mit-han-lab/streaming-llm/issues/51),可以节省KV所需的空间,但是recomputation计算量和长度是平方关系。这个方案也是论文基线方案里效果最接近可用的。
- 针对位置编码的工作,比如线性插值、NTK插值、YaRN,已经被证明是有效的,但是优化的窗口依然是有限的,大概能在原来的训练长度的基础上,提升一倍或者几倍的推理窗口。论文还提到长窗口的一个实践工作,一个关注在NSFW内容(手动滑稽)的模型:https://kaiokendev.github.io/til#extending-context-to-8k
- 针对推理效率,也有一些工作:《Efficiently scaling transformer inference,SmoothQuant: Accurate and efficient post-training quantization for large language models》、《Dynamic context pruning for efficient and interpretable autoregressive transformers》、《Spatten: Efficient sparse attention architecture with cascade token and head pruning》、《H2o: Heavyhitter oracle for efficient generative inference of large language models》,但是在支持的最大长度依然有限制。
- 无损的FlashAttention已经被广泛采用,无论是训练还是推理都很有帮助。
- 其他一些有损的扩展方案,比如Big Bird、Linformer等。

StreamingLLM这篇论文发现在长上下文的情况下,存在attention sink的现象,这也是window attention在上下文超过缓存大小之后效果迅速变差的原因。通过利用attention sink,让模型在超长上下文的情况还可以保持较低的PPL。并且和基线里唯一效果比较好的sliding window with recomputation方案相比,StreamingLLM在速度上有22+倍的提升。

StreamingLLM和其他方案在长上下文上的PPL对比如下图所示,StreamingLLM在各个模型上都能保持稳定较低的PPL。

attention sink

从window attention在上下文超过缓存大小之后的失效,论文发现自回归LLM存在的一个有趣现象:对于输入文本最靠前的少量几个token,无论它们在语义上与语言建模任务的相关性如何,大量的注意力分数都会分配给他们,如下图所示

模型的前两层还能保持attention score更多分配给当前token附近位置的特性,而在其他层,靠前的几个token都会接受到大量的注意力。

论文里把这些token称为attention sink。尽管这些token在语义上很可能并没有什么重要性,但它们却聚集了大量的注意力分数。

出现这个现象的原因就是softmax操作。softmax要求所有上下文token的注意力分数加起来等于1,因此,即使当前token跟前面的其他token都没有匹配的需要,模型仍然需要将多余的注意力值分配到前面的某些token,以使得总和为1。

那么文本最开头的几个初始token就会承担“接收多余的、不需要的注意力”的任务。为什么是初始token来承担这个任务,最简单的原因就是,对于自回归语言建模,初始token对所有后续token都是可见的,这使得它们更容易被训练成attention sink。

上面这个解释还只是猜想,于是论文做了一个实验来验证这个猜想:把初始的4个token都换成没有重要实际语义的换行符号,结果发现模型依然会把大量的注意力分配给这些token,这就说明attention sink这个现象和内容语义无关,而只和这些token所在的位置相关。

这个现象在一些quant的工作里也有发现,比如《SmoothQuant: Accurate and efficient post-training quantization for large language models》和《Quantizable transformers: Removing outliers by helping attention heads do nothing》等。

推理到无限长

基于以上的发现,论文提出了StreamingLLM。StreamingLLM利用了attention sink具有高注意力值的事实,认为保留它们可以保持注意力分数的分布接近正常。

因此,StreamingLLM的方案就是,在window attention的基础上,增加一个策略,保留了attention sink token的KV值,这些token的KV值会和滑动窗口内的token一起构成完整的KV cache,用于进行新token的推理计算,StreamingLLM和其他方案的attention计算如下图所示。

这里有一个问题,要保留多少个初始token才足够。论文在几个不同的模型上做了实验,结果如下表所示。

实验发现,使用4个初始token就可以基本把PPL降下来,而只使用一个或者两个初始token则仍有很大损失。

最终方案是,在window attention的基础上,StreamingLLM把KV cache分成两部分:
- 包含4个初始token的attention sink
- 正常滚动的KV cache

StreamingLLM的KV缓存如下所示

一个重要的点是,StreamingLLM在实际应用上,位置编码中的相对位置,不再直接使用原文中的distance,而是改成使用这些token在cache中的distance。对于像RoPE这样的位置编码,需要缓存引入旋转变换之前token的KV值,解码时再把这些token在缓存内的相对位置加上。这样做使得模型不用处理大于预训练窗口大小的位置编码,而保证了效果。(关于这一点,在下面一篇论文《LM-Infinite》有相关实验)

当前的大模型基本都没有针对attention sink的现象而做针对性的设计,论文提出可以在训练的时候增加一个特殊的token作为attention sink token使用,方便模型把多余的注意力值放在这个特殊token上。

如果不增加一个特殊的token,那另外一个方法就是使用softmax的变体,softmax-off-by-one,替换attention中的softmax。

softmax-off-by-one的公式如下所示

\[\text{SoftMax}_1(x)_i=\frac{e^{x_i}}{1+\sum_{j=1}^Ne^{x_j}}\]

softmax-off-by-one在分母增加了一个常数1,这样就不要求所有attention score的和为1。这相当于有一个K和V全都是0的虚拟token,固定和所有其他token的注意力得分为1,这个虚拟token在一定程度上也起到了attention sink的作用。

为了验证增加特殊token作为attention sink,以及使用softmax-off-by-one的方案效果,论文在相同设置下训练了3个160M参数的模型,效果如下

从结果上看,即使使用softmax-off-by-one(zero sink),模型还是会依赖前几个token作为attention sink。

最后,论文在LLAMA2, MPT, Falcon和Pythia模型验证StreamingLLM在4M长度上的效果。MPT使用的位置编码是ALIBI,其他模型则是RoPE。效果如下图

在4M的长度上,StreamingLLM仍然能保持PPL较为稳定。

理论上StreamingLLM可以把窗口推到无限长,但是由于本质上还是window attention,因此超过cache大小的部分还是会被丢弃。因此StreamingLLM虽然能保持PPL较低,但是对于需要高精度阅读理解和推理的任务,StreamingLLM还是有局限。

LM-Infinite

论文:LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models

时间:2023年8月

长度:在2k或者4k长度上训练,在200M长度上推理保持较低PPL

阶段:无需训练,直接应用在推理阶段

LM-Infinite可以和现有的LLM结合使用,在Passkey Retrieval和Qasper任务上有比较好的效果,且和原始模型相比,在解码速度上有2.7倍提升,显存上节省7.5倍。

三个挑战

LM-Infinite认为,Transformer LLM无法有效地泛化到长上下文的情况,主要是因为会面临3个挑战。

  • 挑战1:challenges in handling unseen distances among tokens

对于使用相对位置编码的模型,两个token的位置对注意力值的影响取决于两个token之间的距离。当推理的长度越来越长,就会出现在训练时没有见过的相对距离,此时attention logits就会倾向于爆炸增长到无穷大,以区分在训练时从未见过的距离。论文对此给出了数学证明。

而从实践上来看,把LLAMA2所有注意力头在ArXiv数据集上8k以内的attention logits抽取出来,其均值和方差如下图所示。

LLAMA2的预训练窗口大小为4k,可以看到4k之后,attention logits的均值和方差开始有明显的上升。

那么为了缓解这个问题,一个自然的想法是将token之间的相对距离值限制在模型预训练期间看到过的最大值,即设置一个距离上限。这样可以解决logits爆炸,但是会导致下面这个问题。

  • 挑战2:attending to unseen numbers of tokens

对于较长的上下文,位置靠后的token需要在更大的上下文长度范围内,分配它的注意力权重。论文提出,如果注意力logits有界,那么随着上下文长度变长,注意力熵(attention entropy)将增长至无穷大。这里原文也给出了数学证明。

直观来说,在注意力权重分配为平均分配的时候,注意力熵最大。而注意力熵增大,表示当前的token不能确定应该将主要注意力放在上文的哪些token,即无法集中注意力。

同样地,我们把模型在8k内的注意力熵画出来,结果如下图所示,随着长度增大,注意力熵也逐渐增大。

这表明我们应该限制attention context size,以确保注意力熵保持在预训练见过的范围内。

如果简单地使用一个window attention,使得每个token只能关注到一定距离内的token,在一定程度上可以一定程度上handle挑战1和挑战2。这样的思路和XPos、Longformer类似。

但是使用window attention,又会引入新的问题。

  • 挑战3:starting tokens occupy a distinct feature space

即使没有显式的绝对位置编码,位置最靠前的几个初始token的注意力输出和其他位置相比,也会占据一个独特的表示空间。

这一结论来自于《The impact of positional encoding on length generalization in transformers》中的定理1,该定理证明了即使没有使用位置编码,单向注意力使得单个注意力层中的token输出也可以隐式地编码绝对位置,而初始token的信号是最强的,且很容易和其他token区分开来。

这个观察和StreamingLLM中指出的attention sink很相似,总之就是最靠前的几个初始token有很大的影响,因此使用window attention把初始token忽略会带来很大的效果损失。

论文把LLAMA2中各层的hidden state output做了PCA,降到2维之后画出来,如下图所示

蓝色的点表示初始token,红色的点表示靠后的token。上图显示了初始token的特征在特征空间中占据了和其他token很不同的区域。这就解释了为什么简单地将注意力窗口限制在最近的token上会效果不好:由于注意力本质上是加权平均,丢弃初始的几个token影响了注意力输出到达它们应占据的空间。

因此,我们需要保留初始token。

解决方案

基于上面的这些观察和发现,论文提出LM-Infinite,给LLM做zero-shot的长度泛化。LM-Infinite设计了两个主要机制,Λ-shaped attention mask和Distance ceiling。

  • Λ-shaped attention mask

Λ-shaped attention mask在概念上和Longformer、LongNet、Big Bird的思路类似。

Λ-shaped attention mask在window attention的基础上,保持了初始token对后续所有token可见。具体来说,每个token能关注到两部分:(1)\(n_{\mathrm{starting}}\) 个初始token(2)\(L_{\mathrm{pretrain}}\) 个最近的token。其中\(n_{\mathrm{starting}}\) 为超参,\(L_{\mathrm{pretrain}}\) 是预训练时的最大窗口大小。除了这两个部分的token之外,其他token都会被直接忽略。

论文对 \(n_{\mathrm{starting}}\) 的数量进行了实验,发现 \(n_{\mathrm{starting}}\in[5,100]\) 时的效果比较好,如下图所示。

注意这里不需要对模型进行任何训练和微调,只需要修改推理的方式即可。

Λ-shaped attention mask可以解决挑战2和挑战3。

  • Distance ceiling

LM-Infinite把最大距离限制在 \(L_{\mathrm{pretrain}}\),这样做基本只影响初始token的位置编码。

具体来说,对于相对位置编码,假设原来的attention logit是 \(w(\mathbf{q},\mathbf{k},d)\),其中 \(d\) 是两个token之间的距离;那么Distance ceiling就是把attention logit的计算变成:

\[\text{attention logits}=w(\mathbf{q},\mathbf{k},d')\]

\[d'=\min(d,L_\text{pretrain})\]

Distance ceiling限制了最大距离,可以解决挑战1。

  • Optionally attending to top-k tokens in the middle

除了Λ-shaped attention mask和Distance ceiling以外,论文还设计了一个可选的机制,用于把“除了当前窗口内的最近token和初始token”以外的中间部分token也加入到attention计算中,这部分token本来是会被直接丢弃的。

具体来说,会从原本被抛弃的中间token里,选择 \(k\) 个attention logits最大的token加入到attention 计算中来。这 \(k\) 个token的距离都会被统一设置为 \(d=\frac12L_\text{pre-train}\)

\(k\) 的取值,论文在Passkey Retrieval的validation set上做了实验,结果如下

最终取了 \(k=5\),并且只在模型>5的层中加入。

加入middle token主要用于优化需要精准获取中间信息的任务,对于没有这个需要的任务,加入过多的middle token不仅无益,甚至有害,原因如前面挑战2所分析的,会增大注意力熵。

最终,LM-Infinite的整体方案如下图所示

实验

论文在LLaMA-7B、LLaMA2-7B、MPT-7B和GPT-J-6B这几个模型验证LM-Infinite的效果。除了MPT-7B使用Alibi之外,其他模型都是用RoPE。

  • Language Modeling

首先是语言建模的评测,使用的数据集是ArXiv和OpenWebText2。

选择的几个基线模型,以及这些模型使用LM-Infinite的推理策略之后,在0-12k长度上的PPL如下图所示

LLAMA2在略长于10K的长度上就输出了NaN,因此其曲线较短。所有原始模型在约32K的时候就出现OOM。

当测试长度超过训练长度时,基线模型的PPL会迅速爆炸。而使用LM-Infinite之后,所有模型都可以泛化到比训练长度长得多,而仍然保持较低PPL。

为了验证LM-Infinite在更大长度上的效果,用LM-Infinite + Llama2,在由ArXiv数据重复抽样拼接构成的200M token长度上的评估PPL。结果如下图所示,在200M长度下,依然可以保持稳定较低的PPL。

  • Passkey Retrieval and Qapser

Passkey Retrieval和Qapser就是需要对整个上下文有精确理解的任务,因此对于这两个任务,加入了top-5个middle token到attention计算中。

结果如下表。原始模型在长度超出训练长度之后,几乎无法输出有效内容,因此Passkey Retrieval上全得0分。而LM-Infinite相比简单的长度裁剪方案,在两个任务上都有更好效果。

  • Ablation study

对于LM-Infinite的两个主要设计,Λ-shaped attention mask和Distance ceiling,论文做了消融实验,如下图所示。

只使用Λ-shaped attention mask或者distance ceiling单一策略的情况下,PPL仍然会出现较明显的上升,说明这两个组件都是必须的。

Transformer-XL

讲下一篇Infini-Transformer之前,先简单回顾下Transformer-XL的思路。

Transformer-XL是2019年6月,由CMU和Google Brain发布的一个工作,目的是赋予在固定长度上训练的模型,在更大长度上的推理能力。

一般来说,transformer的训练和推理如下图所示,训练的时候模型在固定的长度上训练,如果输入文本超过了模型的训练窗口大小,那就把输入切分成多个segment,模型在各个segment分别独立地进行训练;而推理的时候,如果文本超过了模型的推理窗口大小,模型可以使用滑动窗口的方式,逐步进行推理。

在窗口滑动的过程中,模型始终只能看到窗口内的内容,对于超出的部分则无法把信息纳入计算中,这显然对有长距离依赖的任务有损害。

为了让模型看得更远,Transformer-XL把attention计算修改成下面的样子:把输入切分成多个segment,每个的长度为L,对于两个相邻的segment \([\mathbf{s}_{\tau}=x_{\tau,1},\cdots,x_{\tau,L}]\)\(\mathbf{s}_{\tau+1}=[x_{\tau+1,1},\cdots,x_{\tau+1,L}]\),计算如下

\[\begin{aligned}&\widetilde{\mathbf{h}}_{\tau+1}^{n-1}=\left[\mathrm{SG}(\mathbf{h}_\tau^{n-1})\circ\mathbf{h}_{\tau+1}^{n-1}\right]\\&\mathbf{q}_{\tau+1}^n,\mathbf{k}_{\tau+1}^n,\mathbf{v}_{\tau+1}^n=\mathbf{h}_{\tau+1}^{n-1}\mathbf{W}_q^\top,\widetilde{\mathbf{h}}_{\tau+1}^{n-1}\mathbf{W}_k^\top,\widetilde{\mathbf{h}}_{\tau+1}^{n-1}\mathbf{W}_v^\top\\&\mathbf{h}_{\tau+1}^n=\text{Transformer-Layer}\left(\mathbf{q}_{\tau+1}^n,\mathbf{k}_{\tau+1}^n,\mathbf{v}_{\tau+1}^n\right)\end{aligned}\]

其中SG是stop gradient的操作,这样在训练的时候能保持仍然只有固定窗口大小的参数进行梯度计算和更新,而不会扩大显存需求。\(\begin{bmatrix}\mathrm{SG}(\mathbf{h}_\tau^{n-1})\circ\mathbf{h}_{\tau+1}^{n-1}\end{bmatrix}\) 表示在序列长度维度进行拼接。h、q、k、v的上标 \(n\)\(n-1\) 是模型的层序号。

对于当前segment的每一层attention计算,Transformer-XL把上一个segment的信息加入到当前的K和V中再进行注意力计算,这样相当于每层相比上一层,都能往前多关注到一个segment长度的信息,模型的层数越多,模型最终能关注的长度越长,最后一个segment的最终输出依赖第一个segment的第一层输出,如下图所示。

Infini-Transformer

论文:Leave No Context Behind:Efficient Infinite Context Transformers with Infini-attention

时间:2024年4月

长度:实验500k/1M,理论∞

阶段:继续预训练/微调

Google提出的一种叫Infini-attention的注意力计算方式,可以使用有限的显存计算无限长的上下文。

背景

记忆能力是智能的基石,因为它使得针对特定上下文的高效计算成为可能。

但是transformer的注意力计算在时间和空间上都跟输入长度是二次的关系,这就使得在长上下文的场景下很难高效、准确地处理信息。

使用记忆模块可以在大大减少计算量和显存需求的前提下,依然保持长上下文完整信息对模型的可见性。理论上所有输入信息都可以保存在记忆模块中,因此模型具备处理“有长距离依赖,且需要对输入信息进行准确检索”的任务。从这点上看,和StreamingLLM相比,记忆模块的效果上限更高。

针对高效的记忆,目前也有一些相关工作,比如《Metalearned neural memory》/《Enhancing the transformer with explicit relational encoding for math problem solving》所使用的compressive memory,可以用固定大小的记忆模块来编码context,但是现有的记忆模块在效果和设计的simplicity上仍然存在tradeoff,比较难平衡。

Infini-attention

基于此,Infini-Transformer主要的改进点就在记忆模块的设计上。Infini-attention的结构设计如下图所示,每个Transformer block内包括了常规的mask local attention,和储存长期记忆的long term linear attention。

Infini-attention机制复用了标准注意力计算中的所有QKV值,构建了固定大小的注意力模块,用于长期记忆。在处理后续序列时,根据当前的Q从记忆模块中检索相关内容。Infini-attention将长期记忆检索到的值与局部注意力上下文通过concat进行聚合,共同计算最终的输出。

这种对Transformer注意力层的微小但关键的修改,使得现有的LLM能够通过继续预训练和微调自然地扩展到无限长的上下文。

Infini-attention和Transformer-XL的对比如下图所示

和Transformer-XL类似,Infini-Transformer也把输入分成多个segment来处理。不同的是,Transformer-XL每层只能比前一层多看到一个segment的信息,因此需要通过堆叠模型的层数来提升模型处理窗口的大小;而由于层数是有限的,最终窗口的大小还是有限的。

Infini-Transformer则是每层都有对之前所有segment的长期记忆,因此不受模型层数的影响,可以扩展到无限大的窗口。

compressive memory的具体设计,出于simplicity和computational efficiency的考虑,论文参考《Learning associative inference using fast weight memory》,把memory设计成一个associative matrix。

对于记忆模块,最重要的就是记忆的更新和检索。

  • 记忆检索

先看下记忆的检索。假设每个segment的长度为 \(N\),我们使用当前的 \(Q\in\mathbf{R}^{N\times d_{key}}\) 从memory \(M_{s-1}\in\mathbf{R}^{d_{key}\times d_{value}}\) 中检索,计算如下

\[A_{\text{mеm}}=\frac{\sigma(Q)M_{s-1}}{\sigma(Q)z_{s-1}}\]

其中\(\sigma\)是激活函数,使用的是element-wise ELU + 1。

\(z_{s-1}\in\mathbf{R}^{d_{key}}\)是normalization term。这里normalization term参考《Transformers are rnns: Fast autoregressive transformers with linear attention》的做法,使用所有K值的求和值。

  • 记忆更新

完成检索之后,就要更新记忆。需要更新的有normalization term和memory两项,更新方式如下。

\[M_s\leftarrow M_{s-1}+\sigma(K)^TV\]

\[z_s\leftarrow z_{s-1}+\sum_{t=1}^N\sigma(K_t)\]

memory的更新,参考《Metalearned neural memory》和《Learning associative inference using fast weight memory》所使用的delta rule,进一步优化成如下形式

\[M_s\leftarrow M_{s-1}+\sigma(K)^T(V-\frac{\sigma(K)M_{s-1}}{\sigma(K)z_{s-1}})\]

  • 长期记忆和local attention结合

回到当前segment的注意力计算中,需要把长期记忆检索项 \(A_{mem}\) 和local attention state \(A_{dot}\) 结合

\[A=sigmoid(\beta)\odot A_{mem}+(1-sigmoid(\beta))\odot A_{dot}\]

其中 \(\beta\) 是一个可学习的标量参数。

在训练中, \(\beta\) 所学习到的值如下,0则表示这个头只关注当前segment,1则表示值关注长期记忆的内容,介于0和1之间的则是两种的混合,可以看到大部分的注意力头会使用二者的混合结果

整体策略上,和其他包含segment-level memory的模型在方案设计上的比较如下

实验

实验上,每个segment的长度 \(N\) 设为2048,模型的训练窗口大小为32768,即总共有16个segment。

  • 语言建模能力

Infini-Transformer和Transformer-XL/Memorzing Transformer/RMT在PG-19和Arxiv数据集上的language modeling能力对比如下

Infini-Transformer的效果显著更好,而且所需的记忆空间只有Memorizing Transformers的不到1%。

原文指出,把训练的数据长度提升到100k,Infini-Transformer的PPL还能进一步降低。

  • 下游任务

用1B的模型进行轻量级的继续预训练,训练配置如下:
- batch size = 64 - step = 30k - 训练数据长度 > 4k - segment length = 2k

之后在1M长度的passkey retrieval任务上评测,zero-shot和fine-tune模型在各个长度的效果如下表

Infini-Transformer(linear + delta)在1M长度的passkey retrieval任务上做到完全正确。

另外在8B模型上,用8k长度训练30k步,在500k的BookSum任务上评测,结果如下

Infini-Transformer(linear + delta)相比其他模型有一定优势。

小结

  1. Streaming和LM-Infinite思路有些相似,都观察到attention sink对模型效果的影响,并设计了相关的注意力计算机制保留初始token以稳定PPL,而LM-Infinite在这个基础上使用distance ceiling,使得模型在推理时不会使用超出训练所用的token相对距离。虽然这两个工作使得模型可以在M级别的长度仍然保持较低PPL,但是实际大量中间token被丢弃,导致对需要在上下文精准检索的任务效果有损害。
  2. Infini-Transformer使用记忆模块,把输入进行分段,并把靠前的内容通过固定大小的记忆矩阵进行压缩。理论上记忆模块可以提供所有上文的信息,能够应用在更精细和困难的任务。记忆模块的检索和更新的设计是方案的核心,需要考虑复杂度和效果的平衡。

读到这了,来一发点赞收藏关注吧~

博客:http://www.linsight.cn/
知乎:Linsight
微信公众号:Linsight


【往期文章】

MoE模型的前世今生
LLM长上下文的问题
解锁大模型长上下文能力
理解Attention:从起源到MHA,MQA和GQA
Yi技术报告-划重点看细节
transformer中normalization的二三事
从代码实现看normalization-到底做了什么
稀疏注意力计算:sliding window attention
理解LLM位置编码:RoPE
大模型算法题(1)
大模型算法题(2)
大模型算法题(3)
大模型算法题(4)
大模型算法题(5)


Reference

【1】Efficient Streaming Language Models with Attention Sinks https://arxiv.org/abs/2309.17453
【2】Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context https://arxiv.org/abs/1901.02860
【3】Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention https://arxiv.org/abs/2404.07143
【4】LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models https://arxiv.org/abs/2308.16137