论文阅读：ReZero is All Y ou Need: Fast Convergence at Large Depth

0.前言

论文地址：https://arxiv.org/abs/2003.04887v1
代码地址：https://github.com/majumderb/rezero

1.摘要

深度网络已经实现了跨域的显著性能提升，但它们经常遭受消失/爆炸梯度的影响。这尤其适用于深度超过12层的Transformer架构（一种i用于自然语言处理的深度学习模型，详见https://blog.csdn.net/longxinchen_ml/article/details/86533005），如果没有大型数据集和充足的计算资源，很难对其进行训练。一般来说，我们发现低效率的信号传播阻碍了深层网络的学习。在Transformer中，多头自聚焦是造成信号传播不良的主要原因。为了便于深层信号传播，我们提出了ReZero，这是对架构的一个简单更改，它使用每层一个额外的学习参数将任意层初始化为恒等映射（就是说f（x）=x，那么f就是x上的恒等映射）。我们将此技术应用于语言建模，发现我们可以很容易地训练100层以上的ReZero-Transformer网络。当应用于12层Transformer时，ReZero在enwiki8上收敛速度快56%。ReZero也适用于其他剩余网络：可以使深度完全连接网络的收敛速度加快1500%，在CIFAR 10上训练的ResNet-56的收敛速度加快32%。

1.Introduction

深度学习大幅提高了不同领域各种任务的最优表现[1，2，3，4]。神经网络的表达能力通常随深度呈指数增长[5]，因而具有很强的泛化性能，但过深的神经网络通常会导致梯度消失/爆炸和通过模型的不良信号传播[6]。研究人员依赖于谨慎的初始化[7，8]和规范化技术（如BatchNorm[9]和LayerNorm[10]）来缓解这个问题，但这些技术可能代价高昂且作用有限。
在这项工作中，我们提出了ReZero，这是一个小型的体系结构添加，它动态地加快优质梯度和任意深层信号的传播。这个想法很简单：ReZero初始化每个层为恒等映射。对于每一层，我们为输入信号x引入一个残差连接和一个用于调节当前网络层输出F（x）的可训练参数α。即：

在开始训练的时候α=0。这使得在网络训练初期，F这个变换函数中所有参数的梯度都会消失（因为乘以0了嘛，x_i+1=x_i），但在训练的初始阶段会动态地演化到合适的值。我们在图1中演示了体系结构。

ReZero有两个主要优点：
**深层学习：**信号通过深层网络有效传播，这使得我们可以在更深层次的网络上进行训练（以往太深的网络会导致梯度消失或者爆炸）。ReZero成功地训练了10000层完全连接的网络，还可以训练超过100层的Transformer模型， 并且没有使用学习速率热身和 LayerNorm 这些奇技淫巧。与[11]相比，我们发现要在这个深度上获得好的结果，不必增加辅助损失。
**更快的收敛：**我们观察到，与带有标准化的规则残差网络相比，ReZero网络的收敛速度明显加快。当ReZero应用于Transformers时，我们的收敛速度比vanilla Transformer快56%，在enwiki8语言建模基准上达到1.2bpb。当应用于ResNets时，我们在CIFAR 10上获得32%的速度，达到85%的精度。

2.Background and related work

深度为L层、宽度为w的网络通常具有指数级的深度表示能力，但宽度不是[12，5]。在基于梯度的训练方法中，过于深的网络结构往往会遇到麻烦。在深度模型的训练过程中，训练数据中的信号必须从输入层向前传播到输出层，然后，代价函数梯度必须向后传播，以提供有意义的权重更新。如果每一层中扰动的大小被一个因子r改变，信号和梯度都会以r[^L]的速率消失或爆炸，使得许多深网络在实际中不可追踪。
具体来说，考虑一个深度网络：它的输入是宽度为w的输入信号x₀，将输入传播到L层这个过程中，假设每一层网络都通过一个映射函数变换函数F[Wi]：R[^w] → R[^w]，其中w_i表示在层 i = 1……L时的所有参数。那么这整个传播过程可以表示为：

为了改善信号在深层网络中的传播，人们做了很多尝试，它们通常分为三类：初始化方案、标准化层和残差连接。我们在表1中展示了一些将残差网络与标准化相结合的常用方法。

2.1谨慎的初始化

近年来，人们利用平均场理论[13，8，14]对随机初始化的深宽神经网络中的信号传播动力学进行了形式化描述。对于一些深度神经网络，包括完全连接和卷积结构，两个不同信号的余弦距离(x_i * x_j / (||x_i|| ||x_i||))接近于在深度处消失或接近统一的固定点。如果此固定点为1，则网络的行为是稳定的，并且每个输入都映射到相同的输出，从而导致权重更新消失(就是网络收敛了)。如果该固定点为0，则网络的行为是混沌的，甚至类似的输入被映射到非常不同的输出，从而导致爆炸式的权重更新（就是网络无法收敛）。为了了解网络是处于稳定还是混沌阶段，我们考虑了 input-output
Jacobian：

该矩阵的均方奇异值x决定了平均输入信号扰动在网络中传播时的增长/衰减。网络在收敛与发散之间有一个界限——x= 1.训练可以在x=1时有效的进行。（这段我看的迷迷糊糊的，有大神懂的可以评论区留言。）
这种行为在文献[15，6]中得到了认可，这促使了权重的重新标度，使得x≈1和平均信号强度既没有增强也没有减弱。

Pennigton等人。[13，14]认识到input-output Jacobian的单位均方平均值不足以保证可训练性。例如，如果响应于非常大/很小奇异值的j_io奇异向量与数据中的扰动很好地一致，那么训练仍然是低效的。他们提出了动力等距的更强条件[16]，要求所有的J_io奇异值都接近1。这意味着输入信号的所有扰动都能同样地通过网络传播。对于输入信号的一些扰动，ReLU激活函数映射到零，因此，具有ReLU激活的深网络不可能满足动态等距性，这是直观的，如在[13]中严格建立的。对于一些激活函数和网络结构，详细的初始化方案允许网络在初始化时满足动态等距性，这显著改善了训练动态性[17，5，18，19]。

2.2标准化

提高深层神经网络可训练性的另一种方法是将显式提供规范化的层合并起来。已经提出了许多规范化模块，其中最流行的两个模块是BatchNorm[9]和LayerNorm[10]。一般来说，标准化的目的是确保信号最初在通过网络传播时具有零均值和单位方差，从而减少“协变量偏移”[9]。为了简单起见，我们将主要关注与LayerNorm的比较，因为BatchNorm具有额外正则化效果——那与我们的研究不在同一方向。
标准化方法在加速深层网络的训练方面取得了成功，但它们确实会给网络带来计算成本，并且会带来额外的超参数来调整（例如，将标准化放在哪里）。与标准化方法相比，我们提出的方法简单且易于实现。即使在缺乏各种标准化的情况下，仅ReZero也足以训练更深层次的网络。虽然ReZero使标准化对于收敛来说是多余的，但是我们发现BatchNorm的正则化效果是我们方法的补充。

2.3 残差连接

在[2]中引入的身份映射使得在卷积网络的背景下用于图像识别的深度残差学习框架显著增加了可训练的深度。BatchNorm和ResNets[2]的互补使用使得人们能够训练超过100层的卷积神经网络。但对于LayerNorm和Transformer架构，情况并非如此。Yang等人。[18] 研究了残差全连接网络，证明了由于跳跃连接，信号在传播过程中衰减更慢（多项式），从而可以有效地训练更深层的网络。
最近有一篇工作SkipInit[20]，针对ResNet架构，提出了类似ReZero的一种方法来作为BatchNorm的替代方案。作者发现，在没有BatchNorm的深度ResNets中，需要一个标量乘子来保证收敛性。对于[20]中考虑的具体情况，我们得出了类似的结论，并研究了更广泛的信号在跨多个体系结构和超出BatchNorm的较深网络中的传播。

3.ReZero

ReZero 对深度残差网络的结构进行了简单的更改，可促进动态等距（dynamical isometry）并实现对极深网络的有效训练。研究者在初始阶段没有使用那些非平凡函数 F[W_i] 传递信号，而是添加了一个残差连接并通过初始为零的 L 个可学习参数α_i（作者称其为残差权重）来重新缩放该函数。目前，信号根据以下方式进行传递：

在初始阶段，该网络表示为恒等函数并且普遍满足动态等距关系。我们在下面演示一个玩具模型，这个架构可以指数级加速训练。在该架构修改中，即使某一层的 Jacobian 值消失，也可以训练深度网络（正如 ReLU 激活函数或自注意力机制出现这样的状况）。这一技术还可以在现有的已训练网络上添加新层。

3.1 一个实例模型

为了说明ReZero连接如何加速训练，让我们考虑一个由L个单神经元隐藏层描述的深层神经网络的玩具模型，这些层没有偏差，并且都具有相同的权重w和α_i=α ∀i。然后，该网络简单地将输入x0映射到输出：

固定参数α=1将表示全连接残差网络的toy模型，而初始化α=0并将α视为学习参数对应于ReZero网络。input-output Jacobian 由J_io=（1+αw）[^L]给出，表明对于w≈1和α=1的初始化。深度（即L>>1） 网络对输入的任何小扰动都非常敏感，而当α=0时，输入信号的幅度保持不变。虽然这个例子太简单，无法显示收敛，但它确实精确地模拟了深网络中常见的消失和爆炸梯度问题。假设学习率λ和代价函数C，梯度下降根据公式：

对于α=1，初始权值w≈1的梯度下降收敛要求步数不大于1，因此学习速度在深度L上是指数小的：

其中我们只保留了对w和L的参数依赖性，当w>>1时等式6会出现梯度爆炸，而对于w≈-1，等式6又会出现梯度消失。初始化α=0解决了这两个问题：假设一个条件充分的代价函数，梯度下降的第一步将更新残差权重α，使其在保持网络表达能力的同时，避免输出过大，并将参数轨迹保持在条件良好的区域内。残差权重的第一个非平凡步骤计算公式为：

梯度下降收敛于网络深度L的多项式学习率。因此，在这个简单的例子中，ReZero连接允许以比普通残差网络少得多的优化步骤进行收敛。我们在图2中演示了培训动态、代价函数和梯度。

4.加快深度全连接网络的训练

我们现在研究ReZero对deep ReLU网络的影响，并结合表1行中列出的各种有助于深入学习的方法。具体来说，我们将比较普通的完全连接网络（FC，第1行）、具有残差连接的深度网络（FC+Res，第2行）和具有LayerNorm（FC+Norm，第3行）的深度网络，最后我们的建议ReZero（第6行）。我们选择初始权值为Winormally分布，方差为训练最优，即σ_w[^2] =2/W（除了vanilla residual network），其中σ_w[^2] ≈0.25/W，见[6，18]。
作为一个示例任务，我们在CIFAR-10数据集上训练了四种不同的网络结构，用于监督图像分类。我们只对训练动力学感兴趣，并研究模型拟合数据需要多少迭代。
我们在图3中显示了训练损失的演变。在我们的简单实验中，一个32层的网络ReZero架构收敛到比其他技术快7到15倍的速度来拟合训练数据。请注意，与完全连接的纯网络相比，如果没有额外的标准化层，剩余连接会降低收敛速度。我们推测，这是因为在初始化时，信号的方差并不独立于深度，见[18]。（图3论文里是一片空白。。。。反正就是这种方法收敛快就完事儿了。）
随着深度的增加，ReZero架构的优势变得更加明显。为了验证这种体系结构确保了大深度的可训练性，我们成功地训练了完全连接的ReZero网络，在一台笔记本电脑上使用一个gpu，最多可连接10000层，最后直接使得训练过拟合。

5.更快地训练更深层的Transformer

在这一节中，我们研究了信号传播和ReZero在Transformer结构中的应用[21]。Transformer在有监督和无监督的NLP任务中获得了显著的普及和成功[23，11]。Transformer是通过堆叠模块来构建的，这些模块首先是自我注意，然后是点式前馈变换。
原始的Transformer[21]实现可以看作是一个后标准化的残差网络（表1中的第5行）。在一个Transformer模块中，每个子层的输出通过一个残差连接添加，然后通过LayerNorm规范化。

其中子层∈{自我注意，前馈}，如图4的左面板所示。

5.1 Transformer中的信号传播

与原始Transformer层中的信号传播相关的两个关键组件包括LayerNorm[10]和（multi-head）self-attention[21]。我们认为，无论是单分量还是与普通残差连接相结合，都不能对所有输入信号满足所有的动态等距。这一发现促使使用ReZero连接来替换LayerNorm和vanilla残差连接。
层规范化去除了一个给定层的所有神经元的平均值并缩放方差，并引入可学习参数γ和β来重新缩放方差

从这个定义可以清楚地看出，通过纯粹改变输入x的均值或方差的变换来扰动输入x，将使输出保持不变。因此，这些扰动产生了 input-output Jacobian的两个消失奇异值。在Transformer架构[21]中，范数被应用于输入语句的n个元素中的每一个，从而导致每个Transformer层的Jacobian的总共2×n个消失奇异值。
自注意允许模型通过计算输入序列的加权和来关联位于不同位置的内容。具体地说，n×d矩阵x包含一个包含d维嵌入向量的n行输入序列，从中我们可以计算查询、键和值矩阵Q，K，V=x·W[^{Q，K，V]，其中W[}Q，K，V]矩阵是d×d维矩阵。标度的点积关注度由下式给出：

一般来说，这个注意过程的雅可比矩阵的奇异值谱是复杂的。现在，我们只是争论，对于某些输入x和权重WQ，K，V，Jacobian有大量消失的奇异值（下面我们根据经验评估这个声明）。考虑权重或输入，使softmax函数的每个参数都比1小。然后，softmax函数简单地返回一个n×n维矩阵，其中填充了所有近似于1/n的项，这意味着注意函数将输入序列的所有嵌入向量投影到一个对角线方向上。这意味着在n×d的雅可比奇异值中，只有d是不消失的，因此大部分输入信号丢失。残差连接可以恢复一些丢失的信号，但即使这样，一些扰动也会被放大，而另一些则会被衰减。这个例子表明，在深度Transformer网络中，自关注与动态等距和无阻碍的信号传播是不相容的。对于多头注意，很容易验证相同的结论。仔细初始化权重可能会缓解其中的一些问题，但我们不知道任何初始化方案会使Transformer层与动态等距一致。

我们给出了一个理论论点，即vanilla Transformer包含抑制深信号传播的元素。在这里，我们通过在输入序列x的每个n×d项的无穷小变化下评估注意过程的变化来获得注意过程的 input-output Jacobian，从而在实践中验证了这些。我们在图5a中用Xavier统一的初始化权重显示了不同深度的Transformer编码器层的input-output Jacobian。当浅层Transformer的奇异值分布在单位附近达到峰值时，我们清楚地观察到深层结构的Jacobian具有大量的奇异值，这些奇异值会消失到机器精度。
虽然分布因初始化方案的细节而异，但定性语句的适用范围更广。这些结果与普遍观察到的深层Transformer网络的训练具有极大的挑战性相一致。

我们应用ReZero来解决Transformer层中信号传播不良的问题，方法是替换层模型，重新缩放自关注块。具体地说，这将公式（9）修改为

其中α是图4右侧面板中的学习残差权重参数。对于Transformer层中的一对多头自关注前馈网络，我们共享相同的α_i参数。在初始化时，α_i=0，这使得信号传播不受阻碍：input-output Jacobian 的所有奇异值都是1，并且模型平凡地满足动态等距。为了验证模型在整个训练过程中保持接近动态等距，对于更大的α_i，我们在图5b中的WikiText-2[24]上的语言建模玩具任务64层模型训练期间显示雅可比奇异值的直方图。在训练期间，残差连接的权重逐渐增加，允许在保持接近动态等距的信号传播特性的同时，对极其复杂的函数进行建模。

5.2 收敛速度

我们选择enwiki8[25]上的语言建模作为基准，因为强大的语言模型是下游NLP任务性能的良好指标[4]。我们在这些实验中的目的是通过测量12层Transformer在enwiki8上达到每字节1.2位（BPB）所需的迭代次数来测量每种方法的收敛速度。
自从引入Transformers[21]以来，人们已经在Transformer中竞争性地放置了几个LayerNorm以实现更好的收敛性[4，26]。我们用3种Transformer归一化方法进行了实验，并与ReZero Transformer进行了比较。POST范数（表5中的行1）相当于[21]中的vanilla Transformer，在[4]中引入了预规范（表1中的行4）方法，GPT2范数（x_i+1= x_i+ Norm(F(x_i))）被用于GPT2（4）的训练，它成功地训练了多达48层的Transformer。最后，我们用我们提出的将α初始化为0或1的ReZero方法进行了实验。超参数见附录A。

我们的结果（表2）表明，后范数在训练期间发散，而所有其他模型都能够收敛。这并不奇怪，因为最初的Transformer实现需要学习率预热，这在[26]中也得到了证实。为了验证这一点，我们用100步的学习率预热重新运行后范数设置，发现该模型在13690次迭代中能够收敛到1.2bpb。在这种情况下，我们比较了其他分层布局方案与后规范。我们发现，其他位置最初导致更快的收敛，但最终后范数赶上性能，导致相对较慢的前范数和GPT2范数收敛。然而，与后范数相比，其他分层形式的布局有一个优势，即它们不需要学习率预热，因此有较少的超参数可调。α=1的ReZero Transformer与普通Transformer相比没有改进，表明初始化α=0的重要性。在我们提出的α=0的初始化条件下，ReZero的收敛速度比vanilla Transformer快56%。

5.3 深度Transformer

在许多NLP任务中实现最先进性能的转换器模型[23]通常少于24层。在我们的工作中，最深的模型使用了78层[27]，需要256个GPU进行训练。在本节中，我们将扩展到数百个Transformer层之外，并且仍然可以在台式机上进行训练。为了检验我们的方法是否适用于更深层次的Transformer，我们将12层ReZero Transformer从第5.2节扩展到64层和128层，并与普通Transformer（后规范）进行比较。超参数见附录B节。

我们的结果（表3）表明，一个12层的ReZero Transformer在收敛后达到与常规Transformer相同的BPB，这表明我们用ReZero替换LayerNorm并没有丢失模型中的任何表示性。我们发现，尝试训练深vanilla Transformer要么导致收敛困难，要么训练时间慢。当扩展到64层时，vanilla Transformer即使有一个预热计划也无法收敛。初始化为α=1的ReZero Transformer发散，支持我们在α=0时的理论激励初始化。较深的ReZero Transformer比较浅的Transformer能获得更好的性能。

为了进行比较，我们还显示了Character Transformer[11]的结果，它具有类似的设置供参考。然而，字符型Transformer使用的参数较多，附加的辅助损耗也较多，这与我们的工作是正交的。我们的128层Transformer在没有任何中间损耗的情况下实现了类似的性能，使用了一半的参数，并且具有更大的深度。我们没有调整我们的超参数，并且我们的模型可以通过更强的正则化和学习率计划潜在地获得更好的结果。

为了深入探讨我们的模型，我们研究了12层和64层ReZero Transformer在训练过程中的残差权重α的行为（图6）。将|α_i|视为每一层对网络整体信号的贡献量是有用的。我们看到一个有趣的模式出现在浅层和深层的ReZero Transformer。在训练的早期迭代中，残差权值迅速增加到一个峰值，然后在整个训练过程中慢慢衰减到一个小值。在训练早期，高层往往占主导地位（他们的峰值更早），在训练接近尾声时，每一层都得到了相似程度的利用。12层和64层模型训练结束时的平均α_i分别为0.0898和0.0185，约为1/L，其中L是残差层的数量。

6.更快的训练ResNet

在前面的章节中，我们看到了ReZero连接如何使包含具有消失Jacobian奇异值的层的深层网络训练成为可能，例如ReLU激活或自我注意。如果没有 ReZero连接或其他架构更改，这些架构中的一些是不可训练的。在这一节中，我们将ReZero连接应用到用于图像识别的深残差网络[2]。虽然这些网络可以在没有重新归零连接的情况下进行训练，但我们观察到对于ResNet56，CIFAR-10数据集上的模型在将模型中的所有vanilla residual connections变为ReZero 连接后，训练（最多200个阶段）的验证错误显著提高（从（7.37±0.06）%提高到（6.46±0.05）%。实现ReZero后，将验证误差降低到15%以下的次数也减少了（32±14）%。虽然这些结果本身只能提供有限的见解，但它们指向了ReZero连接更广泛的适用性，并推动了进一步的研究。

7.结论

我们引入了ReZero，这是一个简单的结构修改，它有助于信号在深度网络中的传播，并有助于网络保持动态等距。将ReZero应用于各种残差架构（完全连接的网络、Transformer和resnet），我们观察到显著提高了收敛速度。此外，我们还能够有效地培训数百层的Transformer，这对于最初的架构来说是困难的。我们相信更深层次的Transformer将为未来的探索打开大门。
在ReZero训练模型中，我们发现在训练过程中，每一层α_i的残差权重值都存在有趣的模式。这些模式可能暗示某种形式的课程学习，并允许逐步堆叠层，以进一步加快训练[28]。残差权重的模式对于理解这种更深层网络的训练动态是至关重要的，并且对于模型性能可能是重要的，我们将在未来的工作中对此进行探讨。

参考文献
神经网络结构优化：这篇论文让你无惧梯度消失或爆炸，轻松训练万层神经网络