Muon 优化器：矩阵正交化驱动的梯度更新

在前两篇文章中，我们建立了奇异值分解和谱范数的理论基础。SVD 告诉我们矩阵可以分解为旋转和拉伸的组合，谱范数量化了最大拉伸因子，条件数决定了优化的难易程度。现在，我们将这些概念串联起来，理解一个正在改变大模型训练范式的优化器——Muon。

Muon 的核心思想极其简洁：对梯度矩阵做正交化，然后沿正交化后的方向更新参数。这个看似简单的操作，背后连接着 Polar 分解、Newton-Schulz 迭代和幂方法等经典矩阵计算工具。

从动量法到矩阵动量

标准的 SGD 动量法将梯度视为向量，维护一个指数移动平均：

m_t = \beta m_{t-1} + (1 - \beta) g_t

\theta_t = \theta_{t-1} - \eta \cdot m_t

动量的效果是"惯性"——在梯度方向一致时加速，在梯度方向频繁变化时抑制震荡。但这里有一个隐含的假设：梯度是向量，动量是向量的指数移动平均。

然而在神经网络中，梯度实际上是矩阵（或高阶张量）。每一层的权重矩阵 $W \in \mathbb{R}^{m \times n}$ 对应的梯度 $G \in \mathbb{R}^{m \times n}$ 也是一个矩阵。将 $G$ 展平成向量后做动量，忽略了梯度的矩阵结构——尤其是忽略了不同奇异方向上的信息。

Muon 的出发点是：能否利用梯度的矩阵结构，设计更高效的更新规则？

答案指向了正交化。直觉上，正交矩阵是"最均匀"的线性变换——它没有偏袒任何方向，所有方向上的拉伸因子都等于 1。如果我们能将梯度投影到一个正交矩阵，就能确保更新在所有方向上均匀施加，避免某些方向被过度压缩或过度放大。

msign 算子：投影到最近的正交矩阵

Muon 的核心操作是 msign 算子。对于矩阵 $G \in \mathbb{R}^{m \times n}$ ，设其 SVD 为 $G = U\Sigma V^T$ ，msign 定义为

\text{msign}(G) = U_{[:,:r]} V_{[:,:r]}^T

其中 $r = \text{rank}(G)$ ， $U_{[:,:r]}$ 和 $V_{[:,:r]}$ 分别取 $U$ 和 $V$ 的前 $r$ 列。

msign 的几何意义可以通过 Polar 分解来理解。任何矩阵 $G$ 都可以唯一分解为

G = QP

其中 $Q$ 是部分正交矩阵（ $Q^TQ = I$ ）， $P$ 是半正定矩阵。Polar 分解中， $Q$ 正是 $G$ 的"方向"部分， $P$ 是"幅度"部分。msign 提取的正是这个方向： $\text{msign}(G) = Q$ 。

更精确地说，msign 将 $G$ 投影到最近的正交矩阵——这是在 Frobenius 范数意义下的最优近似：

\text{msign}(G) = \arg\min_{Q: Q^TQ = I} \|G - Q\|_F

其证明利用了正交 Procrustes 问题的经典结论。直观地理解，msign 剥离了梯度中的"幅度"信息（奇异值），只保留"方向"信息（奇异向量）。这种操作的效果是：无论梯度在不同方向上的大小如何悬殊，正交化后的更新在所有方向上都具有相同的影响力。

对于向量 $g \in \mathbb{R}^n$ ，msign 退化为 $\text{sign}(g) = g / \|g\|_2$ ——这正是符号梯度（sign SGD）的操作。因此 msign 可以看作 sign 函数从向量到矩阵的自然推广。

msign 的等价形式

msign 还有一个重要的等价表示：

\text{msign}(M) = M(M^T M)^{-1/2}

这个恒等式可以通过 SVD 直接验证：设 $M = U\Sigma V^T$ ，则 $M^T M = V\Sigma^2 V^T$ ， $(M^T M)^{-1/2} = V\Sigma^{-1} V^T$ ，代入得 $U\Sigma V^T \cdot V\Sigma^{-1} V^T = U_{[:,:r]} V_{[:,:r]}^T = \text{msign}(M)$ 。

这个等价形式将 msign 与极分解联系了起来： $M = QP$ 中 $Q = \text{msign}(M)$ ，而 $P = (M^T M)^{1/2}$ ，因此 $M = \text{msign}(M) \cdot (M^T M)^{1/2}$ 。

更重要的是，这个等价形式是 Newton-Schulz 迭代的另一条推导路线的出发点。考虑标量函数 $f(x) = (1-x)^{-1/2}$ 的 Taylor 展开：

(1-x)^{-1/2} = 1 + \frac{1}{2}x + \frac{3}{8}x^2 + \cdots

将其截断并应用于 $(I - X)^{-1/2}$ 的计算，就得到了 Newton-Schulz 迭代的不同阶近似。苏剑林正是从这条路线推导 Newton-Schulz 迭代的，与我们后面采用的标量 Newton 法路线殊途同归。

不同范数约束下的最速梯度下降

在深入 msign 的计算方法之前，我们先从一个更根本的问题出发：Muon 到底在做什么？答案隐藏在"最速梯度下降"的统一框架中。

梯度下降的本质可以表述为一个约束优化问题：在约束 $\|d\| \leq 1$ 下，寻找使目标函数下降最快的方向

d^* = \arg\min_{\|d\| \leq 1} \langle \nabla f, d \rangle

不同的范数约束给出了不同的"最速方向"：

$L_2$ 范数约束： $\|d\|_2 \leq 1$ 。由 Cauchy-Schwarz 不等式， $\langle g, d \rangle \leq \|g\|_2 \|d\|_2$ ，取等条件为 $d = -g / \|g\|_2$ 。不考虑归一化，最速方向就是 $d = -g$ ——这正是标准 SGD。
$L_\infty$ 范数约束： $\|d\|_\infty \leq 1$ 。每个元素 $|d_i| \leq 1$ ，最优解为 $d_i = -\text{sign}(g_i)$ ，即 $d = -\text{sign}(g)$ ——这正是 Signum 优化器。
谱范数约束： $\|d\|_2 \leq 1$ （矩阵的谱范数，即最大奇异值不超过 1）。由 Frobenius 内积与谱范数的关系， $\langle G, D \rangle_F \leq \|G\|_2 \|D\|_F$ （Hölder 不等式的矩阵版本），但更精确的界是

\langle G, D \rangle_F \leq \|G\|_2 \|D\|_*

其中 $\|D\|_*$ 是核范数（奇异值之和）。在谱范数约束 $\|D\|_2 \leq 1$ 下，取等条件为 $D = \text{msign}(G) = U_{[:,:r]} V_{[:,:r]}^T$ ——这正是 Muon 优化器。

苏剑林的核心论点是：Muon 是"谱范数约束下的最速梯度下降"，与 SGD（2-范数约束）和 Signum（∞-范数约束）构成统一框架。三种优化器的区别不在于动量或自适应学习率的技巧，而在于它们对"最速方向"的定义不同——这取决于选择什么范数来约束更新步长。

这个统一框架的深刻之处在于：它将优化器的设计选择归结为范数的选择，而范数的选择反映了对参数空间几何结构的先验假设。 $L_2$ 范数假设参数空间是各向同性的， $L_\infty$ 范数假设每个参数独立，谱范数假设参数具有矩阵结构且方向比幅度更重要。

把这个判断丢到 3D 损失面上检验：在 Himmelblau 函数（4 个等深极小值的非凸损失面）上，从同一个鞍点附近的起点出发，让 SGD、Adam、Muon 三种步几何各自展开 200 步轨迹。同一个梯度场推出三条完全不同的路径——三种步形各自走进不同盆地。

3D 非凸损失面：SGD / Adam / Muon 轨迹分流到不同极小值

Loading 3D scene...

Himmelblau 函数 f(x,y)=(x²+y−11)²+(x+y²−7)² 有 4 个等深极小值（红点），鞍点（白点）位于它们之间。从同一初始点 (−0.27,−0.92)（鞍点附近）出发：红线 SGD（lr=0.005）沿原始梯度向最近的局部最陡方向下降；绿线 Adam（lr=0.02）按动量+方差自适应每个坐标轴；蓝线 Muon（lr=0.04，本图中用 sign(grad) 作为 1D 教学化代理——对 1D「列向量」而言 msign 即 sign，等价于谱范数=1 约束下幅度归一的步——真正的矩阵级 Muon 是 Newton-Schulz 正交化迭代）。可以看到三条轨迹被同一个梯度场推向不同盆地：SGD/Adam 跟随梯度的「方向+幅度」，而 Muon 的恒定幅度步使它跨越鞍脊到达另一组极小值。拖拽旋转、滚轮缩放；用按钮重启动画、随机换起点、或暂停。

const W = container.clientWidth;
const H = container.clientHeight || 560;

// ---------- Himmelblau on [-4,4]^2 ----------
function loss(x, y) {
    const a = x * x + y - 11;
    const b = x + y * y - 7;
    return a * a + b * b;
}
function grad(x, y) {
    const a = x * x + y - 11;
    const b = x + y * y - 7;
    return [2 * a * (2 * x) + 2 * b, 2 * a + 2 * b * (2 * y)];
}
const Z_SCALE = 0.05; // sqrt(loss)*0.05; world z range roughly 0..0.8 vs xy span 8

function zOf(x, y) {
    return Math.sqrt(loss(x, y)) * Z_SCALE;
}
function clamp(v, lo, hi) { return Math.max(lo, Math.min(hi, v)); }

// ---------- Scene ----------
const scene = new THREE.Scene();
scene.background = new THREE.Color(0x0f172a);

const camera = new THREE.PerspectiveCamera(45, W / H, 0.1, 100);
camera.position.set(6.5, 4.5, 6.5);

const renderer = new THREE.WebGLRenderer({ antialias: true });
renderer.setPixelRatio(Math.min(window.devicePixelRatio || 1, 2));
renderer.setSize(W, H);
renderer.domElement.style.display = 'block';
container.appendChild(renderer.domElement);

const controls = new OrbitControls(camera, renderer.domElement);
controls.enableDamping = true;
controls.dampingFactor = 0.08;
controls.minDistance = 4;
controls.maxDistance = 25;
controls.target.set(0, 0.5, 0);

// ---------- Reference grid + axes ----------
const grid = new THREE.GridHelper(8, 16, 0x334155, 0x1e293b);
grid.position.y = 0;
scene.add(grid);

// Axes (X red along plane-x, Y green = world-up = loss-axis, Z blue along plane-y)
const axCol = [0xef4444, 0x22c55e, 0x3b82f6];
const axDir = [
    new THREE.Vector3(1, 0, 0),
    new THREE.Vector3(0, 1, 0),
    new THREE.Vector3(0, 0, 1)
];
axDir.forEach((d, i) => {
    const a = new THREE.ArrowHelper(d, new THREE.Vector3(-4, 0, -4), 1.6, axCol[i], 0.14, 0.08);
    scene.add(a);
});

// ---------- Loss surface (wireframe plane displaced in y by loss) ----------
// PlaneGeometry is in its own xy; rotate -PI/2 around x so its y becomes world-z.
// Map: planeX -> worldX (param x), planeY (rotated) -> worldZ (param y), height -> worldY.
const SEG = 80;
const planeGeom = new THREE.PlaneGeometry(8, 8, SEG, SEG);
planeGeom.rotateX(-Math.PI / 2);
{
 const pos = planeGeom.attributes.position;
 for (let i = 0; i < pos.count; i++) {
 const px = pos.getX(i);
 const pz = pos.getZ(i);
 pos.setY(i, zOf(px, pz));
 }
 pos.needsUpdate = true;
 planeGeom.computeBoundingBox();
}
const surfaceMat = new THREE.MeshBasicMaterial({
 color: 0x60a5fa,
 wireframe: true,
 transparent: true,
 opacity: 0.32
});
const surface = new THREE.Mesh(planeGeom, surfaceMat);
scene.add(surface);

// ---------- Minima markers (4 known Himmelblau minima) ----------
const MINIMA = [
    [3.0, 2.0], [-2.805118, 3.131312], [-3.779310, -3.283186], [3.584428, -1.848126]
];
const minMat = new THREE.MeshBasicMaterial({ color: 0xfacc15 });
const minGeom = new THREE.SphereGeometry(0.14, 18, 18);
MINIMA.forEach(([x, y]) => {
    const s = new THREE.Mesh(minGeom, minMat);
    s.position.set(x, zOf(x, y) + 0.04, y);
    scene.add(s);
});

// ---------- Saddle markers (4 known Himmelblau saddles, approximate) ----------
const SADDLES = [
    [-0.270845, -0.923039],
    [-3.073026, -0.081353],
    [0.086678, 2.884255],
    [3.385154, 0.073851]
];
const sadMat = new THREE.MeshBasicMaterial({ color: 0xffffff });
const sadGeom = new THREE.SphereGeometry(0.09, 14, 14);
SADDLES.forEach(([x, y]) => {
    const s = new THREE.Mesh(sadGeom, sadMat);
    s.position.set(x, zOf(x, y) + 0.04, y);
    scene.add(s);
});

// ---------- Optimizers ----------
const STEPS = 200;

// ---------- Build polyline + animated marker for a trajectory ----------
const OPTS = [
    { name: 'SGD',  color: 0xef4444, run: runSGD,  lr: 0.005 },
    { name: 'Adam', color: 0x22c55e, run: runAdam, lr: 0.02 },
    { name: 'Muon', color: 0x3b82f6, run: runMuon, lr: 0.04 }
];

const trajectoryGroup = new THREE.Group();
scene.add(trajectoryGroup);

let trajectories = []; // [{name, color, pts, line, ball}]
let initPoint = [-0.27, -0.92]; // saddle-adjacent default; SGD->one basin, Adam/Muon->another

let startMarker = null;
rebuildTrajectories(initPoint[0], initPoint[1]);

// ---------- Animation: balls march along their trajectory, loop every ~3s ----------
const LOOP_MS = 3000;
let animStartTs = performance.now();
let paused = false;
let pauseAccum = 0;
let pauseStartTs = 0;

function tickAnimation(now) {
    if (paused) return;
    const elapsed = ((now - animStartTs - pauseAccum) % LOOP_MS) / LOOP_MS;
    // Map elapsed in [0,1] to step index in [0,STEPS]
    const fStep = elapsed * STEPS;
    const k = Math.floor(fStep);
    const t = fStep - k;
    for (const tr of trajectories) {
        const p0 = tr.pts[Math.min(k, STEPS)];
        const p1 = tr.pts[Math.min(k + 1, STEPS)];
        const x = p0[0] + (p1[0] - p0[0]) * t;
        const y = p0[1] + (p1[1] - p0[1]) * t;
        tr.ball.position.set(x, zOf(x, y) + 0.06, y);
    }
}

// ---------- Controls UI ----------
const ctrlBar = document.createElement('div');
ctrlBar.className = 'three-viz-controls';
container.appendChild(ctrlBar);

function mkBtn(label, color, fn) {
    const b = document.createElement('button');
    b.textContent = label;
    if (color !== undefined) {
        const hex = '#' + color.toString(16).padStart(6, '0');
        b.style.borderColor = hex;
        b.style.color = hex;
    }
    b.addEventListener('click', fn);
    ctrlBar.appendChild(b);
    return b;
}

mkBtn('重启', undefined, () => {
    animStartTs = performance.now();
    pauseAccum = 0;
    if (paused) { paused = false; pauseBtn.textContent = '暂停'; }
});

mkBtn('随机初始化', undefined, () => {
    // pick a random init point in [-3.5, 3.5]^2 (avoid the edge clamp)
    const nx = (Math.random() * 2 - 1) * 3.5;
    const ny = (Math.random() * 2 - 1) * 3.5;
    initPoint = [nx, ny];
    rebuildTrajectories(nx, ny);
    animStartTs = performance.now();
    pauseAccum = 0;
    if (paused) { paused = false; pauseBtn.textContent = '暂停'; }
});

// Legend
const legend = document.createElement('div');
legend.className = 'three-viz-hint';
legend.innerHTML = '■ SGD　■ Adam　■ Muon (proxy)　● 极小值　● 鞍点';
container.appendChild(legend);

// ---------- Render loop ----------
let raf = 0;
function animate() {
    raf = requestAnimationFrame(animate);
    tickAnimation(performance.now());
    controls.update();
    renderer.render(scene, camera);
}
animate();

// ---------- Resize + dispose hooks ----------
window['_threeViz_muon-loss-landscape'] = {
    onResize(w, h) {
        if (!w || !h) return;
        camera.aspect = w / h;
        camera.updateProjectionMatrix();
        renderer.setSize(w, h, false);
    },
    dispose() {
        cancelAnimationFrame(raf);
        controls.dispose();
        renderer.dispose();
        scene.traverse(o => {
            if (o.geometry) o.geometry.dispose();
            if (o.material) o.material.dispose();
        });
    }
};

Muon 的完整更新规则为：

G_t = \nabla_W \mathcal{L}

M_t = \beta M_{t-1} + (1 - \beta) G_t

W_t = W_{t-1} - \eta \cdot \text{msign}(M_t) \cdot \max(\|M_t\|_F / d, \epsilon)

其中 $M_t$ 是矩阵动量，最后一步的缩放因子恢复了幅度信息（正交化丢失的奇异值幅度由动量矩阵的 F 范数近似恢复）， $d$ 是矩阵的维度参数。在实践中，这个缩放因子可以简化为 $\|M_t\|_F / \sqrt{\max(m,n)}$ 。

Newton-Schulz 迭代

直接计算 msign 需要对 $G$ 做完整的 SVD，计算量为 $O(\min(m,n) \cdot mn)$ ，对于大模型中的巨型矩阵来说代价过高。Newton-Schulz 迭代提供了一种避免 SVD 的近似方法。

Newton-Schulz 迭代的出发点是求解矩阵符号函数。对于可对角化矩阵 $A = V\Lambda V^{-1}$ ，其符号函数定义为 $\text{sign}(A) = V \cdot \text{sign}(\Lambda) \cdot V^{-1}$ ，其中 $\text{sign}(\Lambda)$ 将正对角元素映射为 $+1$ ，负对角元素映射为 $-1$ 。

Newton-Schulz 迭代公式为：

Z_{k+1} = \frac{1}{2}Z_k(3I - Z_k^2) = \frac{3}{2}Z_k - \frac{1}{2}Z_k^3

从适当的初始值出发， $Z_k$ 会收敛到 $\text{sign}(A)$ 。

这个迭代公式看起来简洁得有些神秘，但其推导并不复杂。考虑标量情形：求解 $\text{sign}(a)$ 等价于求解 $z^2 = 1$ （即 $z = \pm 1$ ）。对 $f(z) = z^2 - 1$ 应用 Newton 法：

z_{k+1} = z_k - \frac{z_k^2 - 1}{2z_k} = \frac{z_k^2 + 1}{2z_k} = \frac{1}{2}\left(z_k + \frac{1}{z_k}\right)

这就是经典的 Newton 迭代求平方根倒数。但这个迭代涉及 $1/z_k$ ，推广到矩阵时需要求逆。为了避免矩阵求逆，我们可以使用一种等价的"逆-free"形式。

注意到 Newton 法的迭代可以改写为 $z_{k+1} = \frac{3}{2}z_k - \frac{1}{2}z_k^3$ 。这个形式不涉及除法（矩阵求逆），可以直接推广到矩阵。具体来说，从 $Z_0 = A / \|A\|_F$ 出发（需要 $\|A\|_2 < 1$ 保证收敛）， $Z_k$ 会收敛到 $\text{msign}(A)$ 。

收敛性的关键在于：每次迭代， $Z_k$ 的奇异值都被"推向" $+1$ 或 $-1$ 。设 $Z_k$ 的奇异值为 $\sigma_i^{(k)}$ ，则迭代对奇异值的作用为

\sigma_i^{(k+1)} = \frac{3}{2}\sigma_i^{(k)} - \frac{1}{2}(\sigma_i^{(k)})^3 = f(\sigma_i^{(k)})

函数 $f(\sigma) = \frac{3}{2}\sigma - \frac{1}{2}\sigma^3$ 在 $\sigma \in (0, 1)$ 上有 $f(\sigma) > \sigma$ 且 $f(\sigma) \to 1$ ，即奇异值逐步趋近于 1。这正是正交化所需要的：所有奇异值收敛到 1 意味着矩阵收敛到正交矩阵。

Newton-Schulz 迭代收敛过程

Loading visualization...

随机初始矩阵的奇异值逐步收敛到 1。左：当前矩阵的奇异值分布。右：奇异值到 1 的偏差随迭代步数的变化。

const W = container.clientWidth, H = container.clientHeight;
const margin = {top: 30, right: 25, bottom: 55, left: 50};
const pW = (W - margin.left - margin.right - 50) / 2;
const pH = H - margin.top - margin.bottom;

const svg = d3.select(container).append("svg").attr("width", W).attr("height", H);
const gMain = svg.append("g").attr("transform", `translate(${margin.left},${margin.top})`);

// 生成随机 4x4 矩阵并归一化
const n = 4;
let A0 = [];
for (let i = 0; i < n; i++) {
 A0[i] = [];
 for (let j = 0; j < n; j++) {
 A0[i][j] = (Math.random() - 0.5) * 2;
 }
}
// 简单归一化: A0 = A / ||A||_F
let fnorm = 0;
for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) fnorm += A0[i][j] * A0[i][j];
fnorm = Math.sqrt(fnorm);
for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) A0[i][j] /= fnorm;

// 预计算 Newton-Schulz 迭代的所有步骤
function matMul(A, B) {
 const C = [];
 for (let i = 0; i < n; i++) {
 C[i] = [];
 for (let j = 0; j < n; j++) {
 C[i][j] = 0;
 for (let k = 0; k < n; k++) C[i][j] += A[i][k] * B[k][j];
 }
 }
 return C;
}

function matSub(A, B) {
  return A.map((row, i) => row.map((v, j) => v - B[i][j]));
}

function matScale(A, s) {
  return A.map(row => row.map(v => v * s));
}

// 用简单幂迭代计算奇异值 (近似)
function singularValues(A, nIter) {
 // A^T A 的特征值 = sigma^2
 const ATA = matMul(transpose(A), A);
 // 幂迭代找特征值
 const evals = [];
 let M = ATA.map(r => [...r]);
 for (let k = 0; k < n; k++) {
 let v = Array(n).fill(0).map(() => Math.random());
 let vnorm = Math.sqrt(v.reduce((s, x) => s + x*x, 0));
 v = v.map(x => x / vnorm);
 for (let it = 0; it < 200; it++) {
 const newV = Array(n).fill(0);
 for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) newV[i] += M[i][j] * v[j];
 vnorm = Math.sqrt(newV.reduce((s, x) => s + x*x, 0));
 v = newV.map(x => x / vnorm);
 }
 const lambda = v.reduce((s, x, i) => s + x * M[i].reduce((ss, xx, j) => ss + xx * v[j], 0), 0);
 evals.push(Math.sqrt(Math.max(0, lambda)));
 // Deflation
 for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) M[i][j] -= lambda * v[i] * v[j];
 }
 return evals.sort((a, b) => b - a);
}

function transpose(A) {
  return A[0].map((_, j) => A.map(row => row[j]));
}

// Newton-Schulz: Z_{k+1} = 1.5 * Z_k - 0.5 * Z_k^3
const maxIter = 20;
const trajectory = [];
let Z = A0.map(r => [...r]);
trajectory.push(singularValues(Z, 50));
for (let iter = 0; iter < maxIter; iter++) {
 const Z2 = matMul(Z, Z);
 const Z3 = matMul(Z2, Z);
 Z = matSub(matScale(Z, 1.5), matScale(Z3, 0.5));
 trajectory.push(singularValues(Z, 50));
}

let currentIter = 0;
let playing = false;
let playTimer = null;

function draw() {
  gMain.selectAll("*").remove();

const svals = trajectory[currentIter];

// === 左图：奇异值柱状图 ===
  const gLeft = gMain.append("g");
  const xBar = d3.scaleBand().domain(d3.range(n).map(i => "σ" + (i+1))).range([0, pW]).padding(0.25);
  const yBar = d3.scaleLinear().domain([0, 1.5]).range([pH, 0]);

gLeft.append("g").call(d3.axisLeft(yBar).ticks(5)).attr("font-size", "10px");
  gLeft.append("g").attr("transform", `translate(0,${pH})`)
    .call(d3.axisBottom(xBar)).attr("font-size", "11px");

// 目标线 σ=1
  gLeft.append("line")
    .attr("x1", 0).attr("x2", pW).attr("y1", yBar(1)).attr("y2", yBar(1))
    .attr("stroke", "#e74c3c").attr("stroke-dasharray", "4,3").attr("stroke-width", 1);

gLeft.selectAll(".bar").data(svals).join("rect")
    .attr("x", (d, i) => xBar("σ" + (i+1)))
    .attr("width", xBar.bandwidth())
    .attr("y", d => yBar(d))
    .attr("height", d => pH - yBar(d))
    .attr("fill", (d, i) => d3.interpolateWarm(0.2 + 0.6 * i / n))
    .attr("opacity", 0.85);

gLeft.append("text").attr("x", pW/2).attr("y", -10)
    .text(`迭代第 ${currentIter} 步`).style("text-anchor", "middle").style("font-size", "13px").style("font-weight", "bold");

// 偏差信息
  const maxDev = Math.max(...svals.map(s => Math.abs(s - 1)));
  gLeft.append("text").attr("x", pW/2).attr("y", pH + 35)
    .text(`最大偏差: ${maxDev.toFixed(4)}`).style("text-anchor", "middle").style("font-size", "11px").style("fill", "#666");

// === 右图：偏差收敛曲线 ===
  const gRight = gMain.append("g").attr("transform", `translate(${pW + 50}, 0)`);
  const xLine = d3.scaleLinear().domain([0, maxIter]).range([0, pW]);
  const yLine = d3.scaleLinear().domain([0, 1.2]).range([pH, 0]);

gRight.append("g").call(d3.axisLeft(yLine).ticks(5)).attr("font-size", "10px");
  gRight.append("g").attr("transform", `translate(0,${pH})`)
    .call(d3.axisBottom(xLine).ticks(5)).attr("font-size", "10px");

gRight.append("text").attr("x", -pH/2).attr("y", -38).text("|σ - 1|")
    .style("text-anchor", "middle").style("font-size", "11px").attr("transform", "rotate(-90)");
  gRight.append("text").attr("x", pW/2).attr("y", pH + 30).text("迭代步数")
    .style("text-anchor", "middle").style("font-size", "11px");

// 画每条奇异值的收敛曲线
 for (let si = 0; si < n; si++) {
 const color = d3.interpolateWarm(0.2 + 0.6 * si / n);
 const pts = trajectory.map((svals, step) => ({
 x: xLine(step), y: yLine(Math.min(Math.abs(svals[si] - 1), 1.2))
 }));
 gRight.append("path").datum(pts)
 .attr("d", d3.line().x(d => d.x).y(d => d.y))
 .attr("fill", "none").attr("stroke", color).attr("stroke-width", 1.5);
 }

// 当前步标记
  gRight.append("line")
    .attr("x1", xLine(currentIter)).attr("x2", xLine(currentIter))
    .attr("y1", 0).attr("y2", pH)
    .attr("stroke", "#333").attr("stroke-dasharray", "3,3").attr("stroke-width", 1);

gRight.append("text").attr("x", pW/2).attr("y", -10)
    .text("奇异值偏差收敛").style("text-anchor", "middle").style("font-size", "13px").style("font-weight", "bold");
}

// 控件
const ctrl = svg.append("g").attr("transform", `translate(${margin.left}, ${H - 25})`);

ctrl.append("text").attr("x", 0).attr("y", 4).text("迭代步:").style("font-size", "13px").style("fill", "#555");

ctrl.append("foreignObject").attr("width", 200).attr("height", 22).attr("x", 55).attr("y", -10)
  .append("xhtml:input")
  .attr("type", "range").attr("min", 0).attr("max", maxIter).attr("step", 1).attr("value", 0)
  .style("width", "200px")
  .on("input", function() { currentIter = +this.value; draw(); });

const stepBtn = ctrl.append("g").style("cursor", "pointer").attr("transform", "translate(355, -8)");
stepBtn.append("rect").attr("width", 60).attr("height", 22).attr("rx", 4).attr("fill", "#50a050");
stepBtn.append("text").attr("x", 30).attr("y", 15).text("→ 步进")
 .style("font-size", "12px").style("fill", "#fff").style("text-anchor", "middle");
stepBtn.on("click", () => {
 if (currentIter < maxIter) { currentIter++; draw(); }
});

draw();

从可视化中可以观察到：Newton-Schulz 迭代的收敛速度很快，通常 5-10 步就足以将奇异值偏差降低到 $10^{-3}$ 以下。收敛速度是三次的（cubic convergence），远快于线性收敛的幂方法。

但 Newton-Schulz 迭代有一个前提：初始矩阵的谱范数必须严格小于 1，否则迭代可能发散。这就是为什么我们需要先将 $G$ 除以 $\|G\|_F$ （或 $\|G\|_2$ ）进行归一化。

优化系数：从理论到实践

上面使用的迭代公式 $Z_{k+1} = \frac{3}{2}Z_k - \frac{1}{2}Z_k^3$ 的系数 $(3/2, -1/2)$ 来源于 Newton 法的理论推导，对任意矩阵都能保证收敛。但在实际应用中，我们可以使用更高阶的迭代公式来加速收敛。

Muon 官方实现使用的是五阶迭代公式，系数为 $(3.4445, -4.7750, 2.0315)$ ：

Z_{k+1} = a Z_k + b Z_k^3 + c Z_k^5

这些系数是通过优化得到的，并非随意选取。其求解思路是：

假设输入的奇异值分布：大模型中梯度矩阵的奇异值近似服从 Marchenko-Pastur 分布（随机矩阵理论中描述 Wishart 矩阵奇异值分布的经典结果）。
固定迭代步数 $T$ ：例如 $T = 3$ 步或 $T = 5$ 步。
以奇异值偏差的平方误差为损失： $\mathcal{L}(a, b, c) = \mathbb{E}_{\sigma \sim \text{MP}} [(\sigma^{(T)} - 1)^2]$ ，其中 $\sigma^{(T)}$ 是经过 $T$ 步迭代后奇异值 $\sigma$ 的终值。
求解最优系数：对 $\mathcal{L}$ 做梯度下降，得到使期望偏差最小的 $(a, b, c)$ 。

苏剑林在《Muon优化器赏析》中详细分析了这些系数的来源和求解方法。核心洞察是：理论系数 $(3/2, -1/2)$ 对所有输入分布都安全，但对特定分布（如 Marchenko-Pastur）不是最优的；优化系数利用了输入分布的先验知识，在相同迭代步数下可以达到更高的精度。

实践中，使用优化系数可以显著提升 msign 的精度——在相同的 3 步迭代下，优化系数的近似误差比理论系数低一个数量级以上。这是"利用问题结构加速收敛"的又一个实例。

流式幂迭代：工程折衷

Newton-Schulz 迭代虽然收敛快，但每步需要计算 $Z^3 = Z \cdot Z \cdot Z$ ，两次矩阵乘法的计算量在大模型场景下仍然不低。更关键的是，每次参数更新都需要跑完整的迭代序列，这在训练循环中引入了显著的开销。

流式幂迭代（Streaming Power Iteration）提供了一种巧妙的工程折衷：将多步迭代平摊到训练的每一步。

核心思想是：在训练的每一步，只做一次幂迭代更新，而不是跑完整个 Newton-Schulz 序列。具体来说，维护两个向量 $u$ 和 $v$ ，每步做：

v \leftarrow \frac{G^T u}{\|G^T u\|}, \quad u \leftarrow \frac{G v}{\|G v\|}

经过 $T$ 步训练后， $u$ 和 $v$ 会收敛到 $G$ 的最大左右奇异向量。msign 的近似为：

\text{msign}(G) \approx u v^T

这恰好是 rank-1 近似。对于完整的 msign，需要对每个奇异方向分别维护一对向量，但实践中 rank-1 近似已经足够——苏剑林的实验表明，更高秩的近似对训练效果的影响有限。

流式幂迭代的优势在于：

零额外计算：每步只需两次矩阵-向量乘法，与标准反向传播相比几乎无额外开销。
自然适配动量： $u$ 和 $v$ 的更新与动量机制天然兼容，梯度动量 $M_t$ 的变化被"流式"地追踪。
隐式正则化：单步幂迭代等价于对 msign 做一步近似，这种不完全正交化反而可能带来正则化效果。

但也需要注意其局限：流式幂迭代的追踪能力取决于 $M_t$ 的变化速度。如果动量矩阵变化过快（ $\beta$ 过小）， $u$ 和 $v$ 可能来不及收敛就被新的梯度"拉走"。

MuP 与谱条件：Muon 的理论优势

理解 Muon 为什么有效，需要回到参数化尺度（Parameterization Scale）的框架。MuP（Maximal Update Parameterization）是 Yang et al. 提出的一套超参数迁移理论，其核心思想是：网络的超参数（学习率、初始化方差等）应该按照特定方式随宽度 $d$ 缩放，使得训练动力学在宽度变化时保持稳定。

在 MuP 框架下，隐藏层权重的更新应该满足 $\Delta W = \Theta(1)$ （不随 $d$ 变化），这要求学习率 $\eta = \Theta(1/d)$ 。但标准的 Adam 优化器对梯度做了逐元素归一化，其更新幅度为 $\Theta(1/\sqrt{d})$ ——比 MuP 要求的小 $\Theta(1/\sqrt{d})$ 倍。这意味着 Adam 在宽网络中的更新偏小，可能需要更高的学习率来补偿。

Muon 的 msign 操作天然满足谱条件。由于 msign 将梯度投影到正交矩阵，其每个元素的量级为 $\Theta(1/\sqrt{d})$ ，而谱范数为 1。考虑缩放因子 $\|M_t\|_F / \sqrt{d}$ ：由于 $\|M_t\|_F \approx \sqrt{d} \cdot \sigma_1(M_t)$ （当矩阵的谱快速衰减时），缩放后的更新幅度约为 $\sigma_1(M_t)$ ，恰好是 $\Theta(1)$ ——满足 MuP 的稳定性要求。

更深层的原因在于正交化消除了梯度中的"宽度偏差"。在宽网络中，随机初始化的权重矩阵接近正交矩阵，梯度矩阵也倾向于各向同性。标准优化器的逐元素归一化破坏了这种各向同性，而 msign 恰好保持了它。

应用：Kimi K2 训练中的 Muon 实践

Kimi K2 是 Moonshot AI 发布的大语言模型，其训练中采用了 Muon 优化器。根据公开的技术报告，Muon 在 K2 的训练中实现了约 2 倍的加速——达到相同的验证损失所需的训练步数减半。

这一加速主要来自两个方面：

更快的收敛：在条件数较大的问题上（如注意力层的权重矩阵），Muon 的正交化更新避免了之字形震荡，收敛速度显著快于 Adam。
更好的超参数迁移：Muon 满足 MuP 条件，从窄模型调好的超参数可以直接迁移到宽模型，省去了大量的调参成本。

Muon vs Adam 收敛对比

Loading visualization...

在不同条件数的二次型问题上，Muon（正交化梯度）与 Adam 的收敛曲线对比。高条件数时 Muon 优势明显。

const W = container.clientWidth, H = container.clientHeight;
const margin = {top: 30, right: 20, bottom: 55, left: 55};
const pW = W - margin.left - margin.right;
const pH = H - margin.top - margin.bottom;

const svg = d3.select(container).append("svg").attr("width", W).attr("height", H);
const g = svg.append("g").attr("transform", `translate(${margin.left},${margin.top})`);

let condNum = 10;

// 二次型 f(x) = 0.5 * x^T A x, A = diag(kappa, 1, ..., 1)
// 2D 简化：f(x,y) = 0.5 * (kappa*x^2 + y^2)
// GD: x <- x - eta * kappa * x, y <- y - eta * y
// Adam: per-element adaptive learning rate
// Muon: msign gradient (in 2D, this is a rotation matrix)

function simulateAdam(kappa, nSteps) {
 const path = [0.5 * (kappa * 9 + 6.25)]; // initial loss
 let x = 3, y = 2.5;
 let mx = 0, my = 0, vx = 0, vy = 0;
 const beta1 = 0.9, beta2 = 0.999, eps = 1e-8;
 const lr = 0.05;
 for (let t = 1; t <= nSteps; t++) {
 const gx = kappa * x, gy = y;
 mx = beta1 * mx + (1 - beta1) * gx;
 my = beta1 * my + (1 - beta1) * gy;
 vx = beta2 * vx + (1 - beta2) * gx * gx;
 vy = beta2 * vy + (1 - beta2) * gy * gy;
 const mxh = mx / (1 - Math.pow(beta1, t));
 const myh = my / (1 - Math.pow(beta1, t));
 const vxh = vx / (1 - Math.pow(beta2, t));
 const vyh = vy / (1 - Math.pow(beta2, t));
 x -= lr * mxh / (Math.sqrt(vxh) + eps);
 y -= lr * myh / (Math.sqrt(vyh) + eps);
 path.push(0.5 * (kappa * x * x + y * y));
 }
 return path;
}

function simulateMuon(kappa, nSteps) {
 const path = [0.5 * (kappa * 9 + 6.25)];
 let x = 3, y = 2.5;
 let mx = 0, my = 0;
 const beta = 0.9;
 const lr = 0.05;
 for (let t = 1; t <= nSteps; t++) {
 const gx = kappa * x, gy = y;
 mx = beta * mx + (1 - beta) * gx;
 my = beta * my + (1 - beta) * gy;
 // msign on 2D vector = normalize
 const mnorm = Math.sqrt(mx * mx + my * my);
 const sx = mnorm > 1e-8 ? mx / mnorm : 0;
 const sy = mnorm > 1e-8 ? my / mnorm : 0;
 // 缩放因子
 const scale = mnorm;
 x -= lr * sx * scale;
 y -= lr * sy * scale;
 path.push(0.5 * (kappa * x * x + y * y));
 }
 return path;
}

const nSteps = 200;

function draw() {
  g.selectAll("*").remove();

const adamPath = simulateAdam(condNum, nSteps);
  const muonPath = simulateMuon(condNum, nSteps);

// 找 y 轴范围
  const allVals = [...adamPath, ...muonPath].filter(v => v > 0.001);
  const yMax = allVals[0];
  const yMin = Math.max(0.001, Math.min(...allVals.filter(v => v > 0)));

const xSc = d3.scaleLinear().domain([0, nSteps]).range([0, pW]);
  const ySc = d3.scaleLog().domain([Math.max(0.01, yMin), yMax * 1.5]).range([pH, 0]);

// 坐标轴
  g.append("g").call(d3.axisLeft(ySc).ticks(5, ".1f")).attr("font-size", "10px");
  g.append("g").attr("transform", `translate(0,${pH})`)
    .call(d3.axisBottom(xSc).ticks(6)).attr("font-size", "10px");

g.append("text").attr("x", -pH/2).attr("y", -42).text("Loss (log)")
    .style("text-anchor", "middle").style("font-size", "11px").attr("transform", "rotate(-90)");
  g.append("text").attr("x", pW/2).attr("y", pH + 35).text("训练步数")
    .style("text-anchor", "middle").style("font-size", "11px");

// Adam 曲线
  const adamPts = adamPath.filter(v => v > 0.01).map((v, i) => ({x: xSc(i), y: ySc(v)}));
  g.append("path").datum(adamPts)
    .attr("d", d3.line().x(d => d.x).y(d => d.y).defined(d => isFinite(d.y)))
    .attr("fill", "none").attr("stroke", "#e07b39").attr("stroke-width", 2);

// Muon 曲线
  const muonPts = muonPath.filter(v => v > 0.01).map((v, i) => ({x: xSc(i), y: ySc(v)}));
  g.append("path").datum(muonPts)
    .attr("d", d3.line().x(d => d.x).y(d => d.y).defined(d => isFinite(d.y)))
    .attr("fill", "none").attr("stroke", "#4a90d9").attr("stroke-width", 2);

// 图例
  const legend = g.append("g").attr("transform", `translate(${pW - 140}, 10)`);
  legend.append("line").attr("x1", 0).attr("x2", 25).attr("y1", 0).attr("y2", 0)
    .attr("stroke", "#4a90d9").attr("stroke-width", 2);
  legend.append("text").attr("x", 30).attr("y", 4).text("Muon")
    .style("font-size", "12px").style("fill", "#4a90d9");
  legend.append("line").attr("x1", 0).attr("x2", 25).attr("y1", 20).attr("y2", 20)
    .attr("stroke", "#e07b39").attr("stroke-width", 2);
  legend.append("text").attr("x", 30).attr("y", 24).text("Adam")
    .style("font-size", "12px").style("fill", "#e07b39");

// 条件数信息
  g.append("text").attr("x", pW/2).attr("y", -10)
    .text(`条件数 κ = ${condNum.toFixed(0)}`)
    .style("text-anchor", "middle").style("font-size", "13px").style("font-weight", "bold");
}

// 控件
const ctrl = svg.append("g").attr("transform", `translate(${margin.left}, ${H - 25})`);

ctrl.append("text").attr("x", 0).attr("y", 4).text("条件数 κ:").style("font-size", "13px").style("fill", "#555");

const condSlider = ctrl.append("foreignObject").attr("width", 200).attr("height", 22).attr("x", 75).attr("y", -10)
  .append("xhtml:input")
  .attr("type", "range").attr("min", 2).attr("max", 100).attr("step", 1).attr("value", condNum)
  .style("width", "200px")
  .on("input", function() { condNum = +this.value; draw(); });

const condLabel = ctrl.append("text").attr("x", 285).attr("y", 4)
  .style("font-size", "13px").style("font-weight", "bold");

const origDraw = draw;
draw = function() {
  condLabel.text(condNum.toFixed(0));
  origDraw();
};

draw();

从对比中可以看到，当条件数较低时（ $\kappa \approx 2$ ），Adam 和 Muon 的收敛速度相近；但当条件数增大时（ $\kappa > 20$ ），Adam 的收敛明显变慢，而 Muon 仍保持较快收敛。这是因为 Adam 的逐元素自适应学习率相当于对角预处理，只能处理坐标对齐的各向异性；而 Muon 的正交化操作可以处理任意方向的各向异性——这在高维问题中尤为重要，因为权重矩阵的主奇异方向通常不与坐标轴对齐。

Shampoo 等价性：两个独立优化器家族的会合

Muon 与 Adam 的对比是"正交化 vs 对角归一化"的故事，但还有一个更令人意外的联系：Muon 与 Shampoo 优化器在无正则化极限下完全等价。

Shampoo 的更新规则为

\theta_{t+1} = \theta_t - \eta (GG^T + \epsilon I)^{-1/4} G (G^T G + \epsilon I)^{-1/4}

其中 $G$ 是梯度矩阵， $\epsilon$ 是正则化常数。当 $\epsilon = 0$ 时，我们可以通过 SVD 证明 Shampoo 与 Muon 等价。设 $G = U\Sigma V^T$ ，则

(GG^T)^{-1/4} = (U\Sigma^2 U^T)^{-1/4} = U\Sigma^{-1/2} U^T

(G^T G)^{-1/4} = (V\Sigma^2 V^T)^{-1/4} = V\Sigma^{-1/2} V^T

代入得

(GG^T)^{-1/4} G (G^T G)^{-1/4} = U\Sigma^{-1/2} U^T \cdot U\Sigma V^T \cdot V\Sigma^{-1/2} V^T = U_{[:,:r]} V_{[:,:r]}^T = \text{msign}(G)

这意味着 Shampoo 在 $\epsilon = 0$ 时恰好退化为 Muon。这个等价性连接了两个独立发展的优化器家族：Muon 从矩阵正交化出发，Shampoo 从 preconditioned gradient 出发，两者在极限处殊途同归。

当 $\epsilon > 0$ 时，Shampoo 的行为发生变化：正则化项 $\epsilon I$ 使得矩阵分数幂的计算不再简单地约去奇异值， $(GG^T + \epsilon I)^{-1/4}$ 的对角元素开始占据主导。此时 Shampoo 逐渐退化为 Adam 类行为——对角近似。因此 $\epsilon$ 可以看作从 Muon（ $\epsilon = 0$ ，完整谱结构）到 Adam（ $\epsilon \to \infty$ ，纯对角结构）的连续插值参数。

小结

Muon 优化器的思想脉络可以概括为：梯度是矩阵而非向量，因此应该利用矩阵的结构来设计更新规则。msign 算子将梯度投影到最近的正交矩阵，等价于剥离奇异值（幅度）只保留奇异向量（方向）。在最速梯度下降的统一框架下，Muon 是谱范数约束的最优选择，与 SGD（2-范数约束）和 Signum（∞-范数约束）构成完整的谱系。Newton-Schulz 迭代提供了计算 msign 的高效方法，优化系数利用输入分布的先验知识进一步加速收敛；流式幂迭代则将计算平摊到训练循环中。在 MuP 框架下，Muon 的正交化更新天然满足稳定性条件，这使得超参数可以从窄模型直接迁移到宽模型。

从矩阵论的视角看，Muon 的本质是：用谱范数替代 F 范数来度量梯度的"大小"。传统优化器将梯度展平为向量来处理，丢失了矩阵的结构信息；Muon 通过正交化保留了梯度的几何结构，让更新方向更加"均匀"。Shampoo 在 $\epsilon = 0$ 时与 Muon 的等价性进一步印证了这一点：完整的谱结构利用（而非对角近似）是高效矩阵优化的关键。这种从向量到矩阵的思维转换，或许是深度学习优化器设计的一个重要方向。

参考文献

苏剑林. Muon优化器赏析：从向量到矩阵的本质跨越. https://kexue.fm/archives/10592
苏剑林. 从谱范数梯度到新式权重衰减的思考. https://kexue.fm/archives/10648
苏剑林. MuP之上：1. 好模型的三个特征. https://kexue.fm/archives/11340
苏剑林. 基于流式幂迭代的Muon实现：1. 初识. https://kexue.fm/archives/11654
Jordan, K., et al. (2024). Muon: An optimizer for hidden layers in neural networks. arXiv preprint.
Gupta, V., Korenda, M., Begue, A., & Barest, T. (2018). Shampoo: Preconditioned Stochastic Tensor Optimization. ICML 2018.
Yang, G., et al. (2022). Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. NeurIPS 2022.
Higham, N. J. (2008). Functions of Matrices: Theory and Computation. SIAM.

从动量法到矩阵动量#

msign 算子：投影到最近的正交矩阵#

msign 的等价形式#

不同范数约束下的最速梯度下降#

Newton-Schulz 迭代#

优化系数：从理论到实践#

流式幂迭代：工程折衷#

MuP 与谱条件：Muon 的理论优势#

应用：Kimi K2 训练中的 Muon 实践#

Shampoo 等价性：两个独立优化器家族的会合#

小结#

相关概念#

参考文献#

相关文章