添加链接
link之家
链接快照平台
  • 输入网页链接,自动生成快照
  • 标签化管理网页链接

摘要: 大型Transformer模型是自然语言处理最新进展的核心。然而,这些模型的训练和推理成本增长迅速,变得昂贵得令人望而却步。在这里,我们的目标是通过寻找更高效的Transformer来降低Transformer的成本。与以前的方法相比,我们的搜索是在更底层的级别上执行的,在定义Transformer TensorFlow程序的原语上执行。我们确定了一个名为Primer的架构,它比原始Transformer和其他用于自回归语言建模的变体具有更小的训练成本。Primer的改进主要归因于两个简单的修改:平方ReLU激活和在自注意的每个Q、K和V投影后添加深度卷积层。实验表明,Primer相对于Transformer的增益随着计算规模的增加而增加,并且在最佳模型规模下质量遵循幂律(d follow a power law with respect to quality at optimal model sizes.)。我们还从经验上验证了Primer可以放入不同的代码库,从而显著加快训练速度,而无需进行额外的调优。例如,在500M的参数大小下,Primer改进了C4自回归语言建模的原始T5架构,将训练成本降低了4倍。此外,降低的训练成本意味着Primer需要更少的计算来达到目标one-shot性能。例如,在类似GPT-3 XL的1.9B参数配置中,Primer使用1/3的训练计算来实现与Transformer相同的one-shot性能。我们在T5中公开了我们的模型和一些比较,以帮助提高再现性: https://github.com/google-research/google-research/tree/master/primer

1    引言

在过去几年中,Transformer[1]在许多NLP进展中被广泛使用(例如[2,3,4,5,6,7])。随着规模的扩大,Transformer的性能越来越好[3,7,8,9],但训练大型模型的成本变得昂贵得令人望而却步。

本文旨在降低Transformer语言模型的训练成本。为此,我们建议通过修改其TensorFlow计算图[10],寻找更有效的Transformer替代方案。给定TensorFlow程序的搜索空间,我们使用evolution[11,12,13,14,15,16,17]来搜索在给定固定训练计算量的情况下实现尽可能低的验证损失的模型。使用TensorFlow程序作为搜索空间的一个优点是,更容易找到优化Transformer的简单底层改进。由于自回归语言建模(LM)的通用性和成功性[18,7,19,20,21],我们将重点放在仅解码器的自回归语言建模(LM)上(注释2:我们在TensorFlow中提供了原语搜索的详细信息,但同样的方法也可以应用于其他深度学习库。)

发现的模型名为Primer(PRIMitives searched transformER),在自回归语言建模方面比常见的transformER变体表现出强大的性能改进。

我们的实验表明,Primer具有以下优点:(1)使用更小的训练成本实现目标质量,(2)在给定固定训练成本的情况下实现更高的质量,(3)使用更小的推理成本实现目标质量。这些优势是强大的,可以跨越模型大小(20M至1.9B参数)、计算规模(10至10^5加速小时)、数据集(LM1B、C4、PG19[22])、硬件平台(TPUv2、TPUv3、TPUv4和V100)、使用默认配置的多个Transformer代码库(Tensor2Tensor、Lingvo和T5)并跨越多个模型族(dense Transformer[1]、sparse mixture-of-experts Switch Transformer[8]和Synthesizers[23])。我们公开这些比较,以帮助我们的结果的再现性。

我们的主要发现是,当控制模型尺寸和质量时,随着训练成本的增加,Primer的计算节省量增加。当使用最佳尺寸的模型时,这些节省遵循关于质量的幂律(follow a power law with respect to quality when using optimally sized models)。为了证明Primer在已建立的训练设置中的节省,我们使用Raffel等人[5]应用于自回归语言建模的精确配置,将500M参数Primer与原始T5架构进行比较。在此设置中,Primer在相同的训练成本下实现了0.9的复杂度改进,并且,在达到与T5基线模型相同的质量时,计算量减小4.2倍。我们进一步证明,在类似于GPT-3 XL[7]的设置中,在1.9B参数下,通过比较Primer与Transformer,将Primer的节约评估变为one-shot评估。在这里,Primer在训练前的困惑和下游的one-shot任务上都达到了与Transformer相似的性能,而训练计算小了3倍。

我们的分析表明,Primer的改进主要归因于两个主要的修改:平方ReLU激活和在自注意的每个Q、K和V投影后添加深度卷积层。这两个修改很简单,可以放入现有的Transformer代码库中,以获得自回归语言建模的显著收益。

2    搜索空间和搜索方法

在TensorFlow程序上搜索:

为了构造Transformer备选方案的搜索空间,我们使用TensorFlow(TF)中的操作。在这个搜索空间中,每个程序定义了自回归语言模型的可堆叠解码器块。给定输入张量X∈ R^{n×d}表示嵌入长度为d、长度为n的序列,我们的程序返回相同形状的张量。堆叠时,其输出表示每个序列位置的下一个token预测嵌入。我们的程序只指定模型架构,没有其他内容。换句话说,输入和输出嵌入矩阵本身以及输入预处理和权重优化不在我们的程序范围内。

图1显示了如何在我们的搜索空间中构造程序。每个程序都是从进化搜索DNA构建的,进化搜索DNA是子程序的索引集合。子程序0是作为执行入口点的MAIN()函数,其他子程序是DNA子程序库的一部分。每个子程序都是一个无长度约束的索引指令数组。指令是具有一组输入参数的操作。操作表示指令执行的函数。每个操作映射到原语词汇表中的TF函数或DNA子程序库中的另一个子程序。原语词汇表由简单的原语TF函数组成,如ADD、LOG和MATMUL(详见附录A.1)。值得强调的是,像自注意这样的高级构建块不是搜索空间中的操作,而是可以从我们的低级操作中构建出来的。DNA的子程序库由附加程序组成,这些程序可以通过指令作为函数执行。每个子程序只能调用子程序库中索引较高的子程序,这样就消除了循环的可能性。

图1:定义解码器模型程序(即自回归语言模型)的DNA概述。每个DNA都有一组子程序,其中子程序0是MAIN()函数的入口点。每个子程序都由指令组成,这些指令被转换为TensorFlow代码行。指令操作映射到原语词汇表中的基本TensorFlow库函数或父DNA子程序之一。使用父指令的参数集填充操作的参数,该参数集包含所有潜在操作参数的值;只会忽略特定操作未使用的参数。


每个指令的参数集包含每个指令操作的潜在参数值列表。参数字段集表示所有操作原语使用的字段的并集:

•    输入1:将用作第一个张量输入的隐藏状态索引。每个隐藏状态的索引是生成它的指令的索引,子程序的输入状态在索引0和1处。SIN是使用此操作的一个示例。

•    输入2:第二个张量输入的索引。这仅用于与张量输入有关的二进制操作。使用此操作的一个示例是ADD。

•    常数:实值常数。使用此选项的操作示例为MAX;C=0的tf.math.maximum(x,C)是表示TransformerReLU激活的方式。

•    维度大小:一个整数,表示使用权重矩阵的转换的输出维度大小。CONV 1X1是使用此功能的一个操作示例,它是Transformer的注意力投射和前馈部分使用的密集投射。关于我们如何利用相对尺寸[13]调整模型尺寸,请参见附录A.2。

通过将每个子程序指令转换为相应的TF代码行(按索引顺序一次转换一行),我们的搜索子程序被转换为TF程序。要创建TF行,将指令操作映射到相应的TF基元函数或DNA子程序,并插入任何相关参数(完整的TF基元词汇表,包括参数映射,请参见附录A.1);其他参数被忽略。最终指令生成的TF张量作为子程序输出。我们不使用TF-Eager,因此构造程序的一个有用特性是,根据TF的原始延迟执行设计,忽略对程序输出没有贡献的无关节点[10]。有关子程序如何转换为TF图的说明,请参见图2,有关TF图如何构造的更多详细信息,包括我们如何处理因果掩蔽,请参见附录A.2。

图2:转换为相应TensorFlow图的程序示例。由于TensorFlow的延迟执行设计,不会执行对程序输出没有贡献的节点。

进化搜索: 我们进化搜索的目标是在搜索空间中找到最高效的架构。为此,我们给每个模型一个固定的训练预算(24 TPUv2小时),并将其适应性定义为Tensor2Tensor[25]中10亿字基准(LM1B)[24]上的困惑。这种方法,我们称之为固定训练预算的隐式效率目标,与以前的架构搜索工作形成对比,这些架构搜索工作的明确目标是在优化效率时减少训练或推理步骤时间[26,27,28,29]。我们的目标不同,因为步长时间和样本效率之间的权衡是隐含的。例如,在我们的搜索中,将步长时间增加一倍,但将采样效率提高三倍的修改是一个很好的修改,因为它最终使架构的计算效率更高。 事实上,我们发现这些修改是最有益的,将ReLUs平方化,并在注意力上增加深度卷积,这些增加了训练步长时间。 然而,它们极大地提高了模型的样本效率,通过大幅减少达到目标质量所需的训练步骤数量,减少了达到目标质量所需的总计算量。

我们使用的搜索算法是带障碍(hurdles)[13]的正则化进化[30]。我们使用第50百分位(hurdles)配置,并将其隔开,以便在每个(hurdles)中投入相等的计算;与具有完整模型评估的相同实验相比,这将搜索成本降低了6.25倍(有关更多详细信息,请参见附录a.3)。此外,我们使用7小时的训练作为全天训练的替代,因为原始Transformer只需7小时的训练就可以解决其24小时训练难题的90%。这将搜索成本进一步降低3.43倍,总计算缩减系数为21.43倍。因此,尽管我们的目标是提高24小时性能,但平均只需1.1小时即可对个体进行评估(更多搜索细节,包括突变细节和超参数,请参见附录A.4)。我们进行搜索∼25K个架构,并对搜索任务中排名前100位的进行再训练,以选择最佳架构。

我们的搜索空间不同于以前的搜索空间(见[31]的架构搜索调查),后者往往存在严重的偏置(编者注:bias),使得随机搜索表现良好(见[32,33,34]的分析)。由于我们的搜索空间没有这种偏置,由于数值不稳定,我们空间中78%的随机程序(长度等于Transformer程序)的训练时间不能超过五分钟。由于这种开放性和大量的退化程序,有必要使用Transformer[13]的副本初始化搜索群体(输入嵌入大小dmodel=512,前馈向上投影大小df=2048,层数L=6)(图3)。要将此初始化应用于搜索空间,我们必须确定如何将Transformer程序划分为子程序。为此,我们按照构成它的机器学习概念进行划分。例如,我们使用常用的实现方法为self-attention、ReLU和layer-norm分别创建一个子程序(完整列表见附录A.5)。我们称这种方法为概念初始化,因为它通过初始化为搜索引入了偏置,同时留下了进化的搜索空间和突变的行动空间。这与之前大量通过搜索空间引入偏置的作品形成了对比。尽管一些研究也探索了类似于我们的微型任务[35]的开放式搜索空间,但我们证明了我们的技术可以扩展到全尺寸的深度学习模式(见第4节)。

图3:比较概念初始化和随机初始化的小规模搜索(10个TPUv2芯片上10小时)。我们的搜索空间是开放的,如果没有强大的初始化,搜索是不可行的。

3    Primer

Primer: 我们将发现的模型命名为Primer,它代表原语(完整程序参见附录图23)。Primer在对搜索任务进行再训练时显示出显著的改进,需要不到一半的Transformer计算就能达到相同的质量(图6)。在第4节中,我们还展示了Primer在迁移到其他代码库、训练机制、数据集和下游one-shot任务时获得同样大的收益。

PrimerEZ: 这项工作的一个核心动机是开发语言建模实践者易于采用的简单技术。为了实现这一点,我们对两个代码库(T5[5]和Tensor2Tensor[25])进行了烧蚀测试, 并确定哪些Primer修改是通用的(附录图26) 。产生最稳健改进的两种方法是平方前馈RELU和向注意力多头投影添加深度卷积(图4)。 我们把只有这两个简单修改的Transformer称为Primer-EZ ;这是我们推荐给对使用Primer感兴趣的语言建模实践者的起点。我们现在解释这些修改,然后测量它们的经验有效性。

图4:给Primer带来大部分收益的两个主要修改:增加注意力多头投射的深度卷积和平方ReLU激活。这些修改很容易实现,并且可以很好地跨代码库迁移。我们把只有这两个简单修改的Transformer称为Primer-EZ。蓝色表示原始Transformer的部分,红色表示我们建议的修改之一。

平方ReLU:

最有效的修改是将Transformer前馈块中的ReLU激活改进为平方ReLU激活。已在神经网络激活函数[36]的背景下研究了不同程度的校正多项式(Rectified polynomials of varying degrees have been studied in the context of neural network activation functions),但并不常用;据我们所知,这是第一次证明这种整流多项式激活在Transformer中有用。有趣的是,高阶多项式[37]的有效性也可以在其他有效Transformer非线性中观察到,例如GLU[38]变体,如ReGLU[39]( y=Ux\odot max(Vx,0)\odot 表示元素相乘element-wise product),以及point-wise激活,像近似GELU[40]( y=0.5x(1+tanh(\sqrt{2/\pi }(x+0.044715x^3) )) )。然而,当x趋近于无穷大时,与最常用的激活功能相比(ReLU、GELU和Swish),平方ReLU具有截然不同的渐近性(图5左侧)。平方ReLU确实与ReGLU有显著重叠,事实上,当ReGLU的U和V权重矩阵相同时,平方ReLU和ReGLU是等效的,并且,平方ReLU是紧接在权重矩阵U线性变换后面。这使我们相信平方ReLU在更简单的同时,也能获得这些GLU变体的好处,无需额外参数,并提供更好的质量(图5右侧)。

图5:左:平方ReLU与其他常见激活函数相比具有截然不同的渐近性。中心:平方ReLU与GLU变体[39]有显著重叠,后者使用ReLU样渐近的激活,如ReGLU和SwiGLU。我们的实验表明,在Transformer语言模型中,平方ReLU优于这些GLU变体。右图:比较在C4自回归LM上训练的Transformer525K步数的不同非线性。

多头深度卷积注意力(MDHA):

另一个有效的修改是在自注意中查询Q、键K和值V的每个多头投影之后添加3x1深度卷积。这些深度卷积是在每个密集投影输出的空间维度上执行的。有趣的是,这种pointwise卷积和深度卷积的顺序与典型的可分离卷积相反,我们在附录A.6中发现典型的可分离卷积顺序不太有效。我们还发现更宽的深度卷积和标准卷积不仅不能提高性能,而且在某些情况下会损害性能。虽然以前的工作[41,42]中,深度卷积已用于Transformer,但在每次密集头部投影之后使用深度卷积并不是最好方法,就我们所知。MDHA类似于卷积注意[43],它使用可分离卷积而不是深度卷积,并且不像我们那样对每个注意头应用卷积运算。

其他修改:

其他不太有效的Primer修改。附录A.5中给出了每次修改的图表,附录A.7中给出了烧蚀研究。我们在此简要介绍这些修改及其有用性:

•    共享Q和K深度表示:Primer共享Q和K的一些权重矩阵。K使用前面描述的MDHA投影创建,而Q是通过Q=KW来创建,W表示可学习权重矩阵W∈ R^{d×d}。我们发现这通常会影响性能。

•    预规范化和后规范化:Transformer的标准实践已将规范化置于自注意和前馈转换之前[44,45]。Primer在自注意之前使用规范化,但在前馈转换之后应用第二个规范化。我们发现这在某些情况下有用,但并非所有情况下都有用。

•    自定义规范化:Primer使用了层规范化的修改版本[46],该版本使用x(x− µ)而不是(x− µ)^2,但我们发现这并不总是有效的。

•    12倍瓶颈投影:发现的模型使用较小的d_model,大小为384(与基线的512相比),较大的d_f f大小4608(与基线的2048相比)。我们发现,这种较大的投影在较小的尺寸下显著改善了结果(∼35M参数),但对于更大的模型效果较差,如前所述[9]。因此,在引用Primer或Primer EZ时,我们不包括此修改。

•    后Softmax空间选通:发现的模型在注意Softmax之后有一组每通道可学习的标量,这改善了固定长度序列的复杂性。然而,这些标量不能应用于可变序列长度,因此我们在实验的Primer中不包括这种修改。

•    无关的修改:有少数额外的修改不会在发现的架构中产生有意义的差异。例如,隐藏状态乘以-1.12。为了验证这些修改既不帮助也不损害质量,我们将它们从正文的讨论中排除,并且在试验Primer时不包括它们。这些无关的修改仍可在附录A.5中找到。

4    结果

在我们的实验中,我们将Primer与三种Transformer变体进行比较:

•    原始Transformer:原始Transformer[1],具有ReLU激活和残差路径外的层规范化[46]。

•    Transformer+GELU:一种常用的vanilla Transformer变体,使用GELU[40]近似激活函数[2,7]。

•    Transformer++:具有以下增强功能的Transformer:RMS归一化[47]、Swish激活[48],以及在前馈反向瓶颈中使用GLU乘法分支[38](SwiGLU)[39]。这些修改经过基准测试,并在T5中证明了是有效的[49]。

我们在三个不同的代码基上进行比较:Tensor2Tensor(T2T)[25]、T5[5]和Lingvo[50]。Tensor2Tensor是我们用于搜索的代码库,因此我们的大多数并排工作都是在T5和Lingvo中完成的,以证明可迁移性。在所有情况下,我们对每个代码基使用默认Transformer hyperparameters,禁用正则化。有关超参数的更多详细信息,请参见附录A.8。

在下面的部分中,我们将在四个主要的自回归语言建模实验中展示我们的结果。首先,我们将展示Primer在搜索任务上优于基线模型。接下来,我们将展示Primer在Transformer上的计算节省与模型质量之间的关系在最佳模型尺寸下遵循幂律。这些节省还可以跨数据集和代码库进行迁移。然后,我们将研究Primer在已建立的训练机制中的收益,并表明它使用完整的compute T5训练,在500M参数大小下可以节省4.2倍的计算。最后,我们将证明这些增益迁移到GPT-3建立的预训练和one-shot下游任务设置[7]。

4.1    搜索任务比较

我们首先分析Primer在搜索任务中的性能:序列长度为64的LM1B语言建模,∼35M型号参数、4096个token批次和24小时训练。我们对比了Tensor2传感器(T2T)[25]和T5[5]以及TPUv2s和V100 GPU上的基线模型。我们根据每种型号达到vanilla Transformer最终质量的速度对其性能进行分级,我们将其称为加速系数。图6显示,在所有情况下,Primer在Transformer上提供1.7倍或更高的加速系数。图6还显示Primer和Primer EZ都可以推广到其他硬件平台和代码库

图6:使用35M参数模型对LM1B搜索任务进行比较。“加速系数”是指每个模型需要达到与24小时训练的普通Transformer同等质量的计算分数。Primer和Primer EZ在所有情况下都能实现超过1.7倍的加速比。请注意,结果在代码库和不同硬件之间迁移。

接下来我们研究了Primer的扩展规律。在这里,我们通过使用这些超参数的排列组合的每一个来训练模型,创建一个从23M到385M的参数范围,从而将Primer与我们的基线进行比较:L∈ {6,9,12}层,d_model∈ {384,512,1024}初始嵌入大小,p∈ {4,8,12}前馈向上投影比率。图7所示的结果证实了之前的观点,即在最佳参数大小下,计算和语言模型质量之间的关系大致遵循幂律[9]。也就是说,验证损失 l 和训练计算 c 之间的关系遵循关系 l=ac^{-k} ,a和k是经验常数。这表示为double log空间中的一条线(图7): \mathrm{log} l=-k\mathrm{log} c+\mathrm{log} a 。但是,对于每种架构,这些线并不相同。这些线大致平行,但上下移动。在附录A.9中,我们表明,给定 \mathrm{log} b^k 的垂直间距,此类平行线表示高级建模的计算节省量 s 也遵循形式为 l=a_1{(1-1/b)}^ks^{-k} ,这背后的直觉是,对于所有的 lb 是常数计算折减系数,因此,与 l 相关的训练计算的幂律也会导致与 l 相关的幂律节约(见附录a.9)。

图7:左:当浏览最优模型尺寸时,Transformer语言模型质量和训练计算之间的关系大致遵循幂律[9]。右图:比较不同Transformer修改的这些幂律线,拟合平滑的MSE损失。模型线彼此平行,这意味着通过使用高级建模计算节省的成本也可以作为具有质量的幂律进行缩放。垂直虚线之间的间距表示计算中的2倍差异。请注意,两个图中的x轴和y轴均为对数。


Primer还具有改进推理的能力,尽管我们的搜索侧重于训练计算。图8显示了使用前馈传递计时作为推断代理时,质量与推断的帕累托前沿比较。我们使用前向传递时间作为推理的代理,因为有多种方法来解码语言模型,每种方法都有不同的计算成本。可以进行更深入的研究,分析Primer在不同解码方法、服务平台、数据集等方面的推理性能,但这超出了本工作的范围。

图8:LM1B上的帕累托最优推断比较。Primer在大多数目标质量上都表现出改进的推理能力。我们观察到这些模型的训练步长和推理时间之间的相关性为0.97。

4.2    向其他代码库、数据集和模型类型的初始可迁移性

我们现在研究Primer在另一个代码库T5中迁移到更大数据集PG19和C4的能力。此外,我们还扩展到一个更高的计算区域,该区域已被先前的研究用作大规模训练的代理[49,5];批次增加到65Ktoken,序列长度更长512,每个解码器有110M参数(dmodel=768,df f=3072,L=12),每个模型训练为∼在4个TPUv3芯片上执行525K步进。我们还将继续对每个模型进行100万步的训练,以研究更大的计算预算对Primer成本节约的影响。图9所示的结果表明,在较大的数据、较高的计算区域中,Primer模型与在较小的LM1B区域中一样强大。与普通基线相比,Primer和Primer EZ在PG19和C4训练结束时的效率至少提高1.8倍。

图9还显示,Primer修改与其他高效模型族兼容,例如Switch Transformer[8]等大型稀疏混合专家和Synthesizer[23]等高效Transformer近似。对于这些实验,我们使用Narang等人提供的T5实现[49]。增加深度卷积和平方ReLUs的Primer EZ技术将Switch Transformer的计算成本降低了1.5倍;当控制计算时,这转化为0.6的复杂性改进(见附录a.10)。向Synthesizer中添加平方ReLUs可将训练成本降低2.0X倍,并在完全训练时将困惑度提高0.7倍。

图9:将Primer转换为较大数据集(C4和PG19)和不同型号系列(Switch Transformer和合成器)在不同的代码库(T5)中的比较,计算量比搜索任务多一个数量级。与普通基线相比,Primer和Primer EZ在PG19和C4训练结束时的效率至少提高1.8倍。在所有情况下,Primer 计算节约的比例都会随着计算投入的增加而增加。Primer EZ改造还改进了Switch Transformer(550M参数)和Synthesizer(145M参数),表明它与其他有效方法兼容。根据每个基线训练525K和1M步所需的时间选择计算预算。具体数字见附录A.10。

4.3    大规模T5自回归语言模型训练

在大规模计算配置中,Primer计算节省率甚至更高。为了证明Primer在已建立的高计算训练设置中的节省,我们将其扩展到完整的T5计算模式,完全复制Raffel等人[5]。这与上一节中的C4配置相同,但使用了批处理∼1Mtoken、64个TPUv3芯片和537M参数(dmodel=1024,df f=8192,L=24)。Primer的计算效率比原始T5模型高4.2倍,比我们的强化Transformer++基线(表1)高2倍。

表1:使用完整的T5计算规模,在537M参数下,C4 LM上达到目标质量的计算使用比较。根据基线模型的最终性能选择目标质量。Primer实现了与原始T5体系结构相同的质量,而计算量小4.2倍。

这里的节省之所以更好,是因为在固定的规模下,更多的计算投资会产生更高的计算节省。图10显示了随着模型训练时间的延长,为了达到与原始T5架构相当的性能,Primer需要的计算的占比会越来越小(编者注:节省量越来越高)(Figure 10 shows how the fraction of compute Primer needs to achieve parity with the original T5 architecture shrinks as the models are trained for longer);这是由于控制和变量复杂度训练曲线的渐近性质造成的。这不同于第A.6节中描述的幂律节约。在这里,我们为每个计算预算使用最佳数量的参数,因此计算节省系数b保持不变。对于固定的模型大小,计算节省因子随着计算投入的增加而增加,这意味着计算节省可能超过幂律估计。注意,这意味着,通过投入比基线模型所需的更多的计算,可以“玩”比较( can be “gamed”),比如这里给出的比较。正是出于这个原因,我们使用了Raffel等人[5]的训练制度的精确复制品:在已经发布的训练配置中展示Primer的节约。

图10:计算C4 LM上Primer与原始T5体系结构随时间的节省。由于两条困惑曲线的渐近性质,在训练中投入的计算越多,节省的成本就越高。Primer达到与原始T5体系结构相同的性能只需小4.2倍的计算量。

4.4    向下游one-shot任务迁移Primer的能力

在我们最后的比较中,我们展示了Primer的改进也适用于预训练和one-shot下游任务迁移机制。语言建模的最新趋势是在大型数据集上训练大型模型,这被称为“预训练”。然后,这些模型被迁移到看不见的数据集和任务中,并且在没有太多或任何额外训练的情况下,展示了在这些“下游”任务中表现良好的能力[2,51]。在我们这里研究的仅解码器的自回归语言建模配置中,GPT-3[7]取得了最令人印象深刻的结果,这表明大型语言模型可以在仅给出一个示例的未知任务上表现出强大的性能,称为“one-shot”学习。在本节中,我们将演示Primer的训练计算节省超出了达到目标预训练的困惑,并确实迁移到下游one-shot任务性能。

为此,我们复制了GPT-3预训练和one-shot评估设置。(注释3:本节中使用的训练数据集和评估管道的开发是其独立的工作。这项工作的全部细节将很快在一份单独的技术报告中公布。)由于GPT-3不是开源的,因此此复制与用于GPT-3的复制不完全相同。因此,这些实验并不意味着直接与GPT-3进行比较,因为存在配置差异。相反,这些实验被用作Transformer和Primer结构的受控比较。我们使用专有的预训练数据集在Lingvo代码库中进行这些实验。下游任务的配置方式与Brown等人[7]所述的相同,每个任务的输入都有一个前缀示例输入到每个模型中。我们将(1)基线1.9B参数Transformer(dmodel=2048,df f=12288,L=24)与GELU激活进行比较,这意味着接近GPT-3 XL架构,以及(2)没有共享QK表示的完整Primer,根据附录A7,这只会损害性能。每一个模型都是使用成批的∼使用512 TPUv4芯片的2Mtoken用于∼140小时(∼71.8K加速器总小时数或∼1M的列车台阶)。我们再次使用T5训练超参数表,无需任何额外调整。

图11显示Primer实现了与Transformer+GELU相同的预训练困惑和one-shot下游性能,同时计算量小3倍。附录中的表6给出了27个评估下游任务中每个任务的准确性能数字。Primer尽管使用的计算量减少了3倍,但在5项任务上优于Transfomer+GELU,在1项任务上表现更差,在其余21项任务上表现相当。同一表格显示,在给定等效计算时,Primer在15项任务上优于Transformer+GELU,在2项任务上表现更差,在其余10项任务上表现相当。这一结果表明,Primer不仅可以改善语言建模的困惑,而且改进还可以迁移到下游NLP任务。

图11:Transformer+GELU和Primer在1.9B参数下的比较,以及下游one-shot任务的不同训练计算预算,类似于GPT-3。当预训练计算量减少3倍时,Primer的性能略好于Transformer,而在相同的预训练计算量下,Primer的性能明显更好。在这里,我们停止在72K TPUv4小时,以大致匹配GPT-3 XL的质量,但如果我们让两个模型运行更长时间,Primer的计算节省将更大(见图10)。注意,这是一个粗略的比较,使用了我们评估的27项one-shot任务的平均分数。关于每个单独任务的准确分数,请参见附录A.11(表6和图27)。

Primer的投入回报率:

这个大规模实验中的计算节省证明了Primer搜索的投入回报。寻找Primer本身就要付出代价∼2.14E+21FLOPs。为本实验训练Transformer的费用∼2.96E+22FLOPs,这意味着Primer为达到相同性能而保存的计算是∼1.98E+22FLOPs。因此,对于这个单一的训练,架构搜索的投入回报大约是9.24倍。请注意,搜索成本是one-shot成本,Primer可以在将来的训练中重用,以节省更多的计算时间。有关能源成本和碳排放估算的更多详细信息,请参见附录A.13。

5    结论

限制:

这项研究有局限性。首先,我们的模型参数扫描大约比Kaplan等人原始研究中进行的扫描小一个数量级[9]。同样,尽管我们的大型模型使用了大量计算,但它们仍然比最先进的模型(如全尺寸GPT-3)小几个数量级[7]。

另一个限制是,我们主要关注仅解码器模型,而仅编码器[2,3,4]和编码器-解码器序列模型[52,1,6]仍然广泛使用。在附录A.12中,我们在T5中进行了编码器-解码器屏蔽语言建模比较,但没有对结果进行深入研究。这里的主要发现是,尽管Primer修改改进了vanilla Transformer,但它们的性能仅与Transformer++一样好。这一结果表明,对于仅适用于解码器的自回归语言模型的架构修改可能不一定对基于编码器的屏蔽语言模型有效。开发一个同样适用于屏蔽语言模型的架构是我们未来研究的主题。

实践讨论:

这项工作的主要动机是对Transformer进行简单实用的修改,使其易于采用。为此,我们回答了从业人员可能提出的一些问题:

•    在所有设置中,Primer训练计算节省是否都相同?不。通过我们自己提供的实验,Primer可以节省不同的计算量。这是因为计算深度取决于硬件细节深度学习lib库操作速度、特定任务的模型采样效率以及其他可能在设置中不同的因素。我们使用T5训练的精确副本来演示在已建立的配置(4.2X)中的节约效果,但预期结果会因配置而异。

•    Primer能否改善BERT[2]?这项工作的重点是自回归语言建模的具体任务,随着GPT-3的发展,它被证明对传统NLP应用程序以及生成应用程序都很重要。我们只简要研究了Primer在屏蔽语言建模和编解码器模型中的应用(附录A.12)。我们的研究表明,虽然Primer改进了vanilla Transformer,但并不明显优于Transformer++。因此,适用于自回归语言建模的修改可能不如适用于掩蔽语言建模的修改有效。未来的工作可能会调查,是否可以以更有效的方式将Primer修改集成到编码器-解码器和仅编码器模型中,从而改进像BERT这样的模型。未来的工作还可以应用本文描述的搜索方法来寻找更好的基于编码器的屏蔽语言模型。

•    是否需要重新调整超参数配置以使用Primer?我们的目的是使Primer修改不需要任何额外的超参数调整。为此,在我们的实验中,我们没有调整任何超参数,而是使用来自已建立库的Transformer超参数。不过,如果进行额外的调整,Primer可能会工作得更好。

•    PrimerEZ是否优于Primer?在我们的对比实验中,我们发现PrimerEZ有时比T5代码库中的Primer更好。然而,在应用于其他代码库(如Lingvo和T2T)时,我们发现完整的Primer可以比PrimerEZ提供更好的性能。因此,我们建议从业者首先尝试使用Primer EZ,以便于实现,如果他们有兴趣实现进一步的收益,则可以继续实现完整的Primer。

建议和未来方向:

我们建议采用Primer和Primer EZ进行自回归语言建模,因为它们具有强大的性能、简单性以及对超参数和代码库更改的鲁棒性。为了证明它们的潜力,我们简单地将它们放入已建立的代码库中,并且在没有任何更改的情况下,表明它们可以显著提高性能。此外,在实践中,额外的调优可以进一步提高它们的性能。

我们也希望我们的工作能鼓励更多研究开发高效Transformer。例如,本研究的一个重要发现是,激活函数的微小变化可以产生更高效的训练。在降低Transformer成本的努力中,对此类简单变化的开发进行更多投入可能是未来探索的一个有希望的领域。

致谢:

We thank Zhen Xu for his help with infrastructure. We also thank Gabriel Bender, Hallie Cramer,Andrew Dai, Nan Du, Yanping Huang, Daphne Ippolito, Norm Jouppi, Lluis-Miquel Munguia, Sharan Narang, Ruoming Pang, David Patterson, Yanqi Zhou, and the Google Brain Team for their help and feedback.

A     Appendix

A.1 TensorFlow Primitives Vocabulary

表2:TensorFlow(TF)原语词汇表。“Name”是搜索空间中操作的名称。“TF函数”是当DNA指令转换为一行TensorFlow代码时,名称映射到的TensorFlow函数。“参数映射”描述了DNA参数集中的值如何映射到相应的TensorFlow函数参数。该词汇表主要由创建Transformer所需的最低级别TF操作构成(见附录A.5)。我们还将这些操作扩展到包括相邻操作;例如,我们将MAX扩展为包含MIN,将RED SUM扩展为包含RED乘积,将CONV 1X1扩展为包含CONV 3X1。我们还添加了常用的数学原语,如SIN和ABS。

A.2    构造张量流图

TensorFlow图是根据DNA程序构建的,如正文第2节所述。在这里,我们提供了附加的实现细节。

相对尺寸:

对于每条指令的“维度大小”参数,我们使用相对维度[13]而不是绝对维度。这允许我们调整模型的大小,以适应我们的参数限制(32M到38M参数)。这些相对维度的词汇是[1,2,4,8,12,16,24,32,48,64]。这个词汇表没有调整。

Values Bank:

对于“常量”和“维度大小”参数字段,我们创建每个指令引用的共享值库。常量组包含2个值,尺寸大小组包含6个值;这些数字没有调整。不是每条指令都拥有各自的参数值,而是拥有这些共享库的索引。这允许多条指令共享相同的值,并在该值更改时同时更改。例如,Q、K和V的每个单独注意力多头投影开始共享相同的输出维度大小,因此如果该值改变,它们都会同时改变。有关这些银行价值如何变化的示例,请参见A.4。

因果掩蔽:

教师强制语言模型训练的一个重要部分是,位置无法“看到”他们试图预测的标志。每个位置只能从以前的位置获取信息,否则在不提供目标的情况下,模型将退化。为了加强这种因果约束,我们为空间移动信息的操作增加了额外的开销,以屏蔽来自未来位置的任何信息。例如,当应用卷积时,我们遵循将输入空间移动(核宽度-1)的标准实践, 因此,每个位置仅接收来自先前位置的信息。

分支:

为了启用Transformer搜索种子的多头功能,我们在指令中添加了一个称为“分支”的元参数。该参数可以接受[1,2,4,8,16]中的任何值,并确定该指令并行执行的次数,结果张量沿嵌入轴连接在一起。分支可以用于任何TensorFlow原语以及任何DNA子程序。这允许我们通过将子程序1(自注意)分支8次(子程序实现见附录A.5)来初始化具有多头自注意的搜索。Primer没有以任何有意义的方式利用这个分支功能,除了使用初始化的多头部注意。

解决维度不匹配:

我们不限制张量维度如何变异,因此程序可能无效,因为它们对大小不兼容的张量执行二进制操作。例如,一个程序可以描述将两个具有不同嵌入大小的张量相加。为了解决这些维度不匹配的问题,我们确定地伪随机设置一个张量维度来匹配另一个张量维度。

A.3    将hurdles减半

我们根据身体状况来设置跨栏(hurdles)[13],使前50%的人通过每一个跨栏。我们以这样一种方式来划分跨栏,即在搜索结束时,用于训练每个跨栏栏带的预期计算量大致相等。也就是说,考虑到我们个人的最大训练计算量为7小时或25200秒,我们在812.9秒、2438.7秒、5690.3秒和12193.5秒的分数上构建跨栏。因此,计算预算的1/5用于训练每个人达到第一个障碍(812.9秒),计算预算的1/5用于训练∼50%的人从第一栏训练到第二栏(2438.7秒− 812.9s=1625.8s),计算预算的1/5用于训练∼25%的人从第二栏训练到第三栏(5690.3秒− 2438.7s=3251.6s),等等。这种配置策略,我们称之为“将跨栏减半”,只需要设置一个超参数,即跨栏数,并且不再需要设置跨栏阈值和比较步骤,正如之前所做的那样[13,35]。我们选择四个栏位,因为五个栏位要求第一个栏位在不到十分钟的训练时间内固定,根据经验我们发现这是一个过于嘈杂的信号。以这种方式使用跨栏将每个模型的平均列车时间减少到4064秒或约1小时8分钟,从而将计算成本减少了一倍∼6.2X。

该策略与bandit算法(如连续减半[53]和Hyperband[54])不同,但我们不使用预先创建的静态个体群体,而是将我们的减半与不断变化的进化群体相结合

图12:将初级搜索的障碍减半。每个点代表进化产生的个体的最终适应度。因为每个跨栏有不同的训练津贴,所以形成了不同的“乐团”。随着时间的推移,所有等级都有所改善,这意味着所有计算津贴的中位数适合度都有所改善。一个模型在不同培训预算下的表现之间的这种相关性使我们能够将总搜索成本降低大约6.2倍。

A.4    进化搜索详细信息

我们使用正则化进化[30],种群规模为100,a tournament selection size为10。这些值没有调整。我们使用的突变如下。

突变:

为了在我们的搜索中创建新的候选者,我们从我们的搜索群体中随机选择一个父代,并对其应用单一突变。我们采用五种不同的突变类型(除非另有规定,否则选择和决定是统一随机执行的):

•    删除 从子程序中删除指令。

•    插入:创建指令并将其插入子程序。

•    删除和插入:先执行删除突变,然后执行插入突变[55]。

•    变异字段:从指令中选择一个字段并更改其值。

•    交换:在随机选择的子程序中交换两条指令的位置。每个指令的输入张量也被交换,因此净效应是切换指令在计算图中的位置。

•    改变组值:改变相应组中相对张量尺寸或常数的值。相对张量尺寸的值从其词汇表中选择(见附录A.2)。常数值根据cnew进行更改:=cprev·10X+Y,用于先前值cprev、新值cnew和随机变量X,Y∼ N(0,1)。

在应用了一个变异之后,我们运行一个轻检查,看看结果候选者的计算图是否与父代的计算图完全等价。如果是,我们进行另一次突变。

A.5    Transformer和Primer程序比较

这里,我们介绍了Transformer种子和发现的Primer模型的程序。表3是将操作名称映射到后续图形的图形符号的键。图13至图22描述了每个模型的子程序,其中以橙色突出显示了Primer更改。图23描述了每个模型的完整计算图,所有子程序解析为其组成原语。图24和25描述了Transformer和Primer的DNA程序,解析了所有子程序,插入了所有指令库值。

表4:使用两个不同硬件平台(TPUv2s和V100 GPU)和两个不同库(Tensor2Tensor和T5)的默认超参数,对LM1B上的搜索任务、自回归语言建模进行比较。此表包含图6的精确数字。“加速”描述了每个模型使用的计算分数,以实现与使用完整计算预算训练的基线相同的结果。尽管底漆是使用TPUv2s在Tensor2传感器中开发的,但它在GPU和T5上显示出强大的性能。对于每个库的默认标记化,都会报告困惑。

A.7    消融和插入研究

这项工作的核心动机之一是开发简单而健壮的Transformer修改。为此,我们研究了每种Primer修改的单独有效性,如正文第3节所述。我们使用插入和消融研究来衡量这种有效性。在插入研究中,我们将每个修改单独添加到普通Transformer中。在烧蚀研究中,我们每次从Primer中去除一种修改。我们对这些修改如何影响搜索库Tensor2Tensor以及其他库的性能感兴趣。因此,我们在另一个库T5中进行了这些插入和消融研究,并将修饰可迁移性作为我们建模建议的关键指导指标。

这些研究的结果如图26所示。“规范化PPLX增量”描述了修改对性能的帮助或影响程度。对于基线复杂度Pb和修改复杂度Pm,“标准化PPLX Delta”定义为Pb−插入研究中的Pm Pb和Pm−Pb用于消融研究。这些定义不同,因此正值总是表示修改是好的,负值总是表示修改是坏的。三种技术在所有场景中都是有益的。第一种是“12X proj”,它在控制参数的同时增加了Transformer前馈向上投影的大小。我们发现这对于较小的模型很有效,但对于较大的模型却没有用处。第二个是MDHA和squared ReLUs,它们是Primer EZ的定义性修改,Primer EZ是一个更简单的模型,捕获了完整Primer的大部分收益。

图26:LM1B上Primer改性的可转移性研究∼35M参数。在“插入研究”中,我们将每个修改插入到一个普通Transformer中。在“烧蚀研究”中,我们去除了Primer上的每一种修改。“标准化PPLX Delta”表示经过处理的模型受这些修改影响的程度;标准化值,使其在代码库中具有可比性,因此正值表示两项研究中的有益技术。同样,负值表示两项研究中的有害技术。

A.8    完整的训练细节

在所有的实验中,我们都使用了以前发布的超参数设置,这些设置是为Transformer进行调整的,正则化被禁用,并且没有为Primer进行额外的调整。在Tensor2传感器(T2T)中,这些是TRANSFORMER_TPU超参数,在T5和Lingvo中,这些是先前T5研究中使用的开源参数[5,49]。它们都指定了一个Adafactor优化器[56],在学习率为0.01的情况下具有10K预热步骤,然后是倒数平方根学习率衰减。T2T使用位置嵌入和子词标记化,而T5和Lingvo使用相对注意[57]和句子片段[58]。

对于LM1B,我们使用T2T默认设置,最大序列长度为64,批次为4096个token;这是合适的,因为LM1B的平均序列长度约为32。对于C4和PG19,我们使用T5默认的最大序列长度512。对于单次预训练,我们使用1024的最大序列长度。在第4.2节中,我们使用65Ktoken的批次,在第4.3节中,我们使用1Mtoken的批次,在第4.4节中,我们使用2Mtoken的批次。

A.9    幂律计算储蓄推导

在正文的第4.1节中,我们再现了Kaplan等人[9]的结果,并表明,在最佳参数大小下,语言模型质量和训练计算之间的关系遵循幂律: l=ac^{-k} ,其中l是验证损失,c是训练计算,a和k是经验常数。这表示为double log space中的一线(图7): \mathrm{log} l=-k\mathrm{log} c+\mathrm{log} a 。然而,对于我们所比较的每种架构,这些线并不相同。这些线大致平行,但上下移动。因此,将两个架构线之间的偏移定义为 \mathrm{log} b^k ,我们可以推导出它们的训练成本关系如下:

表5:较大数据集上的语言建模比较,将Primer转换为T5代码库。在这种转移制度下,Primer改善了所有基线。此外,PrimerEZ不仅达到与Primer同等的水平,而且在某些情况下超过了Primer。Switch Transformer和 Synthesizer 也从EZ修改中受益。计算预算比较点是根据达到525K和1M训练步骤所需的基线时间选择的。关于句子,给出了困惑。此表包含图9的精确数字。

A.11 各个one-shot任务的表现

图27:Transformer+GELU和Primer在1.9B参数下在1/3和完全预训练预算下的下游one-shot任务之间的比较。根据独立t检验,使用5个训练前体重检查点的样本,提供95%的置信区间。当给定1/3的预训练计算时,Primer的性能与Transformer大致相同,而在给定相同的预训练计算时,Primer在大多数任务上的性能更强。具体数字见表6。

A.12    掩蔽语言建模

编码器-解码器风格屏蔽语言建模(MLM)不是这项工作的重点。然而,因为它是原始T5项目的重点,我们在这里包括MLM比较以确保完整性(表7)。具体而言,我们使用了Narang等人[49]使用的精确比较配置,他们对几种Transformer变体进行了基准测试;一个区别是,我们只进行一次模型训练,因为这一制度不是我们研究的重点。对于“Primer EZ解码器”,我们使用Transformer++编码器和Primer EZ解码器。我们的处理表明,Primer EZ修改有能力改进编码器-解码器MLM模型,但与Transformer++相比,改进的程度可能更小。我们相信,这表明仅解码器LM和编码器-解码器MLM从不同的建模决策中受益——这可以在未来的工作中进行研究。我们还相信,直接在编码器-解码器MLM上运行搜索可以产生对该任务更有利的修改。

表7:T5中C4与编码器-解码器样式模型的屏蔽语言建模比较。这些结果与Narang等人[49]的配置完全相同,尽管我们只运行一次模型,因为MLM不是我们工作的重点。*指示从该研究中获取的行

A.13    碳排放估计数

根据Patterson等人[59]的建议,我们发布了我们最大实验的碳排放估算。

为了估算建筑搜索的碳排放量4,5,我们利用Patterson等人的测量结果。他们对建筑搜索的碳排放量估算为1360天TPUv2使用量的3.2 MTCO2e[59]。在这里,我们使用1145.8天的TPUv2计算进行搜索。此外,在我们搜索时,数据中心6的PUE为1.08而不是1.10,其净碳强度平均值为0.336 MTCO2e/MWh而不是0.431 MTCO2e/MWh。7,8因此,我们架构搜索实验的比例排放量估计为3.2 MTCO2e∗ 一千一百四十五点八一三六零∗ 1.08 1.10 ∗ 336431=2.06 MTCO2e。作为比较,从旧金山到纽约的单程旅客往返机票是∼1.2 MTCO2e[59]因此我们的搜索成本大约为1.72张这样的机票。

我们遵循与Patterson等人的测量相同的过程来估算我们的大规模T5实验的排放量。Patterson等人对11B参数T5的排放量估计为10249天TPUv3使用的46.7Tco2e。我们的T5型号更小,因此平均只需要687.5 TPUv3天的训练。我们进行了3次训练(Primer、original T5和T5++),以展示Primer在基线上的改进,共产生2062.5 TPUv3天。当我们进行实验时,数据中心9的PUE为1.10而不是1.12,其净碳强度平均值为0.540 MTCO2e/MWh而不是0.545 MTCO2e/MWh。因此,这些T5模型训练的比例总估算值为46.7 MTCO2e∗ 2062.5 10,249 ∗ 1.10 1.12 ∗ 540 545=8.54 MTCO2e。

为了估算Lingvoone-shot预训练的排放量,我们采用与Patterson等人相同的方式测量系统平均功率[59]。包括内存、网络接口、风扇和主机CPU,每个TPUv4芯片的平均功率为343W。我们使用与Patterson等人相同的方程来计算

我们两次大规模预训练的二氧化碳当量:2∗ 343W∗ 71800h∗ 1.08(PUE)∗ 0.055 MTCO2e/MWh=29.26 MTCO2e.10

我们的大规模T5和one-shot比较的排放成本高于架构搜索本身的成本。我们对这些大规模的比较进行了投资,以证明我们的有效修改的潜在节约。例如,正文第4.4节中所述的Transformer上使用Primer的节约量等于9.75 MTCO2e,仅此一项是∼4.7倍架构搜索的成本。注意,硬件设置的差异会影响这些节省。例如,在俄克拉何马州训练one-shot模型,与进行Primer搜索的格鲁吉亚相比,俄克拉何马州具有良好的MTCO2e/MWh。从触发器的角度来看,为了消除这些硬件特定因素,Primer在one-shot实验中节省的成本是搜索本身成本的9.24倍,如正文第4.4节所述。因此,即使在我们相对较小的比较规模下,架构搜索也会产生投资回报,这大约比全尺寸GPT-3小100倍[7]。

A.14    与Evolved Transformer的比较

这项工作建立在Evolved Transformer[13]的基础上,Transformer还试图通过架构搜索发现改进的序列模型。表4中LM1B上的T5和表5中C4提供了与Evolved Transformer架构的计算效率比较。表8给出了这些实验中与Evolved Transformer结构的样本效率比较。在本节中,我们将讨论这些比较,以及它们如何突出我们的Primer搜索相对于Evolved Transformer搜索的改进。