Layer Normalization & RMSNorm

import numpy as np

rs = np.random.RandomState(0)
x = rs.randn(8) * 2.0 + 3.0  # nonzero mean on purpose
eps = 1e-5
gamma = np.ones_like(x)
beta = np.zeros_like(x)

def layernorm(x, gamma, beta, eps):
    mu = float(x.mean())
    var = ((x - mu) ** 2).mean()
    sigma_eps = np.sqrt(var + eps)
    normalized = (x - mu) / sigma_eps
    output = gamma * normalized + beta
    return {
        "mean": mu,
        "variance": float(var),
        "sigma_eps": float(sigma_eps),
        "normalized": normalized,
        "output": output,
    }

def rmsnorm(x, gamma, eps):
    mean_square = (x ** 2).mean()
    rms_eps = np.sqrt(mean_square + eps)
    inverse_rms = 1.0 / rms_eps
    normalized = x * inverse_rms
    output = gamma * normalized
    return {
        "mean_square": float(mean_square),
        "rms_eps": float(rms_eps),
        "inverse_rms": float(inverse_rms),
        "normalized": normalized,
        "output": output,
    }

ln = layernorm(x, gamma, beta, eps)
rn = rmsnorm(x, gamma, eps)
cos = float(
    (ln["output"] @ rn["output"])
    / (np.linalg.norm(ln["output"]) * np.linalg.norm(rn["output"]))
)

assert x.shape == gamma.shape == beta.shape == (8,)
assert ln["normalized"].shape == rn["normalized"].shape == x.shape
assert np.isclose(rn["rms_eps"] ** 2, rn["mean_square"] + eps)

print("mean(x), std(x):", round(float(x.mean()), 3), round(float(x.std()), 3))
print("mean_square, rms_eps:", round(rn["mean_square"], 3), round(rn["rms_eps"], 3))
print("cos(LN, RMS):", round(cos, 3))
print("LN mean/std:", round(float(ln["output"].mean()), 3), round(float(ln["output"].std()), 3))
print(
    "RMS mean/rms:",
    round(float(rn["output"].mean()), 3),
    round(float(np.sqrt((rn["output"] ** 2).mean())), 3),
)

The residual block in Python

This second program uses only Python's standard library. It computes the analytic reverse pass and one simultaneous SGD step from the fixed arrays, without training a model or calling a service. theta is flattened in W1 row order, b1, W2 row order, b2, gamma, beta order. Its gradients agree with independent central-difference estimates up to floating-point rounding, not bit for bit.

from math import sqrt, tanh, isfinite

X = [[1., -1., .5], [.2, .8, -.6]]
TARGET = [[.3, -.4, .2], [-.2, .5, -.1]]
THETA = [.5, -.3, .2, .4, -.4, .1, .1, -.2,
         .3, -.5, .2, -.2, .1, .4, .05, -.05, .1,
         1., .8, 1.2, 0., .1, -.1]
EPS = 1e-5

def grid(n, m, fn=lambda t, d: 0.):
    return [[fn(t, d) for d in range(m)] for t in range(n)]

def average(values):
    return sum(values) / len(values)

def norm(x, gamma, beta, wrong_axis=False):
    groups = ([[(t, d) for t in range(2)] for d in range(3)]
              if wrong_axis else [[(t, d) for d in range(3)] for t in range(2)])
    u, means, variances, scales = grid(2, 3), [], [], []
    for group in groups:
        mu = average([x[t][d] for t, d in group])
        var = average([(x[t][d] - mu)**2 for t, d in group])
        s = sqrt(var + EPS)
        means.append(mu); variances.append(var); scales.append(s)
        for t, d in group:
            u[t][d] = (x[t][d] - mu) / s
    return dict(input=x, mean=means, variance=variances, scale=scales,
                standardized=u, output=grid(2, 3, lambda t, d: gamma[d]*u[t][d]+beta[d]),
                axis="tokens" if wrong_axis else "features")

def norm_back(trace, upstream, gamma):
    groups = ([[(t, d) for t in range(2)] for d in range(3)]
              if trace["axis"] == "tokens" else [[(t, d) for d in range(3)] for t in range(2)])
    dx, u = grid(2, 3), trace["standardized"]
    for group, s in zip(groups, trace["scale"]):
        p = [upstream[t][d]*gamma[d] for t, d in group]
        correlation = average([q*u[t][d] for q, (t, d) in zip(p, group)])
        for q, (t, d) in zip(p, group):
            dx[t][d] = (q-average(p)-u[t][d]*correlation)/s
    dg = [sum(upstream[t][d]*u[t][d] for t in range(2)) for d in range(3)]
    db = [sum(upstream[t][d] for t in range(2)) for d in range(3)]
    return dx, dg, db

def forward(theta, x, target, mode):
    w1 = [theta[i:i+2] for i in range(0, 6, 2)]
    b1, w2 = theta[6:8], [theta[i:i+3] for i in (8, 11)]
    b2, gamma, beta = theta[14:17], theta[17:20], theta[20:23]
    n = None if mode == "post" else norm(x, gamma, beta, mode == "wrong-axis")
    u = x if n is None else n["output"]
    a = grid(2, 2, lambda t, h: sum(u[t][d]*w1[d][h] for d in range(3))+b1[h])
    h = [[tanh(v) for v in row] for row in a]
    f = grid(2, 3, lambda t, d: sum(h[t][j]*w2[j][d] for j in range(2))+b2[d])
    r = grid(2, 3, lambda t, d: x[t][d]+f[t][d])
    if n is None:
        n = norm(r, gamma, beta)
    y = n["output"] if mode == "post" else r
    loss = sum((y[t][d]-target[t][d])**2 for t in range(2) for d in range(3))/12
    return dict(branchInput=u, preactivation=a, hidden=h, branch=f, residual=r,
                normalization=n, output=y, loss=loss)

def inspect(mode="pre", rate=.1, theta=THETA, x=X, target=TARGET):
    def vector(v, length):
        return isinstance(v, (list, tuple)) and len(v) == length and all(
            type(a) in (float, int) and isfinite(a) and abs(a) <= 4 for a in v)
    if mode not in ("pre", "post", "wrong-axis") or rate not in (.1, 10):
        raise ValueError("Choose a declared placement and SGD rate (.1 or 10)")
    if not vector(theta, 23) or not all(isinstance(a, (list, tuple)) and len(a) == 2
                                      and all(vector(row, 3) for row in a) for a in (x, target)):
        raise ValueError("Expected 23 parameters and two 3-feature rows, finite and in [-4, 4]")
    before = forward(theta, x, target, mode)
    u, h, n = before["branchInput"], before["hidden"], before["normalization"]
    w1 = [theta[i:i+2] for i in range(0, 6, 2)]
    w2, gamma = [theta[i:i+3] for i in (8, 11)], theta[17:20]
    gy = grid(2, 3, lambda t, d: (before["output"][t][d]-target[t][d])/6)
    post = norm_back(n, gy, gamma) if mode == "post" else None
    direct = gy if post is None else post[0]
    gw2 = grid(2, 3, lambda j, d: sum(h[t][j]*direct[t][d] for t in range(2)))
    gb2 = [sum(direct[t][d] for t in range(2)) for d in range(3)]
    ga = grid(2, 2, lambda t, j: sum(direct[t][d]*w2[j][d] for d in range(3))*(1-h[t][j]**2))
    gw1 = grid(3, 2, lambda d, j: sum(u[t][d]*ga[t][j] for t in range(2)))
    gb1 = [sum(ga[t][j] for t in range(2)) for j in range(2)]
    gu = grid(2, 3, lambda t, d: sum(ga[t][j]*w1[d][j] for j in range(2)))
    gn = norm_back(n, gu, gamma) if post is None else post
    through = gn[0] if post is None else gu
    gradients = [v for row in gw1 for v in row]+gb1+[v for row in gw2 for v in row]+gb2+gn[1]+gn[2]
    before["gradients"] = dict(parameters=gradients,
        input=grid(2, 3, lambda t, d: direct[t][d]+through[t][d]), direct=direct,
        throughBranch=through, output=gy, activation=ga, branchInput=gu)
    updated = [v-rate*g for v, g in zip(theta, gradients)]
    return dict(placement=mode, rate=rate, before=before, parameters=updated,
                after=forward(updated, x, target, mode))

if __name__ == "__main__":
    for mode in ("pre", "post", "wrong-axis"):
        r = inspect(mode)
        print(mode, "loss", r["before"]["loss"], "->", r["after"]["loss"])
        print("W1[0,0]", THETA[0], "gradient", r["before"]["gradients"]["parameters"][0],
              "updated", r["parameters"][0])