原文信息
原文题目:《Retentive Network: A Successor to Transformer for Large Language Models》
原文引用:Sun Y, Dong L, Huang S, et al. Retentive Network: A Successor to Transformer for Large Language Models[J]. arXiv preprint arXiv:2307.08621, 2023.
原文链接:https://arxiv.org/pdf/2307.08621.pdfhttps://arxiv.org/pdf/2307.08621.pdf
0.摘要
在这项工作中,我们提出了Retentive Network (RETNET)作为大型语言模型的基础架构,同时实现了训练并行性、低成本推理和良好的性能。我们从理论上推导了循环和注意力之间的联系。然后,我们提出了序列建模的保留机制,支持三种计算范式,即并行、递归和分块递归。具体而言,并行表示允许进行训练并行化。递归表示实现了低成本的O(1)推理,提高了解码吞吐量、延迟和GPU内存,而不会牺牲性能。分块递归表示以线性复杂度实现了高效的长序列建模,每个块都是并行编码的同时递归总结块。语言模型的实验结果表明,RETNET实现了有利的扩展结果、并行训练、低成本部署和高效推理。这些有趣的特性使RETNET成为大型语言模型中Transformer的强有力的继任者。代码将在https://aka.ms/retnet上提供。
图1:相较于Transformer,保留网络(RetNet)在GPU内存、吞吐量和延迟等方面实现了低成本推理、训练并行性和良好的缩放曲线。推理成本的结果以输入长度为8k进行报告。图6展示了不同序列长度的更多结果。
1.引言
Transformer[VSP+17]已成为大型语言模型[BMR+20]的事实上的架构,最初是为了解决循环模型[HS97]的顺序训练问题而提出的。然而,Transformer的训练并行性是以低效推理为代价的,因为每个步骤的复杂度为O(N),而且受限于内存的键值缓存[Sha19],这使得Transformer在部署方面不友好。不断增长的序列长度增加了GPU内存消耗,同时也增加了延迟,并降低了推理速度。
许多努力不断致力于开发下一代架构,旨在保持与Transformer相当的训练并行性和竞争性能,同时具有高效的O(1)推理。同时实现上述目标是具有挑战性的,即所谓的“不可能三角形”,如图2所示。
已经有三个主要的研究方向。首先,线性化注意力[KVPF20]使用核函数ϕ(q) · ϕ(k)来近似标准的注意力分数exp(q · k),以便将自回归推理重写为循环形式。然而,模型的建模能力和性能较Transformer差,这阻碍了该方法的普及。第二个方向是回归模型,用于高效推理,但牺牲了训练并行性。为了弥补这一点,使用逐元素运算符[PAA+23]进行加速,但会损害表示能力和性能。第三个研究方向是探索用其他机制替代注意力,例如S4[GGR21]及其变体[DFS+22,PMN+23]。然而,以前的工作都无法突破“不可能三角形”,与Transformer相比没有明确的胜出者。
在这项工作中,我们提出了保留网络(RetNet),实现了低成本推理、高效的长序列建模、与Transformer可比的性能以及并行模型训练。具体而言,我们引入了一种多尺度保留机制来替代多头注意力,它具有三种计算范式,即并行、循环和分块循环表示。首先,通过并行表示,我们能够充分利用GPU设备进行训练并行性。其次,通过循环表示,我们能够在内存和计算方面实现高效的O(1)推理,从而大大降低了部署成本和延迟。此外,实现过程中无需使用键值缓存技巧,大大简化了实现。第三,分块循环表示可以进行高效的长序列建模。我们同时对每个本地块进行并行编码以提高计算速度,同时对全局块进行循环编码以节省GPU内存。
我们进行了大量实验,将RetNet与Transformer及其变种进行了比较。在语言建模方面的实验结果表明,RetNet在缩放曲线和上下文学习方面始终具有竞争力。此外,RetNet的推理成本与序列长度无关。对于一个7B模型和8k序列长度,RetNet的解码速度比带有键值缓存的Transformer快8.4倍,并节省了70%的内存。在训练过程中,RetNet相比标准的Transformer节省了25-50%的内存,并且比高度优化的FlashAttention[DFE+22]加速了7倍。此外,RetNet的推理延迟对于批量大小不敏感,可以实现巨大的吞吐量。这些有趣的特性使得RetNet成为大型语言模型中强有力的Transformer继任者。
图2:RetNet使得“不可能三角形”成为可能,实现了训练并行性、良好的性能和低推理成本的同时存在。
2.保留网络
保留网络(RetNet)由L个相同的块堆叠而成,其布局与Transformer[VSP+17]中的布局相似(即残差连接和预LayerNorm)。每个RetNet块包含两个模块:多尺度保留(MSR)模块和前馈网络(FFN)模块。我们在以下章节介绍MSR模块。给定输入序列x = x1 · · · x|x|,RetNet以自回归的方式对序列进行编码。首先,将输入向量{xi} |x| i=1打包成X0 = [x1, · · · , x|x|] ∈ R |x|×dmodel,其中dmodel是隐藏维度。然后我们计算上下文化的向量表示Xl = RetNetl(Xl−1),其中l ∈ [1, L]。
2.1.保留机制
在这个部分,我们介绍了一种具有循环和并行性的保留机制。因此,我们可以以并行的方式训练模型,同时循环地进行推理。给定输入X ∈ R |x|×dmodel,我们将其投影到一维函数v(n) = Xn · wV。考虑一个将v(n)映射到o(n)的序列建模问题,其中通过状态sn进行映射。为了简化表示,我们用vn和on表示v(n)和o(n)。我们以循环方式来表达这种映射: 其中†表示共轭转置。该公式在训练实例内很容易进行并行化。总结起来,我们从如式(1)所示的循环建模开始,然后推导出如式(4)所示的并行公式。我们将原始映射v(n) 7→ o(n)视为向量,并得到保留机制如下。
其中Θ是Θ的复共轭,D ∈ R |x|×|x| 将因果屏蔽和相对距离的指数衰减组合为一个矩阵。与自注意力类似,这种并行表示使我们能够高效地使用GPU训练模型。
保留的循环表示 如图3b所示,提出的机制也可以表示为循环神经网络(RNN),这对于推理是有利的。对于第n个时间步,我们循环地获得输出为:
保留的块状循环表示 对于加速训练,特别是对于长序列,可以使用并行表示和循环表示的混合形式。我们将输入序列划分为块。在每个块内,我们按照并行表示(公式(5))进行计算。相反,跨块的信息按照循环表示(公式(6))进行传递。具体而言,设B表示块的长度。我们通过以下方式计算第i个块的保留输出:
图3:RetNet的双重形式。"GN"是GroupNorm的缩写。
2.2.门控多尺度保留
我们在每一层中使用h = dmodel/d保留头,其中d是头的维度。这些头使用不同的参数矩阵WQ,WK,WV ∈ R d×d 。此外,多尺度保留(MSR)为每个头分配不同的γ。为了简化起见,我们使不同层之间的γ相同,并保持固定。此外,我们添加了一个swish门[HG16,RZL17]来增加保留层的非线性。形式上,给定输入X,我们将该层定义为:
其中WG,WO ∈ R dmodel×dmodel是可学习的参数,GroupNorm [WH18]对每个头的输出进行归一化,遵循[SPP+19]中提出的SubLN。请注意,头部使用多个γ尺度,这导致了不同的方差统计。因此,我们单独对头部输出进行归一化。保留的伪代码总结如图4所示。
保留分数的归一化 我们利用GroupNorm的尺度不变性来提高保留层的数值精度。具体来说,在GroupNorm中乘以一个标量值不会影响输出和反向梯度,即GroupNorm(α ∗ headi) = GroupNorm(headi)。我们在公式(5)中实现了三个归一化因子。首先,我们将QK⊺归一化为QK⊺/√d。其次,我们用D˜nm = Dnm/√Pn i=1 Dni替换D。第三,设R表示保留分数R = QK⊺ ⊙ D,我们将其归一化为R˜nm = Rnm/max(|Pn i=1 Rni|,1)。然后,保留输出变为Retention(X) = RV˜。上述技巧不会影响最终结果,同时稳定正向和反向传递的数值流动,因为具有尺度不变性的属性。
图4:保留的三种计算范式的伪代码。
2.3.保留网络的整体架构
对于一个L层的保留网络,我们堆叠多尺度保留(MSR)和前馈网络(FFN)来构建模型。形式上,将输入序列{xi} |x| i=1通过一个词嵌入层转换为向量。我们使用打包的嵌入X0 = [x1, · · · , x|x| ] ∈ R |x|×dmodel作为输入,并计算模型输出XL:
其中LN(·)是LayerNorm [BKH16]。FFN部分的计算方式为FFN(X) = gelu(XW1)W2,其中W1、W2是参数矩阵。
训练:训练过程中我们使用并行(公式(5))和分块递归(公式(7))表示。在序列或块内部进行并行化能够有效利用GPU加速计算。更重要的是,分块递归在长序列训练中非常有用,既节省了计算量也节省了内存消耗。
推理:推理过程中采用了递归表示(公式(6)),这非常适合自回归解码。O(1)复杂度减少了内存和推理延迟,同时实现了相同的结果。
表1:从不同角度比较模型。RetNet实现了训练并行化、恒定的推理成本、线性的长序列内存复杂度和良好的性能。
2.4.与之前方法的关系与区别
表1从不同的角度比较了RetNet和先前的方法。比较结果反映了图2中呈现的“不可能的三角形”。此外,由于分块递归表示,RetNet对于长序列具有线性的内存复杂度。我们还总结了与具体方法的比较如下。
Transformer RetNet的并行表示与Transformer [VSP+17]具有相似的思路。最相关的Transformer变体是Lex Transformer [SDP+22],它使用xPos作为位置嵌入。如公式(3)所述,retention的推导与xPos相一致。与注意力机制相比,retention去除了softmax并启用了递归形式,从而显著提高了推理性能。
S4 与公式(2)不同,如果Qn和Kn不具备内容感知性,公式可以退化为S4 [GGR21],其中O = (QK⊺, QAK⊺, .., QA|x|-1K⊺) * V。
线性注意力:线性注意力变体通常使用各种核函数ϕ(qi)ϕ(kj )/P|x|n=1 ϕ(qi)ϕ(kn)来替代softmax函数。然而,线性注意力在有效编码位置信息方面存在困难,使得模型的性能较低。此外,我们重新审视序列建模的过程,而不是旨在近似softmax。
AFT/RWKV Attention Free Transformer (AFT) 将点积注意力简化为逐元素操作,并将softmax移动到键向量上。RWKV使用指数衰减替换了AFT的位置嵌入,并在训练和推理过程中对模型进行递归运行。相比之下,retention保留了高维状态以编码序列信息,这有助于提高模型的表达能力和性能。
xPos/RoPE 与为Transformer提出的相对位置嵌入方法相比,公式(3)呈现了与xPos [SDP+22]和RoPE [SLP+21]类似的表达式。
Sub-LayerNorm 如公式(8)所示,retention层使用了Sub-LayerNorm [WMH+22]来对输出进行归一化。由于多尺度建模导致头部具有不同的方差,我们用GroupNorm替代了原始的LayerNorm。
3.实验
我们进行了语言建模的实验来评估RetNet。我们使用各种基准测试来评估所提出的架构,包括语言建模性能以及下游任务的零/少样本学习。此外,对于训练和推理,我们还比较了速度、内存消耗和延迟。
表2:语言建模实验中模型的大小和学习超参数。
图5:随着模型规模的扩大,困惑度下降。我们经验证实,当模型大小大于2B时,RetNet往往优于Transformer。
3.1.设置
参数分配:为了进行公平比较,我们重新分配了MSR和FFN中的参数。在这里,为了简单起见,我们用d表示dmodel。在Transformer中,自注意力机制中有大约4d^2个参数,其中WQ、WK、WV、WO∈R^d×d,FFN中有8d^2个参数,其中中间维度为4d。相比之下,RetNet中的retention层有8d^2个参数,其中WQ、WK∈R^d×d,WG、WV∈R^d×2d,WO∈R^2d×d。注意,V的头部维度是Q、K的两倍。扩展的维度通过WO投影回d。为了保持参数数量与Transformer相同,RetNet中FFN的中间维度为2d。同时,我们在实验中将头部维度设置为256,即查询和键的维度为256,值的维度为512。为了公平比较,我们在不同的模型大小中保持γ相同,其中γ=1−e^(linspace(log 1/32,log 1/512,h))∈R^h,而不是公式(8)中的默认值。
语言模型训练:如表2所示,我们从头开始训练各种规模的语言模型(即1.3B、2.7B和6.7B)。训练语料库是The Pile [GBB+20]、C4 [DMI+21]和The Stack [KLBA+22]的精选编译版本。我们添加了标记来表示序列的开始。训练的批量大小为4M个标记,最大长度为2048。我们使用100B个标记,即25k步来训练模型。我们使用AdamW [LH19]优化器,设置β1=0.9,β2=0.98,并将权重衰减设置为0.05。预热步数为375,使用线性学习率衰减。为了保证训练稳定性,参数初始化遵循DeepNet [WMD+22]的方法。实现基于TorchScale [MWH+22]。我们使用512个AMD MI200 GPU来训练模型。
3.2.与Transformer的对比
语言建模:如图5所示,我们针对基于Transformer和RetNet的语言模型在验证集上报告了困惑度。我们使用三种模型大小(即1.3B、2.7B和6.7B)展示了规模曲线。RetNet在性能上与Transformer相当。更重要的是,结果表明RetNet在大小扩展方面更具优势。除了性能,RetNet在我们的实验中训练非常稳定。实验结果表明,对于大型语言模型,RetNet是Transformer的一个强有力的竞争对手。经验证实,当模型大小大于2B时,RetNet开始优于Transformer。我们还在附录B中总结了具有不同上下文长度的语言建模结果。
零样本和少样本在下游任务上的评估。我们还在广泛的下游任务上比较了语言模型。我们使用6.7B模型进行了零样本和4样本的学习评估。如表3所示,数据集包括HellaSwag (HS) [ZHB+19]、BoolQ [CLC+19]、COPA [WPN+19]、PIQA [BZB+20]、Winograd、Winogrande [LDM12]和StoryCloze (SC) [MRL+17]。准确率与图5中的语言建模困惑度一致。在零样本和上下文学习设置中,RetNet在性能上与Transformer相当。
表3:基于Transformer和RetNet进行零样本学习和少样本学习。模型大小为6.7B。 表4:Transformer(Trm)、带有FlashAttention的Transformer(Trm+FlashAttn)和RetNet的训练成本。我们报告了内存消耗和训练吞吐量(每秒处理的单词数;wps)。
3.3.训练代价
如表4所示,我们比较了Transformer和RetNet的训练速度和内存消耗,其中训练序列长度为8192。我们还与FlashAttention [DFE+22]进行了比较,FlashAttention通过重新计算和内核融合来提高速度并减少GPU内存IO。相比之下,我们使用纯粹的PyTorch代码实现了RetNet,并将内核融合或类似FlashAttention的加速留待未来工作。我们使用基于块的递归保留表示,如方程(7)所述。块大小设置为512。我们使用八个Nvidia A100-80GB GPU进行评估,因为FlashAttention对A100进行了高度优化。对于6.7B和13B模型,启用了张量并行。
实验结果表明,在训练过程中,RetNet比Transformer更节约内存,并且具有更高的吞吐量。即使与FlashAttention相比,RetNet在速度和内存成本方面仍然具有竞争力。此外,由于不依赖于特定的内核,RetNet在其他平台上的高效训练变得更加容易。例如,我们在具有良好吞吐量的AMD MI200集群上训练RetNet模型。值得注意的是,RetNet通过高级实现(如内核融合)有进一步降低成本的潜力。
3.4.推断代价
如图6所示,我们比较了Transformer和RetNet在推断过程中的内存消耗、吞吐量和延迟。Transformer重复使用先前解码的令牌的KV缓存。RetNet使用了方程(6)中描述的递归表示。我们在实验中使用A100-80GB GPU对6.7B模型进行评估。图6显示,RetNet在推断成本方面优于Transformer。
内存消耗 如图6a所示,由于KV缓存,Transformer的内存消耗呈线性增长。相比之下,RetNet的内存消耗即使对于长序列也保持一致,需要更少的GPU内存来容纳RetNet。RetNet的额外内存消耗几乎可以忽略不计(约为3%),而模型权重占据了97%的内存。
吞吐量 如图6b所示,随着解码长度的增加,Transformer的吞吐量下降。相比之下,通过利用保留的递归表示,RetNet在解码过程中具有更高且与长度无关的吞吐量。
延迟 延迟是部署中的一个重要指标,它极大地影响用户体验。我们在图6c中报告了解码延迟。实验结果显示,增加批量大小会使Transformer的延迟增加。此外,随着输入长度的增加,Transformer的延迟增长更快。为了使延迟可接受,我们必须限制批量大小,这会损害Transformer的整体推断吞吐量。相比之下,RetNet的解码延迟优于Transformer,并且在不同的批量大小和输入长度下基本保持不变。
图6:使用6.7B模型的Transformer和RetNet的推断成本。RetNet在内存消耗、吞吐量和延迟方面优于Transformer。
(a) Transformer和RetNet的GPU内存消耗。
(b) Transformer和RetNet的吞吐量。
(c) 不同批量大小下的推断延迟。
3.5.与Transformer变体的对比
除了Transformer之外,我们还将RetNet与各种高效的Transformer变体进行了比较,包括线性变换器 [KVPF20]、RWKV [PAA+23]、H3 [DFS+22]和Hyena [PMN+23]。所有模型都有200M的参数,16层,隐藏维度为1024。对于H3,我们将头维度设置为8。对于RWKV,我们使用TimeMix模块替代自注意力层,同时保持FFN层与其他模型一致,以进行公平比较。我们使用0.5M标记的批量大小进行了10k步的训练。大多数超参数和训练语料库与第3.1节中保持一致。
表5报告了领域内验证集和其他领域外语料库(例如,Project Gutenberg 2019-2022 (PG22) [SDP+22]、QMSum [ZYY+21]、GovReport [HCP+21]、SummScreen [CCWG21, SSI+22])上的困惑度数据。总体而言,RetNet在不同的数据集上表现优于先前的方法。RetNet不仅在领域内语料库上取得了更好的评估结果,还在几个领域外数据集上获得了更低的困惑度。这种优越的性能使RetNet成为Transformer的一个强有力的继任者,除了显著降低成本的好处(第3.3和3.4节)。
此外,我们还讨论了比较方法的训练和推断效率。假设d表示隐藏维度,n表示序列长度。在训练过程中,RWKV的令牌混合复杂度为O(dn),而Hyena的复杂度为O(dn log n),使用快速傅里叶变换进行加速。上述两种方法通过使用逐元素操作来减少训练FLOPS,以换取建模能力。与之相比,保留模型的分块递归表示复杂度为O(dn(b + h)),其中b是分块大小,h是头维度,我们通常将b设置为512,h设置为256。对于大型模型(即较大的d)或序列长度,额外的b + h对性能影响可以忽略不计。因此,RetNet的训练效率非常高,同时不会牺牲建模性能。在推断过程中,与其他高效架构相比,Hyena与Transformer具有相同的复杂度(即每步为O(n)),而其他方法可以进行O(1)的解码。
表5:语言建模的困惑度结果。RetNet在领域内评估集和各种领域外语料库上均优于其他架构。
表6:领域内和领域外语料库上的消融实验结果。
3.6.消融研究
我们对RetNet的各种设计选择进行了消融实验,并在表6中报告了语言建模的结果。评估设置和指标与第3.5节相同。
架构 我们消融了方程(8)中描述的swish gate和GroupNorm。表6显示上述两个组件提高了最终性能。首先,门控模块对于增强非线性和改进模型能力至关重要。需要注意的是,在去除门控后,我们使用了与Transformer相同的参数分配。其次,保留中的分组归一化平衡了多头输出的方差,提高了训练稳定性和语言建模结果。
多尺度衰减 方程(8)表明我们使用不同的γ作为保留头部的衰减率。在消融研究中,我们检查了去除γ衰减(即“- γ decay”)和在所有头部上应用相同的衰减率(即“- multi-scale decay”)。具体来说,消除γ衰减等效于γ = 1。在第二种设置中,我们将所有头部的γ设置为127/128。表6表明,衰减机制和使用多个衰减率都可以提高语言建模性能。
头部维度 从方程(1)的循环透视来看,头部维度意味着隐藏状态的记忆容量。在消融研究中,我们将默认头部维度从256降低到64,即查询和键的维度为64,值的维度为128。我们保持隐藏维度dmodel不变,因此头部的数量增加了。表6中的实验结果显示,较大的头部维度能够实现更好的性能。
4.总结
在本文中,我们提出了适用于序列建模的保留网络(RetNet),它能够实现多种表示,包括并行、循环和分块循环。与Transformer相比,RetNet在推理效率(记忆、速度和延迟),训练并行化和性能方面取得了显著的改进。上述优势使RetNet成为大型语言模型的理想替代品,特别是考虑到O(1)推理复杂度带来的部署优势。在未来,我们希望在模型大小[CDH+22]和训练步骤方面扩展RetNet。此外,通过压缩长期记忆,保留可以与结构化提示[HSD+22b]有效配合工作。我们还将使用RetNet作为骨干架构来训练多模态大型语言模型[HSD+22a, HDW+23, PWD+23]。此外,我们对在各种边缘设备上部署RetNet模型(如手机)很感兴趣。
A.超参数
B.不同上下文长度的分组结果
如表8所示,我们报告了不同上下文长度下的语言建模结果。为了使数据可比较,我们使用2048个文本块作为评估数据,并仅计算最后128个标记的困惑度。实验结果表明,RetNet在不同上下文长度下的性能优于Transformer。此外,RetNet可以利用更长的上下文来获得更好的结果。
表8:使用不同上下文长度的RetNet和Transformer的语言建模困惑度。结果显示,RetNet在序列长度方面具有一致的优势。