8 papers
Dion3: Full-Stack Orthogonal Updates
Noah Amsel, Jack Zhang, Kwangjun Ahn +5
The Muon optimizer incurs a significant overhead cost due to its cubic-time Newton-Schulz orthogonalization step. When weights are sharded, communication overhead compounds this co…
Customizing the Inductive Biases of Softmax Attention using Structured Matrices
Yilun Kuang, Noah Amsel, Sanae Lotfi +3
The core component of attention is the scoring function, which transforms the inputs into low-dimensional queries and keys and takes the dot product of each pair. While the low-dim…
The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm
Noah Amsel, David Persson, Christopher Musco +1
Computing the polar decomposition and the related matrix sign function has been a well-studied problem in numerical analysis for decades. Recently, it has emerged as an important s…
Linear Systems and Eigenvalue Problems: Open Questions from a Simons Workshop
Noah Amsel, Yves Baumann, Paul Beckman +36
This document presents a series of open questions arising in matrix computations, i.e., the numerical solution of linear algebra problems. It is a result of working groups at the w…
Compositional Reasoning with Transformers, RNNs, and Chain of Thought
Gilad Yehudai, Noah Amsel, Joan Bruna
It is well understood that different neural network architectures are suited to different tasks, but is there always a single best architecture for a given task? We compare the exp…
Dion2: A Simple Method to Shrink Matrix in Muon
Kwangjun Ahn, Noah Amsel, John Langford
The Muon optimizer enjoys strong empirical performance and theoretical grounding. However, the super-linear cost of its orthonormalization step introduces increasing overhead with…