Longformer:处理长文档的Transformer模型


🕙发布时间:2025-02-19

近日热文:
1. 全网最全的神经网络数学原理(代码和公式)直观解释
2. 大模型进化史:从Transformer到DeepSeek-R1的AI变革之路
3. 2W8000字深度剖析25种RAG变体:全网最全~没有之一
知乎【柏企
公众号【柏企科技说】【柏企阅文

原始的Transformer模型包含一个自注意力组件,其时间和内存复杂度为 $O(n²)$,这里的 $n$ 指的是输入序列的长度。为了应对这一挑战,我们根据一种 “注意力模式” 对完整的自注意力矩阵进行稀疏化处理,这种模式会指定相互关注的输入位置对。与完整的自注意力不同,我们提出的注意力模式随输入序列呈线性扩展,这使得它在处理较长序列时效率更高。

滑动窗口

鉴于局部上下文的重要性,我们的注意力模式采用围绕每个标记的固定大小窗口注意力机制。通过堆叠多个这样的窗口注意力层,可以得到较大的感受野,顶层可以访问所有输入位置,并且有能力构建整合整个输入信息的表示,这与卷积神经网络(CNNs)类似 。

扩张滑动窗口

为了在不增加计算量的情况下进一步扩大感受野,滑动窗口可以进行 “扩张”。这类似于扩张卷积神经网络,其中窗口存在大小为扩张率 $d$ 的间隔。

全局注意力

我们在少数预先选定的输入位置上添加 “全局注意力”。重要的是,我们使这种注意力操作具有对称性:也就是说,具有全局注意力的标记会关注序列中的所有标记,并且序列中的所有标记也会关注它。

全局注意力的线性投影

我们使用两组投影,即 $Q_s$、$K_s$、$V_s$ 来计算滑动窗口注意力的注意力分数,使用 $Q_g$、$K_g$、$V_g$ 来计算全局注意力的注意力分数。额外的投影为对不同类型的注意力进行建模提供了灵活性,我们发现这对于在下游任务中获得最佳性能至关重要。$Q_g$、$K_g$、$V_g$ 均使用与 $Q_s$、$K_s$、$V_s$ 相匹配的值进行初始化。

实现

在常规的Transformer中,昂贵的操作是矩阵乘法 $QK^T$,因为 $Q$ 和 $K$ 都有 $n$(序列长度)个投影。对于Longformer,扩张滑动窗口注意力仅计算 $QK^T$ 的固定数量的对角线元素。

在自回归语言建模中,我们使用扩张滑动窗口注意力。我们在不同层使用不同的窗口大小。具体来说,较低层使用较小的窗口大小,随着层数升高逐渐增大窗口大小。这使得顶层能够学习整个序列的更高级表示,同时让底层捕捉局部信息。

为了最大化底层学习和利用紧邻局部上下文的能力,我们不在底层使用扩张滑动窗口。

训练

在第一阶段,我们从较短的序列长度和窗口大小开始,然后在后续每个阶段,将窗口大小和序列长度加倍,并将学习率减半。这使得训练速度加快,同时将计算量较大的部分留到最后。我们总共分5个阶段训练模型,起始序列长度为2048,最后一个阶段的结束序列长度为23040。

评估

我们使用长度为32256的序列进行评估。我们将数据集分割为大小为32256的重叠序列,步长为512,并报告序列最后512个标记的性能。

预训练

我们使用掩码语言建模(MLM)对Longformer进行预训练。由于MLM预训练成本较高,我们从发布的RoBERTa检查点继续进行预训练,同时仅进行支持Longformer注意力机制所需的最小改动。需要注意的是,我们的注意力模式可以插入到任何预训练的Transformer模型中,而无需更改模型架构。

注意力模式

我们使用窗口大小为512的滑动窗口注意力,因此计算量与RoBERTa相同。

位置嵌入

RoBERTa使用学习到的绝对位置嵌入,最大位置为512。为了支持更长的文档,我们添加了额外的位置嵌入,以支持高达4096的位置。

继续进行MLM预训练

我们使用fairseq在长文档语料库上对Longformer进行预训练。

冻结RoBERTa权重

我们还在冻结所有RoBERTa权重的情况下对Longformer进行预训练,仅训练新的位置嵌入。

Longformer-Encoder-Decoder (LED)

为了便于对长序列进行序列到序列学习建模,我们提出了Longformer的一个变体,它同时具有编码器和解码器Transformer堆栈,但编码器中不使用完整的自注意力,而是采用Longformer高效的局部+全局注意力模式。解码器对整个编码后的标记以及先前解码的位置使用完整的自注意力。

由于预训练LED成本较高,我们从BART初始化LED参数,并在层数和隐藏层大小方面遵循BART的确切架构。唯一的区别是,为了处理更长的输入,我们将位置嵌入扩展到16K个标记(BART为1K个标记)。

论文

Longformer: The Long-Document Transformer 2004.05150

## 推荐阅读
1. DeepSeek-R1的顿悟时刻是如何出现的? 背后的数学原理
2. 微调 DeepSeek LLM:使用监督微调(SFT)与 Hugging Face 数据
3. 使用 DeepSeek-R1 等推理模型将 RAG 转换为 RAT
4. DeepSeek R1:了解GRPO和多阶段训练
5. 深度探索:DeepSeek-R1 如何从零开始训练
6. DeepSeek 发布 Janus Pro 7B 多模态模型,免费又强大!

本文由mdnice多平台发布


柏企科技圈
15 声望5 粉丝