Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Residual Connections › Residual Scaling

Michael BrenndoerferUpdated June 8, 20257 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 103Residual Connections › Residual Scaling

Reading section 6 of 12 · Residual Scaling · 7 minSection 6 of 12 · 7 min

Section 6 of 12 · 7 min read

Residual Connections › Residual ScalingOne section · Sections in this chapter6 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Residual connections solve the vanishing gradient problem beautifully, but they introduce a subtler issue in very deep networks. Recall the unrolled form: $x_L = x_0 + \sum_{l=0}^{L-1} F_l(x_l)$ . We're summing $L$ residuals. If each residual has similar magnitude, the sum grows with depth.

For a 12-layer transformer, this growth is manageable. For a 96-layer model, summing residuals can cause the representation magnitude to explode, destabilizing training even though gradients flow well. The network needs a mechanism to control this accumulation.

The Scaling FactorLink Copied

Residual scaling dampens each residual contribution before adding it to the skip path. Instead of $y = x + F(x)$ , we compute:

y = x + \alpha \cdot F(x)

where:

$\alpha$ : a scaling factor, typically $\alpha < 1$
$F(x)$ : the residual function output
$\alpha \cdot F(x)$ : the dampened residual, smaller in magnitude than the original

The design question is: what value should $\alpha$ take? Several strategies have proven effective:

Fixed scaling: $\alpha = 0.1$ or similar constant, chosen empirically. Simple and effective for moderately deep networks.
Depth-dependent: $\alpha = 1/\sqrt{L}$ where $L$ is total layers. This ensures the variance of the sum $\sum_{l} \alpha F_l(x_l)$ stays bounded as depth increases, based on the statistics of summing independent random variables.
Learned: $\alpha$ as a trainable parameter per layer, allowing the network to discover appropriate scaling during training. More flexible but adds parameters.

Each approach reduces each layer's contribution, preventing the sum of residuals from growing unboundedly. Importantly, scaling doesn't hurt gradient flow: the gradient still flows through the unscaled skip connection, maintaining the gradient highway effect.

In[17]:

Code

class ScaledResidualBlock(nn.Module):
    """Residual block with scaling factor."""

    def __init__(self, dim, scale=0.1):
        super().__init__()
        self.scale = scale
        self.linear1 = nn.Linear(dim, dim)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x):
        residual = self.linear1(x)
        residual = self.relu(residual)
        residual = self.linear2(residual)
        return x + self.scale * residual  # Scaled residual


def compare_scaling_effects(n_layers, scales):
    """Compare output magnitude across different scaling factors."""
    results = {}

    for scale in scales:
        np.random.seed(42)
        torch.manual_seed(42)

        # Stack residual blocks
        blocks = nn.ModuleList(
            [ScaledResidualBlock(64, scale) for _ in range(n_layers)]
        )

        # Forward pass
        x = torch.randn(1, 64)
        magnitudes = [x.norm().item()]

        h = x
        for block in blocks:
            h = block(h)
            magnitudes.append(h.norm().item())

        results[scale] = magnitudes

    return results


scaling_results = compare_scaling_effects(30, [1.0, 0.5, 0.1])

Out[18]:

Visualization

Line plot showing representation magnitude growth across 30 layers for three scaling factors. — Output magnitude across layers with different residual scaling factors. Without scaling (factor 1.0), magnitude grows unboundedly. With aggressive scaling (0.1), magnitude stays controlled but may limit expressiveness.

The visualization reveals the trade-off clearly. Without scaling (α = 1.0), the representation magnitude grows steadily as each layer adds its full residual contribution. The growth appears roughly linear because we're adding vectors of similar magnitude. With moderate scaling (α = 0.5), growth slows substantially. With aggressive scaling (α = 0.1), the magnitude stays nearly constant. Each layer contributes so little that the sum barely increases.

The optimal choice balances two concerns: too little scaling allows magnitude explosion in very deep networks, while too much scaling limits each layer's expressive contribution. A depth-dependent rule like $\alpha = 1/\sqrt{L}$ provides automatic balance: deeper networks use smaller scaling factors.

ReZero: Learning to "Turn On" LayersLink Copied

The fixed and depth-dependent scaling approaches require choosing $\alpha$ before training. But what if the network could discover the right scaling for each layer on its own? The ReZero approach takes this idea to its logical extreme: initialize every scaling factor to zero, and let the network learn when and how much to "turn on" each layer.

y = x + \alpha \cdot F(x), \quad \alpha_{\text{init}} = 0

where:

$\alpha$ : a learnable scalar parameter, initialized to 0
$\alpha_{\text{init}} = 0$ : at the start of training, $\alpha = 0$ , so $y = x$ exactly

At initialization, every residual block is exactly an identity mapping, not approximately, but perfectly. The entire network is just a pass-through. As training progresses, gradient descent increases $\alpha$ values where the residual function $F$ provides useful contributions. Layers that aren't helpful can keep $\alpha$ near zero.

This approach offers elegant properties:

Perfect identity initialization: No approximation errors, no disruption to signal flow at the start of training
Layer-specific adaptation: Important layers develop larger $\alpha$ values; less useful layers remain suppressed
Automatic depth scaling: In very deep networks, the collective $\alpha$ values naturally stay small enough to prevent magnitude explosion

In[19]:

Code

class ReZeroBlock(nn.Module):
    """Residual block with learnable zero-initialized scaling."""

    def __init__(self, dim):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(1))  # Start at zero
        self.linear1 = nn.Linear(dim, dim)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x):
        residual = self.linear1(x)
        residual = self.relu(residual)
        residual = self.linear2(residual)
        return x + self.alpha * residual


# Demonstrate zero initialization
rezero_block = ReZeroBlock(64)
x_test = torch.randn(4, 64)

Out[20]:

Console

ReZero Block Initialization
--------------------------------------------------
Initial alpha: 0.000000
Input equals output: True

The output confirms that ReZero works as intended. With alpha initialized to exactly zero, the block is a perfect identity function: input and output are identical, not merely close. The residual function $F(x)$ still computes something (it has weights and activations), but its contribution is completely suppressed by the zero multiplier.

During training, gradient descent will increase alpha when doing so reduces the loss. Early in training, alpha values remain small, keeping the network close to identity. As the residual functions learn useful transformations, their alpha values grow to let those transformations contribute. This creates a natural curriculum: the network starts simple and gradually increases complexity.

Out[21]:

Visualization

Line plot showing alpha values for 12 layers growing from zero over training steps, with different final values per layer. — Simulated evolution of ReZero alpha values during training for a 12-layer network. Each layer's scaling factor starts at zero and grows as the residual function learns useful transformations. Different layers develop different final alpha values based on their importance.