1000000000！微软改进Transformer一次能记住这么多token了

丰色 发自 凹非寺

量子位 | 公众号 QbitAI

微软亚研院最新研究，有点震撼：

他们开发了一个Transformer变体，居然将其token拓展到了10亿。

什么概念？

目前最强的GPT-4也才最大支持一次处理32k token，相当于50页文字。

而能够只用1分钟看完一本数万字小说的Claude，其token数也不过“才”100k（10万）。

一次性扩展到10亿，并且这个数字理论上其实还是无限的，这不就意味着：

不久的将来，整个语料库甚至互联网都能视为一个序列？

若真如此，简直不敢想象。

如何做到？

大预言模型时代，扩展序列长度已成为一大关键需求。

然而，现有的方法要么难以解决计算复杂性，要么搞不定模型表达力，导致长度很受限。

在此，作者提出一个Transformer变体：LongNet，它应用了一种叫做“膨胀注意力（dilated attention）”的机制，可以随着距离的增长，让注意力场（模型感知范围）呈指数级扩展。

具体而言，dilated attention替代了普通Transformer中的注意力机制的，其一般的设计原则是：

让注意力的分配随着token之间距离的增长，呈指数级下降。

如下图构建块所示，dilated attention由一系列用于建模短期和长期依赖关系的注意力pattern组成。pattern的数量可以根据序列长度进行扩展。

实验证明，dilated attention能够产生线性计算复杂度和token之间的对数依赖性，从而解决了注意力资源有限，但每一个token都可访问的矛盾。

除此之外，作者还引入了多头dilated attention。如下图所示，不同头之间的注意力pattern会因位置的连续变化而不同。

在最终实现中，LongNet可转换为密集的Transformer，从而无缝集成Transformer现有的优化方法（比如内核融合、量化和分布式训练）。

而利用线性复杂度，LongNet还可以跨节点并行训练，利用分布式算法打破计算和内存的限制。

如下图所示，LongNet通过划分序列维度在两个GPU上进行并行化训练。随着设备数量的增长，计算和通信成本几乎仍然是不变的。

因此，这使得LongNet可以在几乎恒定的运行时间内有效地将序列长度扩展到10亿token（同时无需牺牲较短序列的性能）。

相比之下，普通Transformer还在受困于二次复杂度带来的影响。

实验评估

作者将LongNet与普通Transformer和稀疏Transformer进行了比较。

三者之间的体系结构差异只体现在注意力层，模型的序列长度都从2K扩展到32K，同时减少batch size以保持每个batch的token数量不变。

对于LongNet，作者使用w＝{2048,4096,8192,16384,32768}的分段长度，同时扩张比率为r＝{1,2,4,6,12}，并实现了稀疏注意力的固定模式，再将稀疏比进行调整为匹配LongNet的计算flops，从而使实验公平。

（由于计算限制，只能将LongNet扩展到32K序列长度）。

下表为以上三个模型在Stack数据集上的结果，衡量指标是困惑度，测试了每个模型不同的序列长度。

当输入长度超过模型支持的最大长度时，作者就使用一种最先进的语言模型推理外推方法：分块因果注意（BCA）。

此外，还去除了绝对位置编码。

结果表明：

首先，在训练过程中增加序列长度一般会得到效果更好的语言模型。

其次，在长度远大于模型支持的情况下，序列长度外推法（推理时）并不适用。

最后，可以看到，LongNet的每一项数据都优于其他两个基线，证明了它在语言建模方面的有效性。

除此之外，作者还证明，在扩展语言模型上下文长度时，LongNet也是一种更有效的方法。

网友：Memery is all you need

不少人看完这项成果总结道，说来说去，还是“memery is all you need”（手动狗头）

有点厉害！非常希望它是真的且可复现。

不过，也不乏争议。

比如有人吐槽实验只将LongNet扩展到了32k长度，离10亿还差的远呢。

还有不少人称，方法不算很创新，其实就跟稀疏注意力模式差不多。

作者介绍

本论文作者7位，全部来自微软亚研院。

两位共同一作：

Jiayu Ding，微软亚研院实习生，本硕毕业于北京邮电大学计算机科学专业；

马树铭，微软亚研院自然语言计算组研究员，2019年加入，本硕毕业于北京大学。

通讯作者为Furu Wei，微软亚研院自然语言计算组全球研究合伙人，博士毕业于武汉大学，曾任IBM中国研究院研究员、香港理工大学电子计算系研究助理。

论文地址：
https://arxiv.org/abs/2307.02486

参考链接：
[1]https://twitter.com/arankomatsuzaki/status/1676765133362675712
[2]https://twitter.com/AlphaSignalAI/status/1677009901334192143?s=20
[3]https://www.reddit.com/r/LocalLLaMA/comments/14rym30/longnet_scaling_transformers_to_1000000000_tokens/