比如明星大模型GPT-4支持32ktoken,相当于50页的文字;OpenAI前成员创立的Anthropic更是将Claude处理token能力提升到100k,约75000个单词,大概相当于一键总结《哈利波特》第一部。在微软最新的一项研究中,他们这次直接将Transformer扩展到10亿token。
这为建模非常长的序列开辟了新的可能性,例如将整个语料库甚至整个互联网视为一个序列。
作为比较,普通人可以在5小时左右的时间里阅读100,000个token,并可能需要更长的时间来消化、记忆和分析这些信息。
Claude可以在不到1分钟的时间里完成这些。要是换算成微软的这项研究,将会是一个惊人的数字。
- 论文地址:
https://arxiv.org/pdf/2307.02486.pdf
- 项目地址:
https://github.com/microsoft/unilm/tree/master
具体而言,该研究提出了LONGNET,这是一种Transformer变体,可以将序列长度扩展到超过10亿个token,而不会牺牲对较短序列的性能。
文中还提出了dilatedattention,它能指数级扩展模型感知范围。
LONGNET具有以下优势:
1)它具有线性计算复杂性;
2)它可以作为较长序列的分布式训练器;
3)dilatedattention可以无缝替代标准注意力,并可以与现有基于Transformer的优化方法无缝集成。
实验结果表明,LONGNET在长序列建模和一般语言任务上都表现出很强的性能。在研究动机方面,论文表示,最近几年,扩展神经网络已经成为一种趋势,许多性能良好的网络被研究出来。
在这当中,序列长度作为神经网络的一部分,理想情况下,其长度应该是无限的。但现实却往往相反,因而打破序列长度的限制将会带来显著的优势:
- 首先,它为模型提供了大容量的记忆和感受野,使其能够与人类和世界进行有效的交互。
- 其次,更长的上下文包含了更复杂的因果关系和推理路径,模型可以在训练数据中加以利用。相反,较短的依赖关系则会引入更多虚假的相关性,不利于模型的泛化性。
- 第三,更长的序列长度可以帮助模型探索更长的上下文,并且极长的上下文也可帮助模型缓解灾难性遗忘问题。
然而,扩展序列长度面临的主要挑战是在计算复杂性和模型表达能力之间找到合适的平衡。
例如RNN风格的模型主要用于增加序列长度。然而,其序列特性限制了训练过程中的并行化,而并行化在长序列建模中是至关重要的。最近,状态空间模型对序列建模非常有吸引力,它可以在训练过程中作为CNN运行,并在测试时转换为高效的RNN。
然而这类模型在常规长度上的表现不如Transformer。另一种扩展序列长度的方法是降低Transformer的复杂性,即自注意力的二次复杂性。现阶段,一些高效的基于Transformer的变体被提出,包括低秩注意力、基于核的方法、下采样方法、基于检索的方法。
然而,这些方法尚未将Transformer扩展到10亿token的规模(参见图1)。
下表为不同计算方法的计算复杂度比较。N为序列长度,d为隐藏维数。
方法
该研究的解决方案LONGNET成功地将序列长度扩展到10亿个token。具体来说,该研究提出一种名为dilatedattention的新组件,并用dilatedattention取代了VanillaTransformer的注意力机制。
通用的设计原则是注意力的分配随着token和token之间距离的增加而呈指数级下降。该研究表明这种设计方法获得了线性计算复杂度和token之间的对数依赖性。
这就解决了注意力资源有限和可访问每个token之间的矛盾。在实现过程中,LONGNET可以转化成一个密集Transformer,以无缝地支持针对Transformer的现有优化方法(例如内核融合(kernelfusion)、量化和分布式训练)。
利用线性复杂度的优势,LONGNET可以跨节点并行训练,用分布式算法打破计算和内存的约束。最终,该研究有效地将序列长度扩大到1B个token,而且运行时(runtime)几乎是恒定的,如下图所示。
相比之下,VanillaTransformer的运行时则会受到二次复杂度的影响。
该研究进一步引入了多头dilatedattention机制。
如下图3所示,该研究通过对查询-键-值对的不同部分进行稀疏化,在不同的头之间进行不同的计算。
分布式训练
虽然dilatedattention的计算复杂度已经大幅降低到,但由于计算和内存的限制,在单个GPU设备上将序列长度扩展到百万级别是不可行的。有一些用于大规模模型训练的分布式训练算法,如模型并行[SPP+19]、序列并行[LXLY21,KCL+22]和pipeline并行[HCB+19],然而这些方法对于LONGNET来说是不够的,特别是当序列维度非常大时。
该研究利用LONGNET的线性计算复杂度来进行序列维度的分布式训练。
下图4展示了在两个GPU上的分布式算法,还可以进一步扩展到任意数量的设备。
实验
该研究将LONGNET与vanillaTransformer和稀疏Transformer进行了比较。架构之间的差异是注意力层,而其他层保持不变。
研究人员将这些模型的序列长度从2K扩展到32K,与此同时减小batch大小,以保证每个batch的token数量不变。
表2总结了这些模型在Stack数据集上的结果。研究使用复杂度作为评估指标。
这些模型使用不同的序列长度进行测试,范围从2k到32k不等。当输入长度超过模型支持的最大长度时,研究实现了分块因果注意力(blockwisecausalattention,BCA)[SDP+22],这是一种最先进的用于语言模型推理的外推方法。此外,研究删除了绝对位置编码。
首先,结果表明,在训练过程中增加序列长度一般会得到更好的语言模型。其次,在长度远大于模型支持的情况下,推理中的序列长度外推法并不适用。
最后,LONGNET一直优于基线模型,证明了其在语言建模中的有效性。
序列长度的扩展曲线
图6绘制了vanillatransformer和LONGNET的序列长度扩展曲线。该研究通过计算矩阵乘法的总flops来估计计算量。
结果表明,vanillatransformer和LONGNET都能从训练中获得更大的上下文长度。
然而,LONGNET可以更有效地扩展上下文长度,以较小的计算量实现较低的测试损失。这证明了较长的训练输入比外推法更具有优势。实验表明,LONGNET是一种更有效的扩展语言模型中上下文长度的方法。这是因为LONGNET可以更有效地学习较长的依赖关系。
扩展模型规模
大型语言模型的一个重要属性是:损失随着计算量的增加呈幂律扩展。为了验证LONGNET是否仍然遵循类似的扩展规律,该研究用不同的模型规模(从1.25亿到27亿个参数)训练了一系列模型。
27亿的模型是用300B的token训练的,而其余的模型则用到了大约400B的token。图7(a)绘制了LONGNET关于计算的扩展曲线。该研究在相同的测试集上计算了复杂度。
这证明了LONGNET仍然可以遵循幂律。这也就意味着denseTransformer不是扩展语言模型的先决条件。此外,可扩展性和效率都是由LONGNET获得的。
长上下文prompt
Prompt是引导语言模型并为其提供额外信息的重要方法。
该研究通过实验来验证LONGNET是否能从较长的上下文提示窗口中获益。
该研究保留了一段前缀(prefixes)作为prompt,并测试其后缀(suffixes)的困惑度。
并且,研究过程中,逐渐将prompt从2K扩展到32K。为了进行公平的比较,保持后缀的长度不变,而将前缀的长度增加到模型的最大长度。
图7(b)报告了测试集上的结果。它表明,随着上下文窗口的增加,LONGNET的测试损失逐渐减少。这证明了LONGNET在充分利用长语境来改进语言模型方面的优越性。
本文来源:机器学习研究组订阅,原标题:《微软新出热乎论文:Transformer扩展到10亿token》