Which Way Is Down? First-Order Optimizers as a Choice of Norm

We consider the general optimization problem

minπœƒβˆˆΞ˜π’₯οΈ€(πœƒ),

(1)

where Θ is a vector space and π’₯οΈ€ is differentiable. First-order methods solve it iteratively: starting from πœƒ0, they take steps πœƒπ‘‘+1=πœƒπ‘‘+β„Žπ‘‘, where each step β„Žπ‘‘ is computed from the derivative 𝐷π’₯οΈ€(πœƒπ‘‘), and possibly from earlier ones. This post is about how to choose β„Žπ‘‘, and in particular how the choice of a norm on Θ affects it.

1. Derivatives and gradients

The derivative 𝐷π’₯οΈ€(πœƒ) maps a direction β„ŽβˆˆΞ˜ to a number, so it is an element of the dual space Ξ˜βˆ— of linear maps Ξ˜β†’β„, and we write its value at β„Ž as the pairing ⟨𝐷π’₯οΈ€(πœƒ),β„ŽβŸ©. Both Ξ˜βˆ— and the pairing come with Θ itself, so nothing is chosen here. An inner product βŸ¨β‹…,β‹…βŸ©Ξ˜ on Θ, on the other hand, is a choice, for example the Frobenius inner product. With an inner product chosen, the linearization reads

π’₯οΈ€(πœƒ+β„Ž)β‰ˆπ’₯οΈ€(πœƒ)+⟨𝐷π’₯οΈ€(πœƒ),β„ŽβŸ©=π’₯οΈ€(πœƒ)+βŸ¨βˆ‡πœƒπ’₯οΈ€(πœƒ),β„ŽβŸ©Ξ˜.

(2)

The second form holds because every linear map Ξ˜β†’β„ is β„Žβ†¦βŸ¨π‘˜,β„ŽβŸ©Ξ˜ for exactly one π‘˜βˆˆΞ˜, its Riesz representative,11 This is the Riesz representation theorem, as in Batched Backpropagation as a Linear Solve: π‘˜β†¦βŸ¨π‘˜,β‹…βŸ©Ξ˜ is linear and one-to-one, and Θ and Ξ˜βˆ— have the same dimension. and the gradient βˆ‡πœƒπ’₯οΈ€(πœƒ) is the representative of 𝐷π’₯οΈ€(πœƒ). Unlike the derivative, the gradient lies in Θ, but it depends on the inner product: another inner product gives another gradient of the same function. Only the gradient needs an inner product. The derivative and the steps below need none.

ExampleΒ 1: Two inner products, two gradients

Example 1 (Two inner products, two gradients)Β Β  Let Θ=ℝ2 and π’₯οΈ€(πœƒ)=πœƒ12+πœƒ2, so that ⟨𝐷π’₯οΈ€(πœƒ),β„ŽβŸ©=2πœƒ1β„Ž1+β„Ž2.22 This follows from π’₯οΈ€ alone, as the directional derivative d/(d𝑑)π’₯οΈ€(πœƒ+π‘‘β„Ž)|𝑑=0=2πœƒ1β„Ž1+β„Ž2, with no inner product involved. Writing the derivative as the array (2πœƒ1,1) uses the coordinates of Θ=ℝ2, but the value of the pairing does not depend on them. Then

βŸ¨β„Ž,π‘˜βŸ©Ξ˜=β„Ž1π‘˜1+β„Ž2π‘˜2β‡’βˆ‡πœƒπ’₯οΈ€(πœƒ)=(2πœƒ1,1),βŸ¨β„Ž,π‘˜βŸ©Ξ˜=β„Ž1π‘˜1+4β„Ž2π‘˜2β‡’βˆ‡πœƒπ’₯οΈ€(πœƒ)=(2πœƒ1,1/4),

(3)

since both give βŸ¨βˆ‡πœƒπ’₯οΈ€(πœƒ),β„ŽβŸ©Ξ˜=2πœƒ1β„Ž1+β„Ž2.

2. Two surrogates

The linearization π’₯οΈ€(πœƒ)+βŸ¨π‘”,β„ŽβŸ©, with 𝑔=𝐷π’₯οΈ€(πœƒ), is a good model of π’₯οΈ€(πœƒ+β„Ž) only for small β„Ž, and on its own it has no minimizer: unless 𝑔=0, it decreases without bound along any direction with βŸ¨π‘”,β„ŽβŸ©<0. A step therefore needs a second ingredient that says how large β„Ž we allow. The length of β„ŽβˆˆΞ˜ is measured by a norm β€–β‹…β€–, which is something that we again have to choose. There are two natural ways to combine the norm with the linearization: as a constraint, or as a penalty,33 The norm here can be any norm on Θ, not just the one induced by the inner product on Θ.

β„ŽπΆ=arg minβ€–β„Žβ€–β‰€πœ‚βŸ¨π‘”,β„ŽβŸ©,β„Žπ‘ƒ=arg minβ„ŽβˆˆΞ˜βŸ¨π‘”,β„ŽβŸ©+πœ†2β€–β„Žβ€–2.

(4)

The first is a trust-region model: it trusts the linearization only within a radius πœ‚, which caps the length of the step. The second instead penalizes long steps, and as the penalty grows faster than the linearization can decrease, it always has a minimum.44 If π’₯οΈ€ is 𝐿-smooth with respect to β€–β‹…β€–, meaning that π’₯οΈ€(πœƒ+β„Ž)≀π’₯οΈ€(πœƒ)+βŸ¨π‘”,β„ŽβŸ©+𝐿2β€–β„Žβ€–2 for all β„Ž, then for πœ†β‰₯𝐿 the second surrogate is an upper bound on π’₯οΈ€, and β„Žπ‘ƒ decreases π’₯οΈ€ by at least the decrease of the surrogate. Convergence analyses use this; in practice πœ† is often chosen smaller.

Interactive figureΒ 1: The two surrogates over many steps

Interactive figure 1 (The two surrogates over many steps)Β Β  The graph shows both steps for

π’₯οΈ€(πœƒ)=3πœƒ2/4βˆ’5/4logcoshπœƒ,

(5)

with the norm |β„Ž|. In one dimension every norm is a multiple of |β„Ž|, so the choice of norm only matters from two dimensions on. The constrained step always has length πœ‚, since the linearization has no minimum inside the trust region, so close to the minimum it jumps back and forth over it. The penalized step shrinks with the derivative, but with πœ†=𝐿 set by the largest curvature it crawls where π’₯οΈ€ is flat. Constrained methods such as sign descent and Muon therefore let πœ‚ decrease over the iterations, while adaptive norms, later in the post, take longer steps where π’₯οΈ€ is flat.

At iteration π‘˜, the dashed line is the linearization at the constrained iterate πœƒπ‘˜, with the trust region |β„Ž|β‰€πœ‚ shaded, and the dashed parabola is the quadratic model at the penalized iterate πœƒπ‘˜. Drag πœƒ0, change πœ‚ and πœ†, and step through the iterations with π‘˜.

3. The dual norm and the sharp operator

To write the solutions of the constrained and penalized problems in (4) explicitly, we need two things: the dual norm, which measures the size of a derivative in Ξ˜βˆ—, and the sharp operator, which maps a derivative back to an element of Θ. In both cases it is important to note that they depend on the choice of a norm on Θ, rather than the inner product on Θ.

DefinitionΒ 1: Dual norm and sharp operator

Definition 1 (Dual norm and sharp operator)Β Β  For a norm β€–β‹…β€– on Θ, the dual norm of π‘”βˆˆΞ˜βˆ— and the sharp of 𝑔, π‘”β™―βˆˆΞ˜, are

β€–π‘”β€–βˆ—=maxβ€–β„Žβ€–β‰€1βŸ¨π‘”,β„ŽβŸ©,𝑔♯=βˆ’arg minβ„ŽβˆˆΞ˜βŸ¨π‘”,β„ŽβŸ©+12β€–β„Žβ€–2.

(6)

The dual norm is the slope of the linearization in its steepest direction; that is, a step of unit length with respect to the chosen norm changes the linearization by at most β€–π‘”β€–βˆ—, so that |βŸ¨π‘”,β„ŽβŸ©|β‰€β€–π‘”β€–βˆ—β€–β„Žβ€– for all β„ŽβˆˆΞ˜. The sharp operator, the map 𝑔↦𝑔♯, takes 𝑔 back to Θ: βˆ’π‘”β™― is the penalized step for πœ†=1, and its direction and length are set by the dual norm,55 Write β„Ž=𝑑𝑒 with 𝑑β‰₯0 and ‖𝑒‖=1. For a fixed 𝑑, the best 𝑒 has βŸ¨π‘”,π‘’βŸ©=βˆ’β€–π‘”β€–βˆ—, which leaves βˆ’π‘‘β€–π‘”β€–βˆ—+𝑑2/2, smallest at 𝑑=β€–π‘”β€–βˆ—.

βŸ¨π‘”,π‘”β™―βŸ©=β€–π‘”β€–βˆ—2and‖𝑔♯‖=β€–π‘”β€–βˆ—.

(7)

The minus sign is a convention. Like a gradient, 𝑔♯ points where π’₯οΈ€ increases, and a step subtracts it. When the minimizer is not unique, 𝑔♯ is any one of them.

The sharp operator is the map from Ξ˜βˆ— back to Θ that a norm provides, while the Riesz map is the one that an inner product provides.66 The two coincide exactly when the norm comes from the inner product, β€–β„Žβ€–2=βŸ¨β„Ž,β„ŽβŸ©Ξ˜. The naming convention and notation 𝑔♯ follow Kelner et al. [1].77 Bernstein and Newhouse [2] call 𝑔♯/β€–π‘”β€–βˆ— the duality map. Both the constrained and the penalized steps can be written explicitly using the sharp operator. The constrained problem is bounded below by βŸ¨π‘”,β„ŽβŸ©β‰₯βˆ’β€–π‘”β€–βˆ—β€–β„Žβ€–β‰₯βˆ’πœ‚β€–π‘”β€–βˆ—,88 The first inequality is the bound |βŸ¨π‘”,β„ŽβŸ©|β‰€β€–π‘”β€–βˆ—β€–β„Žβ€– from the definition of the dual norm. The second holds because the constraint gives β€–β„Žβ€–β‰€πœ‚ and β€–π‘”β€–βˆ—β‰₯0. which the step below attains, and writing β„Ž=π‘˜/πœ† turns the penalized problem into the definition. For 𝑔≠0,

β„ŽπΆ=arg minβ€–β„Žβ€–β‰€πœ‚βŸ¨π‘”,β„ŽβŸ©=βˆ’πœ‚π‘”β™―β€–π‘”β€–βˆ—,β„Žπ‘ƒ=1πœ†arg minπ‘˜βˆˆΞ˜βŸ¨π‘”,π‘˜βŸ©+12β€–π‘˜β€–2=βˆ’1πœ†π‘”β™―.

(8)

Here 𝑔♯/β€–π‘”β€–βˆ— is the unit vector, in the chosen norm, along 𝑔♯, so β„ŽπΆ goes straight to the border of the trust region. Both steps point along βˆ’π‘”β™―, which depends on the chosen norm, and differ only in length: πœ‚ for β„ŽπΆ and β€–π‘”β€–βˆ—/πœ† for β„Žπ‘ƒ, so they coincide for πœ†=β€–π‘”β€–βˆ—/πœ‚.

ExampleΒ 2: The dual norm and the sharp operator in the plane

Example 2 (The dual norm and the sharp operator in the plane)Β Β  Let Θ=ℝ2 and βŸ¨π‘”,β„ŽβŸ©=3β„Ž1+4β„Ž2. For the Euclidean norm β€–β„Žβ€–2=β„Ž12+β„Ž22,

β€–π‘”β€–βˆ—=maxβ€–β„Žβ€–2≀13β„Ž1+4β„Ž2=5,𝑔♯=βˆ’arg minβ„Žβˆˆβ„23β„Ž1+4β„Ž2+12(β„Ž12+β„Ž22)=(3,4).

(9)

By the Cauchy–Schwarz inequality, 3β„Ž1+4β„Ž2≀5β€–β„Žβ€–2, with equality at β„Ž=(3,4)/5, and in the problem for 𝑔♯, setting the derivative (3+β„Ž1,4+β„Ž2) of the objective to zero gives the minimizer β„Ž=(βˆ’3,βˆ’4), so 𝑔♯=(3,4), the gradient for the Euclidean inner product. For the max-norm β€–β„Žβ€–βˆž=max(|β„Ž1|,|β„Ž2|),

β€–π‘”β€–βˆ—=maxβ€–β„Žβ€–βˆžβ‰€13β„Ž1+4β„Ž2=7,𝑔♯=βˆ’arg minβ„Žβˆˆβ„23β„Ž1+4β„Ž2+12max(|β„Ž1|,|β„Ž2|)2=(7,7).

(10)

Both coordinates can be moved as far as allowed, so the maximum is 3+4 at β„Ž=(1,1). For the minimum, fix 𝑑=max(|β„Ž1|,|β„Ž2|): the linear part is then smallest at β„Ž=(βˆ’π‘‘,βˆ’π‘‘), which leaves βˆ’7𝑑+𝑑2/2, smallest at 𝑑=7. In both cases βŸ¨π‘”,π‘”β™―βŸ©=β€–π‘”β€–βˆ—2 and ‖𝑔♯‖=β€–π‘”β€–βˆ—. The same derivative gives different directions for different norms.99 The same argument gives 𝑔♯=‖𝑔‖1sign(𝑔) for the max-norm, and this map is not linear: (3,4)β™―+(βˆ’1,2)β™―=(7,7)+(βˆ’3,3)=(4,10), while (2,6)β™―=(8,8). So the max-norm is not the norm of any inner product, whose sharp operator, the Riesz map, is linear.

Interactive figureΒ 2: The sharp operator

Interactive figure 2 (The sharp operator)Β Β  The graph shows the definition for 𝑔=(3,4) of the example above and the ℓ𝑝-norms β€–β„Žβ€–π‘=(|β„Ž1|𝑝+|β„Ž2|𝑝)1/𝑝, which give the Euclidean norm for 𝑝=2 and the max-norm for 𝑝=∞. Its contours are those of πœ‘(β„Ž)=βŸ¨π‘”,β„ŽβŸ©+1/2β€–β„Žβ€–2, which is smallest at βˆ’π‘”β™―. The dual norm shows in two places: βˆ’π‘”β™― lies on the border β€–β„Žβ€–=β€–π‘”β€–βˆ— of the ball β€–β„Žβ€–β‰€β€–π‘”β€–βˆ—, and the level set βŸ¨π‘”,β„ŽβŸ©+β€–π‘”β€–βˆ—2=0, a line in the plane, touches the ball there. Only 𝑝=2 comes from an inner product. For every other 𝑝 the sharp operator is nonlinear, and for 𝑝 near 1 or ∞ it jumps between corners of the ball as 𝑔 is dragged, which no linear map can do.

The ball is shaded and the line dashed, with βˆ’π‘”β™― drawn from the origin. Drag 𝑔, change 𝑝 from 1 to ∞, and drag the point β„Ž to check that no point has a smaller value πœ‘(β„Ž) than βˆ’π‘”β™―.

4. The effect of the chosen norm

Once the norm is chosen, the step is fixed. To write the steps down, we need coordinates. Let Θ=β„π‘šΓ—π‘›, so that πœƒ is a matrix, for example the weights of a layer in a neural network. The derivative is then represented by the matrix 𝐺 of partial derivatives πœ•π’₯οΈ€/πœ•πœƒπ‘Žπ‘, with

βŸ¨π‘”,β„ŽβŸ©=⟨𝐺,β„ŽβŸ©πΉ=tr(πΊβŠ€β„Ž),𝑔=𝐷π’₯οΈ€(πœƒ).

(11)

This uses the Frobenius inner product, which makes 𝐺 the Riesz representative of 𝑔: the gradient for the Frobenius inner product.1010 This might seem like a limitation, but as shown in SectionΒ 4.1 every inner product on Θ can be written using the Frobenius inner product. Some norms on Θ, such as the max-norm, only look at the entries of β„Ž, as if the matrix were one long vector of π‘šπ‘› numbers. Others, such as the spectral norm, use that πœƒ is a matrix, a linear map from inputs to outputs, which is what the weights of a layer are. For matrices, β€–β‹…β€–1 and β€–β‹…β€–βˆž act on the entries as for vectors, while β€–β‹…β€–2 is the spectral norm, the largest singular value. Each norm then gives a dual norm and a sharp operator in terms of 𝐺, and with them a known method. Since both the constrained and the penalized steps follow from β€–π‘”β€–βˆ— and 𝑔♯, each example shows only the one that gives the known method. The other step is a rescaled version of it, for example normalized gradient descent in place of gradient descent. The table below collects the norms of this section.1111 The downloadable code computes 𝑔♯ for the Frobenius norm, a diagonal inner product, the max-norm, and the spectral norm by solving the defining problem as a small conic program [3], and prints it next to the closed form.

TableΒ 1: Three norms and the methods they give. For the first row the method is the penalized step βˆ’π‘”β™―/πœ†, and for the last two it is the constrained step βˆ’πœ‚π‘”β™―/β€–π‘”β€–βˆ—.

Norm on Ξ˜β€–π‘”β€–βˆ—π‘”β™―Method
Inner productβ€–π’œοΈ€βˆ’1(𝐺)β€–π’œοΈ€π’œοΈ€βˆ’1(𝐺)Gradient descent
Max-norm‖𝐺‖1‖𝐺‖1sign(𝐺)Sign descent
Spectral‖𝐺‖nuc‖𝐺‖nucπ‘ˆπ‘‰βŠ€Muon

The norms fall into two families. A norm from an inner product has a linear sharp operator, the Riesz map, and its step is a preconditioned gradient step. Every other norm has a nonlinear sharp operator, and its step is not the gradient for any fixed inner product, one that does not depend on 𝑔.

4.1. Norms from an inner product

Every inner product on Θ can be written as

βŸ¨β„Ž,π‘˜βŸ©π’œοΈ€=βŸ¨π’œοΈ€(β„Ž),π‘˜βŸ©πΉ

(12)

where the linear map π’œοΈ€:Ξ˜β†’Ξ˜ is symmetric and positive definite, i.e.,

βŸ¨π’œοΈ€(β„Ž),π‘˜βŸ©πΉ=βŸ¨β„Ž,π’œοΈ€(π‘˜)⟩𝐹andβŸ¨π’œοΈ€(β„Ž),β„ŽβŸ©πΉ>0for allβ„Žβ‰ 0.

(13)

Since βŸ¨π‘”,β„ŽβŸ©=⟨𝐺,β„ŽβŸ©πΉ=βŸ¨π’œοΈ€βˆ’1(𝐺),β„ŽβŸ©π’œοΈ€,1212 We look for the π‘˜βˆˆΞ˜ with βŸ¨π‘˜,β„ŽβŸ©π’œοΈ€=⟨𝐺,β„ŽβŸ©πΉ for all β„Ž. By the definition of βŸ¨β‹…,β‹…βŸ©π’œοΈ€, the left side is βŸ¨π’œοΈ€(π‘˜),β„ŽβŸ©πΉ, so βŸ¨π’œοΈ€(π‘˜)βˆ’πΊ,β„ŽβŸ©πΉ=0 for all β„Ž, which gives π’œοΈ€(π‘˜)=𝐺 and π‘˜=π’œοΈ€βˆ’1(𝐺). The inverse exists because π’œοΈ€ is positive definite. the gradient for βŸ¨β‹…,β‹…βŸ©π’œοΈ€ is βˆ‡π’œοΈ€π’₯οΈ€=π’œοΈ€βˆ’1(𝐺).1313 In Example 1, the second inner product is π’œοΈ€=diag(1,4), and its gradient (2πœƒ1,1/4) is π’œοΈ€βˆ’1 applied to the first. The inverse π’œοΈ€βˆ’1 may look unusual, but it is the familiar preconditioner of scientific computing, where a problem is rescaled before it is solved: a gradient step with βˆ‡π’œοΈ€π’₯οΈ€ is plain gradient descent in the rescaled variables πœ‰=π’œοΈ€1/2(πœƒ).1414 With πœ‰=π’œοΈ€1/2(πœƒ), the derivative of π’₯οΈ€(π’œοΈ€βˆ’1/2(πœ‰)) is represented by π’œοΈ€βˆ’1/2(𝐺), so a gradient step in πœ‰ changes πœ‰ by βˆ’1πœ†π’œοΈ€βˆ’1/2(𝐺), and πœƒ=π’œοΈ€βˆ’1/2(πœ‰) by βˆ’1πœ†π’œοΈ€βˆ’1(𝐺). With π’œοΈ€ close to the curvature of π’₯οΈ€, the rescaled problem is well conditioned, and the step comes close to Newton’s method.

ExampleΒ 3: Preconditioned gradient descent

Example 3 (Preconditioned gradient descent)Β Β  Let β€–β„Žβ€–2=βŸ¨β„Ž,β„ŽβŸ©π’œοΈ€, so that βŸ¨π‘”,β„ŽβŸ©=βŸ¨βˆ‡π’œοΈ€π’₯οΈ€,β„ŽβŸ©π’œοΈ€ with the gradient βˆ‡π’œοΈ€π’₯οΈ€=π’œοΈ€βˆ’1(𝐺). Then

β€–π‘”β€–βˆ—=maxβ€–β„Žβ€–π’œοΈ€β‰€1βŸ¨βˆ‡π’œοΈ€π’₯οΈ€,β„ŽβŸ©π’œοΈ€=β€–βˆ‡π’œοΈ€π’₯οΈ€β€–π’œοΈ€=⟨𝐺,π’œοΈ€βˆ’1(𝐺)⟩𝐹1/2,𝑔♯=βˆ’arg minβ„ŽβˆˆΞ˜βŸ¨βˆ‡π’œοΈ€π’₯οΈ€,β„ŽβŸ©π’œοΈ€+12β€–β„Žβ€–π’œοΈ€2=βˆ‡π’œοΈ€π’₯οΈ€.

(14)

By the Cauchy–Schwarz inequality, βŸ¨βˆ‡π’œοΈ€π’₯οΈ€,β„ŽβŸ©π’œοΈ€β‰€β€–βˆ‡π’œοΈ€π’₯οΈ€β€–π’œοΈ€β€–β„Žβ€–π’œοΈ€, with equality at β„Ž=βˆ‡π’œοΈ€π’₯οΈ€/β€–βˆ‡π’œοΈ€π’₯οΈ€β€–π’œοΈ€. In the problem for 𝑔♯, completing the square gives

βŸ¨βˆ‡π’œοΈ€π’₯οΈ€,β„ŽβŸ©π’œοΈ€+12β€–β„Žβ€–π’œοΈ€2=12β€–β„Ž+βˆ‡π’œοΈ€π’₯οΈ€β€–π’œοΈ€2βˆ’12β€–βˆ‡π’œοΈ€π’₯οΈ€β€–π’œοΈ€2,

(15)

which is smallest at β„Ž=βˆ’βˆ‡π’œοΈ€π’₯οΈ€, so 𝑔♯=βˆ‡π’œοΈ€π’₯οΈ€. The sharp operator is the Riesz map, and the penalized step

β„Žπ‘ƒ=βˆ’1πœ†π’œοΈ€βˆ’1(𝐺)

(16)

is preconditioned gradient descent with the preconditioner π’œοΈ€βˆ’1, and for the Frobenius norm, π’œοΈ€=ℐ︀, plain gradient descent.

Interactive figureΒ 3: A norm from an inner product

Interactive figure 3 (A norm from an inner product)Β Β  The graph shows the constrained step for an inner product on ℝ2, βŸ¨β„Ž,π‘˜βŸ©π’œοΈ€=β„ŽβŠ€π΄π‘˜, where

𝐴=π‘…πœ‘[100πœ…]π‘…πœ‘βŠ€,π‘…πœ‘=[cosπœ‘βˆ’sinπœ‘sinπœ‘cosπœ‘],

(17)

has the eigenvalues 1 and πœ…, with eigenvectors turned by the angle πœ‘. The unit ball is an ellipse, and the step β„ŽπΆ=βˆ’π΄βˆ’1𝑔/β€–π‘”β€–βˆ—, for πœ‚=1, is where the line βŸ¨π‘”,β„ŽβŸ©+β€–π‘”β€–βˆ—=0 touches it. Unlike the steps for 𝑝=1 and 𝑝=∞, it never jumps between corners, and unlike the steps of the other ℓ𝑝-norms, its direction depends linearly on 𝑔, since the sharp operator π΄βˆ’1 is linear.

Currently, 𝐴=

The Euclidean unit circle and the Euclidean step βˆ’π‘”/‖𝑔‖ are dotted.

4.2. Norms not from inner products

Recall that the norm need not come from an inner product. Every method still minimizes the same linearization βŸ¨π‘”,β„ŽβŸ©, with the derivative 𝑔=𝐷π’₯οΈ€(πœƒ)βˆˆΞ˜βˆ—. What changes is the norm that limits the step, and with it the sharp operator, the map from Ξ˜βˆ— back to Θ. For a norm that does not come from an inner product, the sharp operator is nonlinear,1515 If 𝑔↦𝑔♯ were linear, then β€–π‘”β€–βˆ—2=βŸ¨π‘”,π‘”β™―βŸ© would be a quadratic form in 𝑔, so the dual norm would come from an inner product on Ξ˜βˆ—, and then so would the norm on Θ, its dual. so it is not the Riesz map of any fixed inner product, and the step is not a gradient for any inner product that does not depend on 𝑔. The max-norm is one such norm: as in the plane, 𝑔♯=‖𝐺‖1sign(𝐺), and its constrained step β„ŽπΆ=βˆ’πœ‚sign(𝐺) is sign descent. The example below shows the spectral norm, whose constrained step is Muon.

ExampleΒ 4: The spectral norm: Muon

Example 4 (The spectral norm: Muon)Β Β  The spectral norm of a matrix β„Ž is its largest singular value, β€–β„Žβ€–2=𝜎max(β„Ž). Let πœƒβˆˆβ„π‘šΓ—π‘› be the weights of a linear layer 𝑍=πœƒπ‘‹. A step β„Ž changes the outputs by β„Žπ‘‹, and the spectral norm is the largest such change relative to the inputs,

β€–β„Žπ‘‹β€–πΉβ‰€β€–β„Žβ€–2‖𝑋‖𝐹for all𝑋.

(18)

As such, bounding the spectral norm effectively bounds the change of the outputs for every input. The dual norm and the sharp follow from the compact singular value decomposition 𝐺=π‘ˆΞ£π‘‰βŠ€=βˆ‘π‘˜πœŽπ‘˜π‘’π‘˜π‘£π‘˜βŠ€, which keeps only the nonzero singular values,

β€–π‘”β€–βˆ—=maxβ€–β„Žβ€–2≀1tr(πΊβŠ€β„Ž)=‖𝐺‖nuc,𝑔♯=βˆ’arg minβ„ŽβˆˆΞ˜tr(πΊβŠ€β„Ž)+12β€–β„Žβ€–22=‖𝐺‖nucπ‘ˆπ‘‰βŠ€.

(19)

By the cyclic property of the trace,

tr(πΊβŠ€β„Ž)=tr(π‘‰Ξ£π‘ˆβŠ€β„Ž)=tr(Ξ£π‘ˆβŠ€β„Žπ‘‰)=βˆ‘π‘˜πœŽπ‘˜(π‘ˆβŠ€β„Žπ‘‰)π‘˜π‘˜β‰€βˆ‘π‘˜πœŽπ‘˜=‖𝐺‖nuc,

(20)

since each diagonal entry of π‘ˆβŠ€β„Žπ‘‰ is at most its spectral norm, β€–π‘ˆβŠ€β„Žπ‘‰β€–2β‰€β€–β„Žβ€–2≀1. Equality holds at β„Ž=π‘ˆπ‘‰βŠ€, where π‘ˆβŠ€β„Žπ‘‰=𝐼. For the minimum, fix 𝑑=β€–β„Žβ€–2: the linear part is then smallest at β„Ž=βˆ’π‘‘π‘ˆπ‘‰βŠ€, which leaves βˆ’π‘‘β€–πΊβ€–nuc+𝑑2/2, smallest at 𝑑=‖𝐺‖nuc. The nuclear norm, the sum of the singular values, is therefore not chosen but appears as the dual of the spectral norm. The constrained step

β„ŽπΆ=βˆ’πœ‚π‘ˆπ‘‰βŠ€

(21)

is minus the matrix with the singular vectors of 𝐺 and all its nonzero singular values set to πœ‚, while the zero ones stay zero. It is the update of Muon [4], applied to an averaged derivative as in the next section. In a network, the constraint is applied to each block.1616 That is, the norm is maxπ‘–β€–β„Žπ‘–β€–2, as in Bernstein and Newhouse [2]. Its dual norm is βˆ‘π‘–β€–πΊπ‘–β€–nuc, and each block takes the step βˆ’πœ‚π‘ˆπ‘–π‘‰π‘–βŠ€. In practice, Muon computes π‘ˆπ‘‰βŠ€ with a few matrix products by a Newton–Schulz iteration, whose polynomials are improved by The Polar Express Algorithm, and uses π‘ˆπ‘‰βŠ€ only for the weight matrices of the hidden layers. Avoiding the SVD is what makes Muon practical for large networks: an SVD runs poorly on GPUs and needs high precision, while matrix products are fast even in low precision.

Interactive figureΒ 4: The Muon step

Interactive figure 4 (The Muon step)Β Β  Unlike the graphs above, this plane is not Θ, which here is four-dimensional: it holds the vectors a 2Γ—2 block acts on, and each matrix is shown through its image of the unit circle. The block 𝐺 maps the unit circle to an ellipse with the semi-axes 𝜎1𝑒1 and 𝜎2𝑒2. The matrix π‘ˆπ‘‰βŠ€ keeps the singular vectors and sets both singular values to 1: it maps the circle to itself, and its columns π‘ˆπ‘‰βŠ€π‘’1 and π‘ˆπ‘‰βŠ€π‘’2, which take the place of 𝐺𝑒1 and 𝐺𝑒2, are orthonormal.

5. Incorporating the past

The surrogates above have no memory of the past. The linearization is created using only the current iterate πœƒπ‘‘ and its derivative 𝑔𝑑=𝐷π’₯οΈ€(πœƒπ‘‘), and the norm is fixed. Methods with memory use past iterations in two ways: momentum methods change the linearization, and adaptive methods change the norm.

5.1. Momentum

The idea behind momentum is to build the linearization from past derivatives as well. The surrogates keep their form, with a chosen derivative π‘”Μ‚π‘‘βˆˆΞ˜βˆ— in place of 𝑔𝑑,

β„ŽπΆ=arg minβ€–β„Žβ€–β‰€πœ‚βŸ¨π‘”Μ‚π‘‘,β„ŽβŸ©,β„Žπ‘ƒ=arg minβ„ŽβˆˆΞ˜βŸ¨π‘”Μ‚π‘‘,β„ŽβŸ©+πœ†2β€–β„Žβ€–2.

(22)

The derivative 𝑔̂𝑑 enters before the sharp operator, in Ξ˜βˆ—. When the sharp operator is linear, as for a norm from an inner product, averaging the derivatives is the same as averaging the steps they give.1717 A linear sharp operator commutes with averaging, (βˆ‘π‘˜π‘€π‘˜π‘”π‘˜)β™―=βˆ‘π‘˜π‘€π‘˜π‘”π‘˜β™―. A nonlinear one does not: for Muon, the orthogonalized average of the derivatives is not the average of the orthogonalized derivatives.

ExampleΒ 5: Averaged linearizations: Adam and Muon

Example 5 (Averaged linearizations: Adam and Muon)Β Β  Average the past derivatives with weights (1βˆ’π›½)π›½π‘˜ for π›½βˆˆ(0,1), giving

𝑔̂𝑑=(1βˆ’π›½)βˆ‘π‘˜β‰₯0π›½π‘˜π‘”π‘‘βˆ’π‘˜.

(23)

The pairing is linear in 𝑔 for any norm, so βŸ¨π‘”Μ‚π‘‘,β„ŽβŸ© is the same weighted average of the past linearizations βŸ¨π‘”π‘‘βˆ’π‘˜,β„ŽβŸ©. Both steps keep their form with 𝑔̂𝑑 in place of 𝑔𝑑,

β„ŽπΆ=βˆ’πœ‚π‘”Μ‚π‘‘β™―β€–π‘”Μ‚π‘‘β€–βˆ—,β„Žπ‘ƒ=βˆ’1πœ†π‘”Μ‚π‘‘β™―.

(24)

The past derivatives were taken at earlier points, so as a model of π’₯οΈ€ near πœƒπ‘‘ this is worse than the linearization with 𝑔𝑑. But when each 𝑔𝑑 is computed from a minibatch, 𝑔̂𝑑 is a less noisy estimate of the expected derivative. In coordinates, 𝐺̂←𝛽𝐺̂+(1βˆ’π›½)𝐺. Muon takes the step of Example 4 with 𝐺̂ in place of 𝐺, and Adam does the same with the adaptive norm of Example 6.

5.2. Adaptive norms

Past information can also enter the norm. An adaptive method builds an inner product βŸ¨β‹…,β‹…βŸ©π’œοΈ€π‘‘ from past derivatives, so that the norm changes at every step. Each step is still the one of Example 3, with π’œοΈ€π‘‘ in place of π’œοΈ€,

β„Žπ‘ƒ=βˆ’1πœ†π’œοΈ€π‘‘βˆ’1(𝐺).

(25)

Storing a general π’œοΈ€π‘‘ is out of the question for a network, so two cheap structures are used: a diagonal one, which stores as many numbers as πœƒ, and a Kronecker-factored one.

ExampleΒ 6: A diagonal norm: AdaGrad and Adam

Example 6 (A diagonal norm: AdaGrad and Adam)Β Β  AdaGrad [5] sums the squared partial derivatives of all past iterations, and Adam [6] takes a moving average of them instead,

𝑣←𝑣+𝐺2(AdaGrad),𝑣←𝛽2𝑣+(1βˆ’π›½2)𝐺2(Adam),

(26)

where the square, like the square root below, acts entry by entry. So 𝑣 has the same shape as πœƒ, one number per entry. Both measure the step with the weighted norm

β€–β„Žβ€–π‘‘2=βˆ‘π‘Ž(π‘£π‘Ž+πœ€)β„Žπ‘Ž2=βŸ¨π’œοΈ€π‘‘(β„Ž),β„ŽβŸ©πΉ=βŸ¨β„Ž,β„ŽβŸ©π’œοΈ€π‘‘,π’œοΈ€π‘‘(β„Ž)=(𝑣+πœ€)βŠ™β„Ž,

(27)

where π‘Ž runs over all entries of πœƒ, and βŠ™ multiplies two elements of Θ entry by entry, so that π’œοΈ€π‘‘ multiplies each β„Žπ‘Ž by π‘£π‘Ž+πœ€. As ‖⋅‖𝑑 comes from an inner product, Example 3 gives

𝑔♯=π’œοΈ€π‘‘βˆ’1(𝐺),

(28)

which divides each πΊπ‘Ž by π‘£π‘Ž+πœ€, and the penalized step with πœ†=1/πœ‚ is βˆ’πœ‚π’œοΈ€π‘‘βˆ’1(𝐺). Adam also uses the averaged derivative 𝐺̂ of SectionΒ 5.1 in place of 𝐺, so it carries memory in both the linearization and the norm, and its step, without bias correction,1818 Adam starts both averages at zero, so after 𝑑 steps their weights sum to 1βˆ’π›½π‘‘, and early averages are too small. It divides 𝐺̂ by 1βˆ’π›½1𝑑 and 𝑣 by 1βˆ’π›½2𝑑, which scales the step by 1βˆ’π›½2𝑑/(1βˆ’π›½1𝑑) when πœ€ is ignored. The factor approaches 1 on the time scale of 1/(1βˆ’π›½2) steps. is βˆ’πœ‚π’œοΈ€π‘‘βˆ’1(𝐺̂).

Interactive figureΒ 5: The adaptive norm of AdaGrad

Interactive figure 5 (The adaptive norm of AdaGrad)Β Β  The graph shows AdaGrad on the quadratic π’₯οΈ€(πœƒ)=1/2(πœƒ12+πœ…πœƒ22), next to gradient descent with πœ†=𝐿=πœ…. At each iterate, the ellipse is the ball of ‖⋅‖𝑑 through the next iterate, which the step reaches where a line of constant linearization touches the ball. Already the first derivative makes the ball narrow in the steep direction πœƒ2, where the derivatives are large, and wide along the flat coordinate πœƒ1. As 𝑣 accumulates, the ball, and with it the step, shrinks. AdaGrad thereby takes long steps where gradient descent crawls.

Gradient descent is dotted and AdaGrad solid, with thirty steps each. Drag πœƒ0, change πœ‚ and πœ…, and step through the iterations with π‘˜.

ExampleΒ 7: A Kronecker-factored norm: K-FAC and Shampoo

Example 7 (A Kronecker-factored norm: K-FAC and Shampoo)Β Β  Let πœƒβˆˆβ„π‘šΓ—π‘› be the weights of a linear layer 𝑍=πœƒπ‘‹ with inputs π‘‹βˆˆβ„π‘›Γ—π‘. A step β„Ž changes the outputs by β„Žπ‘‹. Measuring this change, weighted on the output side by a symmetric positive definite πΏβˆˆβ„π‘šΓ—π‘š, gives1919 In a network, each layer has its own 𝐿 and 𝑅.

‖𝐿1/2β„Žπ‘‹β€–πΉ2=βŸ¨π’œοΈ€(β„Ž),β„ŽβŸ©πΉ,𝑅=π‘‹π‘‹βŠ€,π’œοΈ€(β„Ž)=πΏβ„Žπ‘…,

(29)

which is positive definite only if π‘‹π‘‹βŠ€ is. Unlike the spectral norm of Example 4, which is the worst case over all inputs, this norm measures the change on the given inputs. Its inverse, π’œοΈ€βˆ’1(𝐺)=πΏβˆ’1πΊπ‘…βˆ’1, needs only an π‘šΓ—π‘š and an 𝑛×𝑛 matrix, not a general one of size π‘šπ‘›Γ—π‘šπ‘›.2020 Stacking the columns of β„Ž into one long vector vec(β„Ž) gives vec(πΏβ„Žπ‘…)=(π‘…βŠ—πΏ)vec(β„Ž), so π’œοΈ€ is the Kronecker product π‘…βŠ—πΏ. This is why such norms and preconditioners are called Kronecker-factored, the KF in K-FAC. Since the inputs change from batch to batch, K-FAC [7] takes moving averages over the iterations of π‘‹π‘‹βŠ€ and of Ξ”Ξ”βŠ€, where Ξ” holds the derivatives of the loss with respect to the outputs 𝑍,2121 In K-FAC, 𝑋 has an extra row of ones for the bias, and Ξ” is computed with targets sampled from the network’s own predictive distribution rather than taken from the data, so that π‘…βŠ—πΏ approximates a block of the Fisher information matrix.

π‘…β†πœŒπ‘…+(1βˆ’πœŒ)π‘‹π‘‹βŠ€,πΏβ†πœŒπΏ+(1βˆ’πœŒ)Ξ”Ξ”βŠ€,

(30)

and its basic step is βˆ’πœ‚πΏβˆ’1πΊπ‘…βˆ’1. K-FAC also adds damping to 𝐿 and 𝑅, and takes the step 𝛼𝑑𝑑+πœ‡β„Žπ‘‘βˆ’1 with 𝑑𝑑=βˆ’πΏβˆ’1πΊπ‘…βˆ’1, where 𝛼 and πœ‡ are chosen at every step by minimizing a quadratic model. This is a momentum term whose coefficients change from step to step. Shampoo [8] instead sums the products of the derivative with itself, starting from 𝐿=πœ€πΌ and 𝑅=πœ€πΌ,

𝐿←𝐿+𝐺𝐺⊀,𝑅←𝑅+𝐺⊀𝐺,

(31)

and takes the step βˆ’πœ‚πΏβˆ’1/4πΊπ‘…βˆ’1/4, so its norm uses the fourth roots of 𝐿 and 𝑅.

6. Summary

TableΒ 2 summarizes the methods of the post by their norm and by where past information enters.

TableΒ 2: The methods of the post as choices of the norm and of where past information enters: the linearization through 𝑔̂𝑑, and the norm. K-FAC also uses the previous step, and implementations of Shampoo often add momentum as well.

MethodNorm𝑔̂𝑑Norm memory
Gradient descentInner product𝑔𝑑–
AdaGradDiagonal𝑔𝑑Sum
AdamDiagonalAverageAverage
K-FACKronecker-factored𝑔𝑑Average
ShampooKronecker-factored𝑔𝑑Sum
Sign descentMax-norm𝑔𝑑–
MuonSpectralAverage–

Cite this post

bibtex
@misc{paltorp2026steepestdescent,
author = {Paltorp, Mikkel},
title = {{Which Way Is Down? First-Order Optimizers as a Choice of Norm}},
year = {2026},
month = oct,
url = {https://mipals.github.io/pdf/2026/10-13-steepest-descent/steepest-descent.pdf},
note = {Blog post}
}

Bibliography

  • [1] J. A. Kelner, Y. T. Lee, L. Orecchia, and A. Sidford, β€œAn Almost-Linear-Time Algorithm for Approximate Max Flow in Undirected Graphs, and its Multicommodity Generalizations,” in Proceedings of the Twenty-Fifth Annual ACM-SIAM Symposium on Discrete Algorithms (SODA), 2014, pp. 217–226. doi: 10.1137/1.9781611973402.16.
  • [2] J. Bernstein and L. Newhouse, β€œModular Duality in Deep Learning.” [Online]. Available: https://arxiv.org/abs/2410.21265
  • [3] P. J. Goulart and Y. Chen, β€œClarabel: An interior-point solver for conic programs with quadratic objectives.” [Online]. Available: https://arxiv.org/abs/2405.12762
  • [4] K. Jordan et al., β€œMuon: An optimizer for hidden layers in neural networks.” [Online]. Available: https://kellerjordan.github.io/posts/muon/
  • [5] J. Duchi, E. Hazan, and Y. Singer, β€œAdaptive Subgradient Methods for Online Learning and Stochastic Optimization,” Journal of Machine Learning Research, vol. 12, no. 61, pp. 2121–2159, 2011, [Online]. Available: https://jmlr.org/papers/v12/duchi11a.html
  • [6] D. P. Kingma and J. Ba, β€œAdam: A Method for Stochastic Optimization,” in International Conference on Learning Representations (ICLR), 2015. [Online]. Available: https://arxiv.org/abs/1412.6980
  • [7] J. Martens and R. Grosse, β€œOptimizing Neural Networks with Kronecker-factored Approximate Curvature,” in Proceedings of the 32nd International Conference on Machine Learning, in Proceedings of Machine Learning Research, vol. 37. 2015, pp. 2408–2417. [Online]. Available: https://proceedings.mlr.press/v37/martens15.html
  • [8] V. Gupta, T. Koren, and Y. Singer, β€œShampoo: Preconditioned Stochastic Tensor Optimization,” in Proceedings of the 35th International Conference on Machine Learning, in Proceedings of Machine Learning Research, vol. 80. 2018, pp. 1842–1850. [Online]. Available: https://proceedings.mlr.press/v80/gupta18a.html