38. Residual Block, and Why It Helps Gradients Flow

Problem

Deep stacks of layers — dozens of Transformer blocks, dozens of SSM blocks — suffer vanishing or exploding gradients without shortcut connections. Wrapping a sub-layer as out = x + f(x) instead of out = f(x) gives the gradient an unimpeded identity path back to the input: d(out)/dx = I + d(f(x))/dx. Even if the branch's Jacobian shrinks toward zero, the I term guarantees the gradient signal reaching x never shrinks below what it would have been without the branch at all.

Implement the block, then be ready to explain the gradient-flow argument precisely — the actual Jacobian decomposition, not just "residuals help training stability" as a slogan.

Source: src/38_residual_block_gradient_flow.py

def residual_block(
    x: torch.Tensor,
    w1: torch.Tensor, b1: torch.Tensor,
    w2: torch.Tensor, b2: torch.Tensor,
) -> torch.Tensor:
    ...

Examples:
>>> x = torch.randn(2, 4)
>>> w1, b1 = torch.randn(4, 8), torch.randn(8)
>>> w2, b2 = torch.randn(8, 4), torch.randn(4)
>>> residual_block(x, w1, b1, w2, b2).shape
torch.Size([2, 4])

Step-by-Step Approach

1. Compute the branch's hidden layer: h = ReLU(x @ w1 + b1), shape (batch, d_ff).
2. Project back down to the model dimension: branch = h @ w2 + b2, shape (batch, d_model), matching x.
3. Add the input back in: out = x + branch. This single + x is the entire residual connection — nothing more structurally is required.
4. Reason about the backward pass symbolically: d(out)/dx = d(x)/dx + d(branch)/dx = I + d(branch)/dx. The identity term I exists regardless of what the branch computes.
5. Operationalize the claim as a test: zero out the branch's second linear layer (w2 = 0, b2 = 0) so out becomes exactly x, then backprop a loss of out.sum() and confirm x.grad is exactly all-ones — the identity path delivering gradient completely unattenuated, independent of whatever w1/b1 happen to be.
6. Generalize mentally to a stack of N such blocks: the gradient at the very first layer's input still contains a product-free identity term (from every block's I), so it can't vanish to zero purely from depth, unlike a plain stack of f(f(f(...x)))) with no skip connections.

The key insight is that the residual connection's gradient contribution isn't approximately good, it's exactly identity — that's why the zero-the-branch test produces a perfectly clean all-ones gradient rather than something merely "close to 1."

Reference solution

def residual_block(x, w1, b1, w2, b2):
    h = F.relu(x @ w1 + b1)          # shape (batch, d_ff)
    branch = h @ w2 + b2             # shape (batch, d_model)
    # The `+ x` here is the entire residual connection. Its backward pass
    # contributes an exact identity Jacobian: d(out)/dx = I + d(branch)/dx,
    # so gradient flowing into `out` always has an unattenuated path to `x`,
    # regardless of how small the branch's own gradient becomes.
    return x + branch

Key Functions & Tricks

• x @ w1 + b1 / F.relu(...) — the branch's own two-layer MLP, unremarkable on its own.
• x + branch — the one line that matters; everything about the "why it helps" argument comes from this addition's backward pass.
• tensor.requires_grad=True + .backward() — used to make the gradient-flow claim empirically checkable rather than purely theoretical.
• Setting w2 = torch.zeros(...) — a deliberate degenerate case that isolates the identity path by forcing the branch's output to exactly zero.
• x.grad after out.sum().backward() — reading the accumulated gradient tensor directly to verify it equals torch.ones_like(x).

How to Recognize This Pattern