Pretraining

这篇文章是关于 LLM 预训练中 MuP / μP 的学习笔记，主要参考原论文 Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer 和苏剑林的博客 《初探MuP：超参数的跨模型尺度迁移规律》。 问题背景 LLM 预训练的成本很高，因此我们通常不希望直接在目标规模的大模型上反复搜索学习率、初始化、权重衰减等超参数。一个自然想法是：先在同架构的小模型上调参，再把超参数迁移到大模型上。 MuP（Maximal Update Parametrization）的核心目标，就是让某些关键超参数在模型宽度变化时尽量保持稳定，从而支持从小模型到大模型的 zero-shot hyperparameter transfer。原论文把这个迁移范式称为 μTransfer。 因为苏老师的博客写得太简洁+深入浅出了，本文也不会重复去讲他讲的很完善的部分，而是对他的内容进行补充和完善。个人在读原博客的时候感觉到一些地方苏老师略过了一些思考过程，导致初次理解时会让人觉得有些跳跃。 正文 MuP问题的出发点很简单，模型是一个黑盒，因此想训练出一个好模型，无法避免地要做大量的尝试（也就是俗称的调参炼丹）。但对于大模型而言，尝试的时间&金钱&人力成本很高。MuP就是对传统的炼丹过程做了一个剪枝，通过数学推导证明了小模型上已经被验证的某些规律可以直接扩展到大模型上。 当然，上述的总结是比较泛化的，具体到实践中，肯定还会问几个问题：模型大小如何界定？哪些规律可以扩展？具体如何扩展？在展开具体方法之前笔者可以回答前两个问题，这里模型的大小用神经网络的宽度/隐藏层维度来量化。可扩展的规律主要指学习率的选择。因此MuP解决的具体问题是“在网络加宽的情况下，学习率应该如何跟随着隐藏层维度改变”。 MuP(朴素版) 这一节的MuP推导没有用到超过大一高数/线代的知识，也没有用到超过机器学习基本常识的知识。用一种非常简单的视角推导了MuP（虽然有些步骤不够严谨） 模型宽度的影响：无法尽善尽美的参数初始化 为什么要先讲参数初始化呢，因为参数初始化提供了一个最基础的视角，来定量描述“宽度影响稳定性”这件事。 前传和反传的最优参数初始化无法兼容。 高维的任意两个向量夹角都是几乎正交的，可以算一下任意向量和单位向量的夹角，这里不赘述。 所以苏剑林老师基于这点给了一个推论： 从$N(0,1/n)$ 中随机选取$n^2$ 个数，组成一个$n×n$ 的矩阵，这个矩阵近似为正交矩阵，且$n$ 越大，近似程度越好。 其实道理是一样的，列向量两两正交，就是正交矩阵。n越大，相当于维度越高，正交概率越大。每个向量里面的元素都是采样出来的，所以每个元素的值大约是sqrt(1/n)，所以整个向量的模长平方就是$n * (sqrt(1/n))^2 = 1$ 正交矩阵有一个好的性质，就是它作用于一个向量时，不改变向量模长。神经网络是对一个输入向量做很多次变换， 得到一个输出向量。我们希望输入向量在变换为输出向量的游走过程中，能一直在一个球面上，也就是模长不变。因为这样 从直觉上可以大幅压缩向量遍历的空间。可以想象一下，在一个完整的高维空间里面找最优解,和在空间内的一个球面 找最优解显然是后者更容易。如果向量变换前后都在同一个球面上或者近似在一个厚度比较薄的球壳上，本文将这种性质称为“稳定性”。 所以最经典的初始化方式是推论里面的采样方式。上述结论也可以通过让变换前后的RMS相等来推导。如果引入了激活函数，初始化的值略有不同，但是推导逻辑类似。 前传和反传区别不大，都是矩阵乘。 $$ \frac{\partial \mathcal{L}}{\partial \boldsymbol{X}} \sim \frac{\partial \mathcal{L}}{\partial \boldsymbol{Y}} \boldsymbol{W}^{\top} $$主要的尺度变化也来自于$W$。输入和输出的维度不相等的时候，就找不到一个两全其美的采样方差。这是一个open的问题，苏老师在原文这里提出这个问题并不是为了直接解决这个问题，而是为了说明模型的宽度和中间层稳定性之间存在着直接的关系。 Loss的稳定性 在苏老师这篇MuP的博客中透露着一个隐含的insight：模型加宽带来的难度就是稳定性下降。（也许这个insight来自于训模型时候的经验）这个稳定性可以是Loss的稳定性，也可以是梯度的稳定性，还可以是每一层输出结果的稳定性。这里考虑了损失增量的稳定性。 文章中需要推导或注释的地方有两点，第一点是公式6如何近似，第二点是公式4如何得到公式7。 公式6的近似： $$ \Delta \mathcal{L}=\mathcal{L}(\boldsymbol{W}+\Delta \boldsymbol{W})-\mathcal{L}(\boldsymbol{W}) $$ 一阶泰勒近似： ...

本文从训练模型要考虑的第一性原理（稳定性和速度）出发，探讨了 LLM 预训练中的初始化设定问题。 参考文献 MuP之上：1. 好模型的三个特征 https://spaces.ac.cn/archives/11605 https://spaces.ac.cn/archives/11647 https://spaces.ac.cn/archives/11729

众所周知MoE有一个Router，也就是门控单元，基于topk选择路由到的专家。训练稳定的基本思想是希望随着宽度d增大，模型的activation、grad、delta loss、delta weight变化都比较小。 \[ \operatorname{RMS}(h_l)=\Theta(1),\quad \operatorname{RMS}(\Delta h_l)=\Theta(1),\quad \Delta L=\Theta(1). \] 不好分析的根源是因为topk不光滑，所以找不到一个有界的全局利普希茨常数。 也就是假设两个输入的x很接近，但是在router这里发生了跳变，导致L也就是$\lVert f(x_2)-f(x_1)\rVert/(x_2-x_1)$趋近于无穷大。 所以说topk没有全局的LipschitzBound。那么为什么lipschitz常数对mup会有影响，原因是对于某个token的一次backward里面参数更新，如果weight更新后导致路由切换，delta y的计算就不只是简单的$\Delta y\approx J_\theta y\Delta\theta$，而是两个expert相减。$\Delta y=F_j(h)-F_i(h)$. 这样可能会导致，$\lVert\Delta x\rVert\to0,\lVert\Delta y\rVert=\Theta(1)$.这个导致的最直接问题就是我们没有办法再去用一阶泰勒近似delta y了，也就是用dense的方式，没有办法用局部梯度去正确描述某个token的路由表随参数更新导致变化的这个事件，也就不能囊括所有的梯度更新类型。 从另一个角度，如果强行进行泰勒展开（当然肯定有各种不严谨，我是工程民科，凑合看看算了），就有$y(\theta+\Delta\theta)-y(\theta)=J_\theta y\Delta\theta+o(\lVert\Delta\theta\rVert)$. 但最大的问题是$\lVert y(\theta+\Delta\theta)-y(\theta)\rVert=\Theta(1)$,$\frac{\lVert y(\theta+\Delta\theta)-y(\theta)\rVert}{\lVert\Delta\theta\rVert}$所以分子分母不同阶。taylor余项也不是无穷小的。 要是想做MoE的mup，就得想办法去控制它由于router切换导致的跳变大小，也就是说必须要建模触发router跳变的条件，以及给跳变以后专家输出的diff加上一个bound。路由跳变对均方特征的贡献可以如下描述： \[ \mathbb{E}\left[\lVert\Delta y_{\mathrm{switch}}\rVert^2\right] \lesssim \int_0^{2\varepsilon}q(u)f_\gamma(u)\,\mathrm{d}u. \] f描述有多少token距离路由边界为u，q表征输出跳变的距离，2sigma代表参数更新能跨越多宽的边界。那么这件事容易吗？ 我们之前只是假设了输入不变，在mup前面的系列里面，mup的目标之一是要让中间的特征发生一些非退化的显著变化，$\Delta h=\Theta(1)$。所以即使router参数冻结了，h的变化也会导致路由切换。这个问题会让特征更新的尺度推导变得困难。 总而言之，传统 Dense μP 能闭环，是因为给定每层的： \[ \operatorname{RMS}(h),\quad \operatorname{RMS}(\delta),\quad \operatorname{RMS}(\Delta W), \] 就能递推出下一步的特征更新： \[ \Delta h\approx J_W h\Delta W, \] 而这个更新尺度仍只由矩阵维度、初始化尺度和学习率决定。 MoE 不闭环，是因为即使这些 RMS 都已知，仍然无法确定输出更新： \[ \Delta y = \Delta y_{\mathrm{smooth}} + \Delta y_{\mathrm{switch}}. \] 第二项还取决于两个额外变量： ...