6 Nov

n个正态随机数的最大值的渐近估计

设$z_1,z_2,\cdots,z_n$是$n$个从标准正态分布中独立重复采样出来的随机数,由此我们可以构造出很多衍生随机变量,比如$z_1+z_2+\cdots+z_n$,它依旧服从正态分布,又比如$z_1^2+z_2^2+\cdots+z_n^2$,它服从卡方分布。这篇文章我们来关心它的最大值$z_{\max} = \max\{z_1,z_2,\cdots,z_n\}$的分布信息,尤其是它的数学期望$\mathbb{E}[z_{\max}]$。

先看结论

关于$\mathbb{E}[z_{\max}]$的基本估计结果是:

设$z_1,z_2,\cdots,z_n\sim\mathcal{N}(0,1)$,$z_{\max} = \max\{z_1,z_2,\cdots,z_n\}$,那么 \begin{equation}\mathbb{E}[z_{\max}]\sim \sqrt{2\log n}\label{eq:E-z-max}\end{equation}

点击阅读全文...

27 Oct

低精度Attention可能存在有偏的舍入误差

前段时间笔者在arXiv上刷到了论文《Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention》,里面描述的实验现象跟我们在训练Kimi K2时出现的一些现象很吻合,比如都是第二层Attention开始出现问题。论文将其归因为低精度Attention固有的有偏误差,这个分析角度是比较出乎笔者意料的,所以饶有兴致地阅读了一番。

然而,论文的表述似乎比较让人费解——当然也有笔者本就不大熟悉低精度运算的原因。总之,经过多次向作者请教后,笔者才勉强看懂论文,遂将自己的理解记录在此,供大家参考。

结论简述

要指出的是,论文标题虽然点名了“Flash Attention”,但按照论文的描述,即便block_size取到训练长度那么大,相同的问题依然会出现,所以Flash Attention的分块计算并不是引起问题的原因,因此我们可以按照朴素的低精度Attention实现来简化分析。

点击阅读全文...

12 Oct

随机矩阵的谱范数的快速估计

《高阶MuP:更简明但更高明的谱条件缩放》的“近似估计”一节中,我们曾“预支”了一个结论:“一个服从标准正态分布的$n\times m$大小的随机矩阵,它的谱范数大致是$\sqrt{n}+\sqrt{m}$。”

这篇文章我们来补充讨论这个结论,给出随机矩阵谱范数的快速估计方法。

随机矩阵论

设有随机矩阵$\boldsymbol{W}\in\mathbb{R}^{n\times m}$,每个元素都是从标准正态分布$\mathcal{N}(0,1)$中独立重复地采样出来的,要求估计$\boldsymbol{W}$的谱范数,也就是最大奇异值,我们以$\mathbb{E}[\Vert\boldsymbol{W}\Vert_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 在推特上指出笔者的方案实际计算中存在误差偏大的问题。本文来具体分析一下这个问题,并给出一个更高效、误差更低的新方案。

点击阅读全文...

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}

点击阅读全文...

2 Jun

等值振荡定理:最优多项式逼近的充要条件

最近在阅读时,遇到了一个关于最优多项式逼近的“等值振荡定理(Equioscillation Theorem)”,证明过程还涉及到无穷范数求导,感觉结论和证明都颇为新奇,特来记录一番。

参考资料:《Notes on how to prove Chebyshev’s equioscillation theorem》《Approximation Theory – Lecture 5》

等值振荡

我们先展示一下结论:

等值振荡定理 设$f(x)$是不超过$n$阶的多项式,$g(x)$是区间$[a,b]$上的连续函数,那么
\begin{equation}f^* = \mathop{\text{argmin}}_f \max_{x\in[a,b]} |f(x) - g(x)|\end{equation}
的充要条件是存在$a\leq x_0 < x_1 < \cdots < x_{n+1} \leq b$以及$\sigma\in\{0,1\}$,使得
\begin{equation}f^*(x_k) - g(x_k) = (-1)^{k+\sigma} \max_{x\in[a,b]} |f^*(x) - g(x)|\end{equation}

点击阅读全文...

11 May

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

在之前的《Muon优化器赏析:从向量到矩阵的本质跨越》《Muon续集:为什么我们选择尝试Muon?》等文章中,我们介绍了一个极具潜力、有望替代Adam的新兴优化器——“Muon”。随着相关研究的不断深入,Muon优化器受到的关注度也在日益增加。

了解过Muon的读者都知道,Muon的核心运算是$\newcommand{msign}{\mathop{\text{msign}}}\msign$算子,为其寻找更高效的计算方法是学术社区的一个持续目标。本文将总结一下它的最新进展。

写在前面

$\msign$的定义跟SVD密切相关。假设矩阵$\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后所得的新矩阵。

点击阅读全文...