1import numpy as np
 2
 3
 4def polar_newton_schulz(a, iterations=5):
 5    """Muon-style semi-orthogonalization via normalized Newton-Schulz."""
 6    x = np.asarray(a, dtype=np.float64).copy()
 7    norm = np.linalg.norm(x, 2)
 8    if norm == 0:
 9        return np.zeros_like(x)
10    x /= norm
11    eye = np.eye(x.shape[1])
12    for _ in range(iterations):
13        x = 0.5 * x @ (3.0 * eye - x.T @ x)
14    return x
15
16
17class BiMaxwellMuon:
18    """Two-relaxation-mode Muon state for one 2-D weight matrix."""
19    def __init__(self, shape, lr=0.025, beta_fast=0.9, beta_slow=0.99,
20                 weight_fast=0.5, ns_iterations=5):
21        if not (0 <= weight_fast <= 1 and 0 <= beta_fast < beta_slow < 1):
22            raise ValueError("require 0 <= beta_fast < beta_slow < 1 and valid weight")
23        self.lr = lr
24        self.beta_fast, self.beta_slow = beta_fast, beta_slow
25        self.weight_fast = weight_fast
26        self.ns_iterations = ns_iterations
27        self.m_fast = np.zeros(shape, dtype=np.float64)
28        self.m_slow = np.zeros(shape, dtype=np.float64)
29
30    def step(self, weight, gradient):
31        self.m_fast = self.beta_fast * self.m_fast + (1-self.beta_fast) * gradient
32        self.m_slow = self.beta_slow * self.m_slow + (1-self.beta_slow) * gradient
33        mixed = self.weight_fast*self.m_fast + (1-self.weight_fast)*self.m_slow
34        update = polar_newton_schulz(mixed, self.ns_iterations)
35        weight -= self.lr * update
36        return weight, update