Innovation-Augmented Polar Decomposition

• Innovation-Augmented Polar Decomposition is a matrix factorization method that uses adaptive, odd polynomial iterations to compute the closest semi-orthogonal matrix, thus enhancing convergence.
• The Polar Express algorithm re-optimizes polynomial coefficients at each iteration based on the current singular value range, achieving worst-case optimal convergence rates without costly inversions.
• The technique incorporates finite-precision safeguards and integrates with the Muon framework, resulting in accelerated training performance and improved deep learning optimization.

Innovation-Augmented Polar Decomposition refers to algorithmic advancements in computing the polar decomposition of matrices, with a focus on addressing the practical demands of modern deep learning optimization frameworks. The Polar Express algorithm exemplifies this concept, providing a GPU-compatible, worst-case optimal, and low-precision-stable approach that advances both the theoretical and applied facets of matrix function computation in machine learning settings (Amsel et al., 22 May 2025).

1. Mathematical Formulation and Classical Methods

The polar decomposition of a real matrix $M \in \mathbb{R}^{m \times n}$ with reduced singular value decomposition (SVD) $M = U \Sigma V^T$ ($U \in \mathbb{R}^{m \times r}$, $V \in \mathbb{R}^{n \times r}$, $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$) is given by

$$M = \underbrace{U V^T}_{\operatorname{polar}(M)} \Sigma$$

where $\operatorname{polar}(M) = U V^T$ is the closest semi-orthogonal matrix to $M$ in spectral or Frobenius norm. For symmetric matrices, the related matrix sign function is defined as

$$\operatorname{sign}(M) = V \operatorname{diag}(\pm 1) V^T$$

Iterative polynomial methods for $\operatorname{polar}(M)$ construct a sequence: $M = U \Sigma V^T$0 where $M = U \Sigma V^T$1 is an odd polynomial of fixed degree $M = U \Sigma V^T$2. The goal is $M = U \Sigma V^T$3.

Classical approaches include:

• Newton–Schulz Iteration: Applies a cubic polynomial $M = U \Sigma V^T$4 to $M = U \Sigma V^T$5. Quadratic convergence is observed once singular values are close to 1, but initial progress is slow if $M = U \Sigma V^T$6.
• Rational-function Methods: For example, Newton’s matrix sign iteration $M = U \Sigma V^T$7 achieves rapid convergence but requires matrix inversion or QR decomposition, making implementation challenging on GPUs and at low precision.

2. Polar Express Algorithmic Innovations

Polar Express enhances the polynomial-iteration paradigm through an "innovation-augmented" minimax optimization at each step. At iteration $M = U \Sigma V^T$8, the current singular value range $M = U \Sigma V^T$9 defines a scalar minimax problem: $U \in \mathbb{R}^{m \times r}$0 where $U \in \mathbb{R}^{m \times r}$1 is the set of odd polynomials of degree $U \in \mathbb{R}^{m \times r}$2. This optimal polynomial is characterized by the Chebyshev equioscillation theorem, yielding a unique optimal solution with $U \in \mathbb{R}^{m \times r}$3 alternating-error points, $U \in \mathbb{R}^{m \times r}$4.

Crucially, the approach re-optimizes $U \in \mathbb{R}^{m \times r}$5 for each iteration’s interval, ensuring that the composite polynomial $U \in \mathbb{R}^{m \times r}$6 is globally worst-case optimal for any fixed number of iterations $U \in \mathbb{R}^{m \times r}$7 and degree $U \in \mathbb{R}^{m \times r}$8. This strategy overcomes slow initial convergence encountered in fixed-polynomial methods.

Each update takes the form: $U \in \mathbb{R}^{m \times r}$9 with $V \in \mathbb{R}^{n \times r}$0. Efficient computation utilizes Horner’s rule to minimize matrix multiplications.

3. Convergence Guarantees and Worst-case Optimality

The greedy minimax procedure used by Polar Express ensures worst-case optimality at every iteration among all polynomial-composition methods of the same degree and depth. The spectral-norm error after $V \in \mathbb{R}^{n \times r}$1 steps with degree $V \in \mathbb{R}^{n \times r}$2 satisfies the super-exponential bound

$V \in \mathbb{R}^{n \times r}$3

providing quadratic convergence for $V \in \mathbb{R}^{n \times r}$4 and cubic convergence for $V \in \mathbb{R}^{n \times r}$5. Since the polynomial is adapted to the current singular value interval, early-phase progress is substantially accelerated compared to fixed-iteration schemes, especially when the minimum singular value is small.

4. Finite-Precision Robustness

The algorithm is specifically designed for low-precision environments such as bfloat16 or float16 on GPUs. Three practical modifications address instabilities:

• Safety Factor: $V \in \mathbb{R}^{n \times r}$6 is evaluated as $V \in \mathbb{R}^{n \times r}$7 to prevent blowup from finite-precision overshoots.
• Cushioning Small Singular Values: The lower bound for the interval is cushioned, using $V \in \mathbb{R}^{n \times r}$8 during polynomial optimization to guarantee $V \in \mathbb{R}^{n \times r}$9.
• Gradient Normalization: Initialization uses $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$0 to stabilize progress when the gradient norm is close to zero.

These measures ensure stable evolution of singular values and guard against the artifacts common in low-precision arithmetic.

5. Integration within the Muon Framework

The Muon optimizer uses "moment-orthogonalized" updates, maintaining a momentum estimate $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$1 (with $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$2), updating model weights via

$$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$3

where $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$4 replaces the standard momentum direction. Previously, fixed-degree Newton–Schulz iterations were used internally. Polar Express is a direct drop-in replacement, employing $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$5 or $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$6 iterations of the optimally-chosen degree-5 polynomials to rapidly approximate $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$7 from the normalized matrix.

6. Empirical Performance and Comparative Analysis

Empirical results demonstrate that Polar Express achieves superior spectral-norm error reductions compared to previous degree-5 polynomial-based methods (including Newton–Schulz, Chen–Chow’s scaled NS, Jordan, and You et al.), halving the number of iterations needed to reach $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$8 error in synthetic tests with singular value gaps as narrow as $$\Sigma = \operatorname{diag}(\sigma_1, \dots, \sigma_r)$$9 to 1. On GPT-2 (124M parameters, 1B FineWeb tokens), Polar Express outperforms AdamW and previous Muon variants in validation loss at all learning rates tested. For instance, best single-epoch, best-tuned losses observed: $$M = \underbrace{U V^T}_{\operatorname{polar}(M)} \Sigma$$0 The convergence improvements are evident in both iteration and wall-clock time metrics, as all compared methods maintain equivalent floating-point operation counts per iteration.

7. The “Innovation-Augmented” Perspective

Polar Express exemplifies "innovation-augmented polar decomposition" by unifying advances in matrix function theory, efficient GPU-based computation, and robust low-precision stability. Its polynomial update adapts at every step to the evolving singular value landscape, maximizing convergence speed in both early and late phases. Unlike rational methods, no inverses or QR decompositions are required, and unlike previous heuristic methods tailored for machine learning, the algorithm is provably optimal among polynomial-based strategies and converges fully to the true polar factor. It is designed as a drop-in replacement for existing Muon/SVIP routines, providing immediate training performance gains on large-scale models.

These characteristics jointly represent an innovation-augmentation over the state of the art: specialized adaptation of classical matrix-function methods to the empirical and hardware-driven constraints of deep learning, emphasizing GPU compatibility, speed over ultimate accuracy, and low-precision stability (Amsel et al., 22 May 2025).