长文本模型近期研究工作梳理

引言

基于Transformer的大语言模型通常具有有限的上下文窗口，在处理超出上下文窗口长度的文本时会导致性能显著下降，这是因为常见的大语言模型位置编码（如RoPE）在处理超过上下文窗口的文本时会出现分布外（OOD）的情况。现有的长文本模型仍然面临诸多问题：

1. 模型在文本长度过长、上下文信息丰富的复杂任务中表现不佳（如代码、多文档问答等）。

2. 较长序列带来的注意力稀释，以及远程衰减问题。

3. 基于Transformer模型的计算复杂度随着序列长度的增长而平方增长，训练和推理开销过大。

SAMBA: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling

https://arxiv.org/pdf/2406.07522

本文试图解决的问题是高效地对具有无限上下文长度的序列进行建模。以往的工作要么面临计算复杂度呈二次方增长的问题，要么在长度泛化能力上受限。为此，本文提出了SAMBA，这是一个简单的混合架构，通过层级混合Mamba、SwiGLU和滑动窗口注意力来实现。SAMBA在保持精确召回记忆能力的同时，选择性地将给定序列压缩成循环隐藏状态。

SAMBA首先利用Mamba层来捕获输入序列的时间依赖性，并通过循环结构和输入选择机制来关注重要的输入信息，从而长期记忆重要信息。随后，SAMBA采用了滑动窗口注意力层来解决Mamba在非马尔可夫依赖性方面的问题，通过在一个滑动窗口上操作来确保计算复杂度与序列长度呈线性关系。此外，SAMBA还结合了多尺度感知（SwiGLU）和多层感知机层，作为非线性变换和事实知识回忆的机制。此外，SAMBA模型通过大规模预训练（3.8B参数规模和3.2T训练tokens）并在特定任务上进行微调以提高性能。

尽管在4K序列长度上进行预训练，SAMBA能够通过零样本外推到1M长度，有效降低了Proof-Pile数据集上的困惑度分数，同时保持线性解码时间复杂度和无限token流。通过指令微调，SAMBA能够在256K上下文长度的Passkey检索任务中实现完美的记忆召回，并且在长上下文摘要任务中表现出更好的性能。

Contextual Position Encoding: Learning to Count What’s Important

https://arxiv.org/pdf/2405.18719

本文提出了一种新的位置编码方法，称为上下文位置编码（Contextual Position Encoding，简称CoPE），旨在解决传统位置编码在大型语言模型中的局限性。CoPE的核心思想是将位置编码与上下文信息结合起来，使得位置的确定更加依赖于上下文内容，而不仅仅是基于标记的顺序。针对关注最后一句话的任务，传统的相对位置编码只能通过衰减注意力的方式近似，依然会关注到倒数第二句，而CoPE则只通过门控机制测量距离，只关注位置0的部分。

CoPE的核心步骤：

1. 基于上下文的门控机制：CoPE首先通过门控机制确定哪些token在测量距离时应该被计入，通过计算每个查询向量q和所有之前的键向量k之间的门控值来实现。

2. 累积求和确定位置：接着，CoPE使用门控值的累积求和来确定前向标记相对于当前标记的位置。位置值可以是分数，能够提供更细致的位置信息。

3. 插值嵌入：由于CoPE的位置值可以是分数，不能直接使用嵌入层来转换位置值。因此，CoPE采用插值方法，通过插值最近整数位置的嵌入向量来计算位置嵌入。

4. 上下文依赖的位置编码：CoPE允许每个查询向量根据其上下文独立地测量距离，这意味着模型可以同时使用多种单位来测量距离，例如同时使用标记位置和句子位置。

CoPE的位置编码矩阵：

最后通过在多个任务（如选择性复制、计数任务、语言建模和代码建模）上的实验，验证了CoPE相较于传统位置编码方法的优势。

Base of RoPE Bounds Context Length

https://arxiv.org/pdf/2405.14591

本文探讨了在大型语言模型中，Rotary Position Embedding (RoPE) 的基值与模型处理上下文长度之间的关系。从语义聚合的期望性质出发，研究了RoPE的底数下界，由此指出更大的训练长度应该选择更大的底数，而不单单是为了配合先在短文本上预训练，后在长文本上微调的训练策略。

目前构建长文本模型的主流做法之一是，先在较小的RoPE基值上（通常为10000）用短文本预训练，然后调大基值并在长文本上微调。这种做法相比对基值不加改动的直接微调，起始损失更小，收敛也更快。本文指出更大的训练长度应该选择更大的底数，与训练策略无关（即使在长文本语料上进行预训练），RoPE的基值对模型能够处理的上下文长度有一个决定性的影响。

本文提出了一个新的理论属性——长期衰减（long-term decay），即模型对相似令牌的注意力分配能力会随着它们之间相对距离的增加而衰减。并推导出RoPE基值的下限，这个下限决定了模型能够处理特定上下文长度的能力。如果基值小于这个下限，模型可能只能保持低困惑度，但失去了从长上下文中检索信息的能力。

Retrieval Head Mechanistically Explains Long-Context Factuality

https://arxiv.org/pdf/2404.15574

论文提出了“检索头”这一概念，这是模型中负责从长文本上下文中检索相关信息的特殊类型的注意力头。这些检索头对于模型能否成功找到并使用输入文本中的信息至关重要。本文定义了检索分数，用于衡量不同的注意力头在自回归解码过程中复制-粘贴行为的频率，并在“大海捞针”（Needle-in-a-Haystack）任务上进行测试。

本文详细研究了检索头的属性，主要包括：（1）普遍性：所有具有长文本能力的模型都有一组检索头（2）稀疏性：只有一小部分（少于5%）的注意力头是检索头。（3）内在性：在预训练时使用短文本的模型已经存在相同的检索头。（4）动态激活：检索头根据上下文动态激活。（5）因果关系：完全剪除检索头会导致模型无法检索相关信息并产生幻觉。本文研究了检索头对下游任务，尤其是需要模型精确回忆输入信息的任务，如抽取式问答和使用链式思考（Chain-of-Thought, CoT）推理任务的影响，发现屏蔽检索头会严重影响模型性能。

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

https://arxiv.org/pdf/2404.07143

本文提出了 一种infini-attention 的注意力技术，可以在有限内存和计算能力的条件下，将 LLM扩展到无限输入长度。infini-attention 将一种压缩式记忆方法引入到传统注意力机制中，并且在单个 Transformer 块中构建了掩码的局部注意力和长期线性注意力机制 。

该方法对标准注意力做了较小的改动，支持即插即用的继续预训练以及长上下文适应。还能够使得 Transformer 通过流式处理极长的输入，在有限的内存和训练资源下无限扩展上下文。infini-attention 通过外加一个检索模块来保存长距离信息，而局部的信息仍然通过传统注意力计算。

infini-attention 的计算流程如下：

1. 将输入拆分成预定长度的segment，逐segment进行计算。

2. 计算当前segment的局部attention。

3. 使用Q与之前片段保存的信息M进行计算，得到全局的attention。

4. 将局部的attention与全局的attention通过一个学习到的门控标量β整合，它允许模型在长期和局部信息流之间进行可学习的权衡。

5. 使用当前segment的KV值更新全局信息M，在下个segment继续参与计算。

本文通过在长上下文语言建模、1M序列长度的密钥检索任务和500K长度的书籍摘要任务上的实验，证明了Infini-attention方法的有效性，并且在内存大小上有114倍的压缩比。