Optimization: Make Each Update Explainable

Optimization changes model parameters to reduce a chosen objective. Your task is to reconstruct an update from its derivatives, preserve the history promised by the API, and distinguish an optimizer defect from a difficult training problem.

Prerequisites: derivatives/partial derivatives, a mean loss, and vector operations. If the derivative of x*x is unfamiliar, return to calculus. You need no knowledge of Adam to solve these two exercises.

One parameter and an exact recurrence

Let f(x)=(x-3)^2, so df(x)=2*(x-3). Gradient descent computes x_next=x-lr*df(x) from the current x. Start at 0 with learning rate 0.25:

Step	Current x	Derivative	Next x	f(next x)
1	0	-6	1.5	2.25
2	1.5	-3	2.25	0.5625
3	2.25	-1.5	2.625	0.140625

gradient_descent_1d(...,n_steps=3) returns final x 2.625 and parameter history [0,1.5,2.25,2.625]. Include the initial value, so the history has n_steps+1 entries. The callable f describes the objective; update arithmetic uses the supplied derivative df. The exercise does not require early stopping or printing.

For f(x)=x^2, the recurrence is x_next=(1-2*lr)*x. At lr=0.1 the magnitude shrinks by 0.8; at 1 it oscillates without shrinking; above 1 it diverges for nonzero initial x. Convexity alone does not make an arbitrary step size converge. A convex objective may also have many minimizers rather than one unique minimum.

Vectors require one consistent snapshot

For f(x,y)=x^2+2*y^2, the gradient is [2*x,4*y]. At [2,-1], it is [4,-4]; with lr=0.25 the next vector is [1,0]. Compute the full gradient before changing either coordinate. In coupled functions, changing one coordinate before calculating another component changes the point being differentiated.

gradient_descent_nd copies the initial list and stores independent snapshots. Reusing one mutable list for every history entry makes the whole history appear to equal the final state. With zero steps, return the initial value and a one-entry history; the caller's list remains unchanged.

Momentum has state and a convention

The exercise uses zero initial velocity and:

v_next = momentum*v + gradient(current_x)
x_next = x - lr*v_next

There is no (1-momentum) multiplier on the gradient in this API. Other optimizer conventions exist; name the one you implement.

On f(x)=x^2, start x=2, lr=0.1, momentum=0.5:

Step	Gradient	New velocity	New x
1	4	4	1.6
2	3.2	0.5*4+3.2=5.2	1.08

Retained velocity can speed movement in a consistent direction; it can also overshoot. It does not guarantee escaping a stationary point: with zero initial velocity and exactly zero gradient, nothing moves. Likewise, stochastic gradients add noise but do not guarantee escape from every bad region.

Derive the regression objective and its factor of two

The second exercise predicts p_i=dot(X[i],w)+b and uses mean squared error, without a factor of 1/2:

L = sum((p_i-y_i)^2)/n
dw_j = (2/n)*sum(X[i][j]*(p_i-y_i))
db = (2/n)*sum(p_i-y_i)

Start with X=[[0],[2]], y=[1,5], zeros for w and b, and lr=0.1. Predictions are [0,0], errors [-1,-5], and loss (1+25)/2=13. Gradients are dw=-10, db=-6. Updating together gives w=1, b=0.6, whose predictions are [0.6,2.6]. The new mean loss is (0.16+5.76)/2=2.96.

The regression API records one pre-update loss per epoch, unlike the optimizer exercise's initial-plus-post-update parameter history. With one epoch it returns weights [1], bias 0.6, and losses [13], not [2.96]. Both histories are useful, but mixing them creates off-by-one diagnostics.

Full-batch training averages across every row. Duplicating all examples must preserve both the mean loss and the next update. Per-row SGD instead updates after each example; order then matters and one epoch contains multiple updates.

Complete the next update yourself

Use f(x)=x^2, x=1, lr=0.25, momentum=0.5, and zero initial velocity. Work two steps and the complete parameter history.

Check the momentum trace

Step 1: gradient 2, velocity 2, x=0.5. Step 2: gradient 1, velocity 0.5*2+1=2, x=0. The history is [1,0.5,0]. Plain gradient descent would give 0.25 after step 2, so forgetting the velocity state is observable.

Independent verification before a longer run

Implement the general optimizer and reproduce a vector trace with a coupled objective such as f(x,y)=(x+y)^2. Then implement the regression exercise and check gradients with central finite differences on moderate inputs. Stanford's CS231n gradient-checking notes explain why local derivative checks and overfitting a tiny dataset answer different debugging questions.

A useful sequence is: verify gradients, verify update direction and history, fit a small problem, then inspect held-out behavior. Falling training loss does not establish generalization. A stationary gradient can describe a minimum, maximum, or saddle. Increasing loss may be an implementation error or an excessive step size; inspect the first numerical discrepancy rather than immediately switching optimizers.

Adaptive methods such as Adam add running gradient statistics and further hyperparameters. They are later extensions, not a substitute for understanding the update above. Next, classification changes the output and loss while retaining the same predict–differentiate–update structure.