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}

点击阅读全文...

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后所得的新矩阵。

点击阅读全文...

26 Sep

脑洞大开:非线性RNN居然也可以并行计算?

近年来,线性RNN由于其可并行训练以及常数推理成本等特性,吸引了一定研究人员的关注(例如笔者之前写的《Google新作试图“复活”RNN:RNN能否再次辉煌?》),这让RNN在Transformer遍地开花的潮流中仍有“一席之地”。然而,目前看来这“一席之地”只属于线性RNN,因为非线性RNN无法高效地并行训练,所以在架构之争中是“心有余而力不足”。

不过,一篇名为《Parallelizing Non-Linear Sequential Models over the Sequence Length》的论文有不同的看法,它提出了一种迭代算法,宣传可以实现非线性RNN的并行训练!真有如此神奇?接下来我们一探究竟。

求不动点

原论文对其方法做了非常一般的介绍,而且其侧重点是PDE和ODE,这里我们直接从RNN入手。考虑常见的简单非线性RNN:
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t)\label{eq:rnn}\end{equation}

点击阅读全文...

4 Apr

数值方法解方程之终极算法

呵呵,做了一回标题党,可能说得夸张了一点。说是“终极算法”,主要是因为它可以任意提高精度、而且几乎可以应付任何非线性方程(至少理论上是这样),提高精度是已知的迭代式上添加一些项,而不是完全改变迭代式的形式,当然在提高精度的同时,计算量也会随之增大。其理论基础依旧是泰勒级数。

我们考虑方程$x=f(y)$,已知y求x是很容易的,但是已知x求y并不容易。我们考虑把y在$(x_0,y_0)$处展开成x的的泰勒级数。关键是求出y的n阶导数$\frac{d^n y}{dx^n}$。我们记$f^{(n)}(y)=\frac{d^n x}{dy^n}$,并且有
$$\frac{dy}{dx}=\frac{1}{(\frac{dx}{dy})}=f'(y)^{-1}$$

点击阅读全文...

6 Mar

(原创)切抛物线法解方程

牛顿法使用的是函数切线的方程的零点来逼近原函数的零点,他所使用的是“切直线”,要是改为同曲率的“切抛物线”,则有更稳定的收敛效果以及更快的收敛速度

设函数$y=f(x)$在$(x_0,y_0)$处有一条“切抛物线”$y=ax^2+bx+c$,则应该有

$a(x_0+\Delta x)^2+b(x_0+\Delta x)+c=f(x_0+\Delta x)$-------(A)
$ax_0^2+bx_0+c=f(x_0)$-------(B)
$a(x_0-\Delta x)^2+b(x_0-\Delta x)+c=f(x_0-\Delta x)$-------(C)

其中$lim_{\Delta x->0}$

点击阅读全文...