科学空间|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}$的特征值全是非负数的那个平方根称为算术平方根。本文要求的矩阵平方根，默认都是指算术平方根。

点击阅读全文...

四个月前，我们发布了Moonlight，在16B的MoE模型上验证了Muon优化器的有效性。在Moonlight中，我们确认了给Muon添加Weight Decay的必要性，同时提出了通过Update RMS对齐来迁移Adam超参的技巧，这使得Muon可以快速应用于LLM的训练。然而，当我们尝试将Muon进一步拓展到千亿参数以上的模型时，遇到了新的“拦路虎”——MaxLogit爆炸。

为了解决这个问题，我们提出了一种简单但极其有效的新方法，我们称之为“QK-Clip”。该方法从一个非常本质的角度去看待和解决MaxLogit现象，并且无损模型效果，这成为我们最新发布的万亿参数模型“Kimi K2”的关键训练技术之一。

问题描述

我们先来简单介绍一下MaxLogit爆炸现象。回顾Attention的定义
\begin{equation}\boldsymbol{O} = softmax(\boldsymbol{Q}\boldsymbol{K}^{\top})\boldsymbol{V}\end{equation}

点击阅读全文...

在文章《Transformer升级之路：20、MLA好在哪里?（上）》中，我们对MLA相比常见MHA、GQA、MQA的一些变化分别做了消融实验，其中的变化包括“增大head_dims”、“Partial RoPE”和“KV共享”，实验的初步结果是这三个变化很可能都是MLA效果优异的原因。

本文我们将从一个更加偏理论的角度出发，来理解MLA的成功之处。

部分旋转

首先，我们把最终的断言放在前面：

在相同训练成本和推理成本下，MLA可能是效果最好的Full Attention变体。

点击阅读全文...

从文章《线性注意力简史：从模仿、创新到反哺》我们可以发现，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}

点击阅读全文...

在中文圈，本站应该算是比较早关注线性Attention的了，在2020年写首篇相关博客《线性Attention的探索：Attention必须有个Softmax吗？》时，大家主要讨论的还是BERT相关的Softmax Attention。事后来看，在BERT时代考虑线性Attention并不是太明智，因为当时训练长度比较短，且模型主要还是Encoder，用线性Attention来做基本没有优势。对此，笔者也曾撰文《线性Transformer应该不是你要等的那个模型》表达这一观点。

直到ChatGPT的出世，倒逼大家都去做Decoder-only的生成式模型，这跟线性Attention的RNN形式高度契合。同时，追求更长的训练长度也使得Softmax Attention的二次复杂度瓶颈愈发明显。在这样的新背景下，线性Attention越来越体现出竞争力，甚至出现了“反哺”Softmax Attention的迹象。

点击阅读全文...