Source-linked AI summary
ReZero is All You Need: Fast Convergence at Large Depth
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W. Cottrell, Julian McAuley
TL;DR
Deep networks can suffer vanishing or exploding gradients and poor signal propagation as depth increases. The paper introduces ReZero, which gates residual transformations with zero-initialized trainable parameters to preserve initial dynamical isometry. ReZero trains extremely deep fully connected networks and Transformers while accelerating convergence on language-modeling and image-recognition benchmarks.
Problem
Deep networks often exhibit vanishing or exploding gradients and poor signal propagation, while existing solutions can involve complex architectures, normalization, or careful initialization.
Method
ReZero adds a skip connection and zero-initialized learnable residual weights α_i that modulate each layer’s non-trivial transformation.
Results
ReZero facilitates efficient training of extremely deep networks and improves convergence across fully connected networks, Transformers, and ResNets, including 56% faster convergence on enwiki8.
Takeaways & Limitations
A simple zero-initialized residual gate supports deep signal propagation and enables training Transformers with hundreds of layers.
Takeaways & Limitations
The observed residual-weight patterns may require further study to determine their role in training dynamics and model performance.
Abstract
from arXiv · showhide
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.
1 Introduction
Deep networks gain expressivity with depth but can suffer poor signal propagation and difficult optimization. ReZero addresses this by zero-initializing a trainable gate on each residual connection, supporting very deep networks and faster convergence.
- Deep networks’ expressivity typically grows exponentially with depth, but deep signal propagation can cause vanishing or exploding gradients.
- ReZero adds a residual connection and one trainable parameter α that modulates each layer’s non-trivial transformation.The parameter is initialized to α = 0, so each layer initially performs the identity operation.
- Widely Applicable: ReZero is simple and architecture agnostic, applying to residual-style architectures with little tuning and few lines of code.
- Deeper learning: 10,000 layers of fully connected networks were successfully trained, and Transformers over 100 layers were trained without learning-rate warm-up, LayerNorm, or auxiliary losses.
- Faster convergence: 56% faster convergence reaches 1.2 BPB on enwiki8 for Transformers, while ResNets achieve a 32% speedup to 85% accuracy on CIFAR 10.
2 Background and related work
Deep networks require mechanisms that preserve signals and gradients across many layers. Related approaches include initialization, normalization, and residual connections, while dynamical isometry provides a stronger criterion for uniformly stable signal propagation.
- Perturbations changed by a factor r per layer cause signals and gradients to vanish or explode at a rate of r^L.
- Deep-network signal-propagation methods commonly use initialization schemes, normalization, or residual connections.
- At the edge of chaos, χ = 1, average signal perturbations are neither enhanced nor attenuated, supporting efficient training.
- Dynamical isometry requires all input-output Jacobian singular values to be close to one, so input perturbations propagate equally well.
- Normalization can accelerate training but adds computational cost and hyperparameters, whereas ReZero is simple, cheap, and sufficient without norms.
- Residual connections improved trainable depth in convolutional networks, but comparable depth increases had not occurred for Transformer architectures.
- Zero-initialized residual-block approaches, including SkipInit, use scalar multipliers to support convergence in deep ResNets without BatchNorm.
3 ReZero
ReZero adds zero-initialized, learnable residual weights so deep networks begin as identity mappings with dynamical isometry, enabling well-conditioned optimization and faster convergence.
- 3 ReZero: ReZero rescales each residual function F[W_i] with learnable residual weights α_i initialized to zero, while preserving a skip connection.At initialization, the network propagates signals through the identity rather than the non-trivial functions.
- 3 ReZero: At initialization, the network represents the identity function and satisfies dynamical isometry.
- 3.1 A toy example: The toy model maps inputs through L shared-weight layers, with input-output Jacobian J_io = (1 + αw)^L.
- 3.1 A toy example: The toy-model trajectories initialized at α = 0 avoid poorly conditioned regions around α ≈1 and converge to αw ≈1.2.
- 3.1 A toy example: For α = 1 and w ≈1, gradient descent requires a learning rate exponentially small in depth L, whereas ReZero converges with a learning rate polynomial in L.
- 3.1 A toy example: Initializing α = 0 updates residual weights toward well-conditioned regions while retaining the network’s expressive power.
4 Training deep fully-connected networks faster
ReZero improves training dynamics in deep fully connected networks: it accelerates fitting on CIFAR-10 and remains trainable at depths reaching 10,000 layers.
- 4 Training deep fully-connected networks faster: ReZero is compared with fully connected, residual, and LayerNorm networks in a 32-layer CIFAR-10 training-dynamics experiment.The networks use identical width and ReLU-based fully connected architectures for the comparison.
- 4 Training deep fully-connected networks faster: ReZero fits the CIFAR-10 training data 7 to 15 times faster than the other tested techniques.
- 4 Training deep fully-connected networks faster: Without an additional normalization layer, the residual connection decreases convergence speed relative to a plain fully connected network.
- 4 Training deep fully-connected networks faster: ReZero fully connected networks were successfully trained to overfit the training set at depths of up to 10,000 layers.The experiment was run on a laptop with one GPU.
Convolutional ResNets faster
ReZero is evaluated in deep convolutional residual networks alongside several signal-propagation methods, with experiments targeting training speed and accuracy.
- Convolutional ResNets faster: ReZero combines the benefits of post-activation residual connections and identity-oriented initialization in a simpler implementation.
- Convolutional ResNets faster: Figure 3 reports cross-entropy loss during training for four 32-layer fully connected network variants, averaged over five runs with 1σ error bands.Training uses Adagrad with learning rate 0.01.
- Convolutional ResNets faster: ReZero is applied to convolutional residual networks because these networks are trainable without modification, yet ReZero is observed to accelerate training and improve accuracies.
- Convolutional ResNets faster: The comparison includes Gated ResNet, zero γ, FixUp, ReZero, and ReZero-modified Pre-Act ResNet on CIFAR-10.The methods are trained with identical hyperparameters and evaluated using validation error and training duration.
- Convolutional ResNets faster: Superconvergence with a ReZero Pre-activation ResNet-18 reduces training time to reach 94% accuracy from 60 epochs to 45 epochs.
6 Training deeper Transformers faster
ReZero applies a learned, zero-initialized residual weight to Transformer layers, preserving dynamical isometry at initialization and keeping signal propagation well behaved during training. This enables faster convergence and effective training of substantially deeper Transformers.
- Signal propagation: ReZero replaces Transformer LayerNorm with a learned residual weight initialized to α_i = 0, making all input-output Jacobian singular values equal to 1 at initialization.The same α_i is shared by the self-attention and feed-forward sublayers within each Transformer layer.
- Signal propagation: 64-layer ReZero Transformers remain much closer to dynamical isometry during training than standard Transformer encoders across the reported depth settings.The Jacobian singular-value histograms use log(λ_io), with dynamical isometry corresponding to mean singular value λ_io ≈1.
- Convergence speed: 56% faster convergence was observed for ReZero than the vanilla Transformer to reach 1.2 BPB on enwiki8.The benchmark measures iterations required for a 12-layer Transformer to reach 1.2 bits per byte on the validation set.
- Convergence speed: Post-Norm diverged without warm-up, whereas ReZero with α = 0 converged 56% faster than the vanilla Transformer; ReZero with α = 1 showed no improvement.With 100 warm-up steps, Post-Norm reached 1.2 BPB in 13,690 iterations.
- Deeper Transformers: During early training, residual weights quickly increase, and their layerwise evolution may help explain training dynamics or support progressive layer stacking.The authors describe this pattern as a possible direction for future work rather than an established mechanism.
7 Conclusion
The paper concludes that ReZero is a simple modification that facilitates signal propagation and helps maintain dynamical isometry across several residual architectures. It also identifies residual-weight patterns as an open direction for understanding and accelerating deep Transformer training.
- Conclusion: ReZero facilitated signal propagation and helped networks maintain dynamical isometry across fully connected networks, Transformers, and ResNets.The conclusion reports significantly improved convergence speeds across these residual architectures.
- Conclusion: The authors observed layerwise residual-weight patterns that may indicate curriculum learning or enable progressive stacking of layers.They state that the role of these patterns in training dynamics and model performance requires future study.
Broader Impact
The broader-impact discussion connects deep-model scaling with training time, compute, carbon, and economic costs, while examining signal-propagation limitations in Transformer components. It explains how normalization and attention can create problematic Jacobian behavior for some inputs.
- Broader impact: Increasing parameter count can improve model performance but also lengthens training, increases GPU-cluster requirements, and contributes to carbon and economic costs.The passage frames these costs as consequences of increasing model capacity.
- Signal propagation: The discussion argues that neither LayerNorm nor self-attention, alone or with a vanilla residual connection, satisfies dynamical isometry for all input signals.This claim is presented as theoretically argued and empirically evaluated through input-output Jacobians.
- Signal propagation: Layer normalization can produce two vanishing Jacobian singular values per normalized input element because mean- and variance-shifting perturbations leave its output unchanged.For a Transformer sentence with n elements, this yields 2 × n vanishing singular values per layer.
- Signal propagation: Self-attention may have many vanishing Jacobian singular values when softmax arguments are small, because attention then projects embeddings toward a single diagonal direction.The passage presents this as an argument for some inputs and weights, evaluated empirically in the paper.
B Convergence speed experimental hyperparameters
The convergence-speed experiments use a controlled 12-layer Transformer setup with fixed architecture, regularization, batch, and optimizer hyperparameters.
- Experimental setup: The experiments fix batch size at 1080, depth at 12 layers, dropout at 20%, hidden and embedding size at 512, and context length at 512.They also use two attention heads and GELU activations in the feed-forward layer.
- Experimental setup: The models use the LAMB optimizer with a fixed learning rate of 0.016, without learning-rate schedules.The passage says schedules were omitted to simplify training.
C Deep Transformers experimental hyperparameters
The paper situates ReZero among residual gates, normalization methods, and initialization schemes for improving signal propagation in deep networks. It emphasizes ReZero’s simple, zero-initialized residual scaling and reports its use in very deep networks.
- D.1 Highway Networks: Highway Networks regulate signal flow with input-dependent transform and carry gates, whereas ReZero uses gates independent of the input signal.
- D.7 ReZero: ReZero is presented as simpler and architecture agnostic, applying to residual-style architectures with little tuning and few lines of code.
- D.7 ReZero: ReZero adds a scalar multiplier initialized at 0 to every residual block.