Training Deep Learning Models with Norm-Constrained LMOs

We are happy that the reviewer found the paper insightful.

"Since you build on top of modded-nanoGPT, I have to ask the question: what is your speedrun result against Muon?"

The speedrun configuration is the 124M parameter model with 512 batchsize in Fig. 1. In this setting Scion consistently across 3 runs has slightly lower validation loss than Muon (with the same wallclock time), which would translate into a slightly better wallclock time if the number of iterations is reduced to exactly match the validation loss of Muon.

The more substantial gain occurs at larger batch sizes. We conduct an additional experiment for the batch size 6144 configuration of Fig. 2, where we simply reduce the number of iterations until Scion matches the larger validation loss of Muon. We find that Scion can achieve the same validation loss as Muon with a 25% smaller wallclock time in this setting.

We thank the reviewer for their feedback. We believe there are important misunderstandings that we address below.

"The algorithmic descriptions in algorithm 1 and 2 are same. Is this a typo?"

This is not a typo: the two algorithms are distinct and solves different problems as detailed in Section 3. Algorithm 1 solves an unconstrained problem, while Algorithm 2 solves a norm-ball constrained problem.

For deep learning this leads to a distinct behaviour of the norm of the weights between Algorithm 1 and 2 as illustrated in Fig. 8-9, since only Algorithm 2 ensures that the norm does not grow beyond the constraint radius $\rho$. This explicit norm control can be particularly important for numerical stability and to avoid overfitting, as e.g. seen in the CIFAR10 experiments were Algorithm 2 (Scion) works without the heuristic weight normalization otherwise present in the Muon implementation for this problem.

"Both algorithm 1 and 2 are simply steepest descent methods with momentum"

Neither of our two proposed algorithms is steepest descent (Please see Section 4 were we explicitly compare with steepest descent). Algorithm 2 is a Conditional Gradient (aka Frank-Wolfe) based algorithm which applies to constrained problems, while with Algorithm 1 we show that the LMO can even be used for unconstrained problems.

One important distinction with steepest descent is that the stepsize $\gamma$ does not have to depend on the Lipschitz constant $L$ (see our main results Lemma 5.3-5.6). Intuitively, this is due to the scale invariance of the LMO, which makes the algorithm more stable (the algorithm is less sensitive to the curvature).

"The experiments on batch size sensitivity does support that the proposed method has better validation loss compares to Muon when batch size is large. Nevertheless, the validation loss curves for all compared algorithms look to be increasing w.r.t. batch size"

The loss is increasing with increasing batch size in the experiments because we keep the total number of tokens fixed. The validation loss as a function of batch size is expected to be worse for any optimizer in this setting – the comparison criterion is instead the rate with which performance deteriorates. As the reviewer points out, Scion have better validation accuracy than the baselines for larger batch sizes.

One way to quantify the validation loss improvement, is to ask how much quicker Scion can achieve the same validation accuracy as one of the baselines. To this end we conduct an additional experiment for the large batch size of 6144, where we simply reduce the number of iterations until Scion matches the larger validation loss of Muon. We find that Scion can achieve the same validation loss as Muon with a 25% smaller wallclock time.

We thank the reviewer for their positive feedback.

"To be verified: In Lemma D.1, I think it is the right constant is $L_2$, and not $L$."

The right constant is $L$, since smoothness is taken w.r.t. the arbitrary norm $\Vert\cdot\Vert$. This is what eventually yields the dependency on $L$ seen in the main results of Section 5.

What might have caused the confusion is that there is a minor typo in the statement of Lemma D.1, since the guarantee of Lemma D.1 should be stated in terms of $\Vert\nabla f(\bar x^n)\Vert_*$ instead of $\Vert\nabla f(\bar x^n)\Vert_2^2$. This is just a typo though, that does not change the final result, since the following Lemma 5.3 is correctly stated in terms of the dual norm. We will correct the typo in the final version.

"One particular line of related research (that is quite old) is natural gradient descent [1]. It also exploits the geometry of the function by defining a pseudo-metric in the parameters space by KL-divergence. It is not very popular nowadays since: 1) Calculating the Fischer matrix and its inverse is quite expensive, and 2) The empirical performance in stochastic setting is not very impressive. However, I think it is worth mentioning because it is among the first papers stepping away from the Euclidean metric and proposing a new metric for optimization, to the best of my knowledge. "

This is a good point. We will remember to cite and discuss natural gradient descent.

Geometry of activation functions

It might be possible to further improve the performance by taking the geometry of activation functions into account, but it seems challenging. One place to maybe look for ideas is the vast literature on initialization.

We thank the reviewer for their feedback and address the remaining concerns below.

"The CIFAR-10 experiments show the transferability of the optimal stepsize, but they do not establish the method's competitiveness in image classification"

To establish competitiveness on image classification we conduct a speedrun to achieve the 94% test accuracy target on the airbench codebase. Scion achieves a 94.08% test accuracy (averaged over 200 runs) in the same wallclock time that the SOTA method Muon uses to achieve the target accuracy. This provides a 13.6% speedup over the tuned baseline of SGD with momentum.

Note that Scion achieves this SOTA performance without the ad-hoc weight normalization otherwise used in the airbench implementation of Muon.

"The paper could benefit from a clearer explanation of the intuition behind why uSCG works for unconstrained problems."

It was indeed also surprising to us initially that we could show convergence for uSCG.

After applying smoothness in the analysis there are mainly two terms to address: i) the direction of the update (given by an inner product), and ii) the magnitude of the update. For simplicity, let us consider the deterministic case where we have

$f(x^{k+1}) \leq f(x^k) + \gamma\rho \langle \nabla f(x^k), \operatorname{lmo}(\nabla f(x^k)\rangle + \tfrac{\gamma^2\rho^2 L}{2}$

The main insight is the relationship between the LMO over a norm-ball and the definition of the dual norm. Specifically, the value that the LMO attains on the boundary of the norm-ball is exactly the dual norm of the gradient. So the direction of the update (the inner product) is still informative and it in fact decreases the function value by the dual norm of the gradient. However, the magnitude of the update (the last term) will be large even close to a solution, but this can be taken care of by the stepsize choice (as made precise by Lemma 5.3 and 5.4).

Notice that we do not need to know the Lipschitz constant $L$ to select the stepsize $\gamma$ as is otherwise the case in e.g. gradient descent. The intuitive reason for this is the scale invariance of the LMO. We will remark on these intuitions in the final version.