Learning Is Optimization: Linear and Logistic Regression, Trained by Gradient Descent

“The model learned to do X” sounds like something happened inside it. What actually happened is much more mechanical. Someone wrote down a number that measures how wrong the model is, and a procedure nudged the model’s parameters to make that number smaller, over and over. Every model in this series, including the ones with hundreds of billions of parameters, is trained that way. The only things that change are the model, the number, and the nudging.

This post, the first of Part I: Learning from data in AI Engineering, from the Inside, builds that idea from nothing, on two small problems taken from this blog itself. It assumes Python, derivatives, and a little linear algebra; the one new tool, eigenvalues, is explained where it’s used. First we predict how many tokens a post costs from its word count. Then we train the first piece of the series’ running example, a router that decides whether a reader’s question belongs with the serving posts. Along the way we meet the three decisions every later post inherits:

• what “wrong” means (the loss);
• how to move downhill (gradient descent);
• how big a step is safe (the learning rate).

Table of contents

Three ways to learn, told by their loss

Machine learning is usually introduced as three families. The cleanest way to tell them apart is to ask what the loss is computed against.

• Supervised learning has a target for every example: a post’s actual token count, or a question’s correct route. The loss compares the model’s prediction to that target. Both problems in this post are supervised.
• Unsupervised learning has no targets. The loss is computed against the data itself: how well a few clusters explain the points, or how well the model can reconstruct its input. Language-model pretraining is often called self-supervised, because the target, the next token, comes free with the text (What “Generative” Means makes this precise).
• Reinforcement learning has no target for each action, only a reward that arrives later, like a game won or an answer a human preferred. The loss has to be built from that delayed signal, which is why RL gets its own post, Reinforcement Learning for Language Models, just before it is needed for RLHF.

In all three, “learning” means the same thing: choose parameters that make the loss small. So the questions are always which loss, and how to minimize it.

A first model: how many tokens does a post cost?

Every post on this blog is eventually read by a language model: a reader’s assistant, a search index, or the Notebook Assistant this series builds. Models are priced and limited in tokens, not words. So a natural first question is how many tokens a post of a given length costs.

Here is the real data: every English post on this blog as of this writing, with its word count and its token count under o200k_base, the tokenizer GPT-4o uses. (How that tokenizer splits text is the subject of Byte-Pair Encoding.)

The counts include everything in the post body, math and code included. That matters, because the ratio column isn’t constant. The post with the most math, transformer-attention-patterns, costs 1.94 tokens per word. LaTeX like \operatorname{softmax} is expensive to tokenize. The table-heavy comparison post is next at 1.85.

A model is a guess at the shape of the relationship, with a few numbers left open. The simplest guess is a straight line:

where $x$ is words (in thousands), $\hat y$ is the predicted tokens (in thousands), and the two parameters are the slope $w$, tokens per word, and the intercept $b$. Learning means choosing $w$ and $b$. To choose, we need to say what makes one line better than another.

Choosing what “wrong” means: L2 and L1 loss

For one post, the error is $e = \hat y - y$. A loss turns all the errors into one number. The two classic choices:

The first is the mean squared error (MSE), the second the mean absolute error (MAE). They look like a matter of taste. They aren’t: they encode different beliefs about what an error costs, and the cleanest way to see that is to make the model as small as possible.

Suppose the model is a single number $c$: every post costs $c$ tokens per word, and the $r_i$ are the observed ratios. Minimize each loss over $c$.

L2. Set the derivative to zero:

Squared error is minimized by the mean.

L1. The derivative of $|r_i - c|$ is $-1$ when $c < r_i$ and $+1$ when $c > r_i$. So the slope of the total is

It is zero exactly when as many points lie below $c$ as above. Absolute error is minimized by the median.

On our twelve ratios, a_loss.py opens with each English post’s word and o200k_base token counts on 2026-09-29 (full MDX body) and converts them to thousands. The rest computes the ratio and both minimizers:

# The simplest possible model: every post costs c tokens per word.
ratio = tokens / words

if __name__ == "__main__":
    mse = lambda c: ((ratio - c) ** 2).mean()
    mae = lambda c: np.abs(ratio - c).mean()
    grid = np.linspace(1.5, 2.0, 50001)
    print(f"L2 minimizer = {grid[np.argmin([mse(c) for c in grid])]:.4f}"
          f"   mean   = {ratio.mean():.4f}")
    lo, hi = np.sort(ratio)[5], np.sort(ratio)[6]   # the two middle values
    print(f"L1 is flat on [{lo:.4f}, {hi:.4f}]:"
          f" MAE = {mae(lo):.5f}, {mae((lo + hi) / 2):.5f}, {mae(hi):.5f}")
    print(f"median = {np.median(ratio):.4f}")

Full script: a_loss.py

L2 minimizer = 1.7378   mean   = 1.7378
L1 is flat on [1.7047, 1.7308]: MAE = 0.07269, 0.07269, 0.07269
median = 1.7178

Two things are worth seeing here.

• The mean is pulled up by the math-heavy post. One ratio of 1.94 moves the L2 answer to 1.738. The L1 answer sits at 1.705–1.731, where the typical post is. Squaring makes a large error count quadratically, so L2 bends toward outliers. L1 charges every error by its size and lets a single strange point be strange.
• With an even number of points, the L1 minimum isn’t a point. Any $c$ between the two middle ratios has exactly six points on each side, so the slope is zero across the whole interval. The printed MAE is identical at both ends and in the middle. The median is just the conventional midpoint of that flat stretch. This is the first hint of something that matters later: a loss can have many equally good answers, and then which one training finds depends on the procedure, not the loss.

Which is right depends on what an error costs. When the budget is a context window and one post blowing past it is a disaster, L2’s fear of large errors is appropriate. When the goal is the typical cost and the outlier is a genuine exception, L1 is. Neither is “the” loss. Choosing one is a modeling decision.

Fitting the line, exactly

Back to the two-parameter line. For L2 there is a closed-form answer. Setting both partial derivatives of the MSE to zero gives two linear equations (the normal equations), whose solution is

On our twelve posts this is $w = 1.672$ tokens per word and $b = 0.206$ thousand tokens, the fit that gradient descent reaches below. In matrix form, with a column of ones appended to the inputs for the bias, the normal equations read $X^\top X\,\theta = X^\top y$. Remember $X^\top X$: it comes back shortly as the curvature of the loss.

For L1 there is no formula. But least absolute deviation is a linear program, and a linear program’s optima sit at vertices, so there is always an L1-optimal line through two of the data points. We can check every pair. The best is $w = 1.747$, $b = -0.110$, running through the BPE post and the vLLM post.

The predictions agree within 15 tokens, but the lines are different lines. The math-heavy post sits 738 tokens above the L2 line, and because it sits near the middle of the word range, it mostly lifts the intercept. L1 ignores it and passes through two ordinary posts instead. The fits aren’t competing answers to one question. Each is the right answer to its own loss.

So why not always solve in closed form? Because it only exists for a few models. The normal equations work because the prediction is linear in the parameters and the loss is quadratic. Put a sigmoid on the output, as we will below, or stack layers, as every neural network does, and setting the gradient to zero no longer gives equations with a solution in closed form. Even when a closed form exists, it costs $O(nd^2 + d^3)$ for $n$ examples and $d$ parameters, and a language model has billions of parameters. We need a procedure that only ever asks “which way is downhill from here?”

Gradient descent

The gradient $\nabla\mathcal{L}$ is the vector of partial derivatives, one per parameter. It points in the direction the loss increases fastest. Gradient descent steps the other way:

where $\theta$ is all the parameters and $\eta$ is the learning rate, the step size. For our line, with $e_i = w x_i + b - y_i$:

Read those in words. The intercept moves against the average error: if we predict too high on average, $b$ goes down. The slope moves against the error weighted by the input: posts with more words pull harder on $w$. Every weight’s gradient in this series has this shape, an error times whatever the weight multiplied. (A bias multiplies 1, so its gradient is the error alone.)

Three steps by hand

Take just three posts: the comparison (1.732k words, 3.209k tokens), vLLM (2.584k, 4.405k) and TensorRT-LLM (3.435k, 6.199k). Start at $w = b = 0$ with $\eta = 0.1$.

At step 0 every prediction is zero, so every error is negative, both gradients are negative, and both parameters go up. The first step overshoots: $w = 2.55$ predicts far too many tokens, so the gradient flips sign and step 2 comes back down, below the answer. The loss still falls every step, from 22.7 to 1.3 in three steps. But $w$ zig-zags, 2.55 → 0.97 → 1.95, and that zig-zag is the learning rate telling us something.

How big a step is safe?

Near a minimum, any smooth loss looks like a bowl. MSE on a linear model is exactly a bowl:

Here $H$ is the Hessian, the matrix of second derivatives. (This $H$ is local to this post. Everywhere else in the series, $H$ is a head dimension.) For our line it is

which is the $X^\top X$ of the normal equations. The gradient of the bowl is $\nabla\mathcal{L} = H(\theta - \theta^*)$, so plugging it into the update rule, the distance to the optimum evolves as

Now the one new tool. $H$ is symmetric, so it has a set of perpendicular eigenvectors: directions along which multiplying by $H$ just scales the vector by a number $\lambda$, its eigenvalue. Along each eigenvector, the update above is a separate one-dimensional problem. Each step multiplies the remaining error in that direction by $(1 - \eta\lambda)$. That single factor explains everything:

• if $0 < \eta\lambda < 1$, the error shrinks without changing sign: a smooth approach;
• if $\eta\lambda = 1$, it vanishes in one step;
• if $1 < \eta\lambda < 2$, it shrinks but flips sign every step: the zig-zag we just saw;
• if $\eta\lambda = 2$, the factor is exactly $-1$, so the error flips sign forever without shrinking;
• if $\eta\lambda > 2$, $|1 - \eta\lambda| > 1$ and the error grows: divergence.

So gradient descent converges only if $\eta < 2/\lambda_\text{max}$. But the direction with the smallest eigenvalue shrinks by only $(1 - \eta \lambda_\text{min})$ per step, and with $\eta$ capped near $2/\lambda_\text{max}$ that factor is at best about $1 - 2/\kappa$. Here $\kappa = \lambda_\text{max}/\lambda_\text{min}$ is the condition number. A large $\kappa$ means a long, thin valley: the step size is set by the steep walls, and progress along the floor crawls.

For our data, with words in thousands:

So the largest stable learning rate is $2/30.72 = 0.065$, and the slow direction still takes hundreds of steps. step_size.py imports the counts from a_loss.py.

import numpy as np
from a_loss import words as x, tokens as y

def gradient_descent(x, y, eta, steps):
    w, b = 0.0, 0.0
    for t in range(steps):
        err = w * x + b - y               # prediction minus target
        w -= eta * 2 * (err * x).mean()   # dL/dw
        b -= eta * 2 * err.mean()         # dL/db
    return w, b, ((w * x + b - y) ** 2).mean()

for eta in [0.01, 0.05, 0.07]:
    w, b, mse = gradient_descent(x, y, eta, steps=400)
    print(f"raw words, eta={eta:<5}: w={w:9.4g}  b={b:9.4g}  MSE={mse:.4g}")

xs = (x - x.mean()) / x.std()             # standardize the feature
for eta in [0.5, 1.0, 1.05]:
    w, b, mse = gradient_descent(xs, y, eta, steps=400)
    print(f"standardized, eta={eta:<4}: w={w:9.4g}  b={b:9.4g}  MSE={mse:.4g}")

Full script: step_size.py

raw words, eta=0.01 : w=    1.654  b=   0.2806  MSE=0.08979
raw words, eta=0.05 : w=    1.672  b=   0.2073  MSE=0.08908
raw words, eta=0.07 : w=-3.273e+24  b=-8.102e+23  MSE=1.746e+50
standardized, eta=0.5 : w=    2.273  b=    6.148  MSE=0.08908
standardized, eta=1.0 : w=-6.963e-13  b=-3.979e-13  MSE=43.05
standardized, eta=1.05: w=-8.198e+16  b=-2.217e+17  MSE=5.588e+34

The first three lines are the theory, running.

• $\eta = 0.01$: stable, but after 400 steps it still hasn’t reached the least-squares fit ($w = 1.654$, not 1.672). That’s the slow floor of the valley.
• $\eta = 0.05$: reaches the fit, to three digits in $w$. Its $b$ is still 0.0015 off, which is the slow direction again.
• $\eta = 0.07$: just past $2/\lambda_\text{max} = 0.065$. It explodes to $10^{24}$.

Now standardize the feature, subtracting its mean and dividing by its standard deviation. The Hessian becomes

because a standardized feature has mean 0 and variance 1. Both eigenvalues are 2, so $\kappa = 1$: the valley is a perfectly round bowl. The last three lines follow exactly.

• $\eta = 0.5$ makes $\eta\lambda = 1$ and lands on the optimum. (The printed $w = 2.273$ is the same fit expressed in standardized units.)
• $\eta = 1.0$ makes $\eta\lambda = 2$, so the factor is exactly $-1$. Every step jumps to the mirror image of the last point. It bounces forever between the start and its reflection, which is why the MSE printed is the starting loss, 43.05.
• $\eta = 1.05$ diverges.

This is the whole case for feature scaling, and it’s more than housekeeping. The same data, the same model and the same optimum take one step or hundreds, depending only on the units we measured in. Notice that standardizing did two things, and both mattered. Scaling each direction by its own curvature alone, dividing by the square root of the diagonal of $H$, only takes $\kappa$ from 128 to 29. Centering removes the rest, because it zeroes the off-diagonal term that couples $w$ and $b$. Later ideas are this fix, grown up:

• Normalization layers rescale and re-center the activations inside a network, so every layer sees well-conditioned inputs. LayerNorm does both steps, and RMSNorm keeps only the rescaling. (Dropout and Normalization covers them.)
• Adaptive optimizers like Adam rescale the step per parameter. That’s a cheap diagonal preconditioner, a rough per-axis stand-in for $H^{-1}$, which is exactly the “scaling” half above.
• Momentum attacks the condition number differently, by building up speed along the flat floor of the valley. For a quadratic like this one, it can cut the number of steps from order $\kappa$ to order $\sqrt{\kappa}$. (Initialization and Optimizers derives it.)

Try it. Toggle between raw and standardized units, pick a learning rate, and watch the path:

The widget starts at $w = 0$, $b = 4$, above the valley and far from the answer. In raw units with $\eta = 0.05$, the first few steps zig-zag across the valley onto its floor. Then the path slides slowly along the floor, and most of the steps before the widget reports convergence (within 0.1% of the best loss) are spent on that slide. That slide is the small eigenvalue at work. Switch to standardized units and $\eta = 0.5$ lands on the ring in a single step, while $\eta = 1.0$ bounces between the start and its mirror image forever.

From numbers to decisions: logistic regression

The second problem is a decision. When a reader asks the Notebook Assistant a question, the first thing it does is route it: which part of the blog should answer? We start with the smallest version, a yes/no router. Does this question belong with the serving posts (KV cache, vLLM, SGLang, TensorRT-LLM) or not?

The data is 24 questions a reader might ask. They are written by hand for this post and are honestly labeled, including a few deliberately hard ones. The one feature is crude on purpose: how many words from a fixed serving vocabulary (“cache”, “batch”, “latency”, “vllm”, “gpu”, “memory”, …) the question contains. Features and Overfitting replaces it with real features, but one number keeps everything visible here.

A straight line would predict any number, and we need a probability. So we pass the line through the sigmoid:

$\sigma$ squashes any real number into $(0, 1)$. Its input $z$ is called the logit, the log-odds $\log\frac{p}{1-p}$, as solving for $z$ shows. The model predicts “serving” when $p \ge 0.5$, that is, when $z \ge 0$. The decision boundary is where $wx + b = 0$.

Why not squared error?

The first instinct is to reuse MSE: $(p - y)^2$, with the label $y \in \{0, 1\}$. It fails in a specific, measurable way. By the chain rule, with $\sigma'(z) = \sigma(z)(1 - \sigma(z))$:

The factor $p(1-p)$ goes to zero when $p$ is near 0 or 1. So when the model is confidently wrong, predicting $p \approx 0$ for a true “serving” question, the gradient nearly vanishes and it barely learns from its worst mistakes. For a true label $y = 1$:

At $z = -5$, badly wrong, the MSE gradient is 75 times weaker than the alternative in the last column. That’s saturation: the sigmoid has flattened out, and squared error inherits its flatness, so a badly wrong model sits on a plateau. There’s a second, separate problem. MSE composed with a sigmoid isn’t convex in the parameters, so it can also have local minima. We’ll watch the first problem stall a real training run below.

The loss that falls out of probability

The right loss comes from asking what $p$ means. If the model says the question is “serving” with probability $p$, the probability it assigns to the label actually observed is

This is $p$ when $y = 1$ and $1-p$ when $y = 0$. Choosing parameters that make the observed labels most probable is maximum likelihood. Products of probabilities underflow and are awkward to differentiate, so take the negative log and average:

This is the log-loss, or binary cross-entropy. Now differentiate. Using $\sigma' = \sigma(1-\sigma)$, the derivative of the bracket with respect to $z$ collapses:

The $p(1-p)$ that crippled MSE cancels exactly. The gradient is prediction minus target, the same shape as linear regression’s error. It stays near $-1$ when the model is confidently wrong, the last column of the table above. By the chain rule, the parameter gradients are

and the log-loss is convex in $(w, b)$. (Its Hessian, $\frac{1}{n}\sum_i p_i(1 - p_i)\,\tilde x_i \tilde x_i^\top$ with $\tilde x_i = (x_i, 1)$, is a sum of outer products with nonnegative weights, so it is positive semidefinite.) So gradient descent can’t get stuck anywhere but the bottom, provided a bottom exists. If some line separates the two classes perfectly, the loss keeps falling as $w \to \infty$ and the weights never settle. Our router’s data isn’t separable, since both labels occur at 0, 1 and 2 hits. The richer features of Features and Overfitting will be, and that’s where regularization enters.

The same argument explains the first half of this post. Suppose a post’s true token count is the line plus Gaussian noise, $y = \hat y + \varepsilon$ with $\varepsilon \sim \mathcal{N}(0, \sigma^2)$. Then

and maximum likelihood is least squares. Laplace noise, $p(\varepsilon) \propto e^{-|\varepsilon|/s}$, gives absolute error the same way. So “they encode different beliefs about what an error costs” can be made exact. Every loss in this post is a negative log-likelihood under some model of the noise: Gaussian for L2, Laplace for L1, Bernoulli for the log-loss. Heavy tails in the noise, like the math-heavy post, favor the Laplace choice.

Training the router

c_router.py holds the 24 labeled questions and the 18 serving words in SERVING: kv, cache, caching, batch, batching, latency, throughput, serve, serving, vllm, sglang, tensorrt, gpu, decode, prefill, paged, quantization and memory. The feature counts those words in each question, and full-batch gradient descent fits the router:

def hits(q):
    return sum(tok in SERVING for tok in re.findall(r"[a-z0-9]+", q.lower()))

x = np.array([hits(q) for q, _ in questions], dtype=float)
y = np.array([label for _, label in questions], dtype=float)
sigmoid = lambda z: 1 / (1 + np.exp(-z))

if __name__ == "__main__":
    w, b, eta = 0.0, 0.0, 0.5
    for step in range(5000):
        p = sigmoid(w * x + b)
        w -= eta * ((p - y) * x).mean()      # the whole gradient: (p - y) x
        b -= eta * (p - y).mean()

    p = sigmoid(w * x + b)
    loss = -(y * np.log(p) + (1 - y) * np.log(1 - p)).mean()
    print(f"w = {w:.3f}, b = {b:.3f}, log-loss = {loss:.4f}")
    for k in range(3):
        print(f"P(serving | {k} keyword hits) = {sigmoid(w * k + b):.3f}")
    print(f"accuracy = {((p >= 0.5) == y).mean():.3f}")

Full script: c_router.py

w = 1.746, b = -1.460, log-loss = 0.4837
P(serving | 0 keyword hits) = 0.188
P(serving | 1 keyword hits) = 0.571
P(serving | 2 keyword hits) = 0.884
accuracy = 0.792

With the router trained, we can test the claim from Why not squared error? on real data. Start the router confidently wrong, $w = -5$, and train it for the same 500 steps with each loss:

import numpy as np
from c_router import x, y, sigmoid

# Start confidently wrong (w = -5) and train with each loss, same steps.
for name in ["log-loss", "squared error"]:
    w, b = -5.0, 0.0
    for step in range(500):
        p = sigmoid(w * x + b)
        grad_z = (p - y) if name == "log-loss" else 2 * (p - y) * p * (1 - p)
        w -= 0.5 * (grad_z * x).mean()
        b -= 0.5 * grad_z.mean()
    p = sigmoid(w * x + b)
    print(f"{name:13s} after 500 steps: w = {w:6.3f}, b = {b:6.3f},"
          f" accuracy = {((p >= 0.5) == y).mean():.3f}")

Full script: wrong_start.py

log-loss      after 500 steps: w =  1.746, b = -1.460, accuracy = 0.792
squared error after 500 steps: w = -4.869, b = -1.480, accuracy = 0.500

The log-loss run recovers completely and finds the optimum. The squared-error run has barely moved: $w$ went from $-5$ to $-4.87$, and the router still gets half the questions wrong. That’s the plateau from the gradient table, made concrete.

The first step is worth tracing. At $w = b = 0$ every prediction is $p = 0.5$, and the loss is $\log 2 = 0.693$, the loss of a coin flip. The classes are balanced, 12 and 12, so the average of $p - y$ is exactly zero and $b$ doesn’t move at all. $w$ gets $-\frac{1}{24}\sum (0.5 - y_i)\,x_i = 0.271$ times $\eta$: questions with more serving words are disproportionately serving questions, so the slope goes up.

After training, the boundary sits at $x = 1.46/1.746 = 0.84$ hits. In practice that means one serving word is enough. The model is right on 19 of 24 questions, and its five mistakes are the interesting part:

• Two serving questions with no serving words. “Why is my chatbot slow to answer long prompts?” and “Why does the first token take so long?” Both are about exactly what the serving posts explain, prefill cost and time to first token, and neither uses the vocabulary.
• Three other questions that use it. “How much GPU memory does training a 7B model need?” (training, not serving), “How does the memory tool in context engineering work?” (agent memory), and “Should I cache the embeddings of every post?” (retrieval).

The model isn’t wrong about its feature. It learned the best possible rule for “count serving words”. The feature is wrong: a keyword count can’t tell GPU memory for training from GPU memory for serving, and it can’t recognize a serving question phrased in plain English. No amount of optimization fixes a feature that doesn’t carry the information. That’s where Features and Overfitting starts.

Was $\eta = 0.5$ safe here? The bowl argument carries over. The log-loss’s curvature is at most $\frac{1}{4} \cdot \frac{1}{n}\sum_i \tilde x_i \tilde x_i^\top$, because $p(1-p) \le \frac{1}{4}$. For our 24 questions that bound puts the largest safe step at about 3.7, so 0.5 is comfortably stable.

One example at a time: stochastic and mini-batch descent

Each step so far used every data point to compute the gradient. That is batch gradient descent. One step costs about $4nd$ floating-point operations for $n$ examples and $d$ parameters: $2nd$ to predict, and $2nd$ more to form the gradient. That’s fine for 24 questions. A language model’s training set has trillions of tokens, so no step can afford to see all of them.

Stochastic gradient descent (SGD) estimates the gradient from one random example, and mini-batch descent from a small random subset. The estimate is noisy but unbiased: on average it points the same way as the full gradient, and it is $n/B$ times cheaper for a mini-batch of size $B$. Train the router with $\eta = 0.5$ both ways, reshuffling the questions each pass, and repeat the whole run with 200 different shuffles:

def log_loss(w, b):
    p = sigmoid(w * x + b)
    return -(y * np.log(p) + (1 - y) * np.log(1 - p)).mean()

best = log_loss(1.746, -1.460)          # the full-batch optimum from above
for batch in [24, 4, 1]:
    finals, needed = [], []
    for seed in range(200):             # 200 different shuffles
        rng = np.random.default_rng(seed)
        w = b = 0.0
        updates, reached = 0, None
        while updates < 5000:
            order = rng.permutation(len(y))
            for start in range(0, len(y), batch):
                i = order[start:start + batch]
                p = sigmoid(w * x[i] + b)
                w -= 0.5 * ((p - y[i]) * x[i]).mean()
                b -= 0.5 * (p - y[i]).mean()
                updates += 1
                if reached is None and log_loss(w, b) < 1.01 * best:
                    reached = updates
        finals.append(log_loss(w, b)); needed.append(reached)
    finals = np.array(finals)
    print(f"{batch:2d} per update: updates to reach 1% of best, median {np.median(needed):.0f};"
          f"  final loss median {np.median(finals):.4f}, worst {finals.max():.4f}")

Full script: minibatch.py

24 per update: updates to reach 1% of best, median 52;  final loss median 0.4837, worst 0.4837
 4 per update: updates to reach 1% of best, median 52;  final loss median 0.4841, worst 0.4921
 1 per update: updates to reach 1% of best, median 44;  final loss median 0.4949, worst 0.7688

All three need about the same number of updates to get close, around 50. But an update on one question costs a 24th of a full-batch update. So SGD gets there after seeing the data about twice, where full-batch descent needs 52 passes. The price is noise, and it shows in the last two columns. With a constant learning rate, SGD never settles. After 5,000 updates its median loss is 0.4949 against the optimum’s 0.4837, and its worst run is at 0.7688, knocked far from the minimum by a few unlucky steps in a row. Mini-batches of 4 average away most of that noise.

The classic fix is a decaying learning rate. Robbins and Monro showed that stochastic steps converge if $\sum_t \eta_t = \infty$ and $\sum_t \eta_t^2 < \infty$: steps big enough to get anywhere, shrinking fast enough to stop. Real LLM training schedules decay the learning rate for the same reason. They also warm it up first, for a different reason, early-training instability, which Initialization and Optimizers covers.

What this model can’t do

Both models here are linear: the prediction, or the log-odds, is a weighted sum of the inputs. That buys a lot. The loss is convex, gradient descent is guaranteed to find the best parameters, and each weight is readable (“1.75 more log-odds per serving word”). It also rules things out.

• Interactions. A linear model adds up each word’s evidence independently: memory adds the same log-odds whether or not training is present. A negative weight on training can approximate “memory means serving, unless training is present”. But no linear model can represent a rule where two words together mean the opposite of what each means alone, the pattern called XOR. That needs handcrafted interaction features, or a model that builds its own features, which is what the hidden layers of a neural network are.
• Features beyond the inputs. It can only weigh the features it’s given, so its ceiling is set by feature engineering.

The neural networks later in this series keep the core of this post: a loss, a gradient, and a step of size $\eta$ downhill. Two things change. The loss is no longer convex, so where training starts and how it steps start to matter. And the gradient of a deep composition of functions needs an algorithm of its own: backpropagation.

Where this leads

• Real features, and the problem they bring. Features and Overfitting replaces the keyword count with every word in the question. That meets the new problem more features create, overfitting, and the regularization that controls it.
• How good is 19 out of 24? Accuracy hides which mistakes were made, and for a router the two kinds of mistake cost different things. Measuring a Classifier Honestly measures the router properly, with precision, recall, cost-aware thresholds and calibration.
• Gradients for any model. The $p - y$ gradient is the base case of backpropagation, which Backpropagation, by Hand derives: it computes the same “error times input” for every weight in a deep network in one backward pass.
• Better steps. The learning-rate ceiling $2/\lambda_\text{max}$ and the condition number $\kappa$ are why deep networks need careful initialization and adaptive optimizers (Initialization and Optimizers) and normalization layers (Dropout and Normalization).
• The same sigmoid, much later. A reward model that learns from “answer A is better than answer B” uses this router’s loss, applied to the difference of two scores. It’s the foundation of RLHF, and DPO turns it back into a logistic loss on log-probability ratios.

Takeaways

• Learning is choosing parameters that make a loss small. The model, the loss and the optimizer are three separate decisions.
• The loss encodes what an error costs. Squared error is minimized by the mean and bends toward outliers. Absolute error is minimized by the median, which can be a whole flat interval of equally good answers.
• Every loss here is a negative log-likelihood: Gaussian noise gives squared error, Laplace noise gives absolute error, and Bernoulli labels give the log-loss.
• Gradient descent steps against the gradient. Every weight’s gradient in this series has the shape “error times what the weight multiplied”.
• On a quadratic bowl, each step multiplies the remaining error by $(1 - \eta\lambda)$ along each curvature direction. So $\eta$ must stay below $2/\lambda_\text{max}$, and a large condition number makes the flat direction crawl. Standardizing our feature took the condition number from 128 to 1, and convergence from hundreds of steps to one. Rescaling alone only reached 29; centering did the rest.
• Mini-batch and stochastic descent trade exact gradients for cheap, noisy, unbiased ones. For the router they needed about the same number of updates but a fraction of the passes over the data. With a constant learning rate, single-example SGD never settles (median loss 0.4949 against 0.4837), which is why learning rates decay.
• For probabilities, maximum likelihood gives the log-loss. Its gradient through a sigmoid is exactly $p - y$, which never vanishes when the model is confidently wrong. Squared error’s gradient is 75× weaker at $z = -5$.

References

To go further, start with Goodfellow, Bengio and Courville, chapters 4 and 5, which cover this post’s material in more depth. Boyd and Vandenberghe §9.3 is the rigorous version of the learning-rate argument. Bottou, Curtis and Nocedal explain why stochastic gradients win at scale.

• Legendre, Nouvelles méthodes pour la détermination des orbites des comètes (1805), and Gauss, Theoria motus corporum coelestium (1809): the method of least squares behind the closed-form fit.
• Lemaréchal, Cauchy and the Gradient Method (Documenta Mathematica, 2012): the history of gradient descent, which Cauchy proposed in 1847.
• Robbins and Monro, A Stochastic Approximation Method (Annals of Mathematical Statistics, 1951): the step-size conditions for stochastic descent.
• Cox, The Regression Analysis of Binary Sequences (JRSS B, 1958): logistic regression, the model behind the router.
• Hastie, Tibshirani and Friedman, The Elements of Statistical Learning (2nd ed., 2009): linear and logistic regression, chapters 3 and 4.
• Boyd and Vandenberghe, Convex Optimization (2004): gradient descent’s convergence on quadratics and the condition number, §9.3.
• Bottou, Curtis and Nocedal, Optimization Methods for Large-Scale Machine Learning (SIAM Review, 2018): why stochastic gradients win at scale.
• Goodfellow, Bengio and Courville, Deep Learning (2016): chapters 4 and 5, numerical optimization and ML basics.