这篇论文让你无惧梯度消失或爆炸，轻松训练一万层神经网络

选自arXiv

作者：Batchlechner等

机器之心编译

机器之心编辑部

深度学习在众多领域都取得了显著进展，但与此同时也存在一个问题：深层网络的训练常常面临梯度消失或梯度爆炸的阻碍，尤其是像 Transformer 这样的大型网络。现在，加州大学圣迭戈分校的研究者提出了一种名为 ReZero 的神经网络结构改进方法，并使用 ReZero 训练了具有一万层的全连接网络，以及首次训练了超过 100 层的 Tansformer，效果都十分惊艳。

深度学习在计算机视觉、自然语言处理等领域取得了很多重大突破。神经网络的表达能力通常随着其网络深度呈指数增长，这一特性赋予了它很强的泛化能力。然而深层的网络也产生了梯度消失或梯度爆炸，以及模型中的信息传递变差等一系列问题。研究人员使用精心设计的权值初始化方法、BatchNorm 或 LayerNorm 这类标准化技术来缓解以上问题，然而这些技术往往会耗费更多计算资源，或者存在其自身的局限。

近日，来自加州大学圣迭戈分校（UCSD）的研究者提出一种神经网络结构改进方法「ReZero」，它能够动态地加快优质梯度和任意深层信号的传播。

论文地址：https://arxiv.org/abs/2003.04887v1

代码地址：https://github.com/majumderb/rezero

这个想法其实非常简单：ReZero 将所有网络层均初始化为恒等映射。在每一层中，研究者引入了一个关于输入信号 x 的残差连接和一个用于调节当前网络层输出 F(x) 的可训练参数α，即：

在刚开始训练的时候将α设置为零。这使得在神经网络训练初期，所有组成变换 F 的参数所对应的梯度均消失了，之后这些参数在训练过程中动态地产生合适的值。改进的网络结构如下图所示：

图 1：ReZero 结构示意图

ReZero 主要带来了以下两个益处：

1. 能够训练更深层神经网络

学习信号能够有效地在深层神经网络中传递，这使得我们能够训练一些之前所无法训练的网络。研究者使用 ReZero 成功训练了具有一万层的全连接网络，首次训练了超过 100 层的 Tansformer 并且没有使用学习速率热身和 LayerNorm 这些奇技淫巧。

2. 更快的收敛速度

与带有标准化操作的常规残差网络相比，ReZero 的收敛速度明显更快。当 ReZero 应用于 Transformer 时，在 enwiki8 语言建模基准上，其收敛速度比一般的 Transformer 快 56%，达到 1.2BPB。当 ReZero 应用于 ResNet，在 CIFAR 10 上可实现 32% 的加速和 85% 的精度。

ReZero (residual with zero initialization)

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

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

实验结果

更快的深层全连接网络训练

图 3 展示了训练损失的演变过程。在一个简单实验中，一个使用了 ReZero 的 32 层网络，拟合训练数据的收敛速度相比其他技术快了 7 到 15 倍。值得注意的是，与常规的全连接网络相比，残差连接在没有额外的标准化层时会降低收敛速度。这可能是因为初始化阶段信号的方差并不独立于网络深度。

随着深度的增加，ReZero 架构的优势更加明显。为了验证该架构可用于深度网络训练，研究者在一台配备 GPU 的笔记本电脑上成功训练了多达 1 万层的全连接 ReZero 网络，使其在训练数据集上过拟合。

图 3：256 宽度和 ReLU 激活的 32 层全连接网络四种变体，在训练过程中的交叉熵损失。

更快的深层 Transformer 训练

研究者提出，常规的 Transformer 会抑制深层信号传递，他们在输入序列 x 的 n x d 个 entry 的无穷小变化下评估其变化，获得注意力处理的输入-输出 Jacobian，从而验证了之前的观点。

图 5a 展示了不同深度中使用 Xavier 统一初始化权重的 Transformer 编码层的输入-输出 Jacobian 值。浅层的 Transformer 表现出峰值在零点附近的单峰分布，可以发现，深层结构中 Jacobian 出现了大量超出计算精度的峰值。虽然这些分布取决于不同初始化方法，但以上量化的结论在很大范围内是成立的。这些结果与普遍认为的相一致，也就是深层 Transformer 很难训练。

图 5：多个输入-输出 Jacobian 矩阵中对数奇异值λ_io 的直方图。（a）层数分别为 4、12、64 层的 Transformer 编码器网络；（b）是 64 层时训练前和训练中的 ReZero Transformer 编码器网络。深层 Transformer 距离动态等距很远，即λ_io ≪ 1，而 ReZero Transformer 更接近动态等距，平均奇异值 λ_io ≈ 1。

能够在多项 NLP 任务中实现 SOTA 的 Transformer 模型通常是小于 24 层的，这项研究中，最深层模型最多使用了 78 层，并且需要 256 个 GPU 来训练。研究者又将这一模型扩展至数百个 Transformer 层，并且仍然可以在台式机上训练。为了检查该方法是否可以扩展至更深层的 Transformer 模型之上，研究者将 ReZero Transformer 拓展到了 64 及 128 层，并与普通 Transformer 进行了对比。

结果显示，收敛之后，12 层的 ReZero Transformer 与常规的 Transformer 取得了相同的 BPB。也就是说，用 ReZero 来替代 LayerNorm 不会失去任何模型表现。训练普通的 Transformer 模型会导致收敛困难或训练缓慢。当达到 64 层时，普通的 Transformer 模型即使用了 warm-up 也无法收敛。ReZero Transformer 在α初始化为 1 时发散，从而支持了α = 0 的初始化理论。深层的 ReZero Transformer 比浅层的 Transformer 表现出了更优越的性能。

表 3：在 enwiki8 测试集上的 Transformers (TX) 对比。

收敛速度比较

选择 enwiki8 上的语言建模作为基准，因为较难的语言模型是 NLP 任务性能的良好指标。在实验中，其目标是通过测量 12 层的 Transformer 在 enwiki8 上达到 1.2 位每字节（BPB）所需的迭代次数，由此来衡量所提出的每种方法的收敛速度。

表二：针对 ReZero 的 12 层 Transformers 归一化后与 enwiki8 验证集上达到 1.2 BPB 时所需的训练迭代比较。

更快的残差网络训练

通过前述部分，看到了 ReZero 的连接是如何使深层网络的训练成为可能的，并且这些深层网络都包含会消失的 Jacobian 奇异值，例如 ReLU 激活或自我注意力。但是，如果没有 ReZero 的连接或者是其他架构的更改，其中某些架构将无法执行训练。在本节中，会将 ReZero 连接应用于深层残差网络从而进行图像识别。

虽然这些网络并不需要 ReZero 连接便可以进行训练，但通过观察发现，在 CIFAR-10 数据集上训练的 ResNet56 model4（最多 200 个 epochs）的验证误差得到了非常明显的提升：从（7.37±0.06）％到（6.46±0.05）％。这一效果是将模型中的所有残差连接转换为 ReZero 连接之后得到的。在实施 ReZero 之后，验证误差降低到 15％以下的次数也减少了（32±14）％。尽管目前这些结果只提供了有限的信息，但它们仍指出了 ReZero 连接拥有更广泛的适用性，从而也推进了进一步的研究。

上手实操

项目地址：

https://github.com/majumderb/rezero

在此提供了自定义的 ReZero Transformer 层（RZTX），比如以下操作将会创建一个 Transformer 编码器：

import torch
import torch.nn as nn
from rezero.transformer 
import RZTXEncoderLayer
encoder_layer = RZTXEncoderLayer(d_model=512, nhead=8)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
src = torch.rand(10, 32, 512)
out = transformer_encoder(src)

创建一个 Transformer 解码器：

import torch
import torch.nn as nn
from rezero.transformer 
import RZTXDecoderLayer
decoder_layer = RZTXDecoderLayer(d_model=512, nhead=8)
transformer_decoder = nn.Transformer
Decoder(decoder_layer, num_layers=6)
memory = torch.rand(10, 32, 512)
tgt = torch.rand(20, 32, 512)
out = transformer_decoder(tgt, memory)

注意确保 norm 参数保留为 None，以免在 Transformer 中用到 LayerNorm。