NLP

最近在阅读Muon is Scalable for LLM Training这篇文章的时候注意到他们使用无权重衰减（weight decay）版本的Muon优化LLM的时候，优化器的收敛优势会随着训练过程逐渐消失，又看到@小明同学在评论区提到的一个细节，很多开源的LLM在技术报告中都提到了使用0.1作为权重衰减的系数，觉得是个比较有意思的发现。结合Kimi的文章中关于bf16的简单陈述，笔者在本文中稍微展开讲下，权重衰减对于LLM的低精度训练中有什么作用。 首先把结论放在前面：除了一般认知中的正则化作用，权重衰减也可能降低精度损失的风险——对于计算机的浮点数而言，绝对值越大，精度越低。对于低精度/混合精度训练而言，使用权重衰减可以控制参数的绝对值范围，从而保证模型参数不落入低精度的数值区间。 浮点数的存储与精度 上述结论主要与浮点数在计算机内的存储形式有关。学过计算机的一些基本课程的读者可能有印象，浮点数的存储是二进制的形式，分为符号位、指数位和尾数位三段。深度学习中常见的浮点数协议（fp32、fp16、bf16、tf32）的区别在于指数位和尾数位的比特数量不同。由于浮点数是一个「指数」的形式，因此它在实数空间的分布是不均匀的。 这里我们考虑规范数的情形（指数位非全0），做一点分析。假设符号位、指数位、尾数位（mantissa）的二进制编码分别是$S$、$E$、$M$，那么对应的浮点数为： $$ \text{value}=(-1)^S\times 1.M\times 2^{E-\text{bias}} $$ 在单精度fp32标准中，$\text{bias}$取${01111111}_{2}=127_{10}$ . fp32浮点数的一个例子 例如在图中的例子中，$S=0$，$E=\underbrace{00\cdots0}_{7\ 0's}1$，$M=\underbrace{00\cdots0}_{22\ 0's}1$，相应的值为 $$ \begin{align} &{-1}^0\times 1.\underbrace{00\cdots0}_{22\ 0's}1_2\times 2^{00000001_2-{01111111}_{2}}\\ &\approx [1.175494490952134\times 10^{-38}]_{10} \end{align} $$ 现在我们来考虑不同的数值范围内的浮点数精度。对于区间范围$[2^{x}, 2^{x+1}],\forall -126\le x\le127$（这里的x已经是经过-bias之后得到的最终指数），我们希望在给定任意浮点数$y\in[2^{x}, 2^{x+1}]$的基础上增加一个最小量$\varepsilon$（即区间内两个浮点数的最小间隔），这个增加的过程是通过操纵二进制编码实现的，那么最小间隔只能是通过在$y$的尾数部分加上 $$ 0.\underbrace{00\cdots0}_{22\ 0's}1 $$ 来实现，对应的最小间隔是 $$ \begin{align} \varepsilon &= 0.\underbrace{00\cdots0}_{22\text{ 0's}}1_2\times 2^{x}\\ &=2^{-23}\times2^{x}=2^{x-23}\\ \end{align} $$ 如果你将上面的公式带入不同的指数位，可以验证与Wikipedia中给出的这张表的Gap是吻合的： 不同指数位下的最小精度，来源：Wikipedia 这个计算可以拓展到其他的精度格式： 对于半精度格式fp16而言，其尾数位有10位，对应的最小间隔是$2^{x-10}$； 对于半精度格式bf16而言，其尾数位有7位，对应的最小间隔是$2^{x-7}$； 从这里也可以看到，bf16相比fp16虽然拓宽了表示范围，但是减少了精度（同样数值范围内的最小间隔更宽了）。 从这里我们得到了一个结论：计算机存储的浮点数之间的最小间隔随着浮点数绝对值数值增加，指数级地增大，换言之，浮点数（绝对值）数值越大，精度越低。并且这个问题对于fp16或bf16格式的浮点数，问题要更加显著。 这个结论的另一个引申的问题是舍入误差，假如一个较大的浮点数和一个较小的浮点数相加，由于浮点数的加法（减法过程相当于取补码后相加，结论是类似的）过程需要先将两个数的指数位对齐，因此绝对值较小的数字的尾数的最后几位数字可能会在加法中丢失。这里我们举一个极端的例子来说明。 假设浮点数存储为fp32格式（8位指数、23位尾数）。 $$ \begin{align} x=(-1)^0\times 1.0_2\times 2^{-1}\\ y=(-1)^0\times 1.0_2\times 2^{-25}\\ \end{align} $$ ...

最近Deepseek比较出圈，连带着里面用到的MLA也被讨论得更多了。MLA无疑是一个非常出色的attention改进，但是由于它的KV cache设计，不能很好地兼容RoPE，因此作者们使用了decoupled RoPE这样的「补丁」来引入位置关系，这无疑也增加了实现的复杂度。 最近，修改注意力KV Cache这一线工作又增添了TPA这个新成员，笔者觉得这篇文章比较有趣，因此希望写一篇简短的导读。之所以叫「导读」，是因为本文不打算太深入文章的formulation和实验，而是从笔者自己的视角出发介绍文章的一些重点贡献。。 MHA的拆解 最标准的多头注意力（Vaswani的版本），大致可以拆解成3个步骤（这里默认讨论自注意力）1： $$ \begin{aligned} \text{step 1 }& \begin{cases} \boldsymbol{q}_i^{(h)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(h)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(h)}\in\mathbb{R}^{d\times d_k} \\ \boldsymbol{k}_i^{(h)} = \boldsymbol{x}_i\boldsymbol{W}_k^{(h)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(h)}\in\mathbb{R}^{d\times d_k}\\ \boldsymbol{v}_i^{(h)} = \boldsymbol{x}_i\boldsymbol{W}_v^{(h)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(h)}\in\mathbb{R}^{d\times d_v} \end{cases} \\ \text{step 2 }& \begin{cases} \boldsymbol{o}_t^{(h)} = \text{Attention}\left(\boldsymbol{q}_t^{(h)}, \boldsymbol{k}_{\leq t}^{(h)}, \boldsymbol{v}_{\leq t}^{(h)}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(h)} \boldsymbol{k}_i^{(h)}{}^{\top}\right)\boldsymbol{v}_i^{(h)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(h)} \boldsymbol{k}_i^{(h)}{}^{\top}\right)} \\ \end{cases} \\ \text{step 3 }& \begin{cases} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(H)}\right] \end{cases} \\ \end{aligned}\\ $$ 这里$\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_l$是输入向量，上标$h\in \{1,\ldots,H\}$表示注意力头，$d,d_k,d_v$分别表示输入维度、key维度、value维度。在第一步中，我们将每个token的表征独立地线性投射到不同的自空间上，第二步是标准的点积注意力，第三步是拼接（后续再做一步线性变换）。 我们重点来审视一下第一个步骤，这里由于每个token的表征是独立操作的（类似FFN），因此我们可以将每个token的表征看做一个样本点，这里做的事情就是将一个$d$维度的向量，转化成3个矩阵$\tilde{\boldsymbol{Q}}\in\mathbb{R}^{H\times d_k},\tilde{\boldsymbol{K}}\in\mathbb{R}^{H\times d_k},\tilde{\boldsymbol{V}}\in\mathbb{R}^{H\times d_v}$，每一个头对应这个矩阵中的一行。接着我们逐行计算每行对应的三组向量的点积注意力，就构成了标准的多头注意力。在标准实现中，这个变换是通过参数矩阵的线性变换+ reshape实现的。 ...

Tied embeddings，即将语言模型中的输入Embeddings权重与输出分类器的权重两组参数共享的操作，一度是语言建模和机器翻译任务的标准配置。在语言模型大规模化之后，这种设计在开源模型中愈发少见了。前几天看到@苏剑林 之前的一篇博客语言模型输出端共享Embedding的重新探索，为tied embeddings的消失提供了一种视角，但也还有值得商榷的地方，本文想从这篇文章出发做一点探讨。 初始Loss的视角 这里先简要概括一下苏老师文章中的阐述框架1。在使用Transformer做语言建模的时候，可能会使用类似DeepNorm等初始化手段，从而使每一个Transformer Block接近于一个恒等映射2，同时由于词元表征是0均值的，因此LayerNorm可以看做与RMSNorm等价3。所以，假设每个残差分支都初始化为0，假设输入中某个位置的初始embedding是$\boldsymbol{w}_i$（对应词表中的第$i$个词，维度是$d$），那么最终得到的表征满足 $$ \frac{\boldsymbol{w}_i}{\Vert\boldsymbol{w}_i\Vert \big/\sqrt{d}} \approx \frac{\boldsymbol{w}_i}{\sigma} $$ 假设在该位置的真实标签是词元$j$，则损失函数可以由如下逼近 $$ \begin{align}\mathcal{L}\triangleq -\log p(j|i) &= \log \sum\limits_k e^{\boldsymbol{w}_i\cdot \boldsymbol{w}_k / \sigma} - \boldsymbol{w}_i\cdot \boldsymbol{w}_j \big/ \sigma \\ &\approx \log \sum_k e^{\boldsymbol{w}_i\cdot \boldsymbol{w}_k / \sigma}\\ &=\log \left(e^{\boldsymbol{w}_i\cdot \boldsymbol{w}_i / \sigma} + \sum\limits_{k|k\neq i} e^{\boldsymbol{w}_i\cdot \boldsymbol{w}_k / \sigma}\right)\\ &\approx\log \left({\color[rgb]{0, 0.5, 0.8}e^{d \sigma}} + (n-1)\right) \end{align} $$ 其中$|n|$是词表大小。在常见的模型维度下，这里的第一项${\color[rgb]{0, 0.5, 0.8}e^{d \sigma}}$是比较大的。我们可以代入几个维度值看下第一项的大小，这里我们假设词表大小是32k，并考虑两种$\sigma$取法，一种是比较常见的初始化超参数$\sigma=0.02$，一种是取$\sigma=1/\sqrt{d}$。可以看到无论是哪种初始化方法，对应的${\color[rgb]{0, 0.5, 0.8}e^{d \sigma}}$都已经远远超过词表大小，响应地初始损失值也处于比较高的水平（按均匀分布的交叉熵是$\log(n)\approx 10.37$）。 不同设定下的「初始损失值」 以上是苏文中给出的关于语言建模中不再共享embedding的一个视角——tied embeddings会使语言模型的初始损失值很大。 ...