ML coding in NumPy: softmax, logistic regression, k-means, AUC, MLP backprop

Before this lesson: the chain rule and backpropagation (Calculus 06, Deep learning 03); gradients of matrix expressions (Linear algebra 16); logistic regression, k-means and ROC curves at the level of a first ML course (for example, the ML fundamentals part of a university course). NumPy broadcasting and the axis= argument. Time: about 75 minutes, plus the notebook. Notation: $n$ examples, $d$ features, $K$ classes, $B$ batch size, $h$ hidden width; $X$ has one example per row.

---

Why this matters #

ML engineer and research engineer interviews often include an "ML coding" round: implement softmax and cross-entropy, logistic regression, k-means, a metric such as AUC, or an MLP with backprop, then answer follow-up questions. Attention and a full training loop come next (lesson 07, lesson 08). Interviewers grade correct math, vectorized code (no Python loop over examples), numerical stability, a check that the code works, and how clearly you explain while you code.

The same details matter in real work. The LLM pretraining loss is a softmax cross-entropy over about $10^5$ tokens, computed with log-sum-exp. Attention uses the same max shift. AUC is the standard score of a detector of rare events. A custom layer or a physics-informed loss needs a hand-written gradient that you check with finite differences.

1. Predict first #

2. The picture #

One routine for every algorithm (times for a 45-minute round):

1. Clarify (3 min). Shapes and dtypes; label coding ($\{0, 1\}$ or $\{-1, +1\}$); mean or sum loss; regularization (never on the bias); what to return. Write the signature with shape comments.
2. Math first (5 min). Write each formula in matrix form as a comment, one line each.
3. Stable form. Mark every exp and log. Use the max shift, log-sum-exp or softplus.
4. Vectorize. One matrix product per layer or per distance table. Loop over iterations, never over examples.
5. Test (5–10 min). A tiny case by hand; a library reference (SciPy, scikit-learn); a property: probabilities sum to 1, the loss decreases, a finite-difference gradient check passes.
6. Talk. Cost in time and memory, and the follow-ups you expect.

Step 1 in practice. The signature and the docstring already prevent most bugs:

def fit_logreg(X, y, lam=0.0, lr=0.1, steps=1000):
    """X: (n, d) float64, y: (n,) in {0, 1}.  Returns w: (d,), b: float.
    Loss: mean(softplus(z) - y * z) + lam/2 * ||w||^2,  z = X @ w + b  (b not penalized)."""
    ...

The simplest case of each idea.

• Softmax of two logits is a sigmoid: $\operatorname{softmax}(z_1, z_2)_2 = \sigma(z_2 - z_1)$, with $\sigma(t) = 1/(1 + e^{-t})$.
• Logistic regression in one dimension without bias has one weight: $p = \sigma(wx)$.
• k-means on a line with $k = 2$ splits the points into a left group and a right group.
• AUC with one positive and one negative is 1 if the positive scores higher, else 0 (and ½ on a tie).
• An MLP with one hidden unit is a chain of scalar functions, and backprop is the chain rule from the end.

Why softmax needs care. A logit is an unnormalized score; softmax turns a vector of logits $z$ into probabilities $p_k = e^{z_k}/\sum_j e^{z_j}$. In float64 the largest number is about $1.8 \times 10^{308}$, so $e^z$ is inf for $z > 709.78$, and $e^z$ rounds to 0 for $z < -745$. Both happen with real logits (for example after a bad initialization, or with large attention scores).

Interactive figure: loading…

3. Definitions and main results #

3.1 Stable softmax, log-softmax and cross-entropy #

For logits $z \in \mathbb{R}^K$: $\operatorname{softmax}(z)_k = e^{z_k} / \sum_{j=1}^K e^{z_j}$.

Shift invariance. $\operatorname{softmax}(z - c\mathbf{1}) = \operatorname{softmax}(z)$ for every number $c$.

Proof. $e^{z_k - c} / \sum_j e^{z_j - c} = e^{-c} e^{z_k} / (e^{-c} \sum_j e^{z_j})$, and $e^{-c}$ cancels. $\blacksquare$

Take $c = m = \max_j z_j$. Then every exponent is $\le 0$, so nothing overflows. The largest term is $e^0 = 1$, so the sum lies in $[1, K]$: no division by 0, and its $\log$ is $\ge 0$.

Log-sum-exp. $\operatorname{LSE}(z) = \log \sum_j e^{z_j} = m + \log \sum_j e^{z_j - m}$. Then $\log \operatorname{softmax}(z)_k = z_k - \operatorname{LSE}(z)$.

Cross-entropy from logits. For the true class $y$, the cross-entropy loss is $\ell = -\log p_y = \operatorname{LSE}(z) - z_y$. Never compute $-\log(\operatorname{softmax}(z)_y)$: a tiny $p_y$ rounds to 0, and the loss becomes $\infty$. Example: $z = (0, -1000)$ with $y = 2$ has $\ell = 1000$, but $-\log$ of the rounded probability gives inf.

Gradient. $\partial \ell / \partial z = p - e_y$, where $e_y$ is the one-hot vector of class $y$.

Proof. $\partial \operatorname{LSE}/\partial z_k = e^{z_k}/\sum_j e^{z_j} = p_k$, and $\partial z_y / \partial z_k = 1$ if $k = y$, else 0. $\blacksquare$

For a mean over a batch of $B$ examples, divide by $B$. The entries of $p - e_y$ sum to 0.

Binary case. Compute the sigmoid $\sigma(t)$ as $1/(1 + e^{-t})$ for $t \ge 0$ and as $e^{t}/(1 + e^{t})$ for $t < 0$, so exp never sees a large positive argument. The binary cross-entropy of a logit $z$ with label $y \in \{0, 1\}$ is $\ell = \operatorname{softplus}(z) - yz$, where the softplus is $\operatorname{softplus}(z) = \log(1 + e^z) = \max(z, 0) + \log(1 + e^{-|z|})$ (np.logaddexp(0, z) computes it).

Why $\operatorname{softplus}(z) - yz$: $-\log \sigma(z) = \log(1 + e^{-z}) = \operatorname{softplus}(z) - z$ and $-\log(1 - \sigma(z)) = \log(1 + e^{z}) = \operatorname{softplus}(z)$. Mix the two with $y$ and $1 - y$.

3.2 Logistic regression #

Data $X \in \mathbb{R}^{n \times d}$, labels $y \in \{0, 1\}^n$, weights $w \in \mathbb{R}^d$, bias $b$. Logits $z = Xw + b$, probabilities $p = \sigma(z)$ (elementwise). The loss with an L2 penalty $\lambda \ge 0$ is

Gradient, step by step.

1. $\frac{d}{dz}\operatorname{softplus}(z) = \frac{e^z}{1 + e^z} = \sigma(z)$, so $\partial \ell_i / \partial z_i = p_i - y_i$.
2. $\partial z_i / \partial w = x_i$ (row $i$ of $X$) and $\partial z_i / \partial b = 1$.
3. Average over $i$ and add the penalty:

The shapes check it: $X^\top$ is $d \times n$ and $p - y$ is $n \times 1$, so the product is $d \times 1$, the shape of $w$.

Convexity and the step size. The Hessian in $w$ is $\frac{1}{n} X^\top S X + \lambda I$ with $S = \operatorname{diag}(p_i(1 - p_i))$. It is positive semidefinite, so $L$ is convex. Since $p(1 - p) \le 1/4$, the gradient is Lipschitz with a constant (the smoothness constant) $\beta \le \lambda_{\max}(X^\top X)/(4n) + \lambda$. (To include $b$, append a column of ones to $X$.) Gradient descent with a step size $\eta \le 1/\beta$ never increases the loss (Optimization 02). Standardize the features first: it makes $\beta$ smaller and the problem better conditioned.

Separable data. The classes are linearly separable if some line (hyperplane) puts all positives on one side and all negatives on the other. Then, with $\lambda = 0$, the loss can come as close to 0 as you like but never reaches it: there is no minimizer. Gradient descent makes $\|w\|$ grow without bound (about like $\log t$), and the direction of $w$ tends to the max-margin direction (Soudry et al., 2018). Use $\lambda > 0$ or stop early.

Interactive figure: loading…

3.3 k-means #

Points $x_1, \dots, x_n \in \mathbb{R}^d$, $k$ centers $\mu_1, \dots, \mu_k$, and an assignment $c_i \in \{1, \dots, k\}$ for each point. The objective, called the inertia, is

Lloyd's algorithm alternates two steps until the assignment stops changing:

• (a) assign: $c_i = \arg\min_j \|x_i - \mu_j\|^2$ (on a tie, keep the current center);
• (b) update: set each $\mu_j$ to the mean of the points with $c_i = j$.

Theorem. $J$ never increases, and Lloyd's algorithm stops after a finite number of steps.

Proof.

1. Step (a) chooses the best center for each point with $\mu$ fixed, so $J$ does not increase. If some $c_i$ changes, that point moved to a strictly closer center (a tie keeps the old one), so $J$ strictly decreases.
2. Step (b) minimizes $J$ over $\mu$ with $c$ fixed. The cost of cluster $j$, $\sum_{i: c_i = j}\|x_i - \mu_j\|^2$, is a convex quadratic in $\mu_j$. Its gradient $2\sum_{i: c_i = j}(\mu_j - x_i)$ is zero at the mean. The old center is one of the candidates, so $J$ does not increase.
3. After step (b), $J$ depends only on the partition (the centers are its means). By steps 1–2, $J$ strictly decreases every time the partition changes, so no partition appears twice.
4. There are at most $k^n$ partitions, so the algorithm stops. $\blacksquare$

It stops at a local optimum, not always the global one (finding the global optimum is NP-hard in general). So run it from several starts (n_init in scikit-learn) and keep the lowest $J$.

Vectorized distances. $D_{ij} = \|x_i - \mu_j\|^2 = \|x_i\|^2 - 2x_i^\top\mu_j + \|\mu_j\|^2$. With the centers as the rows of $M$ ($k \times d$), this needs one matrix product $XM^\top$, of shape $(n, d)(d, k) \to (n, k)$, plus two broadcast vectors. Clip $D$ at 0, because round-off can make tiny values negative. The cost per iteration is $O(nkd)$.

k-means++ seeding. Pick the first center uniformly from the data. Pick each next center at random with probability proportional to $D(x)^2$, the squared distance from $x$ to its nearest chosen center. Then the expected $J$ is within a factor $O(\log k)$ of the optimum (Arthur and Vassilvitskii, 2007).

Empty cluster. If no point is assigned to center $j$, its mean is $0/0$. Re-seed it, for example at the point that is farthest from its own center.

Interactive figure: loading…

3.4 ROC AUC #

Scores $s_i$ (higher means "more likely positive") and labels $y_i \in \{0, 1\}$, with $n_+$ positives and $n_-$ negatives. The ROC curve plots the true positive rate (the fraction of positives above the threshold) against the false positive rate (the fraction of negatives above the threshold) while the threshold moves down from $+\infty$. AUC is the area under it.

Theorem (AUC as a probability). For a random positive with score $S^+$ and a random negative with score $S^-$:

This is the Mann–Whitney $U$ statistic divided by $n_+ n_-$.

Why this is true. When the threshold passes one negative, the curve moves right by $1/n_-$. The height at that moment is the fraction of positives that score above this negative. So the area added is $\frac{1}{n_-} \times$ (fraction of positives above it). Summed over all negatives, this is the fraction of (positive, negative) pairs in the right order. When a positive and a negative share a score, the curve moves diagonally, and the triangle under that step adds exactly half a pair.

Rank formula, $O(n \log n)$. Sort all scores in ascending order and give rank 1 to the smallest; tied scores get the average of their ranks. Let $R_+$ be the sum of the ranks of the positives. Then

Proof (no ties).

1. A positive with rank $r$ has $r - 1$ scores below it: some are negatives, some are positives.
2. Over all positives, the number of (positive, positive) pairs with one below the other is $0 + 1 + \dots + (n_+ - 1) = n_+(n_+ - 1)/2$.
3. So the number of (positive, negative) pairs with the positive above is the total $\sum (r - 1) = R_+ - n_+$ minus the (positive, positive) pairs: $R_+ - n_+ - n_+(n_+ - 1)/2 = R_+ - n_+(n_+ + 1)/2$.
4. Divide by the $n_+ n_-$ pairs. $\blacksquare$

With average ranks, a tied (positive, negative) pair adds $\tfrac12$, as in the theorem.

Interactive figure: loading…

Properties. AUC does not change under a strictly increasing transform of the scores (only the order matters). $\mathrm{AUC}(-s) = 1 - \mathrm{AUC}(s)$. A random scorer has AUC 0.5. AUC uses rates within each class, so it does not change when the class ratio changes. Under heavy imbalance it can look good while the precision is poor, so also report the area under the precision–recall curve (PR-AUC, or average precision). With only one class present, AUC is undefined: raise an error. The $O(n_+ n_-)$ count over all pairs is a good test oracle for your $O(n \log n)$ version.

3.5 MLP backprop #

A 2-layer MLP on a batch $X \in \mathbb{R}^{B \times d}$ with integer labels $y$, hidden width $h$ and $K$ classes. $Y \in \{0, 1\}^{B \times K}$ is the one-hot label matrix, $\odot$ is the elementwise product, $\operatorname{relu}(u) = \max(u, 0)$, and $dW$ means $\partial L / \partial W$.

Forward	Shape	Backward	Shape
$H = XW_1 + b_1$	$B \times h$	$G = (P - Y)/B$	$B \times K$
$A = \operatorname{relu}(H)$	$B \times h$	$dW_2 = A^\top G$, $db_2 = \sum_{\text{rows}} G$	$h \times K$, $K$
$Z = AW_2 + b_2$	$B \times K$	$dA = G W_2^\top$	$B \times h$
$P = \operatorname{softmax}(Z)$, by rows	$B \times K$	$dH = dA \odot \mathbf{1}[H > 0]$	$B \times h$
$L = \frac1B\sum_i \big(\operatorname{LSE}(Z_i) - Z_{i, y_i}\big)$	scalar	$dW_1 = X^\top dH$, $db_1 = \sum_{\text{rows}} dH$	$d \times h$, $h$

Read the backward column from the top: it runs from the loss back to the input.

• $G = (P - Y)/B$ is the cross-entropy gradient of section 3.1, one row per example, divided by $B$ because $L$ is a mean.
• For $Z = AW_2 + b_2$: $dW_2 = A^\top G$ and $dA = GW_2^\top$ (Linear algebra 16 derives these rules).
• ReLU passes the gradient where $H > 0$ and blocks it elsewhere.

Two rules catch most bugs: every gradient has the shape of its parameter, and $dW$ = (layer input)$^\top$ × (upstream gradient). The bias gradient sums over the batch, because $b$ is added to every row.

Gradient check. Compare entries with the central difference $\big(L(\theta + \varepsilon e_j) - L(\theta - \varepsilon e_j)\big)/(2\varepsilon)$, in float64 with $\varepsilon \approx 10^{-6}$; $e_j$ is the unit vector of entry $j$. Look at the relative error $|g_{\text{num}} - g| / (|g_{\text{num}}| + |g|)$: below about $10^{-6}$ means correct; about $10^{-2}$ or more means a bug. Check a tiny network, or a few random entries of a big one: a full check costs two forward passes per parameter (Calculus 07).

Answers to section 1. (1) [nan, nan]: both exponentials overflow to inf, and inf/inf is nan. The true value is $(\sigma(-1), \sigma(1)) \approx (0.2689, 0.7311)$: shift by $m = 1001$ to get $e^{-1}/(e^{-1} + 1)$ and $1/(e^{-1} + 1)$. (2) The positives score $\{0.9, 0.7\}$ and the negatives $\{0.8, 0.6\}$. Of the 4 pairs, 3 put the positive higher ($0.9 > 0.8$, $0.9 > 0.6$, $0.7 > 0.6$, but $0.7 < 0.8$), so $\mathrm{AUC} = 3/4 = 0.75$.

4. Worked example #

Softmax cross-entropy. $z = (2, 1, 0)$, and the true class is the first one (index 0 in code).

1. $m = 2$, so $z - m = (0, -1, -2)$ and $\sum_j e^{z_j - m} = 1 + 0.3679 + 0.1353 = 1.5032$.
2. $\operatorname{LSE}(z) = 2 + \log 1.5032$, and $\log 1.5032 = 0.4076$, so $\operatorname{LSE}(z) = 2.4076$.
3. Divide $(1, 0.3679, 0.1353)$ by $1.5032$: $p = (0.6652, 0.2447, 0.0900)$.
4. $\ell = \operatorname{LSE}(z) - z_1 = 2.4076 - 2 = 0.4076$.
5. $\partial \ell / \partial z = p - e_1 = (-0.3348, 0.2447, 0.0900)$. The entries sum to 0.

Logistic regression, one step. One feature, no bias: $x = (1, 2, -1)$, $y = (1, 1, 0)$, start at $w = 0$.

1. $z = 0$, so $p = \sigma(0) = (0.5, 0.5, 0.5)$, and $L(0) = \log 2 = 0.6931$.
2. $p - y = (-0.5, -0.5, 0.5)$, so $x^\top(p - y) = 1 \cdot (-0.5) + 2 \cdot (-0.5) + (-1) \cdot 0.5 = -2$, and $\nabla L = \frac13 x^\top(p - y) = -0.6667$.
3. With $\eta = 1$: $w = 0 - 1 \cdot (-0.6667) = 0.6667$, so $z = wx = (\tfrac23, \tfrac43, -\tfrac23)$. Use $\operatorname{softplus}(z) - z = \operatorname{softplus}(-z)$ for the two examples with $y = 1$. The three terms $\operatorname{softplus}(z_i) - y_i z_i$ are $\operatorname{softplus}(-\tfrac23) = 0.4144$, $\operatorname{softplus}(-\tfrac43) = 0.2340$ and $\operatorname{softplus}(-\tfrac23) = 0.4144$ (here $y = 0$). So $L = \frac13(0.4144 + 0.2340 + 0.4144) = 0.3542$.
4. The step was safe: $\beta \le x^\top x/(4n) = 6/12 = 0.5$, so any $\eta \le 2$ works.
5. The data are separable ($x > 0 \iff y = 1$). If you keep going, $w$ grows without end: about 2.16 after 10 steps, 4.24 after 100 and 6.51 after 1000, while the loss falls toward 0 (0.077, 0.0096, 0.00099).

k-means on a line. $x = (0, 1, 4, 5, 6)$, $k = 2$, start at $\mu = (0, 1)$.

Iteration	Clusters after (a)	$J$ after (a)	$\mu$ after (b)	$J$ after (b)
1	$\{0\}$, $\{1, 4, 5, 6\}$	50	$(0, 4)$	14
2	$\{0, 1\}$, $\{4, 5, 6\}$	6	$(0.5, 5)$	2.5
3	unchanged: stop	2.5	$(0.5, 5)$	2.5

Check row 1. After (a), the point 1 is a center itself, and 4, 5, 6 are closer to 1 than to 0: $J = 0 + 0 + 3^2 + 4^2 + 5^2 = 50$. After (b), the second center is the mean $(1 + 4 + 5 + 6)/4 = 4$: $J = 0 + (3^2 + 0^2 + 1^2 + 2^2) = 14$. $J$ never goes up, as the theorem says.

AUC with a tie. Scores $s = (0.9, 0.7, 0.7, 0.4, 0.2)$, labels $y = (1, 0, 1, 1, 0)$, so $n_+ = 3$ and $n_- = 2$.

1. Ascending ranks: $0.2 \to 1$, $0.4 \to 2$, the two scores $0.7$ share $(3 + 4)/2 = 3.5$, $0.9 \to 5$.
2. $R_+ = 5 + 3.5 + 2 = 10.5$ and $n_+(n_+ + 1)/2 = 6$, so $\mathrm{AUC} = (10.5 - 6)/6 = 0.75$.
3. Pair check: $0.9$ beats both negatives (2); the positive $0.7$ ties the negative $0.7$ (½) and beats $0.2$ (1); $0.4$ beats only $0.2$ (1). Total 4.5 of 6 pairs. ✓

5. In code #

NumPy, SciPy and scikit-learn check every number above. Softmax and cross-entropy:

import numpy as np
from scipy.special import softmax, logsumexp

z = np.array([1000.0, 1001.0])
with np.errstate(over="ignore", invalid="ignore"):
    print(np.exp(z) / np.exp(z).sum())        # [nan nan]
print(softmax(z))                             # [0.26894142 0.73105858]

z = np.array([2.0, 1.0, 0.0])
p = softmax(z)
print(p.round(4), round(logsumexp(z) - z[0], 4))   # [0.6652 0.2447 0.09  ] 0.4076
print((p - np.eye(3)[0]).round(4))                 # [-0.3348  0.2447  0.09  ]

with np.errstate(divide="ignore"):
    print(-np.log(softmax([0.0, -1000.0])[1]), logsumexp([0.0, -1000.0]) + 1000.0)   # inf 1000.0

One gradient step of logistic regression:

from scipy.special import expit                    # the sigmoid

x = np.array([1.0, 2.0, -1.0]); y = np.array([1.0, 1.0, 0.0])
loss = lambda w: np.mean(np.logaddexp(0, w * x) - y * w * x)   # softplus(z) - y z
grad = lambda w: np.mean((expit(w * x) - y) * x)
print(round(loss(0.0), 4), round(grad(0.0), 4))     # 0.6931 -0.6667
w1 = 0.0 - 1.0 * grad(0.0)                          # one step with eta = 1
print(round(w1, 4), round(loss(w1), 4))             # 0.6667 0.3542

k-means and AUC:

from sklearn.cluster import KMeans
from sklearn.metrics import roc_auc_score
from scipy.stats import mannwhitneyu

X = np.array([[0.0], [1.0], [4.0], [5.0], [6.0]])
km = KMeans(n_clusters=2, init=np.array([[0.0], [1.0]]), n_init=1).fit(X)
print(km.cluster_centers_.ravel(), km.labels_, km.inertia_)   # [0.5 5. ] [0 0 1 1 1] 2.5

print(roc_auc_score([1, 0, 1, 0], [0.9, 0.8, 0.7, 0.6]))      # 0.75
yy = np.array([1, 0, 1, 1, 0])
s = np.array([0.9, 0.7, 0.7, 0.4, 0.2])
U = mannwhitneyu(s[yy == 1], s[yy == 0]).statistic           # 4.5 winning pairs
print(roc_auc_score(yy, s), U / (3 * 2))                       # 0.75 0.75

The MLP table, checked with central differences on every parameter:

rng = np.random.default_rng(0)
B, d, h, K = 4, 3, 5, 2
Xb = rng.normal(size=(B, d)); yb = np.array([0, 1, 1, 0])
P0 = {"W1": rng.normal(size=(d, h)), "b1": np.zeros(h), "W2": rng.normal(size=(h, K)), "b2": np.zeros(K)}

def forward(P):
    H = Xb @ P["W1"] + P["b1"]; A = np.maximum(H, 0)
    Z = A @ P["W2"] + P["b2"]
    L = np.mean(logsumexp(Z, axis=1) - Z[np.arange(B), yb])
    return L, (H, A, Z)

def backward(P, cache):
    H, A, Z = cache
    G = (softmax(Z, axis=1) - np.eye(K)[yb]) / B        # dL/dZ, shape (B, K)
    dA = G @ P["W2"].T
    dH = dA * (H > 0)
    return {"W2": A.T @ G, "b2": G.sum(0), "W1": Xb.T @ dH, "b1": dH.sum(0)}

L, cache = forward(P0)
g = backward(P0, cache)
eps, worst = 1e-6, 0.0
for name, Wt in P0.items():
    for idx in np.ndindex(Wt.shape):
        old = Wt[idx]
        Wt[idx] = old + eps; Lp = forward(P0)[0]
        Wt[idx] = old - eps; Lm = forward(P0)[0]
        Wt[idx] = old
        num = (Lp - Lm) / (2 * eps)
        worst = max(worst, abs(num - g[name][idx]) / max(1e-12, abs(num) + abs(g[name][idx])))
print({k: v.shape == P0[k].shape for k, v in g.items()})   # {'W2': True, 'b2': True, 'W1': True, 'b1': True}
print(f"largest relative error: {worst:.1e}")               # about 1e-8: correct

6. Where it shows up #

• The logistic loss is the last layer of every binary classifier. Platt scaling (calibration) is one-dimensional logistic regression on the scores. RLHF reward models train with the Bradley–Terry loss $-\log\sigma(r_{\text{chosen}} - r_{\text{rejected}})$: the logistic loss on a score difference (lesson 11).
• Softmax cross-entropy is the LLM pretraining loss over a vocabulary of about $10^5$ tokens. Fused kernels compute the LSE without storing the $B \times V$ probabilities. Attention uses the same max shift, and FlashAttention computes it online, block by block (ML systems 04).
• AUC scores detectors of rare events: fraud, earthquake detection, extreme-weather or flood warnings. With rare positives, report PR-AUC as well.
• k-means builds codebooks: vector quantization, the discrete speech targets of HuBERT (k-means on acoustic features), weather-regime clustering of pressure fields, and the initialization of Gaussian mixture models.
• Hand-written backprop is needed for custom layers and fused kernels, for adjoints of simulators, and to debug a gradient when autograd differentiates a function you did not intend.

7. Common mistakes #

• Overflow and $\log 0$. A naive softmax or sigmoid; cross-entropy computed from probabilities. Work with logits: the max shift, LSE and softplus.
• Broadcasting (n, 1) - (n,) gives an $(n, n)$ matrix without any error. Keep p and y both of shape (n,), and assert the shapes.
• Mean versus sum. A mean loss with a summed gradient makes the step $n$ times too large.
• Penalizing the bias, or not standardizing the features before gradient descent.
• Expecting convergence on separable data with $\lambda = 0$. There is no minimizer; $\|w\|$ grows for ever.
• k-means. An empty cluster gives nan; only one random start; a Python loop for the distances; negative squared distances from round-off.
• AUC. Ignoring ties (the answer then depends on the sort order); a fixed grid of thresholds instead of exact ranks; no check for a single class.
• Backprop. Not summing the bias gradient over the batch; dividing by $B$ twice (in $G$ and again in $dW$); a gradient check in float32 or with $\varepsilon = 10^{-12}$ (round-off then dominates).
• Changing the caller's data. X -= X.mean(0) modifies the input array in place; use X = X - X.mean(0).

8. Check yourself #

1. Show that $\operatorname{softmax}(z + c\mathbf 1) = \operatorname{softmax}(z)$. Why choose $c = -\max_j z_j$?
   Answer

   $e^{z_k + c}/\sum_j e^{z_j + c} = e^c e^{z_k}/(e^c \sum_j e^{z_j})$, and $e^c$ cancels. With $c = -\max_j z_j$ every exponent is $\le 0$ (no overflow), and the largest term is 1, so the sum is $\ge 1$ (no division by 0, and its log is $\ge 0$).
2. Logistic regression with labels $y_i \in \{-1, +1\}$ and $z_i = x_i^\top w$: write the mean loss and its gradient.
   Answer

   $L = \frac1n\sum_i \log(1 + e^{-y_i z_i})$ and $\nabla_w L = -\frac1n\sum_i y_i\,\sigma(-y_i z_i)\,x_i$, because $\frac{d}{dz}\log(1 + e^{-yz}) = -y\,\sigma(-yz)$.
3. Why can the update step (b) of k-means not increase $J$?
   Answer

   With the assignment fixed, the cost of cluster $j$ is a convex quadratic in $\mu_j$; its gradient $2\sum_{i: c_i = j}(\mu_j - x_i)$ is zero at the mean, so the mean minimizes it. The old center is one of the candidates, so the cost cannot go up.
4. All scores are equal. What is the AUC? What if you negate all scores of a model with AUC 0.8?
   Answer

   0.5: every (positive, negative) pair is a tie and counts ½. Negating the scores reverses every strict order and keeps the ties, so the AUC becomes $1 - 0.8 = 0.2$.
5. In the 2-layer MLP, what are the shapes of $G$, $dA$ and $db_1$, and how do you compute $db_1$?
   Answer

   $G$ is $B \times K$, $dA = GW_2^\top$ is $B \times h$, and $db_1$ has shape $(h,)$, like $b_1$: $db_1 = \sum_{\text{rows}} dH$, the sum of $dH = dA \odot \mathbf 1[H > 0]$ over the batch.
6. Your logistic regression on separable data with $\lambda = 0$ never converges: the loss keeps falling and $\|w\|$ keeps growing. Bug or expected? What do you do?
   Answer

   Expected: there is no finite minimizer, and the direction of $w$ tends to the max-margin direction. Add $\lambda > 0$, or stop early (for example on a validation loss).

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Why does subtracting $\max_j z_j$ from the logits not change the softmax?

   Show answer

2. Write log-sum-exp in its stable form.

   Show answer

3. What is the gradient of the softmax cross-entropy $\ell = \operatorname{LSE}(z) - z_y$ with respect to the logits?

   Show answer

4. What is the gradient of the mean logistic loss with the penalty $\frac{\lambda}{2}\|w\|^2$?

   Show answer

5. What happens to logistic regression on linearly separable data with $\lambda = 0$?

   Show answer

6. Why does Lloyd's algorithm never increase the k-means objective?

   Show answer

7. What is AUC as a probability?

   Show answer

8. Compute AUC from ranks.

   Show answer

9. How do you get the weight gradient of a linear layer $Y = XW + b$ from the upstream gradient $dY$?

   Show answer

10. Which relative error in a central-difference gradient check means "correct"?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

Which tool for which wording.

The problem says …	Try …
"probabilities from scores", "softmax"	subtract the row maximum, axis=-1, keepdims=True
"loss from logits", "cross-entropy"	$\operatorname{LSE}(z) - z_y$; binary: np.logaddexp(0, z) - y * z
"pairwise distances", "nearest center"	$\|x\|^2 - 2XM^\top + \|\mu\|^2$, clip at 0, argmin(axis=1)
"AUC without scikit-learn"	average ranks (scipy.stats.rankdata), rank formula; pair count as the test
"backprop by hand"	the shape table; $dW = (\text{input})^\top (\text{upstream})$; sum the bias gradient
"verify the gradient"	central differences in float64, relative error
"too much data for memory"	mini-batches: SGD for logistic regression, mini-batch k-means

The routine to say aloud. "Inputs are … with shapes …; I return … The loss in matrix form is … The stable form uses … I vectorize over examples with one matrix product. I test a tiny case by hand, compare with scikit-learn, and run a gradient check. The cost is …"

Typical questions and short model answers.

• "Implement logistic regression." Write the loss with softplus, the gradient $X^\top(p - y)/n + \lambda w$, and gradient descent with a fixed step $\eta \le 1/\beta$ after standardizing. Follow-ups: why not squared error on $\sigma(z)$ (it is not convex in $w$, and its gradient vanishes when $\sigma$ saturates); multiclass (softmax, gradient $X^\top(P - Y)/n$); Newton's method (a few steps, each solving a $d \times d$ system).
• "Implement k-means." Lloyd's algorithm with vectorized distances, k-means++ seeding, an empty-cluster rule, and several starts. Follow-ups: choosing $k$ (elbow of $J$, silhouette score); clusters that are not round (Gaussian mixtures, spectral clustering); cost $O(nkd)$ per iteration; k-means as the zero-variance limit of EM for a Gaussian mixture.
• "Compute AUC without scikit-learn." Sort, average ranks for ties, rank formula, $O(n \log n)$; test against the $O(n_+ n_-)$ pair count. Follow-ups: imbalance (PR-AUC); calibration (AUC ignores it, because only the order matters).
• "Write softmax." Subtract the max along the class axis. Follow-ups: float16 (keep the shift and the sum in float32); a temperature $T$, $\operatorname{softmax}(z/T)$; the online version used by FlashAttention.
• "Backprop through a 2-layer MLP." Give the shape table, then $G = (P - Y)/B$ and the four gradients. Follow-ups: how you verify it (gradient check); the memory of the stored activations ($H$ or the ReLU mask, and $A$).

Typical follow-ups on any of them. What is the time and memory cost? What breaks with float32? How do you test it without a library? What changes for a sum loss instead of a mean?

10. Practice #

The notebook for this lesson runs in Colab. Each exercise has a check cell and a folded solution:

1. log_softmax and cross_entropy from logits, checked against scipy.special on random and extreme logits.
2. fit_logreg by gradient descent, checked against sklearn.linear_model.LogisticRegression with the matching penalty $C = 1/(n\lambda)$.
3. kmeans with k-means++ seeding and vectorized distances; its inertia must be within 1% of scikit-learn's on the same data.
4. roc_auc with ties by the rank formula, checked against sklearn.metrics.roc_auc_score and the $O(n_+ n_-)$ pair count.
5. mlp_backward for the 2-layer MLP, checked with your own central-difference gradient check.

Timed practice (20 minutes each, in an empty file, no help): write each of the five again and test it against the library. Public problem collections such as Deep-ML (sections on machine learning and deep learning) have short versions of the same tasks: softmax, logistic regression, k-means, a single neuron with backpropagation.

Say it aloud (3 minutes): explain to an imaginary interviewer why you subtract the maximum in softmax, why you never compute $-\log(\operatorname{softmax}(z)_y)$, and what the gradient with respect to the logits is.

11. Watch and read #

ROC and AUC, Clearly Explained!StatQuest with Josh Starmer · 16 min

How the ROC curve is built by moving the threshold, and what the area means.

StatQuest: K-means clusteringStatQuest with Josh Starmer · 9 min

The assign and update steps on a small example, and why you try several starts.

StatQuest: Logistic RegressionStatQuest with Josh Starmer · 9 min

The model and the S-shaped curve; this lesson adds the loss, the gradient and the code.

Building makemore Part 4: Becoming a Backprop NinjaAndrej Karpathy · 1 h 55 min

Manual backprop through an MLP, tensor by tensor, with a check against autograd; the cross-entropy part matches section 3.5.

• Stanford CS231n notes: Linear classification (softmax, cross-entropy and the numerical-stability trick) and Backpropagation, intuitions (gradients of matrix products by their shapes).
• Andrew Ng and Tengyu Ma, CS229 lecture notes: logistic regression and its gradient, Newton's method, and k-means.
• T. Fawcett, An introduction to ROC analysis, Pattern Recognition Letters (2006): ROC curves, AUC and its link to the Wilcoxon (Mann–Whitney) statistic, efficient computation with ties.
• D. Arthur and S. Vassilvitskii, k-means++: the advantages of careful seeding (SODA 2007): the seeding rule and the $O(\log k)$ guarantee.
• D. Soudry et al., The implicit bias of gradient descent on separable data (2018): why $\|w\|$ grows like $\log t$ and the direction tends to the max-margin solution.
• scikit-learn, LogisticRegression (the parameter $C$ is the inverse penalty strength) and roc_auc_score.
• Andrej Karpathy, Neural Networks: Zero to Hero: the code of the lecture series, including the backprop exercise.