More Math Lessons

This is a mixed-practice lesson. It covers three habits that separate code that works on a toy example from code that survives real training runs: keeping exponentials from overflowing, predicting shapes before running them, and expressing contractions with einsum. None of them is deep mathematics; all three cause daily bugs.

The idea

Numerical stability

Floating point has limits. In float32 the largest finite value is about $3.4 \times 10^{38}$, so exp(100) overflows to inf even though the true value is only $2.7 \times 10^{43}$. The log-sum-exp trick removes the danger:

Subtracting the max leaves exponents at most $0$, so every exponential is in $(0, 1]$ and cannot overflow. The answer is mathematically identical, and it also avoids underflow when all $x_i$ are very negative. Softmax uses the same shift: $\text{softmax}(x)_i = e^{x_i - m} / \sum_j e^{x_j - m}$. For losses, prefer logaddexp or a fused cross_entropy(logits, target) over log(softmax(x)), which loses precision twice.

Broadcasting

Two arrays combine when their shapes line up from the right. Dimensions match if they are equal or one of them is 1; a missing dimension is treated as 1. Then both are stretched to the common shape without copying data.

• (3, 1) with (1, 4) gives (3, 4).
• (5, 1) with (4,) gives (5, 4) — the (4,) is treated as (1, 4).
• (5,) with (5,) gives (5,), not (5, 5).

The trap is (N,) versus (N, 1): the first broadcasts a length-$N$ vector against rows to produce an $N \times N$ array, which can silently blow up memory. Add an explicit axis when the intent is a column.

Einsum

einsum names each axis with a letter and states which axes survive in the output. Repeated letters are contracted; letters that appear only in inputs are summed; letters in the output are kept.

• "i,i->" — dot product, a scalar.
• "ij,jk->ik" — matrix product.
• "ij,ij->" — Frobenius inner product.
• "bij,bjk->bik" — batched matrix product.
• "ii->" — trace.
• "ij->ji" — transpose.

Worked example

Logits $x = (1000, 1001, 1002)$. Naively, $e^{1002}$ overflows. With $m = 1002$: $e^{-2} = 0.1353$, $e^{-1} = 0.3679$, $e^{0} = 1$. The sum is $1.5032$, giving probabilities $(0.0900, 0.2447, 0.6652)$ that sum to 1. In exact arithmetic the unshifted softmax gives the same numbers; the shift only changes which rounding errors you get.

For broadcasting and einsum, take $A$ of shape $(2,3)$ and $B$ of shape $(3,2)$. Then einsum("ij,jk->ik", A, B) has shape $(2,2)$ and agrees with A @ B to the last bit.

In code

import numpy as np

x = np.array([1000.0, 1001.0, 1002.0])
m = x.max()

def softmax(v):
    e = np.exp(v - v.max())
    return e / e.sum()

print(softmax(x))                  # [0.09  0.2447 0.6652]

print((np.arange(3)[:, None] * np.ones((1, 4))).shape)   # (3, 4)

A = np.ones((2, 3)); B = np.ones((3, 2))
print(np.allclose(np.einsum("ij,jk->ik", A, B), A @ B))  # True
print(np.einsum("ij,ij->", A, A))  # 6.0, Frobenius inner product

Check yourself

1. Why does subtracting the maximum in log-sum-exp not change the result?
2. What shape does (5, 1) broadcast against (4,) produce, and why?
3. Write the einsum string for a batched matrix product and for a trace.

Key takeaways

• Shift by the max before exponentiating; use log-space ops in losses.
• Broadcasting aligns shapes from the right and stretches size-1 axes.
• einsum makes contractions explicit and turns shape bugs into string bugs.