6 papers · 1 filter
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…
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…
On the Benefits of Rank in Attention Layers
Noah Amsel, Gilad Yehudai, Joan Bruna
Attention-based mechanisms are widely used in machine learning, most prominently in transformers. However, hyperparameters such as the rank of the attention matrices and the number…