Posts

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会使语言模型的初始损失值很大。 ...