Tied embeddings，即将语言模型中的输入Embeddings权重与输出分类器的权重两组参数共享的操作，一度是语言建模和机器翻译任务的标准配置。在语言模型大规模化之后，这种设计在开源模型中愈发少见了。前几天看到@苏剑林 之前的一篇博客语言模型输出端共享Embedding的重新探索，为tied embeddings的消失提供了一种视角，但也还有值得商榷的地方，本文想从这篇文章出发做一点探讨。

初始Loss的视角

这里先简要概括一下苏老师文章中的阐述框架¹。在使用Transformer做语言建模的时候，可能会使用类似DeepNorm等初始化手段，从而使每一个Transformer Block接近于一个恒等映射²，同时由于词元表征是0均值的，因此LayerNorm可以看做与RMSNorm等价³。所以，假设每个残差分支都初始化为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会使语言模型的初始损失值很大。

但这个问题实际上可以用一个rescale来解决，我们可以简单地将输出端乘以$1/\sqrt{d}$，则损失函数可以做如下近似 $$ \begin{align}\mathcal{L} &= \log \sum\limits_k e^{\boldsymbol{w}_i\cdot \boldsymbol{w}_k / (\sigma\sqrt{d})} - \boldsymbol{w}_i\cdot \boldsymbol{w}_j \big/(\sigma\sqrt{d}) \\ &\approx\log \left({\color[rgb]{0, 0.5, 0.8}e^{\sigma\sqrt{d}}} + (n-1)\right) \end{align} $$

这时候，对于常见的$\sigma=1/\sqrt{d}$或者$\sigma=0.02$，${\color[rgb]{0, 0.5, 0.8}e^{\sigma\sqrt{d}}}$这一项相对于词表大小都可以忽略不计了。

事实上，早期共享embeddings的预训练模型T5的实现⁴中就使用了这个技巧：

    if self.shared_embedding_and_softmax_weights:
        logits = mtf.einsum(
            [x * (self.model_dim.size ** -0.5), embedding_weights],
            reduced_dims=[self.model_dim])

另外，在一些公开的预训练实现中，一般残差项的初始化不会使用特别小的值，例如在OLMo-2中，就是对残差分支中的各个参数矩阵就是直接用了标准差为0.02的truncated normal初始化。

我们可以使用Llama的模型结构做一个简单的实验，我们使用常见的正态分布初始化，在固定层数为12的情况下，测试不同embedding维度下的初始loss值，结果如下表所示。

可以看到，

1. 无论是否应用tied embeddings，初始loss都有略高于$\log(n)$的情况；
2. 在输出端应用rescale技巧，可以将初始loss控制在$\log(n)$左右。

寻根溯源

笔者认为，初始Loss虽然是一个非常好的视角，但是不能解释当前tied embeddings的式微。讨论tied embeddings的应用，还得稍微追溯学术史，先看看他们是为何被提出的。

在语言建模中引入tied embeddings技巧可以追溯到LSTM-LM时代的两篇工作：Inan 2016.和Press and Wolf 2017.。其中，Inan 2016.通过类似KD的框架构造出一种soft label

$$ \begin{aligned} \boldsymbol{u}_t &= \boldsymbol{L}\boldsymbol{y}^{*}_t \\ \tilde{\boldsymbol{y}}_t &= \text{softmax}(\frac{\boldsymbol{L}^\top \boldsymbol{u}_t}{\tau}) \end{aligned} $$

这里$\boldsymbol{L}, \boldsymbol{y}^{*}_t$分别表示embedding权重和第$t$个位置的目标词元。作者论证了在一定的假设下，tied embeddings设定的语言模型相当于在隐式地学习这个soft label（而不是一般的one-hot目标）。

Press and Wolf 2017.则是通过一系列实验论证了如下几个结论⁵：

1. RNNLM使用tied embeddings时，embeddings的演进方式更接近与untied版本中输出端的embeddings；
2. 使用tied embeddings可以有效降低语言建模中的PPL（PTB数据集），无论是否使用dropout均成立；
3. 在不使用dropout的情况下，在输出embedding之前添加一个额外的投影$P$，并对$P$添加正则化loss，可以进一步降低PPL指标。

这篇文章还提出在机器翻译模型中，对于en-fr这样比较相似的语言，可以在两个语言的语料合集上联合训练一个tokenizer，共享encoder与decoder的embeddings（即encoder的输入、decoder的输入与输出共享一个参数矩阵），后来的Transformer(Vaswani 2017.)也沿用了这一做法。笔者认为这就是初期的很多预训练模型都不约而同地沿用tied embeddings的设定的原因。

但是如今回顾这两篇文章的时候，我们注意到几点：

1. 当时的语言模型一般基于浅层的RNN，输入与输出的embeddings参数在模型中占比很大；
2. 当时的实验基于PTB和WikiText数据集，相对于如今的预训练语料规模，可谓是非常小了，尤其是前者。

笔者认为，tied embeddings的有效性与数据和模型规模离不开关系。当数据与模型的规模比较有限时，tied embeddings可以作为一种很好的正则化手段（显著降低参数数量），从Press and Wolf 2017.的实验来看，在PTB这样的小数据集上，tied embeddings的语言模型在训练集上的PPL并不占优势，这表明它的作用可能有部分来自于过拟合风险的降低。

现在的LLM模做规模化主要是通过加大隐藏层维度和模型层数，non-embedding部分的参数量按$\mathcal{O}(Ld^2)$的级别增长，而embeddings的参数量只随着隐藏层维度线性增长，因此现有的LLM的embeddings所占参数比例已经非常小了，通过tied embeddings减少参数量的作用非常有限。另外，现在的预训练语料的词元规模也通常在万亿这个量级，与PTB这种训练集不到一百万词的数据集已经不能同日而语了。

训练的不稳定、工程的限制

前面我们提到，tied embeddings是源于数据与模型规模都较小的LSTMLM时期的一种正则化方法，逐渐成为一项标准设定，在预训练的早期也被沿用了下来。如今在数据与模型规模化的趋势下，正则化的强问题意识已经逐渐不成立了，这种强正则甚至可能成为训练的负担。例如，在OLMo的talk中作者提到，tied embeddings在7B的模型中会造成训练的不稳定。

除此之外，在语言模型规模化以后，模型的训练越来越依赖于各种跨节点并行计算方法。而使用tied embeddings实际上对并行方法的选择也有一定的限制。例如，使用流水线并行（Pipeline Parallelism）要求将模型纵向拆分部署在多个节点上，那么此时如果将输入与输出层看做两个不同的层，部署在不同的节点上，则首先这两部分参数共享不会节约任何的存储，还需要付出额外的通信成本来同步两个层的梯度。不过笔者觉得这个原因是次要的，如果收益是正向的，那么额外的同步步骤也是值得的。

结语

本文从语言模型输出端共享Embedding的重新探索中的初始loss视角出发，拓展讨论了在语言建模规模化之后，tied embeddings操作不再作为标准设定的原因：模型与数据规模的变化使得正则化的问题意识不再，且从一些公开的实验来看，tied embeddings可能引发训练的不稳定⁶，此外tied embeddings也对并行方法的选型有一定限制。

拓展阅读

• Inan 2016. Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling
• Press and Wolf 2017. Using the Output Embedding to Improve Language Models
• 语言模型输出端共享Embedding的重新探索