在看apple 最近发布的OpenELM 模型,其论文中提到 block-wise scaling 模型结构优化方法,(论文见: OpenELM: An Efficient Language Model Family with Open-source Training and Inference Framework),这里记录下DeLighT论文中的 block-wise scaling,翻译整理下以便对照代码实现,了解背景和原理。DeLighT论文中的实验任务主要是在两个标准的序列建模任务上评估了DeLighT的性能:机器翻译(machine translation)任务 encoder-decoder architecture 和 语言建模( language modeling)decoder architecture,论文中机器翻译任务未对En-Zh(英文译中文)进行实验,可以作为一个复现练习,根据源码实操一下论文中的实验;而语言建模可以作为openELM的来源延伸~ 结合cornet进行复现(也有mxl示例,mxl针对Apple Silicon 硬件进行的优化深度学习框架)。

论文主作者:Sachin Mehta

论文地址:https://arxiv.org/pdf/2008.00623

论文代码: https://github.com/sacmehta/delight (基于当时facebook的 fairseq seq2seq工具库开发)

该论文研究是在作者以前的DeFINE: DEep Factorized INput Token Embeddings for Neural Sequence Modeling 进行改进,模型结构引入更多的GLTs,来学习更宽的权重,并且减少了参数数量。

摘要

我们介绍了一种深度且轻量级的Transformer,名为DeLighT,它在参数数量显著减少的情况下,提供了与标准基于Transformer的模型相似或更好的性能。DeLighT更有效地在每个Transformer块内部(通过DeLighT变换)以及跨块(通过块级缩放)分配参数,允许输入端使用较浅较窄的DeLighT块,输出端使用较宽较深的DeLighT块。总体而言,DeLighT网络比标准Transformer模型深2.5到4倍,但参数和运算量更少。在基准机器翻译和语言建模任务上的实验表明,DeLighT在平均参数数量减少2到3倍的情况下,达到或提高了基线Transformer的性能。

1 引言

基于注意力的Transformer网络(Vaswani et al., 2017)广泛用于序列建模任务,包括语言建模和机器翻译。为了提高性能,模型通常通过增加隐藏层的维度(变得更宽)或堆叠更多的Transformer块(变得更深)来进行扩展。例如,T5(Raffel et al., 2019)使用了65K的维度,而GPT-3(Brown et al., 2020)使用了96个Transformer块。然而,这种扩展显著增加了网络参数的数量(例如,T5和GPT-3分别拥有110亿和1750亿参数),并使学习过程复杂化,即这些模型要么需要非常大的训练语料库(Raffel et al., 2019; Devlin et al., 2019; Brown et al., 2020),要么需要仔细的正则化(Hinton et al., 2012; Wan et al., 2013; Merity et al., 2018a)。在本文中,我们介绍了一种新的参数高效注意力基础架构,可以轻松地扩展为宽和深。 我们的深度轻量级Transformer架构,DeLighT,扩展了Vaswani et al.(2017)的Transformer架构,并在参数和运算量显著减少的情况下提供了相似或更好的性能。DeLighT的核心是DeLighT变换,它使用Mehta et al.(2018)的组线性变换(GLTs)和扩展-缩减策略,有效地变化DeLighT块的宽度和深度。由于GLTs本质上是局部的,DeLighT变换使用特征洗牌,类似于卷积网络(Zhang et al., 2018)中的通道洗牌,以在不同组之间共享信息。这种宽和深的表示促进了用单头注意力和轻量级前馈层替代Transformer中的多头注意力和前馈层,从而减少了总网络参数和运算量。重要的是,与Transformer不同,DeLighT变换将深度和宽度与输入大小解耦,允许我们通过在输入端使用较浅较窄的DeLighT块,在输出端使用较深较宽的DeLighT块,更有效地跨块分配参数。 我们证明了DeLighT模型在两个常见的序列建模任务上,即(i)机器翻译和(ii)语言建模,实现了与Transformer模型相似或更好的性能,同时显著减少了参数和运算量。在资源较少的WMT’16 En-Ro机器翻译数据集上,DeLighT使用2.8倍更少的参数达到了Transformer的性能。在资源充足的WMT’14 En-Fr数据集上,DeLighT以1.8倍更少的参数提供了比基线Transformer更好的性能(+0.4 BLEU分数)。同样,在语言建模方面,DeLighT在WikiText-103数据集上与Transformer-XL(Dai et al., 2019)的性能相当,但参数数量减少了1.5倍。

2 相关工作

改进Transformer(Improving transformers):已经引入了几种方法来改进Transformer架构。第一条研究线索解决了在长输入序列上计算自注意力的挑战(Child et al., 2019; Kitaev et al., 2020; Beltagy et al., 2020)。这些方法可以与我们的架构结合使用。第二条研究线索集中在解释多头注意力(Raganato and Tiedemann, 2018; Brunner et al., 2020)。他们表明,增加Transformer头的数量可能会导致冗余表示(Voita et al., 2019a; Michel et al., 2019),并且使用具有预定义模式的固定注意力头(Raganato et al., 2020)或合成注意力矩阵(Tay et al., 2020)可以提高性能。第三条研究线索集中在通过学习更好的表示来改进Transformer(Wu et al., 2019; 2020; So et al., 2019)。这些工作旨在通过使用不同的变换来提高Transformer的表达能力——例如,使用卷积(Wu et al., 2019; Gehring et al., 2017)、门控线性单元(Dauphin et al., 2017)或多分支特征提取器(So et al., 2019; Wu et al., 2020)。我们的工作属于这一类。与之前的工作不同,我们展示了如何有效地在块级别使用DeLighT变换以及跨块使用块级缩放来分配参数。

模型缩放(Model scaling):模型缩放是提高序列模型性能的标准方法(Vaswani et al., 2017; Raffel et al., 2019; Lan et al., 2020; Devlin et al., 2019; Shoeybi et al., 2019; Tan and Le, 2019; Brown et al., 2020)。在宽度缩放中增加模型维度(Vaswani et al., 2017; Devlin et al., 2019),而在深度缩放中堆叠更多的块(例如,Transformer块)(Shoeybi et al., 2019; Brown et al., 2020; Wang et al., 2019)。在这两种情况(及其组合)中,网络中每个块内的参数是相同的,这可能导致次优解决方案。为了进一步提高序列模型的性能,本文引入了块级缩放,允许不同大小的块并有效地在网络中分配参数。我们的结果表明,(1)输入端的DeLighT块较浅较窄,输出端的DeLighT块较深较宽,能提供最佳性能,以及(2)与单独的模型缩放相比,结合块级缩放的模型缩放实现了更好的性能。我们注意到,卷积神经网络(CNNs)也在输入端学习较浅较窄的表示,在输出端学习较深较宽的表示。与CNNs(例如,He et al. 2016的ResNet)在每个卷积层执行固定数量的操作不同,所提出的块级缩放在每层和块中使用可变数量的操作。

改进序列模型(Improving sequence models):最近还有大量关于改进序列模型的相关方法的工作,包括(1)使用更好的Token级表示提高准确性——例如,使用BPE(Sennrich et al., 2016)、自适应输入(Baevski and Auli, 2019)和输出(Grave et al., 2017a),以及DeFINE(Mehta et al., 2020),以及(2)提高效率——例如,使用压缩(Chen et al., 2018; Sun et al., 2020)、剪枝(Han et al., 2016; Voita et al., 2019b)和蒸馏(Hinton et al., 2015; Sanh et al., 2019)。与我们的工作最接近的是DeFINE变换,它也使用扩展-缩减策略学习表示。DeFINE变换(图1c)和DeLighT变换(图1d)之间的关键区别在于,DeLighT变换在扩展和缩减层中更有效地分配参数。与DeFINE不同,后者在组线性变换中使用较少的组来学习更宽的表示,DeLighT变换使用更多的组以更少的参数学习更宽的表示。DeLighT变换实现了与DeFINE变换相当的性能,但参数数量显著减少。

3 DeLighT: 深度轻量级Transformer

image-20240428123509253
图1. (a, b)标准Transformer块与DeLighT块之间的块级别比较。在DeLighT中,计算注意力的操作数量减少了一半,而FFN中的参数数量(和操作)减少了$16\times$。具有可学习参数的变换(Linear(天空蓝) 和 DeLighT(地上草绿色))显示为彩色。线性变换的形状指示其操作(扩展、减少等)。(c, d) 将DeFINE变换与DeLighT进行了比较。与DeFINE变换相比,DeLighT使用了更多组的组线性变换(group linear transformations GLTs)来学习更宽的权重,并减少了参数数量。不同颜色用于显示GLTs中的组。为简单起见,(d)中未显示特征洗牌。

标准的Transformer块(图1a)由多头注意力组成,使用查询-键-值分解来建模序列标记之间的关系,并且由一个前馈网络(FFN)来学习更宽的表示。多头注意力通过对输入应用三个投影来获取查询$\mathbf{Q}$、键$\mathbf{K}$和值$\mathbf{V}$,每个投影由$h$个线性层(或头)组成,将$d_m$维输入映射到$d_h$维空间,其中$d_h=d_m/h$是头的维度。FFN由两个线性层组成,其中第一个将维度从$d_m$扩展到$d_f$,第二个将维度从$d_f$减少到$d_m$。Transformer块的深度为4,包括(1)用于query、key和value的三个并行分支,(2)将多个头的输出组合的融合层,以及(3)FFN中的两个顺序线性层。一般来说,基于Transformer的网络通过顺序堆叠Transformer块来增加网络的容量和深度。

本文扩展了Transformer架构,引入了一个深度轻量级的Transformer,DeLighT。我们的模型使用深度轻量级的扩展-减少变换,DeLighT(第3.1节),有效地实现了学习更宽的表示(representations)。它还允许用单头注意力和轻量级FFN(第3.2节)替换多头注意力和前馈网络(FFN)层。DeLighT将注意力维度与深度和宽度解耦,使我们能够使用块级别缩放而不是均匀堆叠Transformer块来有效地学习表示(第3.3节)。

3.1 DeLighT

DeLighT将一个$d_m$维输入向量映射到一个高维空间(扩展),然后使用Mehta等人(2018)的$N$ 层组变换将其减少到一个$d_o$维输出向量(缩减),如图1 d所示。在这些扩展和缩减阶段,DeLighT变换使用组线性变换(GLTs),因为它们通过从输入的特定部分派生输出来学习局部表示,并且比线性变换更高效。为了学习全局表示,DeLighT变换使用特征洗牌,在组线性变换中不同组之间共享信息,类似于卷积网络中的通道洗牌(Zhang等人,2018)。

增加Transformer的表达能力和容量的标准方法是增加输入维度$d_m$。然而,线性增加$d_m$也会线性增加标准Transformer块中多头注意力的操作数量($\mathcal{O}(n^2 d_m)$,其中$n$是序列长度)(图1 a)。相比之下,为了增加DeLighT块的表达能力和容量,我们通过扩展和减少阶段增加其中间DeLighT变换的深度和宽度。这使得我们可以使用较小的维度计算注意力,从而减少操作数量。

形式上,DeLighT由五个配置参数控制:(1)GLT层数$N$,(2)宽度倍增器$w_m$,(3)输入维度$d_m$,(4)输出维度$d_o$,以及(5)GLT中的最大组数$g_{max}$。在扩展阶段,DeLighT使用$\lceil \frac{N}{2} \rceil$层将$d_m$维输入线性投影到高维空间$d_{max}=w_m d_m$。在减少阶段,DeLighT使用剩余的$N - \lceil \frac{N}{2} \rceil$ GLT层将$d_{max}$维向量线性投影到$d_o$维空间。数学上,我们定义每个GLT层$l$的输出$\mathbf{Y}$为:

$$ \mathbf{Y}^l = \left\{ \begin{array}{ll} \mathcal{F} \left(\mathbf{X}, \mathbf{W}^l, \mathbf{b}^l, g^l\right), & l = 1 \\ \mathcal{F} \left( \mathcal{H}\left(\mathbf{X}, \mathbf{Y}^{l-1}\right), \mathbf{W}^l, \mathbf{b}^l, g^l\right), & \text{Otherwise} \end{array} \right. \label{eq:glt} $$

其中 $\mathbf{W}^l = \lbrace \mathbf{W}^l_1, \cdots, \mathbf{W}^l_{g^l}\rbrace$ 和 $\mathbf{b}^l = \lbrace \mathbf{b}^l_1, \cdots, \mathbf{b}^l_{g^l}\rbrace$ 是第$l$层GLT的可学习权重和偏置,其中$g^l$是第$l$层GLT的组数。简言之,$\mathcal{F}$函数接受输入$\mathbf{X}$(或$\mathcal{H}\left(\mathbf{X}, \mathbf{Y}^{l-1}\right)$)并将其分成$g^l$个不重叠的组,使得$\mathbf{X} = \lbrace\mathbf{X}_1,\cdots,\mathbf{X}_{g^l}\rbrace$ 。

然后,$\mathcal{F}$ 函数线性地变换每个$\mathbf{X}_i$,得到输出$\mathbf{Y}^l_i = \mathbf{X}_i \mathbf{W}_i^l + \mathbf{b}_i^l$。每个组的输出$\mathbf{Y}^l_i$然后被串联起来,产生输出$\mathbf{Y}^l$。函数$\mathcal{H}$首先对$\mathbf{Y}^{l-1}$中每个组的输出进行洗牌,然后使用Mehta等人(2020)的输入混合器连接将其与输入$\mathbf{X}$组合起来,以避免梯度消失问题。图2 可视化了DeLighT变换中的扩展阶段,其中包括组线性变换、特征洗牌和输入混合器连接。

在DeLighT中,第$l$个GLT的组数计算为: $$ \begin{equation} g^l = \left\{ \begin{array}{lr} \text{min}(2^{l-1}, g_{max}), & 1 \leq l \leq \lceil{N/2}\rceil \\ g^{N-l}, & \text{Otherwise} \end{array} \right. \end{equation} $$ 在我们的实验中,我们使用$g_{max} = \lceil \frac{d_m}{32} \rceil$,以便每个组至少有32个输入元素。

image-20240428132940873
图2.示例说明了DeLighT变换中的扩展阶段,该阶段使用GLTs、特征洗牌和输入混合器连接,以有效地学习更深和更宽的表示(representations)。为了说明,我们使用了相同的输入和输出维度。

3.2 DeLighT Block/Layer

图1 b展示了我们如何将DeLighT集成到Transformer块中以提高其效率。首先将$d_m$维输入馈送到DeLighT中,产生$d_o$维输出,其中$d_o < d_m$。然后将这些$d_o$维输出馈送到单头注意力中,接着是一个轻量级FFN来建模它们之间的关系。

DeLighT层和单头注意力(DeLighT layer and single head attention): 假设我们有一个长度为$n$的输入标记序列,每个标记的维度为$d_m$。首先将这$n$个$d_m$维输入馈送到DeLighT中,产生$n$个$d_o$维输出,其中$d_o < d_m$。然后使用三个线性层同时对这些$n$个$d_o$维输出进行投影,得到$d_o$维的查询$\mathbf{Q}$、键$\mathbf{K}$和值$\mathbf{V}$。然后我们使用缩放点积注意力(公式3)对这些n个Token之间的上下文关系进行建模。为了启用残差连接(He等人,2016),这个注意力操作的$d_o$维输出被线性投影到一个$d_m$维的空间。 $$ \begin{equation} \text{Attention}(\mathbf{K}, \mathbf{Q}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_o}}\right) \mathbf{V} \label{eq:espda} \end{equation} $$ 我们假设DeLighT学习更宽的表示(representations)的能力使我们能够用单头注意力替换多头注意力。标准Transformer中计算注意力的计算成本为$\mathcal{O}(d_m n^2)$,而DeLighT块中为$\mathcal{O}(d_o n^2)$,其中$d_o < d_m$。因此,DeLighT块将计算注意力的成本减少了一个因子$d_m/d_o$。在我们的实验中,我们使用$d_o=d_m/2$,因此与Transformer架构相比,需要$2\times$更少的乘加操作。

轻量级(Light-weight) FFN :与Transformer中的FFN类似,这个块也由两个线性层组成。由于DeLighT块已经通过DeLighT集成了更宽的表示,它允许我们反转Transformer中FFN层的功能。第一层将输入的维度从$d_m$减少到$d_m/r$,而第二层将维度从$d_m/r$扩展到$d_m$,其中$r$是缩减因子(见图1b)。我们的轻量级FFN通过因子$r d_f/d_m$减少了FFN中的参数和操作数量。在标准Transformer中,FFN的维度扩展了4倍(Transformer-base使用$d_m$=512和$d_f$=2048,而Transformer-large使用$d_m$=1024和$d_f$=4096)。 在我们的实验中,我们使用$r=4$。因此,轻量级FFN将FFN中的参数数量减少了$16\times$。

block/layer深度:DeLighT块堆叠了(1)具有 $N$个GLT的DeLighT ,(2)三个并行的用于Key、Query和Value的线性层,(3)一个投影层,和(4)轻量级FFN的两个线性层。因此,DeLighT块的深度为$N+4$。与标准Transformer块(深度为4)相比,DeLighT块更深。

3.3 Block/Layer wise scaling

提高序列模型性能的标准方法包括增加模型维度(宽度缩放)、堆叠更多的块(深度缩放)或两者兼有。然而,这种缩放在小型数据集上并不是非常有效。例如,当在WMT’16 En-Ro语料库上将Transformer-Base($d_m$ = 512)网络替换为Transformer-Large($d_m$ = 1024)时,参数数量增加了大约4倍,而性能并没有显著变化(BLEU:34.28 vs. 34.35)。我们假设这是因为模型宽度和深度的缩放是在块之间均匀分配参数,这可能导致学习冗余参数。为了创建深而宽的网络,我们将模型缩放扩展到块级(见图3)。

image-20240428134732165
图3: 块级缩放(block-wise scaling) 有效地分配参数和操作到各个块中,导致接近输入的DeLighT块更浅更窄,接近输出的DeLighT块更深更宽。在(b)中,具有统一缩放($N=N_{min}=N_{max}=8$)和块级缩放($N_{min}=4$,$N_{max}=8$)的DeLighT网络拥有约16.7百万个参数并执行35亿次操作(针对序列长度$n=30$计算),然而具有块级缩放的DeLighT网络交付了更好的困惑度,提高了2个点。

缩放DeLighT块: DeLighT块使用DeLighT学习深度和宽度可变的表示,其深度和宽度由两个配置参数控制:GLT层的数量$N$和宽度乘数$w_m$(图3a)。这些配置参数允许我们独立于输入$d_m$和输出$d_o$维度增加DeLighT块内的可学习参数数量。这种校准在标准Transformer块中是不可能的,因为它们的表达能力和容量是输入的函数(输入维度=头数$\times$头维度)。在这里,我们介绍了块级别的缩放,它创建了一个具有可变大小DeLighT块的网络,将浅层和窄层的DeLighT块分配给输入附近,而将深层和宽层的DeLighT块分配给输出附近。

为此,我们引入了两个网络范围的配置参数:最小$N_{min}$和最大$N_{max}$ GLT的数量。对于第$b$个DeLighT块,我们使用线性缩放计算如下公式中的GLT数量$N^b$和宽度乘数$w_m^b$。通过这种缩放,每个DeLighT块具有不同的深度和宽度(见图3a)。 $$ N^{b} = N_{min} + \frac{(N_{max} - N_{min})\ b}{\mathcal{B}-1}, \quad w_m^{b} = w_m + \frac{(N_{max} - N_{min})\ b}{N_{min}(\mathcal{B}-1)}, \quad 0 \le b \le \mathcal{B}-1 $$ 这里,$\mathcal{B}$表示网络中DeLighT块的数量。我们在GLT层数$N$和宽度乘数$w_m$上添加上标$b$,以表示这些参数是针对第$b$个块的。

网络深度: Transformer块的深度是固定的,即4。因此,先前的工作(Raffel et al., 2019; Brown et al., 2020; Wang et al., 2019)已经将基于Transformer的网络的深度与Transformer块的数量联系起来。在DeLighT中,我们提出了一种不同的视角来学习更深层次的表示,其中每个块的大小是可变的。为了计算网络深度,我们使用不同领域(包括计算机视觉,例如He等人2016年的ResNet,以及理论机器学习,例如Telgarsky, 2016)的标准定义。这些工作将网络深度定义为连续可学习层的数量(例如,卷积、线性或组线性)。类似地,具有$\mathcal{B}$个块的DeLighT和Transformer网络的深度分别为$\sum_{b=0}^{\mathcal{B}-1} (N^b + 4)$和$4\mathcal{B}$。

4 实验结果

我们在两个标准的序列建模任务上评估了DeLighT的性能:(1)机器翻译(第4.1节)和(2)语言建模(第4.2节)。

4.1 机器翻译(machine translation)

数据集和评估:我们在四个数据集上对DeLighT模型进行了基准测试:(1)IWSLT’14 德语-英语(De-En),(2)WMT’16 英语-罗马尼亚语(En-Ro),(3)WMT’14 英语-德语(WMT’14 En-De),以及(4)WMT’14 英语-法语(WMT’14 En-Fr)。对于IWSLT’14 De-En数据集,我们复制了Wu等人(2019)和Edunov等人(2018)的设置,使用160K/7K/7K句子对进行训练、验证和测试,以及大约10K个令牌的联合BPE词汇表。对于WMT’14 En-De数据集,我们遵循了Vaswani等人(2017)的设置。数据集包含3.9M/39K/3K句子对,分别用于训练、验证和测试,以及44K的联合BPE词汇表大小。对于WMT’14 En-Fr数据集,我们复制了Gehring等人(2017)的设置,使用36M/27K/3K句子对进行训练、验证和测试,以及44K的联合BPE词汇表大小(我们使用的是与Tensor2Tensor库(Vaswani等人,2018)兼容的训练和验证数据,以便与最近的工作(例如,进化Transformer)进行公平比较)。性能是根据测试集上的BLEU分数(Papineni等人,2002)(分数越高越好)来评估的。我们遵循Wu等人(2019)的束搜索相关超参数设置。

架构:我们采用了Vaswani等人(2017)的对称编码器-解码器架构,并使用了正弦位置编码。编码器和解码器都包含了$\mathcal{B}$个DeLighT块。解码器块与编码器块相同(见图1b),不同之处在于它们在轻量级FFN之前有一个额外的源-目标单头注意力单元。在源-目标单头注意力单元中,键和值是对编码器输出的投影(详细内容见附录A)。在我们的实验中,我们对WMT’16 En-Ro、WMT’14 En-De和WMT’14 En-Fr使用了$w_m=2$, $N_{min}=4$, 和 $N_{max}=8$ ;这导致了222层深的DeLighT网络。对于IWSLT’14 De-En,我们使用了$w_m=1$, $N_{min}=3$, 和 $N_{max}=9$;这导致了289层深的网络。为了简化,我们设置了$\mathcal{B}=N_{max}$。我们使用了一个可学习的查找表,将词汇表中的每个Token映射到一个128维的向量。我们使用Fairseq(Ott等人,2019)实现了我们的模型,并使用了他们提供的脚本进行数据预处理、训练和评估。

训练:对于IWSLT’14 De-En模型,我们遵循了Wu等人(2019)的设置,并在单个NVIDIA GTX 1080 GPU上,以4K令牌的批量大小对所有模型进行了50K次迭代的训练。对于WMT’16 En-Ro,我们遵循了Ghazvininejad等人(2019)的训练设置,并在16个NVIDIA Tesla V100 GPU上进行了100K次迭代的训练,有效批量大小为64K令牌。对于WMT’14 En-De和WMT’14 En-Fr,我们遵循了Wu等人(2019)的训练设置,并在16个V100 GPU上分别进行了30K和50K次迭代的训练。我们使用Adam(Kingma和Ba,2015)在训练期间最小化交叉熵损失,并使用了标签平滑值为0.1。为了公平比较,我们训练基线Transformer模型时使用了相同的训练设置。

4.1.1 结果

与基线Transformer的比较:表1比较了DeLighT与Vaswani等人(2017)在不同语料库上的基线Transformer的性能。DeLighT在不同语料库上都以更少的参数实现了更好的或相似的性能。具体来说,在资源较少(WMT’16 En-Ro)和资源丰富(WMT’14 En-De & WMT’14 En-Fr)的语料库上,DeLighT分别以2.8倍和1.8倍更少的参数实现了相似或更好的性能。当参数数量增加时,DeLighT的性能超过了Transformer。例如,在WMT’14 En-Fr数据集上,DeLighT比Transformer深3.7倍,提高了1.3个BLEU分数,但参数数量少了1300万,操作数量少了30亿(见表2)。

小型语料库结果

WMT’14 En-De WMT’14 En-Fr
模型 参数数量 比率 BLEU BLEU变化 参数数量 比率 BLEU BLEU变化
Transformer (Vaswani et al., 2017) 34.4$^\dagger$ 62 M 34.3$^\ddagger$
Transformer (我们的实现) 42 M $1.0\times$ 34.3 62 M $1.0\times$ 34.3
DeLighT 14 M $0.3\times$ 33.8 -0.5 22 M $0.35\times$ 34.3 0.0
DeLighT 30 M $0.7\times$ 35.3 +1.0 53 M $0.85\times$ 34.7 +0.4

大型语料库结果

WMT’14 En-De WMT’14 En-Fr
模型 参数数量 比率 BLEU BLEU变化 参数数量 比率 BLEU BLEU变化
Transformer (Vaswani et al., 2017) 62 M 27.3 62 M 38.1
Transformer (我们的实现) 67 M $1.0\times$ 27.7 67 M $1.0\times$ 39.2
DeLighT 37 M $0.55\times$ 27.6 -0.1 37 M $0.55\times$ 39.6 +0.4
DeLighT 54 M $0.80\times$ 28.0 +0.3 54 M $0.80\times$ 40.5 +1.3

表1:与基线Transformer在机器翻译语料库上的比较。DeLighT模型需要显著较少的参数来达到类似的性能。这里,$^\dagger$和$^\ddagger$分别表示Wu et al. (2019) 和 Ghazvininejad et al. (2019),中报告的最佳Transformer基线。


模型 深度 参数数量 MAC数 BLEU
Transformer 60 67 M 11.1 B 39.2
DeLighT 222 37 M 5.6 B 39.6
DeLighT 222 54 M 8.1 B 40.5

表2:与Transformer相比,DeLighT网络是深度、轻量级且高效的。BLEU分数是在WMT’14 En-Fr数据集上报告的。为了计算网络深度,我们统计了网络中顺序层的数量(第3.3节)。我们使用了20个源Token和20个目标Token来计算乘加操作(MACs)。详细信息见附录C。


特别值得关注的是在WMT'14 En-De语料库上(如图4),DeLighT与Vaswani等人(2017年)的基准Transformer以及其神经搜索变体——So等人(2019年)的Evolved Transformer在两种不同参数设置下的性能比较。对于小型模型(小于1千万参数),DeLighT模型提供了更好的性能,并且要达到与这些模型相同的性能水平,DeLighT模型需要更少的参数

image-20240428183808019
图4:在WMT'14 En-De语料库上,DeLighT与Transformer和Evolved Transformer在两种不同设置下的比较:(1)参数数量相同,(2)性能相同。

与最先进方法的比较:大多数最先进的方法都在WMT’14 En-De上评估了性能,一些也在IWSLT’14 De-En上进行了评估。表3比较了DeLighT与这些语料库上最先进方法的性能。DeLighT实现了与现有方法相似或更好的性能。重要的是要注意,现有方法通过不同的设计选择改进了基线Transformer——例如,不对称的编码器-解码器结构(Wang et al., 2019)和神经架构搜索(So et al., 2019)。我们相信,DeLighT将来也会从这些设计选择中受益。

IWSLT’14 De-En
模型 参数数量 BLEU
Transformers (Vaswani et al., 2017) 42 M 34.3
Variational Attention (Deng et al., 2018) 33.1
Dynamic convolutions (Vaswani et al., 2017) 43 M 35.2
Lite Transformer$^\ddagger$ (Wu et al., 2020) 33.6
DeLighT(我们的模型) 30 M 35.3
WMT’14 En-De
模型 参数数量 BLEU
Transformer (Vaswani et al., 2017) 62 M 27.3
DLCL (Wang et al., 2019) 62 M 27.3
Evolved Transformer $^\dagger$ (So et al., 2019) 46 M 27.7
Lite Transformer$^\ddagger$ (Wu et al., 2020) 26.5
DeLighT(我们的模型) 37 M 27.6

表3:与机器翻译语料库上的最新方法进行比较。DeLighT模型在参数较少的情况下提供了与最新模型相似或更好的性能。这里,$^\dagger$表示网络使用神经架构搜索(NAS),$^\ddagger$ 表示未报告完整的网络参数。

DeLighT模型的扩展:图5 显示了随着网络参数增加,DeLighT模型的性能提高;表明它们能够在不同语料库上学习表示,包括资源较少的语料库。

image-20240428184459900
图5: DeLighT模型的规模扩展。随着网络参数数量的增加,DeLighT在不同语料库(包括低资源语料库WMT'16 En-Ro)上的性能都有所提高。

4.2 语言建模(language modeling)

数据集和评估:我们在WikiText-103数据集上进行评估(Merity et al., 2017),该数据集有103M/217K/245K个Token用于训练、验证和测试。它有一个大约260K个Token的词级词汇表。按照最近的工作(Baevski和Auli,2019;Dai等人,2019),我们在测试集上以困惑度(越低越好)报告性能。

架构:我们使用Baevski和Auli(2019)的基于Transformer的解码器架构,带有B个DeLighT块。我们使用$w_m$=$2$, $N_{min}$=$4$, 和 $N_{max}$=$12$。我们使用自适应输入(Baevski和Auli,2019)作为查找表,并将自适应输出(Grave et al., 2017a)作为分类层,具有一个头(头维度为128)和两个尾部(尾部维度分别为64和32)。我们还共享了输入和输出层之间的权重。

训练:我们遵循Baevski和Auli(2019)的训练设置,只是我们所有的模型都在8个NVIDIA Tesla V100 GPU上进行了100K次迭代的训练,上下文长度为512,有效批量大小为64K令牌。我们使用Adam进行训练,并在测试期间使用480的上下文长度。

结果:表4b比较了DeLighT与之前方法在WikiText-103上的性能。表4a绘制了DeLighT和Transformer-XL(Dai et al., 2019)的困惑度随参数数量变化的图表——后者超过了其他基于Transformer的实现(例如,Baevski和Auli 2019)。两个表格都表明,DeLighT以更少的参数提供了比最先进的方法(包括Transformer-XL)更好的性能,并且使用了更短的上下文长度,表明DeLighT变换有助于学习强大的上下文关系。

image-20240428193117942
表4:在WikiText-103数据集上的结果。与Transformer-XL相比,DeLighT在参数较少的情况下提供了相似或更好的性能(更低的困惑度)。† 对于Transformer-XL,我们使用官方源代码复现结果。对于使用上下文长度为480评估Transformer-XL,我们在官方评估脚本中将mem_len超参数设置为480。

5 计算效率的分析与讨论

训练时间和内存消耗:表5比较了DeLighT与基线Transformer在训练时间和内存消耗方面的情况。为了进行公平比较,我们实现了一个没有NVIDIA专用CUDA内核的Transformer单元,并在16个NVIDIA V100 GPU上对全精度网络进行了30K次迭代的训练。Transformer和DeLighT模型的训练时间分别约为37小时和23小时,GPU内存消耗分别约为12.5 GB和14.5 GB(R1 vs. R2)。当我们启用APEX库3为多头注意力提供的专用CUDA内核时,Transformer模型的训练时间从37小时减少到16小时,而我们没有观察到内存消耗的显著变化。受到这一观察结果的启发,我们为GLT中的分组和反分组函数实现了专用的CUDA内核(见附录E)。这些更改使得DeLighT的训练时间和GPU内存消耗分别减少了约4小时和3 GB。我们强调,分组、线性变换、特征洗牌和反分组可以通过单个CUDA内核高效实现。未来,我们期望为这些操作提供专用的CUDA内核,这将进一步减少内存消耗以及训练/推理时间。

Row # Model # Params (百万) BLEU (WMT'14 En-Fr) 训练时间 内存 (GB)
R1 Transformer (未优化) 67 M 39.2 37小时 12.5
R2 DeLighT(未优化) 54 M 40.5 23小时 14.5
R3 Transformer (使用Apex优化) 67 M 39.2 16小时 11.9
R4 DeLighT(使用优化的分组) 54 M 40.5 19小时 11.5

表5:与基线Transformer在训练速度和内存消耗方面的比较。在R4中,我们仅为分组和取消分组功能实现了CUDA内核(见附录E)。我们预计DeLighT将更有效率,因为它有一个专用的CUDA内核用于分组、转换、特征洗牌和取消分组。内存消耗是在一台单个的NVIDIA GP100 GPU(16GB内存)上测量的,每个批次最多有4096个标记,并且没有任何梯度积累。


正则化:表6显示,DeLighT在参数数量较少的情况下,与基线Transformer相比,需要更少的正则化就能达到相似的性能。这表明,使用更好的变换函数学习表示可以减少对dropout的需求。

Model Dropout BLEU
Transformer (62 M) 0.10 27.3
Transformer (62 M) 0.30 27.7
DeLighT(37 M) 0.05 27.6

表6:与基线Transformer相比,DeLighT需要更少的正则化(数据集:WMT'14 En-De)

6 结论

本文介绍了一种深度且轻量级的Transformer架构,DeLighT,它在DeLighT块内部以及跨DeLighT块有效地分配参数。与最先进的Transformer模型相比,DeLighT模型(1)深度且轻量级,并且(2)提供了相似或更好的性能。未来,我们计划将DeLighT应用于其他任务,包括语言模型预训练、问答和语言生成。

致谢:本研究得到了ONR N00014-18-1-2826、DARPA N6600119-2-403、NSF (IIS-1616112, IIS1252835)以及Allen杰出研究者奖的支持。作者还要感谢华盛顿大学UW-NLP和H2Lab的成员提供的宝贵反馈和评论。

参考文献

  1. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, 和 Illia Polosukhin. “Attention is all you need.” 在《神经信息处理系统进展》中,第5998–6008页,2017年。

  2. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, 和 Peter J. Liu. “Exploring the limits of transfer learning with a unified text-to-text transformer.” arXiv预印本 arXiv:1910.10683, 2019.

  3. Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, 和 Dario Amodei. “Language models are few-shot learners.” arXiv预印本 arXiv:2005.14165, 2020.

  4. Jacob Devlin, Ming-Wei Chang, Kenton Lee, 和 Kristina Toutanova. “BERT: Pre-training of deep bidirectional transformers for language understanding.” 在《北美计算语言学协会2019年会议论文集:人类语言技术》中,第1卷(长论文和短论文),2019年。

  5. Geoffrey E. Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, 和 Ruslan R. Salakhutdinov. “Improving neural networks by preventing co-adaptation of feature detectors.” arXiv预印本 arXiv:1207.0580, 2012.

  6. Li Wan, Matthew Zeiler, Sixin Zhang, Yann LeCun, 和 Rob Fergus. “Regularization of neural networks using dropconnect.” 在《国际机器学习会议》中,第1058–1066页,2013.

  7. Stephen Merity, Nitish Shirish Keskar, 和 Richard Socher. “Regularizing and optimizing LSTM language models.” 在《学习表示国际会议》中,2018a. URL: https://openreview.net/forum?id=SyyGPP0TZ.

  8. Sachin Mehta, Rik Koncel-Kedziorski, Mohammad Rastegari, 和 Hannaneh Hajishirzi. “Pyramidal recurrent unit for language modeling.” 在《自然语言处理的实证方法会议》中,2018年。

  9. Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, 和 Jian Sun. “Shufflenet: An extremely efficient convolutional neural network for mobile devices.” 在《IEEE计算机视觉与模式识别会议》中,第6848–6856页,2018.

  10. Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, 和 Ruslan Salakhutdinov. “Transformer-xl: Attentive language models beyond a fixed-length context.” 在《计算语言学协会》中,2019年。

  11. Rewon Child, Scott Gray, Alec Radford, 和 Ilya Sutskever. “Generating long sequences with sparse transformers.” arXiv预印本 arXiv:1904.10509, 2019.

  12. Nikita Kitaev, Lukasz Kaiser, 和 Anselm Levskaya. “Reformer: The efficient transformer.” 在《学习表示国际会议》中,2020年。

  13. Iz Beltagy, Matthew E. Peters, 和 Arman Cohan. “Longformer: The long-document transformer.” arXiv:2004.05150, 2020.

  14. Alessandro Raganato 和 Jörg Tiedemann. “An analysis of encoder representations in transformer-based machine translation.” 在《2018 EMNLP BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP》研讨会论文集中,2018年。

  15. Gino Brunner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, 和 Roger Wattenhofer. “On identifiability in transformers.” 在《学习表示国际会议》中,2020年. URL: https://openreview.net/forum?id=BJg1f6EFDB.

  16. Elena Voita, Rico Sennrich, 和 Ivan Titov. “The bottom-up evolution of representations in the transformer: A study with machine translation and language modeling objectives.” 在《2019年自然语言处理和计算语言学第九届国际联合会议(EMNLP-IJCNLP)》论文集中,2019a。

  17. Paul Michel, Omer Levy, 和 Graham Neubig. “Are sixteen heads really better than one?” 在《神经信息处理系统进展》中,第14014–14024页,2019年。

  18. Alessandro Raganato, Yves Scherrer, 和 Jörg Tiedemann. “Fixed encoder self-attention patterns in transformer-based machine translation.” arXiv预印本 arXiv:2002.10260, 2020.

  19. Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, 和 Che Zheng. “Synthesizer: Rethinking self-attention in transformer models.” arXiv预印本 arXiv:2005.00743, 2020.

  20. Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, 和 Michael Auli. “Pay less attention with lightweight and dynamic convolutions.” 在《学习表示国际会议》中,2019年。

  21. Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, 和 Song Han. “Lite transformer with long-short range attention.” 在《学习表示国际会议》中,2020年。

  22. David So, Quoc Le, 和 Chen Liang. “The evolved transformer.” 在《第36届国际机器学习会议论文集》中,第5877–5886页,2019年。

  23. Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, 和 Yann N. Dauphin. “Convolutional sequence to sequence learning.” 在《第34届国际机器学习会议论文集-第70卷》中,第1243–1252页。JMLR. org, 2017.

  24. Yann N. Dauphin, Angela Fan, Michael Auli, 和 David Grangier. “Language modeling with gated convolutional networks.” 在《第34届国际机器学习会议论文集-第70卷》中,第933–941页。JMLR. org, 2017.

  25. Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, 和 Radu Soricut. “Albert: A lite bert for self-supervised learning of language representations.” 在《学习表示国际会议》中,2020年。

  26. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, 和 Bryan Catanzaro. “Megatron-lm: Training multi-billion parameter language models using gpu model parallelism.” arXiv预印本 arXiv:1909.08053, 2019.

  27. Mingxing Tan 和 Quoc V. Le. “Efficientnet: Rethinking model scaling for convolutional neural networks.” 在Kamalika Chaudhuri和Ruslan Salakhutdinov编辑的《第36届国际机器学习会议论文集》中,第9-15页,2019年6月9-15日,加利福尼亚州长滩,美国。

  28. Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F. Wong, 和 Lidia S. Chao. “Learning deep transformer models for machine translation.” 在《第57届计算语言学协会年会论文集》中,2019年。

  29. Kaiming He, Xiangyu Zhang, Shaoqing Ren, 和 Jian Sun. “Deep residual learning for image recognition.” 在《IEEE计算机视觉与模式识别会议》中,第770–778页,2016年。

  30. Rico Sennrich, Barry Haddow, 和 Alexandra Birch. “Neural machine translation of rare words with subword units.” 在《第54届计算语言学协会年会论文集》第1卷(长论文)中,8月2016年。

  31. Alexei Baevski 和 Michael Auli. “Adaptive input representations for neural language modeling.” 在《学习表示国际会议》中,2019年。

  32. Édouard Grave, Armand Joulin, Moustapha Cissé, David Grangier, 和 Hervé Jégou. “Efficient softmax approximation for GPUs.” 在《国际机器学习会议》中,2017a。

  33. Sachin Mehta, Rik Koncel-Kedziorski, Mohammad Rastegari, 和 Hannaneh Hajishirzi. “DeFINE: Deep Factorized Input Token Embeddings for Neural Sequence Modeling.” 在《学习表示国际会议》中,2020年。

  34. Patrick Chen, Si Si, Yang Li, Ciprian Chelba, 和 Cho-Jui Hsieh. “Groupreduce: Block-wise low-rank approximation for neural language model shrinking.” 在《神经信息处理系统进展》中,2018年。

  35. Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, 和 Denny Zhou. “Mobilebert: a compact task-agnostic bert for resource-limited devices.” 在《计算语言学协会(ACL)》中,2020年。

  36. Song Han, Huizi Mao, 和 William J. Dally. “Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding.” 在《表示学习国际会议》中,2016年。

  37. Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, 和 Ivan Titov. “Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned.” 在《第57届计算语言学协会年会论文集》中,2019b。

  38. Geoffrey Hinton, Oriol Vinyals, 和 Jeff Dean. “Distilling the knowledge in a neural network.” 在NIPS深度学习和表示学习研讨会中,2015。

  39. Victor Sanh, Lysandre Debut, Julien Chaumond, 和 Thomas Wolf. “Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.” 在第5届能源效率机器学习与认知计算NeurIPS研讨会中,2019。

  40. Matus Telgarsky. “Benefits of depth in neural networks.” COLT, 2016.

  41. Sergey Edunov, Myle Ott, Michael Auli, David Grangier, 和 Marc’Aurelio Ranzato. “Classical structured prediction losses for sequence to sequence learning.” 在《北美计算语言学协会2018年会议论文集:人类语言技术》第1卷(长论文)中,2018年。

  42. Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, 和 Jakob Uszkoreit. “Tensor2tensor for neural machine translation.” CoRR, abs/1803.07416, 2018. URL: http://arxiv.org/abs/1803.07416.

  43. Kishore Papineni, Salim Roukos, Todd Ward, 和 Wei-Jing Zhu. “BLEU: a method for automatic evaluation of machine translation.” 在计算语言学协会第40届年会论文集中,第311–318页,计算语言学协会,2002。

  44. Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, 和 Michael Auli. “Fairseq: A fast, extensible toolkit for sequence modeling.” 在NAACL-HLT 2019: Demonstrations会议论文集中,2019年。

  45. Marjan Ghazvininejad, Omer Levy, Yinhan Liu, 和 Luke Zettlemoyer. “Mask-predict: Parallel decoding of conditional masked language models.” 在2019年自然语言处理和计算语言学第九届国际联合会议(EMNLP-IJCNLP)论文集中,第6114–6123页,2019年。

  46. Diederik P. Kingma 和 Jimmy Ba. “Adam: A method for stochastic optimization.” 在学习表示国际会议中,2015年。

附录A DeLighT架构用于语言建模和机器翻译

DeLighT架构用于语言建模和机器翻译的示意图展示在图6中。对于语言建模,我们遵循了Baevski和Auli(2019)的架构;而对于机器翻译,我们遵循了Vaswani等人(2017)的架构。

**语言建模:**图6a展示了语言建模的架构。该架构堆叠了$\mathcal{B}$个DeLighT块,每个块的配置是使用块级缩放确定的。每个块有三个子层。第一层是DeLighT变换,它在高维空间中学习表示。第二层是单头注意力,它编码上下文关系。第三层是位置感知的轻量级前馈网络。与Vaswani等人(2017)相似,我们采用了残差连接(He等人,2016)。与先前的工作(Baevski和Auli,2019;Dai等人,2019)一样,我们使用绑定的自适应输入(Baevski和Auli,2019)和自适应softmax(Grave等人,2017a)将Token映射到向量和将向量映射到Token。

**机器翻译:**图6b展示了机器翻译的架构。编码器堆叠了B个DeLighT块,每个块的配置也是使用块级缩放确定的。与语言建模类似,每个编码器块也有三个子层。第一层是DeLighT变换,它在高维空间中学习表示。第二层是单头注意力,它编码上下文关系。第三层是位置感知的轻量级前馈网络。与Vaswani等人(2017)相似,我们采用了残差连接(He等人,2016)。我们使用可学习的查找表将Token映射为向量。与编码器类似,解码器也堆叠了$\mathcal{B}$个块。解码器块与编码器块相同,不同之处在于它们在轻量级FFN之前有一个额外的源-目标单头注意力单元。源-目标单头注意力单元中的键和值是对编码器输出的投影。我们使用标准的可学习查找表将Token映射为向量,并使用线性分类层将向量映射为Token。

image-20240428195422419
图6:使用DeLighT进行序列建模。这里,绿色六边形代表了DeLighT转换。

附录B 带有输入混合器连接的组线性变换

组线性转换(group linear transformation GLT)$\mathcal{F}$ 将一个 $d_m$ 维输入 $\mathbf{X}$ 分成 $g$ 个不重叠的组,使得 $\mathbf{X} = \text{Concat}(\mathbf{X}_1, \cdots, \mathbf{X}_g)$,其中 $\mathbf{X}_i$ 是一个 $\frac{d_m}{g}$ 维的向量。然后,使用 $g$ 个线性变换 $\mathbf{W}_i \in \mathbf{R}^{\frac{d_m}{g} \times \frac{d_o}{g}}$ 同时对 $\mathbf{X}_i$ 进行变换,得到 $g$ 个输出 $\mathbf{Y}_i = \mathbf{X}_i \mathbf{W}_i$。然后,将 $\mathbf{Y}_i$ 进行串联以产生最终的 $d_o$ 维输出 $\mathbf{Y} = \text{Concat}(\mathbf{Y}_1, \cdots, \mathbf{Y}_g)$。

图7a展示了DeLighT变换扩展阶段中GLT的一个例子。为了说明,我们在这个例子中使用了相同的维度。回想一下,当我们在扩展阶段深入时,组的数量会增加。在这个例子中,第一层有一组,第二层有两组,第三层有四组。GLT学习特定于组的表示,并且是局部的。为了允许GLT学习全局表示,我们使用了特征洗牌。图7b展示了带有特征洗牌的GLT的一个例子。此外,仅通过堆叠线性或组线性(无论是否带有特征洗牌)来训练深度神经网络是具有挑战性的,因为梯度消失问题。He等人(2016)引入的残差连接缓解了这个问题,并有助于训练深度神经网络。然而,当输入和输出维度不同时(例如,在DeLighT变换的扩展和缩减阶段),就不能使用这样的连接。为了稳定训练并学习更深层次的表示,我们使用了Mehta等人(2020)的输入混合器连接。图7c展示了带有特征洗牌和输入混合器连接的GLT的一个例子。

image-20240428195422419
图 7:本图展示了DeLighT转换中使用的不同变体的组线性转换(group linear transformations GLTs)。

附录C DeLighT中的乘加操作

DeLighT块是使用线性变换、组线性变换(GLTs)和缩放点积注意力构建的。网络中的总乘加操作(MACs)数量是这些单独操作的累积。

让 $n$ 表示源token的数量,$m$ 表示目标token的数量,$d_m$ 表示输入维度,$d_o$ 表示输出维度,$g$ 表示 GLT 中的组数。下面描述了每个操作计算 MAC 数的过程。

组线性变换(GLT):GLT $\mathcal{F}$ 有 $g$ 个可学习矩阵 $\mathbf{W}_i \in \mathbf{R}^{\frac{d_m}{g} \times \frac{d_o}{g}}$。因此,GLT 学习了 $\frac{d_md_o}{g}$ 个参数,并执行了 $\frac{d_md_o}{g}$ 个 MAC 操作,将 $d_m$ 维输入转换为 $d_o$ 维输出。遵循标准做法,例如 He等人(2016)的 ResNet,我们将加法和乘法视为一次操作而不是两次操作,因为这些操作可以在最新的硬件中融合。

重要的是,当 $g=1$ 时,GLT 与线性变换相同。

DeLighT中的自注意力:DeLighT中的缩放点积自注意力定义为: $$ \begin{equation} \text{Attention}(\mathbf{K}, \mathbf{Q}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_o}}\right) \mathbf{V} \end{equation} $$

其中$\mathbf{Q} \in \mathbb{R}^{n \times d_o}$, $\mathbf{K} \in \mathbb{R}^{n \times d_o}$, $\mathbf{V} \in \mathbb{R}^{n \times d_o}$分别表示查询、键和值。

注意机制涉及两个点积。第一个点积是在 $\mathbf{Q}$ 和 $\mathbf{K}$ 之间进行的,而第二个点积是在第一个点积的输出和 $\mathbf{V}$ 之间进行的。两个点积都需要 $d_on^2$ 个乘加操作(MACs)。因此,计算缩放点积自注意力的总MAC数量为 $2d_on^2$。

在源-目标注意力中(比如在机器翻译中),$\mathbf{K}$ 和 $\mathbf{V}$ 来自源(编码器),而 $\mathbf{Q}$ 逐步解码(一次一个令牌)。因此,给定 $n$ 个源令牌,解码 $m$ 个目标令牌所需的MAC数量为 $\sum\limits_{k=1}^{m} 2knd_o$

附录D WikiText-103 数据集的消融(ablation)研究

表7 研究了 WikiText-103 数据集上 DeLighT块参数的影响,即 (1) GLT 的最小数量 $N_{min}$,(2) GLT 的最大数量 $N_{max}$,(3) 宽度倍增器 $w_m$,以及 (4) 模型维度 $d_m$(见图 1b)。图8、图9 和图10 显示了DeLighT、特征洗牌和轻量级 FFN 的影响。表 8 显示了DeLighT在 DeLighT块中位置的影响,而图12 显示了缩放 DeLighT网络的效果。我们选择了 WikiText-103 数据集进行消融研究,因为与其他数据集相比,它的词汇量非常大(267K 对比 30-40K),这使得我们能够测试在大词汇量的情况下的能力。性能以验证集上的困惑度(越低越好)来报告。在我们的消融(ablation)研究中,我们使用了与第 4.2 节相同的训练设置,只是我们只训练了 50K 次迭代。

DeLighT 块:总的来说,表7 显示了使用DeLighT和块级缩放(block-wise scaling)来缩放深度和宽度的效果。我们得出以下观察结果:

  1. 块级缩放(R4,R5)相比统一缩放(R1-R3)提供了更好的性能。例如,具有 $N_{min}=4$ 和 $N_{max}=8$ 的 DeLighT(R4)比具有 $N_{min}=8$ 和 $N_{max}=8$ 的 DeLighT更浅 $1.25\times$,但性能更好,而且参数和操作数量相似。缩放 $w_m$ 可以提高性能(R2 vs. R3),但改进幅度明显低于块级缩放模型(R3 vs. R5)。这表明跨块的参数的非均匀分布使网络能够学习更好的表示。
  2. 不同的 $N_{max}$ 和 $N_{min}$ 之间的比率产生不同的结果。当比率大于或等于两倍时,我们观察到性能显著提升。例如,当我们将 $\frac{N_{max}}{N_{min}}$ 从 2 缩放到 3 时(R6 vs. R8),困惑度提高了约 5 个点,而网络参数仅适度增加。另一方面,当 $\frac{N_{max}}{N_{min}}$ 接近 1 时(R6 vs. R7),性能变化不明显。这可能是因为跨块的参数分配接近于均匀(Eq. \ref{eq:mw})。这与我们之前的观察一致。
  3. 在输入附近学习更深更宽的表示,以及在输出附近学习更浅更窄的表示,可以实现更好的性能。例如,当我们将 $N_{max}$ 从 8 缩放到 12 并且 $N_{min}=4$ 时(R6,R8),与具有 $N_{min}=6$ 的模型相比,DeLighT以相似数量的参数交付了更好的性能(R7,R9)。这可能是因为当 $N_{min}=4$ 时,$N_{max}$ 和 $N_{min}$ 的比率较高,有助于更有效地分配每个块的参数。
  4. 在输入附近学习更深更宽的表示,而在输出附近学习更浅更窄的表示会损害性能(R13 vs. R16)。
  5. 使用 $w_m$ 和 $d_m$ 缩放宽度可以提高性能(R10-R15),但它们的影响是不同的。例如,当我们将 $w_m$ 和 $d_m$ 缩放两倍时,与 $w_m$ 相比,参数和操作数量的增加速度更快。 DeLighT以不同的方式学习更宽的表示可能在选择应用程序特定模型时很有用。
Row # $N_{min}$ $N_{max}$ $w_m$ $d_m$ Depth Parameters MACs Perplexity
R1 4 4 2 256 43 14.1 M 2.96 B 56.19
R2 8 8 2 256 115 16.6 M 3.49 B 48.58
R3 8 8 4 256 115 22.1 M 4.64 B 45.10
R4 4 8 2 256 92 16.7 M 3.51 B 46.30
R5 4 12 2 256 158 21.0 M 4.41 B 41.18
R6 4 8 2 256 92 16.7 M 3.51 B 46.30
R7 6 8 2 256 102 16.5 M 3.46 B 46.68
R8 4 12 2 256 158 21.0 M 4.41 B 41.18
R9 6 12 2 256 172 20.0 M 4.20 B 42.26
R10 4 12 2 256 158 21.0 M 4.41 B 41.18
R11 4 12 3 256 158 23.8 M 4.99 B 39.92
R12 4 12 4 256 158 27.1 M 5.69 B 39.10
R13 4 12 2 256 158 21.0 M 4.41 B 41.18
R14 4 12 2 384 158 29.9 M 6.28 B 35.14
R15 4 12 2 512 158 43.8 M 9.20 B 30.81
R16 12 4 2 256 158 21.0 M 4.41 B 43.10

表7 展示了对DeLighT块的不同方面进行的消融研究,包括统一缩放与块级缩放、深度缩放和宽度缩放。部分突出显示的行具有相同的配置(重复以说明结果)。我们的实验设置类似于第 4节,但我们将模型训练了 50K 次迭代。乘法和加法操作(MACs)是针对 20 个时间步进行计算的。

DeLighT变换的影响:我们将DeLighT块中的DeLighT变换替换为(1)DeFINE变换和(2)线性层堆叠。图8显示,与DeFINE单元和线性层相比,DeLighT变换以显著更少的参数提供了相似的性能。

image-20240428220057115
图 8: 不同变换的影响。DeLighT 变换比 DeFINE 和线性变换更具参数效率。更低的困惑度值意味着更好的性能。

特征洗牌:图9显示特征洗牌通过1-2个困惑度点提高了DeLighT的性能。这里,我们使用了与R13-R15(表7)相同的设置。

image-20240428215843596
图 9: 特征重排的影响。特征重排使我们能够从全局信息中学习表示,并提高了性能。更低的困惑度值意味着更好的性能。

轻量级FFN:图10展示了变化轻量级FFN中缩减因子r的影响。我们使用了与R13(表7)相同的设置。我们没有观察到性能有任何显著下降,直到$r = 4$。当r大于4时,我们看到了性能的下降(困惑度增加了约2点)。在这种情况下,轻量级FFN的内部维度非常小,影响了性能。值得注意的是,当$r = 2^{-2}$时,轻量级FFN的性能与$r = 2^{-2}$时相同,但参数数量少了1.28倍。在$r = 2^{-2}$时,轻量级FFN与Vaswani等人(2017)中的FFN相同。这表明DeLighT变换能够高效地在高维空间中学习表示,使我们能够减轻FFN的计算负担。

image-20240428215610711
图 10: 轻量级FFN中减少因子$r$的影响。DeLighT变换在高维空间中有效学习表示的能力使我们能够减少FFN的计算负担。更低的困惑度值意味着更好的性能。

我们还测试了去除轻量级FFN,虽然它减少了$\sim$0.5-1 M的参数,但性能下降了约2-3个困惑度点。

均匀缩放与块级缩放:图11比较了DeLighT在均匀缩放和块级缩放下的性能。对于给定的模型维度dm,具有块级缩放的DeLighT模型提供了更好的性能。

image-20240428215259755
图 11: 均匀(uniform)尺度与分块(block-wise)尺度对比。(a)对比了均匀尺度和分块尺度方法。(b)比较了WikiText-103数据集上DeLighT与均匀尺度和分块尺度方法的结果。带有分块尺度的DeLighT网络在不同设置下提供更好的性能。更低的困惑度值意味着更好的性能。

DeLighT变换的位置:我们在WikiText-103验证集上研究了DeLighT变换的三种配置:(1)DeLighT变换后跟单头注意力和轻量级FFN,(2)单头注意力后跟DeLighT变换,以及(3)单头注意力后跟DeLighT变换和轻量级FFN。对于相似数量的参数,我们发现(2)和(3)显著降低了(1)的性能。这表明更深和更宽的表示有助于学习更好的上下文表示;使我们能够用单头注意力替代多头注意力。

扩展DeLighT:图12展示了在变化DeLighT变换的配置参数($N_{min}$={4, 6}, $N_{max}$={8, 12}, $w_m$={2, 3, 4}, and $d_m$={256, 384, 512})后获得的DeLighT模型的结果。我们可以看到,同时保持其他配置参数不变(例如,$N_{min}$, $N_{max}$, 和 $w_m$),只扩展一个配置参数(例如,$d_m$)就能一致地提高性能。

image-20240428215036835
图 12: DeLighT的规模扩展。在保持其他配置参数恒定(例如,$N_{\text{min}}$、$N_{\text{max}}$、$w_m$)的情况下,扩展一个配置参数(例如 dm)始终会提高性能。每个柱子顶部的数字表示网络参数(以百万为单位)。更低的困惑度值意味着更好的性能。

这项工作通过手动调整参数 $N_{\text{min}}$、$N_{\text{max}}$、$w_m$ 和 $d_m$,探究它们之间的关系。我们相信,采用更加原理性的方法,例如 Tan 和 Le (2019) 提出的复合缩放方法,建立这些参数之间的关系,将会产生更有效和准确的模型。

附录E 组线性变换的源代码

在PyTorch中实现组线性变换(Group Linear Transformation, GLT)的源代码如下所示。为了在GLT中高效地实现分组功能,我们用专用的CUDA内核替换了PyTorch中简单的reshape和transpose操作,从而减少了内存占用并加快了训练速度。

列表1: “在Pytorch中GLT的简单实现”

import torch

def glt_function(x, n_groups, weights, bias=None):
    ''''
    :param x: 输入张量,大小为 [B x N],其中B是批量大小,N是输入维度
    :param n_groups: GLT中的组数
    :param weights: glt权重 [g x N/g x M/g]
    :param bias: GLT偏置(可选),大小为 [g x 1 x M/g]
    :return: 输出张量,大小为 [B x M]
    '''
    bsz = x.size(0)

    # 分组函数:将 [B x N] 张量转换为 [g x B x N/g]
    x = x.contiguous().view(bsz, n_groups, -1)  # [B x N] --> [B x g x N/g]
    x = x.transpose(0, 1)  # 转置,使组成为第一个维度

    # 变换函数:将 N/g 维空间变换为 M/g 维空间
    x = torch.bmm(x, weights)  # 与权重相乘

    # 如果提供了偏置,则添加偏置
    if bias is not None:
        x = torch.add(x, bias)

    # 重新分组函数:将 [g x B x M/g] 张量转换回 [B x M]
    x = x.transpose(0, 1)  # 再次转置,使批量大小成为第一个维度
    x = x.contiguous().view(bsz, -1)  # [B x g x M/g] --> [B x M]
    return x

列表2: “CUDA中的分组内核”

/* 分组内核:将输入从 [B x N] 转换为 [g x B x N/g] */
template<typename scalar_t>
__global__ void grouping_kernel_forward(const scalar_t* input,
                                         const int groups,
                                         const int total_elements,
                                         const int input_features,
                                         const int group_features,
                                         const int batch_size,
                                         scalar_t* output) {
    const int index = IMUL(blockIdx.x, blockDim.x) + threadIdx.x;
    if (index >= total_elements) {
        return;
    }
    const int b_idx = index / group_features;
    const int g_f_idx = index % group_features;
    int in_offset, out_offset;
    #pragma unroll
    for (int g = 0; g < groups; g++) {
        in_offset = (b_idx * input_features) + (g * group_features) + g_f_idx;
        out_offset = ((g * batch_size + b_idx) * group_features) + g_f_idx;
        output[out_offset] = input[in_offset];
    }
}