I've finally solved steepest descent on Finsler-structured (matrix) manifolds more generally. This generalizes work by me, @jxbz, and @Jianlin_S on Muon, Orthogonal Muon, & Stiefel Muon. --- The general solution turned out to be much simpler than I thought. And it should generalize to any combination of (underlying manifold, Finsler norm) and any number of extra constraints on the updates so long as the feasible set for each constraint is convex. --- I now consider this class of problems as sufficiently solved (by my definition of 'solved') and thus I'm moving on to other things I'm interested about.

This too!

This is not a new result btw! See:

Here's the tracker issue: github.com/marin-community/m…

I recommend reading this work done by this awesome team. They also helped me test the claim that larger batch sizes widens the gap between optimizers like Muon and AdamW. Here's an early result (see pic). And I think this would be a decent interview question: can you explain why this happens and how this speeds up larger scale pretraining?

Amazing work by the Pytorch team!

This is a much-needed boost to my self-confidence atm. Thank you so much @_arohan_ !! Can vouch that these folks are looking at things in a different way. And I hope I could live up to expectations 🙏

Triton implementation by @WhyPhy_Labs !

Also check out @Jianlin_S's solution:

If you're interested in reading more about this problem, check out @jxbz's writeup on Stiefel Muon here: docs.modula.systems/algorith…

This turned out to have a rather simple solution taking advantage of the fact that the feasible sets of the constraints are convex: the primal-dual approach!

Here's the problem we're trying to solve (see pic) More intuitively, given a "raw gradient" G, we want to find the optimal update A that is (1) maximal according to some norm and (2) in the tangent space at the current point W in the manifold.

Finsler manifolds are a generalization of Riemannian manifolds: we just need to remove the requirement of the 'measure of distance' having a bilinear form and viola Think of it this way: Imagine some lovecraftian horror pulls you into its realm. And with you is a ruler. It 'looks' the same in any direction, but its ticks recalibrate as you rotate the ruler. That's what it's like to be in a Finsler geometry. en.wikipedia.org/wiki/Finsle…

But then... 1. What if we want to do steepest descent under other norms? 2. What if we want to constraint the weights to "live in" a manifold to prevent them from blowing up to infinity or from degenerating to 0? Now we have to do optimization on Finsler manifolds.

First, here's a minimal construction of the Muon optimizer: 1. Take R^{m x n} 2. Equip the tangent spaces with the spectral norm 3. Do first-order optimization on the resulting manifold 4. Add momentum and viola, you get Muon. --- What I like about this approach is that we have solid reasons for doing each step: 1. Linear weights are matrices 2. arxiv.org/abs/2310.17813 3. Why do second-order optimization when first-order would suffice? 4. Momentum accelerates optimization + handles gradient noise

I have so many interests I find it hard to focus on any of them I wanna study algebraic topology, category theory, optimization on finsler manifolds but also, I wanna build. I can build the entire AI infra of an AI SaaS, even the UI. I've done it before yet here I am, working on something I'm not an expert at yet

Finally, here's the code: colab.research.google.com/dr… I challenge you to try to beat my record! Rules: - You must reach a median 95% evaluation accuracy in (at least) 64 random seeds - Your model has to be ~1-Lipschitz for fair comparison - Active parameters per token must be at most how many is in my current record I'll setup another track for transformer models. Have fun!

Also check out this article by @Jianlin_S on solving the steepest descent problem on stiefel manifold. It's been fun going back-and-forth with him on the topic! It does seem that the performance gap between our methods is negligible at scale, although one can easily find adversarial examples at much smaller scales that widens the gap via e.g. backpropagation. x.com/Jianlin_S/status/19536…

If you wanna read more about our paper on training transformers with enforced lipschitz bounds, please check out this awesome thread by @LakerNewhouse These fellas were awesome collaborators 💙