[OFT] Chebyshev-Optimized Newton-Schulz (CANS): Faster and Better

[Koratahiu]

The Issue

In OFT, we currently orthogonalize the weights using two different methods:

• Default: Truncated Cayley-Neumann
• Exact Solver: Exact matrix math

The exact solver is typically excluded from practical use because it is computationally slow and scales poorly. On the other hand, the Cayley-Neumann method exhibits a relatively high orthogonalization error. While the exact solver achieves an error of around $10^{-6}$, Cayley-Neumann's error ranges between $0.1$ and $0.5$. It is also unstable for matrices with higher norms (as noted in #1492) and converges poorly in those cases. Which makes it scale variant and prone to error.

Standard alternative approximations, such as standard Newton-Schulz, were evaluated but did not resolve these issues.

Our Solution (CANS)

CANS is a variant of the Newton-Schulz (NS) algorithm designed to achieve strict orthogonality.

Note: Standard NS flattens singular values but is not explicitly optimized to reach true orthogonality.

A known limitation of CANS is that it requires tuning a lower bound parameter to converge optimally. However, for the OFT formulation ($I + Q$), we can define this lower bound simply as:
1 / Frobenius norm of I+Q

Using this bound makes CANS highly suitable for OFT.

---

Convergence compared to current methods

The plot above shows the performance on a random matrix. In this test case, CANS (red) achieves lower orthogonalization error than both Cayley-Neumann and the exact solver.

---

Block-size invariant & Number of iterations - matmuls required

It only requires 7 iterations (14 matmuls) to fully converge in FP32.

• Note: It has 2 matmuls per iteration, unlike standard NS, which has 3 matmuls per iteration. This makes 7 CNS iterations faster than 5 standard iterations.

Test plan

• pre-commit run --all-files passes
• Launched the affected UI or script and exercised the change
• Tested with at least one real preset / config when relevant (note which: ____)

AI assistance

• No AI involvement

Sources:

Accelerating Newton-Schulz Iteration for Orthogonalization via Chebyshev-type Polynomials

[Koratahiu]

This is ready for testing and review.
I tested it myself, and it was the first time I have seen OFT reach the higher norms of 20+ (using CANS + spectral scaling). This pushed it to its limits, which might results in an expected +50% to +80% increase in expressiveness for the same block size.

It is also optimal for DOFT #1335, since the orthogonalization error is very small (1e-6 compared to 0.1-0.01 of Cayley-Neumann, assuming FP32).

[dxqb]

ready for merge?

[Koratahiu]

ready for merge?

Yeah, but #1556 seems to be superior. It has a similar ortho error and speed, but it works in a different geometry (Matrix Exponential instead of Cayley). This is arguably better for standard optimizers because in the Matrix Exponential the relationship between the norm of Q and the rotation angle is linear. If you increase the norm of Q by 0.1, it will yield the same rotation increase regardless of Q's current norm, whereas Cayley shrinks the growth of rotations down to zero as the norm of Q grows (up to infinity)

[dxqb]

row 4 in between row 2