QK-Clip:让Muon在Scaleup之路上更进一步
By 苏剑林 | 2025-07-12 | 20608位读者 | 引用四个月前,我们发布了Moonlight,在16B的MoE模型上验证了Muon优化器的有效性。在Moonlight中,我们确认了给Muon添加Weight Decay的必要性,同时提出了通过Update RMS对齐来迁移Adam超参的技巧,这使得Muon可以快速应用于LLM的训练。然而,当我们尝试将Muon进一步拓展到千亿参数以上的模型时,遇到了新的“拦路虎”——MaxLogit爆炸。
为了解决这个问题,我们提出了一种简单但极其有效的新方法,我们称之为“QK-Clip”。该方法从一个非常本质的角度去看待和解决MaxLogit现象,并且无损模型效果,这成为我们最新发布的万亿参数模型“Kimi K2”的关键训练技术之一。
问题描述
我们先来简单介绍一下MaxLogit爆炸现象。回顾Attention的定义
\begin{equation}\boldsymbol{O} = softmax(\boldsymbol{Q}\boldsymbol{K}^{\top})\boldsymbol{V}\end{equation}
Transformer升级之路:21、MLA好在哪里?(下)
By 苏剑林 | 2025-07-10 | 13459位读者 | 引用MoE环游记:5、均匀分布的反思
By 苏剑林 | 2025-05-16 | 26703位读者 | 引用如果说Meta的LLAMA系列为Dense模型确立了标准架构,那么DeepSeek或许就是MoE标准架构的奠基者。当然,这并非指DeepSeek首创了MoE,也不是说它的MoE不可超越,而是指DeepSeek对MoE所提的一些改进,很可能都是效果增益比较显著的方向,从而逐渐成为MoE的标配。这其中,包括我们在《MoE环游记:3、换个思路来分配》介绍的Loss-Free负载均衡方案,还有本文将要介绍的Shared Expert、Fine-Grained Expert策略。
说到负载均衡,它无疑是MoE一个极为重要的目标,本系列的第2~4篇,可以说都在围绕着它展开。然而,已有读者逐渐意识到,这里边有个尚未回答的本质问题:抛开效率上的需求不谈,均匀分布就一定是效果最好的方向吗?本文就带着这个疑问,去理解Shared Expert、Fine-Grained Expert。
共享专家
让我们再次回顾MoE的基本形式
\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i\end{equation}
Transformer升级之路:20、MLA好在哪里?(上)
By 苏剑林 | 2025-05-04 | 46572位读者 | 引用自从DeepSeek爆火后,它所提的Attention变体MLA(Multi-head Latent Attention)也愈发受到关注。MLA通过巧妙的设计实现了MHA与MQA的自由切换,使得模型可以根据训练和推理的不同特性(Compute-Bound or Memory-Bound)选择最佳的形式,尽可能地达到效率最大化。
诚然,MLA很有效,但也有观点认为它不够优雅,所以寻找MLA替代品的努力一直存在,包括我们也有在尝试。然而,经过一段时间的实验,我们发现很多KV Cache相同甚至更大的Attention变体,最终效果都不如MLA。这不得不让我们开始反思:MLA的出色表现背后的关键原因究竟是什么?
接下来,本文将详细介绍笔者围绕这一问题的思考过程以及相关实验结果。
观察
MLA提出自DeepSeek-V2,本文假设读者已经熟悉MLA,至少了解之前的博客《缓存与效果的极限拉扯:从MHA、MQA、GQA到MLA》所介绍的内容,因此MLA自身的细节将不会过多展开。
MoE环游记:4、难处应当多投入
By 苏剑林 | 2025-03-28 | 23755位读者 | 引用前两篇文章我们都在讨论负载均衡,其中在《MoE环游记:3、换个思路来分配》介绍Loss-Free方案时,笔者留了一个悬念:它引入的Bias项有一个冗余的自由度,这个自由度可以用来做另外有趣的事情。这篇文章我们就来讨论这件事。
我们知道,MoE是为每个Token只选择最匹配的$k$个Expert来进行计算,从而在增大参数量的同时还节省了计算量。然而,当我们仔细思考就会发现,这个策略实际上有明显的可改进之处:直观来看,每个Token的难度并不一样,所以更合理的方案应该是难的Token分配更多的计算资源,简单的token分配更少的资源,这样或许能在同样有限的资源下将效果最大化。
而刚才提到的Bias的额外自由度,恰好可以用来简单地实现这个目标。
为什么梯度裁剪的默认模长是1?
By 苏剑林 | 2025-01-02 | 65735位读者 | 引用我们知道,梯度裁剪(Gradient Clipping)是让模型训练更加平稳的常用技巧。常用的梯度裁剪是根据所有参数的梯度总模长来对梯度进行裁剪,其运算可以表示为
\begin{equation}\text{clip}(\boldsymbol{g},\tau)=\left\{\begin{aligned}&\boldsymbol{g}, &\Vert\boldsymbol{g}\Vert\leq \tau \\
&\frac{\tau}{\Vert\boldsymbol{g}\Vert}\boldsymbol{g},&\Vert\boldsymbol{g}\Vert > \tau
\end{aligned}\right.\end{equation}
这样一来,$\text{clip}(\boldsymbol{g},\tau)$保持跟$\boldsymbol{g}$相同的方向,但模长不超过$\tau$。注意这里的$\Vert\boldsymbol{g}\Vert$是整个模型所有的参数梯度放在一起视为单个向量所算的模长,也就是所谓的Global Gradient Norm。
不知道大家有没有留意到一个细节:不管是数百万参数还是数百亿参数的模型,$\tau$的取值在很多时候都是1。这意味着什么呢?是单纯地复用默认值,还是背后隐含着什么深刻的原理呢?
从谱范数梯度到新式权重衰减的思考
By 苏剑林 | 2024-12-25 | 26280位读者 | 引用在文章《Muon优化器赏析:从向量到矩阵的本质跨越》中,我们介绍了一个名为“Muon”的新优化器,其中一个理解视角是作为谱范数正则下的最速梯度下降,这似乎揭示了矩阵参数的更本质的优化方向。众所周知,对于矩阵参数我们经常也会加权重衰减(Weight Decay),它可以理解为$F$范数平方的梯度,那么从Muon的视角看,通过谱范数平方的梯度来构建新的权重衰减,会不会能起到更好的效果呢?
那么问题来了,谱范数的梯度或者说导数长啥样呢?用它来设计的新权重衰减又是什么样的?接下来我们围绕这些问题展开。
基础回顾
谱范数(Spectral Norm),又称“$2$范数”,是最常用的矩阵范数之一,相比更简单的$F$范数(Frobenius Norm),它往往能揭示一些与矩阵乘法相关的更本质的信号,这是因为它定义上就跟矩阵乘法相关:对于矩阵参数$\boldsymbol{W}\in\mathbb{R}^{n\times m}$,它的谱范数定义为
从Hessian近似看自适应学习率优化器
By 苏剑林 | 2024-11-29 | 25942位读者 | 引用这几天在重温去年的Meta的一篇论文《A Theory on Adam Instability in Large-Scale Machine Learning》,里边给出了看待Adam等自适应学习率优化器的新视角:它指出梯度平方的滑动平均某种程度上近似于在估计Hessian矩阵的平方,从而Adam、RMSprop等优化器实际上近似于二阶的Newton法。
这个角度颇为新颖,而且表面上跟以往的一些Hessian近似有明显的差异,因此值得我们去学习和思考一番。
牛顿下降
设损失函数为$\mathcal{L}(\boldsymbol{\theta})$,其中待优化参数为$\boldsymbol{\theta}$,我们的优化目标是
\begin{equation}\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathcal{L}(\boldsymbol{\theta})\label{eq:loss}\end{equation}
假设$\boldsymbol{\theta}$的当前值是$\boldsymbol{\theta}_t$,Newton法通过将损失函数展开到二阶来寻求$\boldsymbol{\theta}_{t+1}$:
\begin{equation}\mathcal{L}(\boldsymbol{\theta})\approx \mathcal{L}(\boldsymbol{\theta}_t) + \boldsymbol{g}_t^{\top}(\boldsymbol{\theta} - \boldsymbol{\theta}_t) + \frac{1}{2}(\boldsymbol{\theta} - \boldsymbol{\theta}_t)^{\top}\boldsymbol{\mathcal{H}}_t(\boldsymbol{\theta} - \boldsymbol{\theta}_t)\end{equation}
最近评论