Backpropagation & Tied Weights

Chain Rule & Tied Gradients

Analyze analytical weight gradients dL/dW_j = (p_j - y_j) * x and why unobserved tokens require balanced regularization.

Grad rule: dL/dWj = (pj - yj) * x

Plain Language Intuition

Backpropagation answers: 'How much did each of our 185K weights contribute to the prediction error, and in which direction should we nudge them?' It works backward through the chain rule. Residual skip connections are critical because their derivative contains the identity matrix (+I), letting gradients flow across layers without vanishing.

Production Real-World Context

PyTorch handles backprop automatically with loss.backward(). The manual analytical derivations shown here are what autograd calculates under the hood. Understanding weight gradients explains subtle training failures like unobserved token collapse and the AdamW weight decay necessity.

Prerequisites:

Training Loop (Step 8)

Transformer Block (Step 6)

The Chain Rule (Backpropagation)

Rumelhart, Hinton & Williams 1986

Intuition: To determine how a single parameter affects loss, trace backward through every mathematical transformation - multiplying local Jacobian matrices along the path.

Residual Connection Gradient Highway

Derived from Pre-LN Transformer

Intuition: The identity matrices (I) in the derivative mean gradients travel directly through skip connections without diminishing, even if attention or FFN gradients vanish.

Tied Weight Gradient Derivation

Tied Embeddings Derivation

Intuition: For unobserved tokens (y_j = 0), the gradient is strictly p_j * x. In tied models, this pulls unused token embeddings toward average layer activations.

AdamW Optimizer with Decoupled Weight Decay

Loshchilov & Hutter 2019, arXiv:1711.05101

Intuition: AdamW maintains momentum (m_t) and adaptive per-parameter learning rates (v_t), decoupling weight decay (lambda) to prevent tied weight drift.

Unobserved Tokens (Control IDs 0–9)

In byte-level vocabularies, ASCII control codes 0 to 9 never appear in natural Shakespeare or Hafez text. Because their target $y_j = 0$ permanently, their gradients are strictly $p_j \cdot x$. In inspections, these unobserved rows exhibit parallel vertical striping reflecting mean residual activations.

W_j^(t+1) = W_j^(t) - η * (p_j * x + λ * W_j)

Printable Subwords (IDs ≥ 32)

In contrast, printable ASCII and Persian UTF-8 subwords receive active updates $(p_j - 1) \cdot x$ whenever they appear in text. This breaks symmetry and pushes each subword embedding into its unique semantic direction in 64D space, eliminating vertical striping.

Var(W_print) > 0, Uniformity > 0.94

Analytical vs Autograd Verification - PyTorch Codepython

28 lines

Key Takeaways & Core Rules

The chain rule decomposes dL/dw into a product of local Jacobian matrices backwards through the network.
Residual connections create a gradient highway: the identity matrix (+I) allows error signals to skip layers.
Analytical output weight gradient is dL/dW_j = (p_j - y_j) * x.
AdamW decouples weight decay from adaptive gradient scaling, stabilizing training across 50 epochs.

Try This Experiment:

Run the verification code snippet above in Python. Notice how the analytical formula matches autograd down to machine precision (< 1e-7).