全部文章0

科学空间苏剑林··访问 1

流形上的最速下降:6. Muon + 双旋转

原网页
原始内容

我们知道,用Adam、Muon等优化器更新矩阵参数时,奇异值和左右奇异向量都会随之变化,它们通常都是耦合在一起。也正是因为这种耦合性,我们无法简单地调控矩阵参数的奇异值,因此在奇异值出现异常增长时,我们无法简单有效地阻止它,这可能导致训练失败。

本文在《Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation》(后面简称Pion)的启发下,提出了一种单独更新矩阵左右奇异向量的Muon变体——“旋转Muon(MuonR)”,它能维持矩阵的奇异值分布不变,从而保证训练稳定性。

前文回顾

由于左右奇异向量组成的矩阵一定是正交矩阵,所以我们先来简单回顾一下正交约束下的Muon。设参数$\boldsymbol{W}\in\mathbb{R}^{n\times n}$,满足$\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$,设更新量为$\Delta\boldsymbol{W}=-\eta \boldsymbol{\Phi}$,我们希望参数更新之后依然满足正交性,那么对应的谱范数最速下降问题就是
\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\quad(\boldsymbol{W} - \eta \boldsymbol{\Phi})^{\top}(\boldsymbol{W} - \eta \boldsymbol{\Phi})=\boldsymbol{I}\end{equation}

[...]