为什么LONGNET可以处理10亿个令牌

Vaswani等人在2017年引入的Transformers是LLM和许多其他机器学习任务的基石。它成功的主要原因之一是自注意力机制。但是随着模型的扩大这种机制成为计算瓶颈，特别是对于长序列。

本文中将深入研究一种名为“扩张注意力”(Dilated Attention)的创新方法，它是LONGNET(微软开发)的核心，于2023年7月19日发布，它号称可以有效处理数十亿个令牌的序列!作为比较，GPT-4可以处理的最长序列长度是32,768个令牌(相当于大约52页的文本)。

LONGNET的主要的结果是将自注意力的计算复杂度从二次降低到线性。10亿个令牌大约是400万页的文本。

传统的自注意力机制

在传统的Transformers中，给定输入Q (query)、K (key)和V (value)，自注意力机制计算输出O为:

每个查询Q都会关注所有键K和值V，这使得最终计算是序列长度的二次复杂度，这样处理长序列效率低下。

稀疏注意力:一个部分解

为了克服传统的自注意力方法的不足，产生了稀疏注意方法。它限制查询只能访问键和值的子集。稀疏注意力的输出表示为:

这里1S是稀疏关注模式，用于确定查询可以关注哪些键和值。例如，Sparse Transformer利用了固定模式，如局部模式和跨步模式。

扩张注意力

如图下所示，扩张的注意力带来了范式的转变。它包括将输入分割成段，然后通过以特定间隔选择行来稀疏每个段。

给定输入段Q、K、V，分段计算可表示为:

这里的w为段长度，r为区间(扩张率)。这些片段会被并行地输入到注意力机制中。然后将结果输出连接起来，如下所示:

扩张注意力可以通过特定的收缩和分散操作转化为普通注意力，从而重用像Flash attention这样的优化。与传统的自注意力相比，显著降低了计算成本。

扩张注意力的实践

在计算效率和注意力的全局性之间取得平衡是至关重要的。扩张注意力是通过实现具有不同分段大小和膨胀率{ri, wi}^k的膨胀注意力的混合来实现的。最终输出是一个加权的注意力和，其中权重是根据softmax的分母动态计算的。

为了确保有效地计算局部和全局注意力，还会逐渐增加每个注意力的片段长度，同时在几何序列中设置扩张率以获得指数级的注意力域。

扩张注意力的优点

1、并行计算:

{O∣ri,wi}k表示注意力机制在不同扩张率(ri)和段大小(wi)下的输出。这些输出的计算是并行的。因为不同膨胀率和段大小的计算之间没有相互依赖关系。也就是说每一种扩张的注意力模式都是独立的，并且不依赖于另一种模式。

2、动态权重vs.固定权重:

论文通过研究发现，使用动态权值(根据注意softmax的分母计算)比可学习的固定权值效果更好。这意味着该模型受益于根据上下文调整注意力权重，而不是依赖于静态的预定权重。

3、分散后注意力的混合:

在不同注意力计算中处理键，整合方法将这些注意力组合在一起，就像从不同的段收集键一样，然后计算它们的softmax。这使注意力模式可以用不同的方式扩张，然后再把它们结合起来获得更丰富的上下文理解。

4、局部vs.全局注意力:

局部注意力需要精确的计算，因为它更关注即时或附近的令牌，而这些令牌通常具有更直接的关系。

全局注意力:可以近似，因为它捕获了更广泛、不那么直接的关系。

为了平衡这一点，更大的段尺寸(wi)与更大的膨胀率(ri)配对。当模型获得更大的数据块，则会对数据点之间的间隔更加宽容。

5、分段大小逐渐增加:

每个注意力模式的片段大小(wi)逐渐增加，片段大小继续增长，直到达到最大序列长度或注意力模式总数。

6、多头扩张注意力:

采用多头注意机制，并且每个头使用不同的扩张注意模式。每个注意头都集中在输入序列中的不同位置。这样模型可以同时捕获数据中的各个方面和模式。

7、膨胀率:

膨胀率(ri)也遵循一个级数，每一个速率都比前一个速率大。

8、分段大小和扩张率的几何序列

段大小(w)和膨胀率(r)都被设置为几何序列，这意味着每一项都是前一项的固定倍数。这种方法确保了指数增长模式，提供了一种快速增加关注领域和捕捉不同层次上下文的方法。

结果

以下是论文的实验和发现的简要说明:

模型范围从2K到32K的序列长度。相应地调整批大小以保持令牌的一致性。

LONGNET的段长度:w = {2048, 4096, 8192, 16384, 32768}

膨胀比:r = {1,2,4,6,12}

稀疏Transformers:对不同的子块采用固定的模式注意，有多个头。

传统的Transformers:利用密集和完全连接的注意力层，由于计算限制序列长度上限为32K。

框架:所有模型都基于FlashAttention平台，并针对不同的注意力配置进行了定制。

论文的主要成果:

性能提升:更长的训练序列总是产生更好的语言模型。

外推限制:如果推断序列长度大大超出模型的设计，则外推不理想。

在Stack数据集的测试中，LONGNET超越了它的竞争对手，标志着它在语言建模方面的实力。

随着LLM的快速发展，LONGNET的创新方法强调了该领域的巨大潜力和未来的可能性。

LONGNET代表了Transformer架构的重大进步。通过解决传统子注意力的的二次复杂度挑战，它为有效和可扩展的序列建模铺平了道路。

以下是论文地址，有兴趣的看看

https://arxiv.org/pdf/2307.02486.pdf

作者：Preetham