[论文解读] ReZero is All You Need: Fast Convergence at Large Depth

Thomas Bachlechner, Bodhisattwa Prasad Majumder+3 more|arXiv (Cornell University)|Mar 10, 2020

Domain Adaptation and Few-Shot Learning参考文献 39被引用 150

一句话总结

本文提出 ReZero，一种零初始化的残差门控，能够实现初始动态等距性，从而训练极深的网络并在 FC nets、ResNets 与 Transformer 上实现更快的收敛。

ABSTRACT

Deep networks often suffer from vanishing or exploding gradients due to inefficient signal propagation, leading to long training times or convergence difficulties. Various architecture designs, sophisticated residual-style networks, and initialization schemes have been shown to improve deep signal propagation. Recently, Pennington et al. used free probability theory to show that dynamical isometry plays an integral role in efficient deep learning. We show that the simplest architecture change of gating each residual connection using a single zero-initialized parameter satisfies initial dynamical isometry and outperforms more complex approaches. Although much simpler than its predecessors, this gate enables training thousands of fully connected layers with fast convergence and better test performance for ResNets trained on CIFAR-10. We apply this technique to language modeling and find that we can easily train 120-layer Transformers. When applied to 12 layer Transformers, it converges 56% faster on enwiki8.

研究动机与目标

• 说明训练极深网络的困难，原因是梯度消失/梯度爆炸。
• 提出一个简单、与架构无关的修改，以实现初始时的动态等距性。
• 展示在多种体系结构（FCN、CNN/ResNet、Transformer）下的大深度训练能力和更快的收敛。
• 在 CIFAR-10 和语言建模基准如 enwiki8 上展示经验收益。

提出的方法

• 在每一层添加一个残差门控 alpha，并将 alpha 初始化为零： x_{i+1}=x_i+alpha_i F(x_i).
• 证明这种初始化使网络在起始时表现为单位映射，从而实现动态等距性。
• 将 ReZero 应用到全连接网络、卷积型 ResNet 和 Transformer 架构。
• 与标准残差和基于归一化的方法相比，比较收敛速度和最终性能。
• 分析雅可比矩阵谱来说明训练过程中信号传播的保持。

实验结果

研究问题

• RQ1 一个最小的、零初始化的残差门控是否能在不使用归一化的情况下实现极深网络的稳定训练？
• RQ2 ReZero 如何影响跨越 FCN、CNN/ResNet 与 Transformer 的收敛速度和测试性能？
• RQ3 将 alpha 初始化为零在训练动力学上的影响程度，与初始化为 1 相比有多大？
• RQ4 在现代架构中应用 ReZero 时，深度的实际极限及深度相关的性能提升有哪些？
• RQ5 ReZero 能否减少或消除深层模型对 warm-up 或重度归一化的需求？

主要发现

• ReZero 使得可以训练数千层的网络（例如 10,000 层的全连接网络）。
• 使用 ReZero 的 Transformer 超过 100 层也可以在没有 LayerNorm 或 warm-up 的情况下训练。
• 在 enwiki8 上，使用 ReZero 的 12 层 Transformer 收敛速度比未改动的 Transformer 提前 56%，达到 1.2 BPB。
• 在 CIFAR-10 上，带 ReZero 的 ResNet 变体实现更快的收敛和更好的测试准确率（例如带 ReZero 的 ResNet-110 显示出改进的指标）。
• ReZero 在跨多个 ResNet 深度的情况下，将达到目标准确度的迭代次数加速约 32%。
• 在更深的 Transformers（64、128 层）中，ReZero 比 Post-Norm/GPT2-Norm 放置更能维持动态等距性，并且在普通 Transformer 发散的情况下也能实现训练。