[论文解读] LongNet: Scaling Transformers to 1,000,000,000 Tokens

Jiayu Ding, Shuming Ma+6 more|arXiv (Cornell University)|Jul 5, 2023

Topic Modeling被引用 34

一句话总结

LongNet 引入扩张注意力，使 Transformer 的序列长度可扩展到超过 10 亿个 token，具有线性级别的复杂度，支持分布式训练并在较短序列上保持性能。

ABSTRACT

Scaling sequence length has become a critical demand in the era of large language models. However, existing methods struggle with either computational complexity or model expressivity, rendering the maximum sequence length restricted. To address this issue, we introduce LongNet, a Transformer variant that can scale sequence length to more than 1 billion tokens, without sacrificing the performance on shorter sequences. Specifically, we propose dilated attention, which expands the attentive field exponentially as the distance grows. LongNet has significant advantages: 1) it has a linear computation complexity and a logarithm dependency between any two tokens in a sequence; 2) it can be served as a distributed trainer for extremely long sequences; 3) its dilated attention is a drop-in replacement for standard attention, which can be seamlessly integrated with the existing Transformer-based optimization. Experiments results demonstrate that LongNet yields strong performance on both long-sequence modeling and general language tasks. Our work opens up new possibilities for modeling very long sequences, e.g., treating a whole corpus or even the entire Internet as a sequence.

研究动机与目标

• 推动在 Transformer 中将序列长度扩展到超出当前极限的必要性。
• 提出一种扩张注意力机制，以指数级扩展注意力场，同时保持效率。
• 实现跨多设备的分布式训练，以处理极端长的序列。
• 证明 LongNet 在显著扩展上下文长度的同时，仍能保持对短序列的性能。
• 提供与现有 Transformer 优化（内核融合、量化等）的实用集成。

提出的方法

• 用扩张注意力替换普通自注意力，在序列上对 Q、K、V 进行基于段的稀疏化。
• 使用混合的扩张注意力模式，随着段大小和扩张率的增加，达到在 N 上近线性 FLOPs 的目标，同时实现基于对数的 token 依赖。
• 引入多头扩张注意力，每个头使用平移的稀疏模式来捕捉多样的局部/全局信息。
• 将总体输出表述为对多种扩张注意力的加权组合，其权重来自 softmax 分母。
• 描述一种分布式训练算法，将序列维在设备间分区，通过 KV 的 all-gather 以及局部 Q 与全局 KV 的跨注意力实现。
• 提供与 dense Transformer 变体的实现兼容性，兼容标准优化（FlashAttention、内核融合、量化）。

实验结果

研究问题

• RQ1 扩张注意力是否能在没有高昂计算或内存成本的前提下，将 Transformer 的序列长度扩展到 1B tokens？
• RQ2 在扩展上下文长度的同时，LongNet 是否能在较短序列上保持或提升性能？
• RQ3 提出的分布式训练方法是否适用于超过单 GDP 内存极限的极长序列？
• RQ4 相较于密集/其他稀疏 Transformer，LongNet 的上下文长度如何影响语言建模性能和增长规律？

主要发现

模型	长度	批次	Github	2K	8K	32K
Transformer	2K	256	Yes	4.24	5.07	11.29
Sparse Transformer	8K	64	Yes	4.39	3.35	8.79
LongNet (ours)	2K	3.24	3.36
Sparse Transformer	16K	32	Yes	4.85	3.73	19.77
LongNet (ours)	16K	3.26	3.31
Sparse Transformer	32K	16	Yes	5.15	4.00	3.64
LongNet (ours)	32K	3.01	3.33	3.01

• LongNet 在序列长度上的计算复杂度接近线性，且 token 依赖度为 O(log N)。
• 扩张注意力使序列可扩展到 1B tokens，通过分布式的序列并行方法，在设备间 KV 通信成本保持常量的条件下实现。
• 在语言建模任务中，LongNet 在较长序列长度（2K、8K、16K、32K）下持续优于密集 Transformer 和稀疏 Transformer 基线，困惑度更低。
• LongNet 在短序列上保持强劲表现的同时，显著扩展了上下文长度，并遵循与标准 Transformer 相似的基于计算的扩展规律。
• 实验表明在训练中扩展上下文长度是可行的，在相同或相近的计算量下，更长的上下文能带来更好的语言建模结果。