.

Author: AntonOresten

src/rules.jl

@@ -50,6 +50,7 @@ function apply!((; opt, eta, mu, lambda, fallback)::Muon, state, x::AbstractArra
         # Nesterov update fed to NS5: U ← β m + (1-β) g
         U = @. μ * state + (1 - μ) * dx
         # orthogonalize
+        @.. U = U / (norm(U) + T(1e-6))
         Ot = newtonschulz5(U)
         # post shape factor √max(1, r/c)
         r, c... = size(x)
@@ -59,46 +60,50 @@ function apply!((; opt, eta, mu, lambda, fallback)::Muon, state, x::AbstractArra
     end
 end
 
-const NS5_COEFFS = (; a = 3.4445, b = -4.7750, c = 2.0315)
-
-function _newtonschulz5!(X::AbstractMatrix{T}, ::Val{true}) where T
-    n = size(X, 1)
+# Applies `X = a*X + b*X*X'*X + c*X*X'*X*X'*X` five times,
+# with two methods based on X*X' and X'*X respectively
+function newtonschulz5!(X::AbstractMatrix{T}) where T
+    n = minimum(size(X))
     A = similar(X, n, n)
     B = similar(X, n, n)
     x = similar(X) # mirror of X
-    (; a, b, c) = NS5_COEFFS
-    for _ in 1:5
-        mul!(A, X, X')
-        B .= A; mul!(B, A, A, c, b)
-        x .= X; mul!(X, B, x, true, a)
+    a, b, c = 3.4445, -4.7750, 2.0315
+    if size(X, 1) <= size(X, 2)
+        for _ in 1:5
+            mul!(A, X, X')
+            B .= A; mul!(B, A, A, c, b)
+            x .= X; mul!(X, B, x, true, a)
+        end
+    else
+        for _ in 1:5
+            mul!(A, X', X)
+            B .= A; mul!(B, A, A, c, b)
+            x .= X; mul!(X, x, B, true, a)
+        end
     end
     return X
 end
 
-function _newtonschulz5!(X::AbstractMatrix{T}, ::Val{false}) where T
-    n = size(X, 2)
-    A = similar(X, n, n)
-    B = similar(X, n, n)
-    x = similar(X) # mirror of X
-    (; a, b, c) = NS5_COEFFS
-    for _ in 1:5
-        mul!(A, X', X)
-        B .= A; mul!(B, A, A, c, b)
-        x .= X; mul!(X, x, B, true, a)
+function newtonschulz5(X::AbstractMatrix{T}) where T
+    a, b, c = T(3.4445), T(-4.7750), T(2.0315)
+    if size(X, 1) <= size(X, 2)
+        for _ in 1:5
+            A = X * X'
+            B = c * A + b * A * A
+            X = a * X + B * X
+        end
+    else
+        for _ in 1:5
+            A = X' * X
+            B = c * A + b * A * A
+            X = a * X + X * B
+        end
     end
     return X
 end
 
-_newtonschulz5!(G, c::Bool) = _newtonschulz5!(G, Val(c)::Union{Val{true},Val{false}})
-
-function newtonschulz5(G::AbstractMatrix{T}) where T
-    X = G / (norm(G) + T(1e-7))
-    return _newtonschulz5!(X, size(G, 1) <= size(G, 2))
-end
-
-# Applies `a*X + b*X*X'*X + c*X*X'*X*X'*X` five times,
-# with two methods based on X*X' and X'*X respectively
-newtonschulz5(G::AbstractArray) = reshape(newtonschulz5(reshape(G, size(G,1), :)), size(G))
+newtonschulz5!(X::AbstractArray) = reshape(newtonschulz5!(reshape(X, size(X,1), :)), size(X))
+newtonschulz5(X::AbstractArray) = reshape(newtonschulz5(reshape(X, size(X,1), :)), size(X))
 
 adjust(r::Muon, η::Real) = adjust(r, eta = η, opt = adjust(r.opt, eta = (r.opt.eta / r.eta) * η))