import numpy as np
import matplotlib.pyplot as plt
import torch
plt.style.use("../figure_style.mplstyle")
torch.manual_seed(13)
width, depth = 128, 30
def make_layers(depth, width, std):
layers = []
for _ in range(depth):
layer = torch.nn.Linear(width, width, bias=False)
with torch.no_grad():
layer.weight.normal_(0, std)
layers.append(layer)
return layers
def run_plain(layers, x, readout):
h = x
activations = []
for layer in layers:
h = torch.tanh(layer(h))
h.retain_grad() # Keep this intermediate gradient for inspection.
activations.append(h)
loss = (h @ readout.to(h)).sum()
loss.backward()
act_stds = [a.detach().std().item() for a in activations]
grad_norms = [a.grad.norm().item() for a in activations]
return act_stds, grad_normsDemo: Gradient Scale by Layer
A deep network repeatedly transforms both activations and gradients. We will measure those quantities layer by layer to see why initialization, normalization, and residual connections make depth trainable.
Question 1: How Does Initialization Affect Signals Across Layers?
Three plain tanh networks differ only in their weight scale. On a logarithmic axis, a straight slope means that each layer contributes a similar multiplicative factor.
x = torch.randn(256, width)
readout = torch.randn(width) / width ** 0.5
configs = [
("Too small ($0.5/\\sqrt{n_{\\mathrm{in}}}$)", 0.5 / width ** 0.5, "--"),
("Well-scaled ($1/\\sqrt{n_{\\mathrm{in}}}$)", 1.0 / width ** 0.5, "-"),
("Too large ($3/\\sqrt{n_{\\mathrm{in}}}$)", 3.0 / width ** 0.5, ":"),
]Compare the activation and gradient scales at each layer.
fig, axes = plt.subplots(1, 2, figsize=(9, 3.2))
layer_idx = np.arange(1, depth + 1)
for label, std, linestyle in configs:
torch.manual_seed(13) # Scale the same random weights in each comparison.
layers = make_layers(depth, width, std)
act_stds, grad_norms = run_plain(layers, x, readout)
axes[0].plot(layer_idx, act_stds, linestyle=linestyle, linewidth=1.5, label=label)
axes[1].plot(layer_idx, grad_norms, linestyle=linestyle, linewidth=1.5, label=label)
rho = (grad_norms[0] / grad_norms[-1]) ** (1 / (depth - 1))
print(f"std {std:.4f}: act std {act_stds[0]:8.3g} -> {act_stds[-1]:8.3g} | "
f"grad {grad_norms[-1]:8.3g} (output) -> {grad_norms[0]:8.3g} (input) | rho = {rho:.2f}")
axes[0].set_yscale("log")
axes[0].set_xlabel("Layer")
axes[0].set_ylabel("Activation Standard Deviation")
axes[1].set_yscale("log")
axes[1].set_xlabel("Layer")
axes[1].set_ylabel("Gradient Norm")
fig.legend(*axes[0].get_legend_handles_labels(), loc="upper center",
bbox_to_anchor=(0.5, 1.12), ncol=3, fontsize=9)
plt.show()std 0.0442: act std 0.419 -> 7.31e-10 | grad 15.3 (output) -> 3.5e-08 (input) | rho = 0.50
std 0.0884: act std 0.628 -> 0.131 | grad 15.3 (output) -> 4.35 (input) | rho = 0.96
std 0.2652: act std 0.863 -> 0.836 | grad 15.3 (output) -> 1.28e+05 (input) | rho = 1.37
Answer. Small weights make activations and gradients decay. Overly large weights distort their scale in the other direction. A well-scaled initialization reduces this repeated amplification or contraction at the start of training.
Question 2: What Do Residual Connections and Layer Normalization Change?
We return to the too-small initialization. Layer normalization explicitly resets activation scale, while a residual block adds an identity path to both the forward and backward computations.
std_small = 0.5 / width ** 0.5
def run_residual(layers, x, readout):
h = x
activations = []
for layer in layers:
h = h + torch.tanh(layer(h))
h.retain_grad()
activations.append(h)
loss = (h @ readout.to(h)).sum()
loss.backward()
activation_stds = [a.detach().std().item() for a in activations]
gradient_norms = [a.grad.norm().item() for a in activations]
return activation_stds, gradient_norms
def run_layernorm(layers, x, readout):
ln = torch.nn.LayerNorm(width, elementwise_affine=False)
h = x
activations = []
for layer in layers:
h = ln(torch.tanh(layer(h)))
h.retain_grad()
activations.append(h)
loss = (h @ readout.to(h)).sum()
loss.backward()
activation_stds = [a.detach().std().item() for a in activations]
gradient_norms = [a.grad.norm().item() for a in activations]
return activation_stds, gradient_norms
torch.manual_seed(13)
act_plain, grad_plain = run_plain(make_layers(depth, width, std_small), x, readout)
torch.manual_seed(13)
act_residual, grad_residual = run_residual(make_layers(depth, width, std_small), x, readout)
torch.manual_seed(13)
act_layernorm, grad_layernorm = run_layernorm(make_layers(depth, width, std_small), x, readout)
runs = [
("No fix", act_plain, grad_plain, "--"),
("Residual connections", act_residual, grad_residual, "-"),
("Layer normalization", act_layernorm, grad_layernorm, ":"),
]Compare the three forward computations using the same initial weights.
fig, axes = plt.subplots(1, 2, figsize=(9, 3.2))
layer_idx = np.arange(1, depth + 1)
for label, acts, grads, linestyle in runs:
axes[0].plot(layer_idx, acts, linestyle=linestyle, linewidth=1.5, label=label)
axes[1].plot(layer_idx, grads, linestyle=linestyle, linewidth=1.5, label=label)
axes[0].set_yscale("log")
axes[0].set_xlabel("Layer")
axes[0].set_ylabel("Activation Standard Deviation")
fig.legend(*axes[0].get_legend_handles_labels(), loc="upper center",
bbox_to_anchor=(0.5, 1.12), ncol=3, fontsize=9)
axes[1].set_yscale("log")
axes[1].set_xlabel("Layer")
axes[1].set_ylabel("Gradient Norm")
plt.show()
for label, acts, grads, _ in runs:
print(f"{label:22s} act std {acts[0]:8.3g} -> {acts[-1]:8.3g} | "
f"grad {grads[-1]:8.3g} (output) -> {grads[0]:8.3g} (input)")
No fix act std 0.419 -> 7.31e-10 | grad 15.3 (output) -> 3.5e-08 (input)
Residual connections act std 1.08 -> 3.68 | grad 15.3 (output) -> 70.1 (input)
Layer normalization act std 1 -> 1 | grad 15.3 (output) -> 27 (input)Answer. Layer normalization keeps activation scale controlled. Residual connections let information bypass the small nonlinear branch and prevent the severe gradient decay seen in the plain network. The two mechanisms address related but distinct failures.
Question 3: What Happens at One Hundred Layers?
The final comparison increases the depth while holding the width and small initialization fixed. Double precision keeps the plain network’s tiny values measurable.
depth_new = 100 # change me!
# Double precision prevents the unfixed activations and gradients from underflowing.
x2 = torch.randn(256, width).double()
torch.manual_seed(13)
plain_layers = make_layers(depth_new, width, std_small)
for layer in plain_layers:
layer.double()
_, g_plain = run_plain(plain_layers, x2, readout)
torch.manual_seed(13)
residual_layers = make_layers(depth_new, width, std_small)
for layer in residual_layers:
layer.double()
_, g_res = run_residual(residual_layers, x2, readout)Compare how each gradient changes from the output layers back toward the input.
fig, ax = plt.subplots(figsize=(7, 3))
layer_idx = np.arange(1, depth_new + 1)
ax.plot(layer_idx, g_plain, linestyle="--", linewidth=1.4, label="No fix")
ax.plot(layer_idx, g_res, linewidth=1.4, label="Residual connections")
ax.set_yscale("log")
ax.set_xlabel("Layer")
ax.set_ylabel("Gradient Norm")
ax.legend(frameon=False, loc="center left")
plt.show()
Answer. The plain network’s gradient decays exponentially with depth. The residual path remains usable because every block includes an identity contribution instead of forcing all information through another small Jacobian.