Turning an equation into tensor code

Setup#

pip install torch

Verified with torch 2.5.1 and numpy 1.26.4 on CPU. Runs instantly.

Algorithm 1, transcribed#

Kingma and Ba's Adam is ten lines of pseudocode. Here it is as ten lines of numpy, with the paper's names kept exactly, and then checked against PyTorch's implementation on the same problem.

import numpy as np
import torch

def adam_from_paper(grad_fn, theta, alpha=0.001, beta1=0.9, beta2=0.999,
                    eps=1e-8, steps=100):
    """Algorithm 1 of Kingma and Ba (2015), line for line."""
    m = np.zeros_like(theta)
    v = np.zeros_like(theta)
    for t in range(1, steps + 1):                  # t starts at 1, not 0
        g = grad_fn(theta)
        m = beta1 * m + (1 - beta1) * g            # first moment
        v = beta2 * v + (1 - beta2) * g * g        # second moment, elementwise
        m_hat = m / (1 - beta1 ** t)               # bias correction
        v_hat = v / (1 - beta2 ** t)
        theta = theta - alpha * m_hat / (np.sqrt(v_hat) + eps)
    return theta

scales = np.array([1.0, 100.0, 0.01])
mine = adam_from_paper(lambda th: 2 * scales * th, np.ones(3), alpha=0.01, steps=100)

w = torch.ones(3, dtype=torch.float64, requires_grad=True)
opt = torch.optim.Adam([w], lr=0.01)
s = torch.tensor(scales, dtype=torch.float64)
for _ in range(100):
    opt.zero_grad()
    (s * w**2).sum().backward()
    opt.step()

print("from the paper :", mine)
print("torch.optim.Adam:", w.detach().numpy())
print("largest disagreement:", np.abs(mine - w.detach().numpy()).max())

from the paper : [0.22444605 0.22444604 0.22444637]
torch.optim.Adam: [0.22444605 0.22444604 0.22444637]
largest disagreement: 5.551115123125783e-17

A disagreement of 5.6e-17 after a hundred steps is floating-point dust — the two computations are the same. That single line of output is worth more than any amount of staring at the code.

Note what this test does not need: a dataset, a model, a GPU, or a training run. It needs one function whose gradient you can write down.

The walkthrough#

float64 on the PyTorch side is deliberate. numpy defaults to double precision and PyTorch to single. Compare a float32 run against a float64 run and you get disagreements around 1e-7, which are indistinguishable from a real bug in the last decimals. Match the dtypes first, then interpret the difference.

t starts at 1. The paper's loop is one-based, and $t$ appears inside $\beta_1^t$. Start at 0 and the first bias correction divides by $1 - \beta_1^0 = 0$. This is the single most common transcription error in this algorithm.

m and v live outside the loop. They are state carried between steps. Recreating them inside the loop turns Adam into something that looks similar and behaves nothing like it — and the loss will still go down, which is why this bug survives.

g * g is elementwise, not a dot product. Whenever a paper squares a vector, decide whether it means elementwise or an inner product. Here the shapes tell you: v must match theta, so the square is elementwise.

A quadratic is the right test problem. Its gradient is exact and cheap, the minimum is known, and it exercises the per-coordinate behaviour with three different scales. Testing an optimiser against a real network first would mean debugging two things at once.

When there is no reference to compare against#

Most methods you reimplement will not exist in a library. Three substitutes, in order of strength.

Compare a fast version against a slow one. Write the plain nested-loop version first, then the vectorised version, and assert they agree on small inputs. The loop version is too slow to use and too simple to be wrong.

Check gradients numerically. If the paper defines a loss and its derivative, verify the derivative against finite differences — see verifying gradients with gradcheck.

Check the properties the maths promises. Probabilities sum to one. A symmetric kernel produces a symmetric matrix. A normalised embedding has length one. An attention row sums to one. Each of these is one assertion and catches a whole class of errors.

Common mistakes#

Rewriting the algorithm "more cleanly" during transcription. Fusing two lines, reordering updates, dropping a term that looks redundant. Every simplification is an unverified change. Transcribe, verify, then refactor.

Comparing against a library with different defaults. PyTorch's Adam defaults to eps=1e-8 outside the square root, matching Algorithm 1. Other frameworks place epsilon differently, and some apply weight decay by a different route. Read the reference implementation's own formula before calling a mismatch a bug in your code.

Testing only on the happy path. Try a zero gradient, a huge gradient, and a single parameter. Division-by-zero and broadcasting bugs surface at the edges, not in the middle.

Declaring success from a loss curve. A curve going down is compatible with a dozen wrong implementations. Insist on a numerical comparison somewhere in the chain, even if it is against your own slow version.

Try it yourself#

Break the transcription on purpose in four ways: start t at 0 with the divide-by-zero guarded, delete both bias corrections, move m = np.zeros_like(theta) inside the loop, and change g * g to g @ g. Predict which of the four still runs, which still converges, and which produces a disagreement small enough to miss.

What to learn next#

• A staged plan for reimplementing a paper — scaling this up from one equation to a whole method.
• Verifying gradients with gradcheck — the reference check when no library version exists.
• einsum in PyTorch — the safest route from index notation to tensor code.