Chapter 11: Matrix Calculus. In the chapter mind map this icon labels Taylor & Matrix Functions: 2nd-order Taylor Weight Pruning Newton-Schulz Matrix Ortho. The discussion below is excerpted and lightly edited from § Example: Another Newton-Schulz Iteration for Orthogonal Matrix (Polar Decomposition) in Mathematics for AI and Machine Learning. Related material also appears in Chapter 6 ($e^A$ & Neumann).

Beyond matrix inversion, Newton-Schulz iteration has a powerful application in computing the orthogonal factor of the polar decomposition (see the matrix chapter). This is particularly useful in contexts where we need to find the "closest" orthogonal matrix to a given matrix, such as in certain QR decomposition algorithms or in optimization methods like the Muon optimizer.

Given a matrix $G \in \mathbb{R}^{N \times N}$ with polar decomposition $G = Q S$ where $Q$ is orthogonal and $S$ is positive semidefinite, we seek to compute $Q$ without performing SVD. We can formulate this as finding the zero of the function $f : \mathbb{R}^{N \times N} \to \mathbb{R}^{N \times N}$ defined as: where $X \in \mathbb{R}^{N \times N}$ is our current estimate of the orthogonal matrix $Q$, and we want $X^\top X = I$ (orthogonality condition).

However, a more direct approach uses the fact that $Q = G (G^\top G)^{-1/2}$ in the polar decomposition. The Newton-Schulz iteration for computing the orthogonal factor is:

Derivation: Starting from $Q_t$, we want to "correct" it to be more orthogonal. Define the error $E_t = I - Q_t^\top Q_t$. If $Q_t$ is close to orthogonal, then $E_t$ is small. The correction factor $(Q_t^\top Q_t)^{-1/2} = (I - E_t)^{-1/2}$ can be approximated using the Taylor expansion: where $o(\|E_t\|^2)$ denotes the remainder term that satisfies $\lim_{\|E_t\| \to 0} \frac{o(\|E_t\|^2)}{\|E_t\|^2} = 0$, meaning the remainder grows slower than $\|E_t\|^2$ as $E_t \to 0$.

The iteration converges __quadratically__ to an orthogonal matrix when the initial guess $Q_0$ has spectral norm $\|Q_0\|_2 < \sqrt{3}$. In practice, we typically initialize with $Q_0 = G / \|G\|_2$, which normalizes the input matrix $G$ by its spectral norm, ensuring that $\|Q_0\|_2 = 1$ and thus satisfying the convergence condition.

Connection to QR Decomposition: While QR decomposition typically uses Householder reflections or Gram-Schmidt, the Newton-Schulz iteration provides an alternative method to compute the orthogonal factor. Given a matrix $A$, if we apply Newton-Schulz to a normalized version of $A$, we obtain an orthogonal matrix that can be used in certain QR-like factorizations. This approach is particularly valuable when: