Seeking Peer Review: Math Derivations for GPT-2 Backpropagation in llm.mojo

General
1

Hi Everyone,

In honor of Mojo v1.0.0’s imminent arrival, I am working on implementing backpropagation and autograd for an updated version of llm.mojo, which is a pure-Mojo implementation of GPT-2/3 inspired by Andrej Karpathy’s llm.c and Martin Dudeck’s llm.mojo.

To ensure the kernels are correct, I have spent the last few weeks writing out the complete mathematical derivations for the backward pass of every operation in the GPT-2 forward pass.

While my background is in mathematics, it has been some time since I last performed formal derivations at this level of granularity. I would appreciate any peer review or feedback from the community to confirm that these derivations are mathematically correct before I publish this project to the community showcase.

Derivation Methodology

To bridge the gap between matrix calculus and actual code, I wrote all derivations in index form (element-by-element) using Kronecker deltas rather than standard matrix notation. This approach:

  • Reduces multi-dimensional tensor operations to scalar calculus.
  • Uses independent probe indices to evaluate gradients coordinate-by-coordinate.
  • Directly maps to nested loop structures and indices used in writing the Mojo kernels.

I’m open to changing the representation to standard matrix notation, but the index form has a clearer analog to the actual kernels.

What is Covered

The companion document, The Mathematics of Backpropagation, includes complete step-by-step derivations for the following:

  1. Cross-Entropy Loss (with one-hot targets)
  2. Softmax (row-wise log-sum-exp stabilization)
  3. Linear Layer (Matmul)(including parameter gradients for weights and biases)
  4. GELU (approximate formulation and its derivative)
  5. LayerNorm (deriving statistics and dependencies row-wise)
  6. Multi-Head Attention (causal score matrix, softmax, attention probability, and value projections)
  7. Embeddings & Encoder(handling token and positional embeddings)
  8. The Full Backward Pass(integrating all components into a coherent backpropagation sequence)

Where to Find It

The LaTeX source file backprop.tex(and compiled assets) are available in the repository:

If anyone has time to look through the LaTeX document or has run into similar edge cases when deriving backward passes from scratch, I would love to hear your thoughts. Specifically, if you spot any notation issues or mathematical steps that could be simplified or made more rigorous for clarity, please let me know!

Thanks in advance,

Evan