"""
Implementation of Chapters 5-9 of:
"The Simple Mathematics of Large Language Models" (Breeden, 2026)
Chapter 5 - The Core Mechanism: Learned Weighted Averaging
Chapter 6 - Multiple Relations
Chapter 7 - Depth: Stacking Multiple Layers
Chapter 8 - Position Information
Chapter 9 - From Representations to Probabilities
Variable names are kept identical to the document wherever possible:
WA, WB, WC, WO, ai, bj, cj, rij, alpha (influence weight), hi, h,
W1, W2, b1, b2, pi, uv, zv, k, d, H, L.
"""
import numpy as np
# ---------------------------------------------------------------------------
# Chapter 5: The Core Mechanism - Learned Weighted Averaging
# ---------------------------------------------------------------------------
def project(e, W):
"""W e in R^k -- 'a set of learned features extracted from the
original representation.'"""
return W @ e
def influence_score(ai, bj, k):
"""rij = ai . bj -- 'how well what position j offers aligns with
what position i is receptive to.' Scaled by sqrt(k) to prevent the
dot products from becoming very large."""
return (ai @ bj) / np.sqrt(k)
def influence_weights(rij_vector):
"""alpha_ij = e^rij / sum_m e^rim -- the inverse multinomial logit.
'Larger inputs receive exponentially more weight.'"""
r = rij_vector - np.max(rij_vector) # numerical stability, does not change ratios
exp_r = np.exp(r)
return exp_r / np.sum(exp_r)
def core_mechanism(h_prev, WA, WB, WC, i):
"""The Complete Operation (Section 5.8):
1. ai = WA ei for the single position i being updated
2. bj = WB ej, cj = WC ej for each context position j
3. rij = ai . bj / sqrt(k)
4. alpha_ij = softmax(rij)
5. hi = sum_j alpha_ij cj -- 'a new representation for position i
that incorporates information from the entire context.'
"""
n = h_prev.shape[0]
k = WA.shape[0]
ai = project(h_prev[i], WA) # what position i is receptive to
bj = np.stack([project(h_prev[j], WB) for j in range(n)]) # what position j offers
cj = np.stack([project(h_prev[j], WC) for j in range(n)]) # content to aggregate from j
rij = np.array([influence_score(ai, bj[j], k) for j in range(n)])
alpha = influence_weights(rij) # influence weight
hi = np.sum(alpha[:, None] * cj, axis=0) # weighted average
return hi, alpha
# ---------------------------------------------------------------------------
# Chapter 6: Multiple Relations
# ---------------------------------------------------------------------------
def multiple_relations(h_prev, WA_h, WB_h, WC_h, WO, i):
"""'Rather than learning a single set of transformation matrices,
learn H sets in parallel... These are concatenated and projected
back to the original dimension.'
WA_h, WB_h, WC_h: lists of length H (one matrix set per relation).
"""
H = len(WA_h)
o_list = []
for h in range(H):
o_i_h, _ = core_mechanism(h_prev, WA_h[h], WB_h[h], WC_h[h], i)
o_list.append(o_i_h) # o^(h)_i
concatenated = np.concatenate(o_list) # [o^(1)_i ; ... ; o^(H)_i]
hi = WO @ concatenated
return hi
# ---------------------------------------------------------------------------
# Chapter 7: Depth - Stacking Multiple Layers
# ---------------------------------------------------------------------------
def relu(x):
"""phi(x) = max(0, x) -- 'simply replace negative values with zero.'"""
return np.maximum(0, x)
def nonlinear_transformation(hi, W1, b1, W2, b2, phi=relu):
"""hi <- W2 . phi(W1 hi + b1) + b2 -- a nonlinear function applied
element-wise, position by position."""
return W2 @ phi(W1 @ hi + b1) + b2
def residual_connection(hi, f_hi):
"""hi <- hi + f(hi) -- 'each layer makes an incremental additive
adjustment to the representation.'"""
return hi + f_hi
def normalize(hi, eps=1e-8):
"""Normalize each vector to have mean 0 and variance 1 across its
coordinates -- 'a form of standardization that keeps the numbers
in a well-behaved range.'"""
mean = np.mean(hi)
var = np.var(hi)
return (hi - mean) / np.sqrt(var + eps)
def transformer_layer(h_prev, WA_h, WB_h, WC_h, WO, W1, b1, W2, b2):
"""One full layer: multi-relation aggregation + feedforward,
each wrapped in a residual connection, followed by normalization.
h^(l)_i = Normalize(h^(l-1)_i + h~_i + W2 . max(0, W1(h^(l-1)_i + h~_i)))
"""
n = h_prev.shape[0]
h_out = np.zeros_like(h_prev)
h_tilde = np.stack([multiple_relations(h_prev, WA_h, WB_h, WC_h, WO, i)
for i in range(n)])
h_mid = h_prev + h_tilde # residual connection
for i in range(n):
f_hi = nonlinear_transformation(h_mid[i], W1, b1, W2, b2)
h_res = residual_connection(h_mid[i], f_hi) # incremental refinement
h_out[i] = normalize(h_res)
return h_out
def stacked_layers(h0, layer_params):
"""h^(0)_i = e_i (initial word vectors)
h^(l)_i = Layer_l(h^(l-1)_1, ..., h^(l-1)_n)_i
layer_params: list of dicts, one per layer, each with keys
WA_h, WB_h, WC_h, WO, W1, b1, W2, b2.
"""
h = h0
for params in layer_params:
h = transformer_layer(h, **params)
return h
# ---------------------------------------------------------------------------
# Chapter 8: Position Information
# ---------------------------------------------------------------------------
def add_absolute_position(ei, pi):
"""ei <- ei + pi -- early approaches added a position vector
directly to each word embedding."""
return ei + pi
def sinusoidal_position_encoding(position, d):
"""pi defined using mathematical functions such as sines and
cosines at different frequencies (Vaswani et al., 2017)."""
pi_vec = np.zeros(d)
for idx in range(0, d, 2):
freq = 1.0 / (10000 ** (idx / d))
pi_vec[idx] = np.sin(position * freq)
if idx + 1 < d:
pi_vec[idx + 1] = np.cos(position * freq)
return pi_vec
def rope_rotate(x, position, base=10000.0):
"""Rotary Position Embedding: 'rotates the vectors ai and bj by
angles proportional to their positions,' so that the dot product
ai . bj depends on the difference i - j, not on i and j individually."""
d = x.shape[0]
x_rot = np.zeros_like(x)
for idx in range(0, d, 2):
theta = position / (base ** (idx / d))
cos_t, sin_t = np.cos(theta), np.sin(theta)
x0, x1 = x[idx], x[idx + 1] if idx + 1 < d else 0.0
x_rot[idx] = x0 * cos_t - x1 * sin_t
if idx + 1 < d:
x_rot[idx + 1] = x0 * sin_t + x1 * cos_t
return x_rot
def alibi_bias(rij, i, j, slope):
"""ALiBi: 'adds a penalty to influence scores that increases
linearly with distance between positions.'"""
return rij - slope * abs(i - j)
# ---------------------------------------------------------------------------
# Chapter 9: From Representations to Probabilities
# ---------------------------------------------------------------------------
def token_score(uv, h_L_n):
"""zv = uv . h^(L)_n -- a score for each token v in the vocabulary."""
return uv @ h_L_n
def next_token_probabilities(U, h_L_n):
"""P(w_{n+1} = v | w1, ..., wn) = e^zv / sum_v' e^zv'
-- the inverse multinomial logit applied to vocabulary scores.
U: array of shape (|V|, d), one uv row per vocabulary token.
"""
zv = U @ h_L_n # score for each token v in the vocabulary
zv = zv - np.max(zv) # numerical stability
exp_zv = np.exp(zv)
return exp_zv / np.sum(exp_zv)
# ---------------------------------------------------------------------------
# Minimal smoke test
# ---------------------------------------------------------------------------
if __name__ == "__main__":
rng = np.random.default_rng(0)
n, d, k, H, L, vocab_size = 5, 16, 8, 4, 2, 20
h0 = rng.normal(size=(n, d))
for pos in range(n):
h0[pos] = add_absolute_position(h0[pos], sinusoidal_position_encoding(pos, d))
layer_params = []
for _ in range(L):
WA_h = [rng.normal(size=(k, d)) * 0.1 for _ in range(H)]
WB_h = [rng.normal(size=(k, d)) * 0.1 for _ in range(H)]
WC_h = [rng.normal(size=(k, d)) * 0.1 for _ in range(H)]
WO = rng.normal(size=(d, k * H)) * 0.1
W1 = rng.normal(size=(d, d)) * 0.1
b1 = np.zeros(d)
W2 = rng.normal(size=(d, d)) * 0.1
b2 = np.zeros(d)
layer_params.append(dict(WA_h=WA_h, WB_h=WB_h, WC_h=WC_h, WO=WO,
W1=W1, b1=b1, W2=W2, b2=b2))
h_L = stacked_layers(h0, layer_params) # h^(L)_i for all positions
U = rng.normal(size=(vocab_size, d)) * 0.1
probs = next_token_probabilities(U, h_L[-1]) # P(w_{n+1} = v | ...)
print("influence weights sum to 1:", np.isclose(probs.sum(), 1.0))
print("next-token probability distribution:\n", probs)
www.astronews.com
