Table of Contents
Fetching ...

Iterative Orthogonalization Scaling Laws

Devan Selvaraj

TL;DR

The paper investigates potential scaling limitations of muon's iterative orthogonalization as model sizes grow, showing that singular-value distributions shift with scale and may hamper large-scale deployments. It combines empirical experiments on randomly initialized matrices with Marchenko-Pastur-based theory to characterize how singular values shrink with dimension and how NS iterations affect this behavior. A key contribution is articulating a scaling law for singular values with dimension and discussing tuning requirements for NS iterations at scale, without proposing concrete remedies. The findings alert practitioners to possible scaling challenges in very large models and motivate further work on robust orthogonalization strategies and precision considerations in large-scale settings.

Abstract

The muon optimizer has picked up much attention as of late as a possible replacement to the seemingly omnipresent Adam optimizer. Recently, care has been taken to document the scaling laws of hyper-parameters under muon such as weight decay and learning rate. However, at much larger scales the iterative orthogonalization procedure present in muon may suffer a possible issue as the singular values of random matrices shrink with scale. This paper shows this scaling behavior theoretically and empirically on random matrices but does not suggest what to do about it.

Iterative Orthogonalization Scaling Laws

TL;DR

The paper investigates potential scaling limitations of muon's iterative orthogonalization as model sizes grow, showing that singular-value distributions shift with scale and may hamper large-scale deployments. It combines empirical experiments on randomly initialized matrices with Marchenko-Pastur-based theory to characterize how singular values shrink with dimension and how NS iterations affect this behavior. A key contribution is articulating a scaling law for singular values with dimension and discussing tuning requirements for NS iterations at scale, without proposing concrete remedies. The findings alert practitioners to possible scaling challenges in very large models and motivate further work on robust orthogonalization strategies and precision considerations in large-scale settings.

Abstract

The muon optimizer has picked up much attention as of late as a possible replacement to the seemingly omnipresent Adam optimizer. Recently, care has been taken to document the scaling laws of hyper-parameters under muon such as weight decay and learning rate. However, at much larger scales the iterative orthogonalization procedure present in muon may suffer a possible issue as the singular values of random matrices shrink with scale. This paper shows this scaling behavior theoretically and empirically on random matrices but does not suggest what to do about it.
Paper Structure (8 sections, 4 theorems, 9 equations, 2 figures)

This paper contains 8 sections, 4 theorems, 9 equations, 2 figures.

Key Result

Theorem 1

The singular values of a gradient normalized by it's Frobenius norm scale $1/\sqrt{in\_d}$. Proof.Assume the gradient entries are i.i.d with mean $0$ finite variance $\sigma^2$. Assume the $in\_d/out\_d$ ratio remains constant with scale ApplyFrobenius Norm Scaling Lawthm:frob-scaling and $\nabla W$ ApplySingular Value Scaling Lawthm:sval-scaling-1

Figures (2)

  • Figure 1: Muon's chosen polynomial (from Keller's Blog)muon
  • Figure :

Theorems & Definitions (6)

  • Theorem 1: Normalized Singular Value Scaling Law
  • Lemma 1: Frobenius Norm Scaling Law
  • proof
  • Theorem 2: Marchenko-Pastur Theorem
  • Lemma 2: Singular Value Scaling Law
  • proof