宇宙驿站感谢国家天文台LAMOST项目之“宇宙驿站”提供网络空间和数据库资源! 感谢国家天文台崔辰州博士等人的多方努力和技术支持!

版权科学空间致力于知识分享,所以欢迎您转载本站文章,但转载本站内容必须遵循 署名-非商业用途-保持一致 的创作共用协议。

参与科学空间

为了保证你的利益,推荐你注册为本站会员。同时欢迎通过邮件或留言进行交流、建议或反馈科学空间的问题。
会员注册 会员登录 查看全站文章归档页

10 Jul

Transformer升级之路:21、MLA好在哪里?(下)

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

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

部分旋转

首先,我们把最终的断言放在前面:

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

点击阅读全文...

1 Jul

“对角+低秩”三角阵的高效求逆方法

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

点击阅读全文...

23 Jun

通过msign来计算奇异值裁剪mclip(下)

前面我们在《通过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 在推特上指出笔者的方案实际计算中存在误差偏大的问题。本文来具体分析一下这个问题,并给出一个更高效、误差更低的新方案。

点击阅读全文...

23 Jun

矩阵符号函数mcsgn能计算什么?

《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}

点击阅读全文...

20 Jun

线性注意力简史:从模仿、创新到反哺

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

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

点击阅读全文...

13 Jun

msign的导数

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

两种定义

本文依然假设大家已经对$\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后所得的新矩阵。

点击阅读全文...

7 Jun

通过msign来计算奇异值裁剪mclip(上)

前面我们用了两篇文章《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}

点击阅读全文...

5 Jun

msign算子的Newton-Schulz迭代(下)

在上文《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}

点击阅读全文...