Chapter 31

Teaching the Reader with Calculus

We have a function to minimize. Its input is a reader, represented by 650 stored numbers, and its output is the cost of that reader's answers on the training set, together with a penalty for large parameters. We could evaluate this function at any point we like. How should we choose the points?

Before inventing a new method, let's ask what our familiar method would do.

31.1Start with critical points

For a differentiable function $𝑓 (𝑥, 𝑦)$ , an interior local minimum must satisfy

𝑓 𝑥 = 0, 𝑓 𝑦 = 0 .

We solve these equations to find the critical points, then investigate which ones are minima. For instance, if $𝑓 (𝑥, 𝑦) = 𝑥 2 + 2 𝑦 2$ , the equations are $2 𝑥 = 0$ and $4 𝑦 = 0$ . There is just one solution, the origin, and the formula itself tells us it is the global minimum: both squares are nonnegative, and they vanish together only there.

The equations alone did not establish that last conclusion. The function $𝑥 2 - 𝑦 2$ has exactly the same critical point, but along the $𝑥$ -axis it increases away from the origin, while along the $𝑦$ -axis it decreases. The critical point is a saddle.

Figure 31.1 Change the coefficient in

𝑓 (𝑥, 𝑦) = 𝑥 2 + 𝑎 𝑦 2

. For

𝑎 \neq 0

, the critical-point equations always have the same unique solution. The kind of critical point changes when the sign of

𝑎

changes. At

𝑎 = 0

, notice what happens to both the equations and the graph.

Nothing prevents us from writing the same equations for the reader:

𝜕 𝐽 𝜕 Θ 1 = 0, \dots, 𝜕 𝐽 𝜕 Θ 650 = 0 .

But writing them down is much easier than solving them. Each probability contains exponentials of scores involving many parameters, divided by a sum of ten such exponentials. Each partial derivative then adds contributions from thousands of images. These are coupled nonlinear equations: a parameter appearing in one score affects the probabilities of all ten digits.

The problem is not simply that 650 is a large number. A function like $Θ 21 + \dots + Θ 2650$ gives 650 equations we can solve immediately. Here it is the way the unknowns occur together that defeats our usual algebraic approach. Critical points still characterize the answer; we need another way to reach one.

31.2A minimum must be a minimum in every direction

There is a second issue for general functions. Finding a point with zero gradient does not tell us that we should stop searching for lower values. A saddle can hide if we inspect too few directions.

Figure 31.2 Slice the graph of

𝑓 (𝑥, 𝑦) = 𝑥 2 - 𝑦 2

by lines through the origin. Some slices show a minimum, others a maximum, and two directions give a constant slice. One reassuring slice is not enough.

This becomes especially easy to miss with more coordinates. Consider

𝑓 (𝑥 1, \dots, 𝑥 𝑛) = 𝑥 21 + \dots + 𝑥 2 𝑛 - 1 - 𝑥 2 𝑛 .

The origin is a minimum on the entire subspace $𝑥 𝑛 = 0$ , which contains $𝑛 - 1$ independent coordinate directions. Yet moving in the remaining direction lowers $𝑓$ . More generally, a diagonal quadratic with both positive and negative coefficients has a saddle at the origin. There are many possible mixed-sign arrangements as the dimension grows. This observation does not assign a probability to encountering a saddle; it explains why checking a few directions becomes less reassuring.

We will return to the shape of our reader's particular objective. First let's develop a way to seek lower values of a general differentiable function without solving all of its critical-point equations at once.

31.3Let the gradient choose a direction

Suppose we are at $⃗ 𝑎$ and can move a small distance. For a unit vector $⃗ 𝑢$ , the directional derivative is

𝐷 ⃗ 𝑢 𝑓 (⃗ 𝑎) = \nabla 𝑓 (⃗ 𝑎) \cdot ⃗ 𝑢 = ‖ \nabla 𝑓 (⃗ 𝑎) ‖ c o s 𝜃 .

Its smallest value occurs when $⃗ 𝑢$ points opposite the gradient. At a point where the gradient is nonzero, the steepest unit direction of decrease is therefore

⃗ 𝑢 = - \nabla 𝑓 (⃗ 𝑎) ‖ \nabla 𝑓 (⃗ 𝑎) ‖ .

Once again, a sum of partial derivatives has become a dot product, and the dot product has supplied the geometry. This is a local statement: it describes the rate of change as we begin moving. It does not tell us how far we can continue in that direction before the graph turns upward again.

Figure 31.3 Rotate a direction at a fixed point of

𝑓 (𝑥, 𝑦) = (𝑥 2 + 4 𝑦 2) / 2

. The arrows are in the two-dimensional domain; the surface above shows the resulting heights. Find the direction giving the most negative directional derivative.

Now imagine continually turning to face downhill as we move. We can describe the resulting path $⃗ 𝑟 (𝑡)$ by asking its velocity to be the negative gradient at its current location:

⃗ 𝑟' (𝑡) = - \nabla 𝑓 (⃗ 𝑟 (𝑡)) .

This equation describes a gradient flow. We have chosen a speed as well as a direction: the speed is $‖ \nabla 𝑓 ‖$ , so the motion slows where the graph is flatter. We do not need a course in differential equations to check why this motion goes downhill. Apply the chain rule to the height along the path:

𝑑 𝑑 𝑡 𝑓 (⃗ 𝑟 (𝑡)) = \nabla 𝑓 (⃗ 𝑟 (𝑡)) \cdot ⃗ 𝑟' (𝑡) = - ‖ \nabla 𝑓 (⃗ 𝑟 (𝑡)) ‖ 2 \leq 0 .

Wherever the gradient is nonzero, the height decreases. At a critical point this prescribed velocity is zero. That includes saddles, so this calculation does not promise that a flow for every function finds a minimum.

31.4A computer takes finite steps

How do we turn this moving point into a computation? Use the Fundamental Strategy of Calculus: over a short interval, replace the changing motion by its linear approximation. A computer can evaluate the gradient, move a short distance using that velocity, and evaluate it again. Over a short time interval $𝜂$ , the velocity approximation gives

⃗ 𝑟 (𝑡 + 𝜂) \approx ⃗ 𝑟 (𝑡) + 𝜂 ⃗ 𝑟' (𝑡) = ⃗ 𝑟 (𝑡) - 𝜂 \nabla 𝑓 (⃗ 𝑟 (𝑡)) .

This suggests the iteration

⃗ 𝑎 𝑚 + 1 = ⃗ 𝑎 𝑚 - 𝜂 \nabla 𝑓 (⃗ 𝑎 𝑚) .

We call it gradient descent, and $𝜂 > 0$ is its step size or learning rate. Each iteration consists of a gradient evaluation and a vector subtraction.

For our stretched bowl, $\nabla 𝑓 = (𝑥, 4 𝑦)$ , so the update becomes

𝑥 𝑚 + 1 = (1 - 𝜂) 𝑥 𝑚, 𝑦 𝑚 + 1 = (1 - 4 𝜂) 𝑦 𝑚 .

We can see the role of the step size without running anything. With $𝜂 = 0.2$ , both coordinates shrink. With $𝜂 = 0.4$ , the $𝑦$ coordinate alternates signs but still shrinks. With $𝜂 = 0.6$ , it is multiplied by $- 1.4$ at each step and grows. Every step began by pointing downhill; the steps were simply too long.

Figure 31.4 Compare the exact flow with finite descent steps on the same surface. Increase the step size until the discrete path fails to descend. Then turn on step halving and inspect the accepted size.

One practical remedy is to try a step and evaluate its result. If it does not decrease the function enough, halve the step size and try again. We will use the concrete test

𝑓 (⃗ 𝑎 - 𝜂 ⃗ 𝑔) \leq 𝑓 (⃗ 𝑎) - 𝑐 𝜂 ‖ ⃗ 𝑔 ‖ 2, ⃗ 𝑔 = \nabla 𝑓 (⃗ 𝑎), 𝑐 = 10 - 4 .

The right side asks for a small fraction of the decrease predicted by the directional derivative. Why should halving eventually work? The linear approximation gives

𝑓 (⃗ 𝑎 - 𝜂 ⃗ 𝑔) = 𝑓 (⃗ 𝑎) - 𝜂 ‖ ⃗ 𝑔 ‖ 2 + 𝑜 (𝜂) .

When $⃗ 𝑔 \neq 0$ and $𝜂$ is sufficiently small, the error is smaller than the allowance between the full predicted decrease and the fraction $𝑐$ . This is a local justification for the search, not a promise that large steps work. Near a minimum, finite precision also limits how small a decrease a computer can detect.

31.5Now differentiate the reader

We have a method. To apply it, we need the gradient of $𝐽$ with respect to the stored parameters.

The bias is easiest to handle by writing

̃ ⃗ 𝑥 = (1, 𝑥 1, \dots, 𝑥 64), 𝜽 𝑘 = (𝑏 𝑘, 𝑤 𝑘, 1, \dots, 𝑤 𝑘, 64) .

Then $𝑠 𝑘 = 𝜽 𝑘 \cdot ̃ ⃗ 𝑥$ . This is bookkeeping for differentiation: the image still has 64 varying coordinates. The new first entry is always one.

Start with one labeled image and its loss

ℓ = l o g (9 \sum 𝑟 = 0 𝑒 𝑠 𝑟) - 𝑠 𝑦 .

Temporarily treat the scores as ten independent variables. Differentiating with respect to $𝑠 𝑘$ gives

𝜕 ℓ 𝜕 𝑠 𝑘 = 𝑒 𝑠 𝑘 \sum 9 𝑟 = 0 𝑒 𝑠 𝑟 - 𝟏 𝑘 = 𝑦 = 𝑝 𝑘 - 𝟏 𝑘 = 𝑦 .

The symbol $𝟏 𝑘 = 𝑦$ is one when $𝑘$ is the known digit and zero otherwise. Thus the derivative is $𝑝 𝑦 - 1$ for the known digit and $𝑝 𝑘$ for each other digit.

Next, $𝜃 𝑘, 𝑗$ occurs in score $𝑠 𝑘$ with coefficient $̃ 𝑥 𝑗$ , and it occurs in no other score. The chain rule therefore gives

𝜕 ℓ 𝜕 𝜃 𝑘, 𝑗 = (𝑝 𝑘 - 𝟏 𝑘 = 𝑦) ̃ 𝑥 𝑗 .

That is the whole derivative for one image: ten residual numbers, each multiplying the same image vector.

There is a useful check at the zero reader. All probabilities equal $0.1$ . A known three contributes $- 0.9 ̃ ⃗ 𝑥$ to the gradient of the three's parameter vector and $0.1 ̃ ⃗ 𝑥$ to each of the other nine. Subtracting the gradient moves the three's vector toward this image direction and moves the others away. For this single image, its own score increases while the competing scores decrease.

Figure 31.5 Choose a training image and one parameter coordinate. At the zero reader, check the

- 0.9

and

0.1

factors. Then move toward the trained reader. The last column adds every training image's contribution and the parameter penalty; it need not have the same sign as the selected image's contribution.

For the full objective, add the derivatives of all the image losses. The penalty contributes $𝜆 𝜃 𝑘, 𝑗$ , so

𝜕 𝐽 𝜕 𝜃 𝑘, 𝑗 = 𝑁 \sum 𝑖 = 1 (𝑝 𝑘 (⃗ 𝑥 (𝑖); Θ) - 𝟏 𝑘 = 𝑦 𝑖) ̃ 𝑥 (𝑖) 𝑗 + 𝜆 𝜃 𝑘, 𝑗 .

These 650 numbers are $\nabla 𝐽 (Θ)$ . Images can disagree about which way a parameter should move. The sum combines their requests before we take a step. Training by the full gradient is therefore different from making a correction for one image and forgetting the others.

31.6Does this objective have saddles?

For our chosen $𝐽$ , we can say much more than we could for a general function. Look along an arbitrary line in parameter space,

Θ (𝑡) = Θ 0 + 𝑡 𝑉,

where $𝑉$ is a fixed, nonzero parameter direction. For one image, each score has the form $𝑠 𝑘 (𝑡) = 𝑎 𝑘 + 𝑡 𝑐 𝑘$ : a dot product with an affine parameter list is affine in $𝑡$ .

Write $𝑝 𝑘 (𝑡)$ for the resulting probabilities and $¯ 𝑐 = \sum 𝑘 𝑝 𝑘 𝑐 𝑘$ . Differentiate the image loss twice along this line. The first derivative is $\sum 𝑘 𝑝 𝑘 𝑐 𝑘 - 𝑐 𝑦$ . Differentiating once more gives

ℓ ″ (𝑡) = \sum 𝑘 𝑝 𝑘 𝑐 2 𝑘 - (\sum 𝑘 𝑝 𝑘 𝑐 𝑘) 2 = \sum 𝑘 𝑝 𝑘 (𝑐 𝑘 - ¯ 𝑐) 2 \geq 0 .

To check the middle step, the quotient rule gives $𝑝' 𝑘 = 𝑝 𝑘 (𝑐 𝑘 - ¯ 𝑐)$ . To check the last equality, expand the square and use $\sum 𝑘 𝑝 𝑘 = 1$ . No new machinery is hiding in this calculation.

The penalty has second derivative $𝜆 ‖ 𝑉 ‖ 2$ , so the complete restriction satisfies

𝑑 2 𝑑 𝑡 2 𝐽 (Θ 0 + 𝑡 𝑉) = 𝑁 \sum 𝑖 = 1 9 \sum 𝑘 = 0 𝑝 (𝑖) 𝑘 (𝑐 (𝑖) 𝑘 - ¯ 𝑐 (𝑖)) 2 + 𝜆 ‖ 𝑉 ‖ 2 > 0 .

No matter which way we draw a line through parameter space, the function along it curves upward. This property is called strict convexity. If a point has zero gradient, then every line through it has derivative zero there and strictly increasing derivative. Along each line it is the unique minimum. Consequently it is the unique global minimum of $𝐽$ , and there can be no saddle points.

A minimum also exists. Each image loss is nonnegative, so $𝐽 (Θ) \geq 𝜆 ‖ Θ ‖ 2 / 2$ , which grows without bound as we go farther from the origin. Choose a large ball outside which this lower bound exceeds $𝐽 (0)$ . Continuity gives a minimum on the closed ball, and nothing outside can do better.

Figure 31.6 Evaluate both parts of the second derivative along an actual parameter direction. The plotted line is one example; the sum-of-squares argument establishes the result for every line.

This is a welcome feature of our model. The difficult coupled equations have a unique solution, and we can approach it by taking downhill steps. It does not follow that every model used in machine learning has this geometry. Adding nonlinear layers would change the objective we just analyzed.

Figure 31.7 Follow saved checkpoints from the run that produced our demonstration reader. The objective falls while the ten parameter vectors change. The mistake count need not improve on every step: it measures a different function from the one whose gradient we follow.

31.7Extension: carry some motion forward

Ordinary gradient descent forgets its previous motion. In a narrow valley, this can waste many steps crossing back and forth while making slow progress along the bottom.

Try retaining a fraction of the previous displacement. With $⃗ 𝑣 0 = 0$ , define

⃗ 𝑣 𝑚 + 1 = 𝛽 ⃗ 𝑣 𝑚 - 𝜂 \nabla 𝑓 (⃗ 𝑎 𝑚), ⃗ 𝑎 𝑚 + 1 = ⃗ 𝑎 𝑚 + ⃗ 𝑣 𝑚 + 1, 0 \leq 𝛽 < 1 .

This is gradient descent with momentum. The vector $⃗ 𝑣$ records the displacement, and $𝛽$ says how much of it we carry into the next step. Setting $𝛽 = 0$ recovers ordinary descent.

If successive gradients push consistently along the valley, their contributions accumulate. If they alternate across it, their contributions tend to cancel. This suggests a possible gain, but the retained motion can also take us uphill or past a turn. The chain-rule decrease for a gradient flow does not apply to this new iteration.

Figure 31.8 Compare both methods on

𝑓 = (𝑢 2 + 20 𝑣 2) / 2

, where

𝑢 = (𝑥 + 𝑦) / \sqrt 2

and

𝑣 = (𝑥 - 𝑦) / \sqrt 2

. The surface, domain paths, and function values show the same computations. Ordinary descent uses the well-chosen constant rate

2 / 21

; momentum has its own displayed rate and retention factor. Compare equal numbers of gradient evaluations.

Why is $2 / 21$ a fair ordinary rate here? In the rotated coordinates, the update multiplies $𝑢$ by $1 - 𝜂$ and $𝑣$ by $1 - 20 𝜂$ . To make the larger of their absolute contraction factors as small as possible, balance them: $1 - 𝜂 = 20 𝜂 - 1$ . This gives $𝜂 = 2 / 21$ . Momentum's gain in the default comparison does not come from choosing an unnecessarily tiny ordinary step.

Figure 31.9 Now use

𝑓 (𝑥, 𝑦) = (1 - 𝑥) 2 + 10 (𝑦 - 𝑥 2) 2

. The bottom of the valley bends toward

(1, 1)

. Both methods start at

(- 1, 1)

and initially use the same step size. Increase momentum's rate or retention and look for overshoot as well as faster progress.

Momentum costs one additional stored vector and some vector arithmetic per step. It still uses one gradient evaluation. Its usefulness depends on the function and the chosen parameters; the figures are experiments we can explain, not a guarantee of universal acceleration. Our first complete digit-reader implementation will use ordinary descent with step halving, so that the mathematical route from a derivative to an accepted update remains visible.

31.8Things to try

For $𝑓 (𝑥, 𝑦) = (𝑥 2 + 4 𝑦 2) / 2$ , find all constant step sizes that make every ordinary descent orbit converge to the origin. Check the endpoints separately.
Derive the formula $𝑝' 𝑘 = 𝑝 𝑘 (𝑐 𝑘 - ¯ 𝑐)$ along a parameter line, and use it to verify the second derivative of one image loss.
For one labeled image at the zero reader, compute the change in each score after one unpenalized descent step. Express the answer using $‖ ̃ ⃗ 𝑥 ‖ 2$ .
In the narrow-valley figure, set $𝛽 = 0$ and give momentum the same rate as ordinary descent. Explain why the paths must agree.