AdamW的Weight RMS的渐近估计(下)
By 苏剑林 | 2025-11-17 | 5645位读者 | 引用在博客《AdamW的Weight RMS的渐近估计(上)》中,我们推导了AdamW训练出来的模型权重的RMS渐近表达式。不过,那会我们假设了Weight Decay和学习率在整个训练过程中是固定的,这跟实际训练并不完全吻合,所以这篇文章我们将之前的结论推广成动态版。
所谓动态版,即允许Weight Decay和学习率都随着训练步数的增加而变化,比如经典的Cosine Decay、WSD(Warmup Stable Decay)等,从而让结论更为通用。
步骤之一
我们的出发点还是AdamW的定义:
\begin{equation}\text{Adam}\color{skyblue}{\text{W}}:=\left\{\begin{aligned}
&\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\
&\boldsymbol{v}_t = \beta_2 \boldsymbol{v}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t^2\\
&\hat{\boldsymbol{m}}_t = \boldsymbol{m}_t\left/\left(1 - \beta_1^t\right)\right.\\
&\hat{\boldsymbol{v}}_t = \boldsymbol{v}_t\left/\left(1 - \beta_2^t\right)\right.\\
&\boldsymbol{u}_t =\hat{\boldsymbol{m}}_t\left/\left(\sqrt{\hat{\boldsymbol{v}}_t} + \epsilon\right)\right.\\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}})
\end{aligned}\right.\end{equation}
DiVeQ:一种非常简洁的VQ训练方案
By 苏剑林 | 2025-10-08 | 28430位读者 | 引用对于坚持离散化路线的研究人员来说,VQ(Vector Quantization)是视觉理解和生成的关键部分,担任着视觉中的“Tokenizer”的角色。它提出在2017年的论文《Neural Discrete Representation Learning》,笔者在2019年的博客《VQ-VAE的简明介绍:量子化自编码器》也介绍过它。
然而,这么多年过去了,我们可以发现VQ的训练技术几乎没有变化,都是STE(Straight-Through Estimator)加额外的Aux Loss。STE倒是没啥问题,它可以说是给离散化运算设计梯度的标准方式了,但Aux Loss的存在总让人有种不够端到端的感觉,同时还引入了额外的超参要调。
幸运的是,这个局面可能要结束了,上周的论文《DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick》提出了一个新的STE技巧,它最大亮点是不需要Aux Loss,这让它显得特别简洁漂亮!
AdamW的Weight RMS的渐近估计(上)
By 苏剑林 | 2025-10-01 | 20372位读者 | 引用在《为什么Adam的Update RMS是0.2?》中,我们用平均场近似估计了Adam的Update RMS。不久后,读者 @EIFY 指出相同的结果已经出现在论文《Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks》中。阅读后,笔者发现其中不仅包含了Update RMS的估计,还包含了Weight RMS的估计。
也就是说,AdamW训出来的模型,其权重的RMS是可以事先估计出来一个渐近结果的。大家会不会觉得这个结论有点意外?反正笔者第一次看到它是颇为意外的,直觉上权重模长是模型根据训练集自己学出来的,结果它告诉我这已经隐藏在优化器的超参中,可谓很反直觉了。
这篇文章我们还是用平均场近似方法,来复现对Weight RMS的渐近估计。
为什么Adam的Update RMS是0.2?
By 苏剑林 | 2025-09-02 | 36575位读者 | 引用众所周知,我们很早就开始尝试将Muon用于大规模LLM的训练。特别地,在《Muon续集:为什么我们选择尝试Muon?》中,我们提出了“Match Adam Update RMS”的技巧,以便快速从Adam迁移到Muon上,这个技巧同样用到了Kimi K2的训练中。该技巧是指将Muon的Update RMS统一成0.2,这使得我们复用Adam的学习率和权重衰减率。
这一技巧的背后,是我们观察到Adam的Update RMS约等于0.2,并且这一现象是稳定且可复现的。这便引发了一个有趣的问题:为什么Adam的Update RMS是0.2?我们可以从理论上解释它吗?
问题引入
首先描述一下现象:从实验中我们观察到,大致上在Warmup结束、模型进入正式训练后,Adam的Update RMS几乎都保持在0.2~0.3之间,并且不同尺寸的模型也呈现出相似的规律。这些模型的共同点是都用Adam训练,参数是$\beta_1=0.9,\beta_2=0.95$。由于共性很明显,所以这大概率不是巧合,因此笔者尝试分析背后的原理。
重新思考学习率与Batch Size(一):现状
By 苏剑林 | 2025-09-01 | 27578位读者 | 引用在之前的文章《当Batch Size增大时,学习率该如何随之变化?》和《Adam的epsilon如何影响学习率的Scaling Law?》中,我们从理论上讨论了学习率随Batch Size的变化规律,其中比较经典的部分是由OpenAI提出的展开到二阶的分析。然而,当我们要处理非SGD优化器时,这套分析方法的计算过程往往会相当复杂,有种无从下手的感觉。
接下来的几篇文章,笔者将重新整理和思考上述文章中的相关细节,尝试简化其中的一些推导步骤,给出一条更通用、更轻盈的推导路径,并且探讨推广到Muon优化器的可能性。
方法大意
首先回顾一下之前的分析方法。在《当Batch Size增大时,学习率该如何随之变化?》中,我们介绍了多种分析学习率与Batch Size规律的思路,其中OpenAI在《An Empirical Model of Large-Batch Training》提出的二阶近似分析占了主要篇幅,本文也是沿用同样的思路。
这篇文章我们来推导$\newcommand{msign}{\mathop{\text{msign}}}\msign$算子的求导公式。如果读者想要像《Test-Time Training Done Right》一样,将TTT和Muon结合起来,那么本文可能会对你有帮助。
两种定义
本文依然假设大家已经对$\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后所得的新矩阵。
SVD(Singular Value Decomposition,奇异值分解)是常见的矩阵分解算法,相信很多读者都已经对它有所了解,此前我们在《低秩近似之路(二):SVD》也专门介绍过它。然而,读者是否想到,SVD竟然还可以求导呢?笔者刚了解到这一结论时也颇感意外,因为直觉上“分解”往往都是不可导的。但事实是,SVD在一般情况下确实可导,这意味着理论上我们可以将SVD嵌入到模型中,并用基于梯度的优化器来端到端训练。
问题来了,既然SVD可导,那么它的导函数长什么样呢?接下来,我们将参考文献《Differentiating the Singular Value Decomposition》,逐步推导SVD的求导公式。
推导基础
假设$\boldsymbol{W}$是满秩的$n\times n$矩阵,且全体奇异值两两不等,这是比较容易讨论的情形,后面我们也会讨论哪些条件可以放宽一点。接着,我们设$\boldsymbol{W}$的SVD为:
\begin{equation}\boldsymbol{W} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}\end{equation}
通过梯度近似寻找Normalization的替代品
By 苏剑林 | 2025-04-02 | 32017位读者 | 引用不知道大家有没有留意到前段时间的《Transformers without Normalization》?这篇论文试图将Transformer模型中的Normalization层用一个Element-wise的运算DyT替代,以期能提高速度并保持效果。这种基础架构的主题本身自带一点吸引力,加之Kaiming He和Yann LeCun两位大佬挂名,所以这篇论文发布之时就引起了不少围观,评价也是有褒有贬。
无独有偶,上周的一篇新论文《The Mathematical Relationship Between Layer Normalization and Dynamic Activation Functions》从梯度分析和微分方程的视角解读了DyT,并提出了新的替代品。个人感觉这个理解角度非常本质,遂学习和分享一波。
写在前面
DyT全称是Dynamic Tanh,它通过如下运算来替代Normalization层:
\begin{equation}\mathop{\text{DyT}}(\boldsymbol{x}) = \boldsymbol{\gamma} \odot \tanh(\alpha \boldsymbol{x}) + \boldsymbol{\beta}\end{equation}








最近评论