科学空间|Scientific Spaces

上一篇文章《矩阵平方根和逆平方根的高效计算》中，笔者从$\newcommand{mcsgn}{\mathop{\text{mcsgn}}}\mcsgn$算子出发，提出了一种很漂亮的矩阵平方根和逆平方根的计算方法。比较神奇的是，该方案经过化简之后，最终公式已经看不到最初$\mcsgn$形式的样子。这不禁引发了更深层的思考：该方案更本质的工作原理是什么？是否有推广到任意$r$次方根的可能性？

沿着这个角度进行分析后，笔者惊喜地发现，我们可以从一个更简单的角度去理解之前的迭代算法，并且在新角度下可以很轻松推广到任意$r$次方根和逆$r$次方根的计算。接下来我们将分享这一过程。

前情回顾

设$\boldsymbol{G}\in\mathbb{R}^{m\times n}$是任意矩阵，$\boldsymbol{P}\in\mathbb{R}^{n\times n}$是任意特征值都在$[0,1]$内的矩阵，上一篇文章给出：
\begin{gather}
\boldsymbol{G}_0 = \boldsymbol{G}, \quad \boldsymbol{P}_0 = \boldsymbol{P} \notag\\[6pt]
\boldsymbol{G}_{t+1} = \boldsymbol{G}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2) \label{eq:r2-rsqrt}\\[6pt]
\boldsymbol{P}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2)^2\boldsymbol{P}_t \label{eq:r3-rsqrt}\\[6pt]
\lim_{t\to\infty} \boldsymbol{G}_t = \boldsymbol{G}\boldsymbol{P}^{-1/2}\notag
\end{gather}

点击阅读全文...

设$\boldsymbol{P}\in\mathbb{R}^{n\times n}$是一个特征值都是非负实数的$n$阶方阵，本文来讨论它的平方根$\boldsymbol{P}^{1/2}$和逆平方根$\boldsymbol{P}^{-1/2}$的计算。

基本概念

矩阵$\boldsymbol{P}$的平方根，指的是满足$\boldsymbol{X}^2=\boldsymbol{P}$的矩阵$\boldsymbol{X}$。我们知道正数都有两个平方根，因此不难想象矩阵平方根一般也不唯一。不过，“算术平方根”是唯一的，一个正数的算术平方根是正的那个平方根，类似地，我们将$\boldsymbol{P}$的特征值全是非负数的那个平方根称为算术平方根。本文要求的矩阵平方根，默认都是指算术平方根。

点击阅读全文...

从文章《线性注意力简史：从模仿、创新到反哺》我们可以发现，DeltaNet及其后的线性Attention模型，基本上都关联到了逆矩阵$(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-)^{-1}$。本文就专门来探讨一下这类具有“对角+低秩”特点的三角矩阵的逆矩阵计算。

基本结果

我们将问题一般地定义如下：

给定矩阵$\boldsymbol{Q},\boldsymbol{K}\in\mathbb{R}^{n\times d}$和对角矩阵$\boldsymbol{\Lambda}\in\mathbb{R}^{n\times n}$，满足$n\gg d$，定义 \begin{equation}\boldsymbol{T} = \boldsymbol{\Lambda} + \boldsymbol{Q}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-\end{equation} 其中$\boldsymbol{M}^-=\boldsymbol{M} - \boldsymbol{I}$，矩阵$\boldsymbol{M}$定义为 \begin{equation}M_{i,j} = \left\{\begin{aligned} &1, &i \geq j \\ &0, &i < j\end{aligned}\right.\end{equation} 现在要求逆矩阵$\boldsymbol{T}^{-1}$，并且证明其复杂度是$\mathcal{O}(n^2)$。

点击阅读全文...

前面我们在《通过msign来计算奇异值裁剪mclip（上）》讨论了奇异值裁剪$\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$的数值计算，核心思路来自 @leloykun 的文章《Numerically Stable Spectral Clipping Via Newton-Schulz Iteration》（现已重新修订和改名），通过寻找基于$\newcommand{msign}{\mathop{\text{msign}}}\msign$的表达式来避免另外寻找Newton-Schulz迭代，在文章中笔者提出了一个计算量更低的嵌套$\msign$方案。

不过前两天，@leloykun 在推特上指出笔者的方案实际计算中存在误差偏大的问题。本文来具体分析一下这个问题，并给出一个更高效、误差更低的新方案。

点击阅读全文...

在《msign的导数》一文中，我们正式引入了两种矩阵符号函数$\newcommand{msign}{\mathop{\text{msign}}}\msign$和$\newcommand{mcsgn}{\mathop{\text{mcsgn}}}\mcsgn$，其中$\msign$是Muon的核心运算，而$\mcsgn$则是用来解Sylvester方程。那么$\mcsgn$除了用来解Sylvester方程外，还能干些什么呢？本文就来整理一下这个问题的答案。

两种符号

设矩阵$\boldsymbol{M}\in\mathbb{R}^{n\times m}$，我们有两种矩阵符号函数
\begin{gather}\msign(\boldsymbol{M}) = (\boldsymbol{M}\boldsymbol{M}^{\top})^{-1/2}\boldsymbol{M}= \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2} \\[6pt]
\mcsgn(\boldsymbol{M}) = (\boldsymbol{M}^2)^{-1/2}\boldsymbol{M}= \boldsymbol{M}(\boldsymbol{M}^2)^{-1/2}
\end{gather}

点击阅读全文...

这篇文章我们来推导$\newcommand{msign}{\mathop{\text{msign}}}\msign$算子的求导公式。如果读者想要像《Test-Time Training Done Right》一样，将TTT和Muon结合起来，那么本文可能会对你有帮助。

两种定义

本文依然假设大家已经对$\msign$有所了解，如果还没有，可以先移步阅读《Muon优化器赏析：从向量到矩阵的本质跨越》和《msign算子的Newton-Schulz迭代（上）》。现设有矩阵$\boldsymbol{M}\in\mathbb{R}^{n\times m}$，那么
\begin{equation}\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V}^{\top} = \text{SVD}(\boldsymbol{M}) \quad\Rightarrow\quad \msign(\boldsymbol{M}) = \boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top}\end{equation}
其中$\boldsymbol{U}\in\mathbb{R}^{n\times n},\boldsymbol{\Sigma}\in\mathbb{R}^{n\times m},\boldsymbol{V}\in\mathbb{R}^{m\times m}$，$r$是$\boldsymbol{M}$的秩。简单来说，$\msign$就是把矩阵的所有非零奇异值都变成1后所得的新矩阵。

点击阅读全文...

前面我们用了两篇文章《msign算子的Newton-Schulz迭代（上）》和《msign算子的Newton-Schulz迭代（下）》讨论了矩阵的$\newcommand{msign}{\mathop{\text{msign}}}\newcommand{sign}{\mathop{\text{sign}}}\newcommand{clip}{\mathop{\text{clip}}}\newcommand{mclip}{\mathop{\text{mclip}}}\msign$算子的数值计算，这篇文章我们来关注“奇异值裁剪（Singular Value Clipping）”运算，它最近在 @_arohan_ 的推特上引起了热议，我们此前在《高阶MuP：更简明但更高明的谱条件缩放》也提到过，接下来我们简称为$\mclip$。

基本概念

对于标量$x$，$\clip$运算定义为
\begin{equation}\clip(x) = \max(\min(x, 1), -1) = \left\{\begin{aligned}1, &\quad x\geq 1 \\
x, &\quad x\in(-1, 1)\\
-1, &\quad x\leq -1
\end{aligned}\right.\end{equation}

点击阅读全文...

在上文《msign算子的Newton-Schulz迭代（上）》中，我们试图为$\mathop{\text{msign}}$算子寻找更好的Newton-Schulz迭代，以期在有限迭代步数内能达到尽可能高的近似程度，这一过程又可以转化为标量函数$\mathop{\text{sign}}(x)$寻找同样形式的多项式迭代。当时，我们的求解思路是用Adam优化器端到端地求一个局部最优解，虽然有效但稍显粗暴。

而在几天前，arXiv新出了一篇论文《The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm》，作者运用了一系列精妙的数学结论，以优雅且硬核的方式给出了更漂亮的答案。本文让我们一起欣赏和学习一番这篇精彩的论文。

问题描述

相关背景和转化过程我们就不再重复了，直接给出我们要求解的问题是
\begin{equation}\mathop{\text{argmin}}_f d(f(x),1)\end{equation}

点击阅读全文...