​语言模型输出端共享Embedding的重新探索

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

预训练刚兴起时，在语言模型的输出端重用 Embedding 权重是很常见的操作，比如 BERT、第一版的 T5、早期的 GPT，都使用了这个操作，这是因为当模型主干部分不大且词表很大时，Embedding 层的参数量很可观，如果输出端再新增一个独立的同样大小的权重矩阵的话，会导致显存消耗的激增。

不过随着模型参数规模的增大，Embedding 层的占比相对变小了，加之《Rethinking embedding coupling in pre-trained language models》[1] 等研究表明共享 Embedding 可能会有些负面影响，所以现在共享 Embedding 的做法已经越来越少了。

本文旨在分析在共享 Embedding 权重时可能遇到的问题，并探索如何更有效地进行初始化和参数化。尽管共享 Embedding 看起来已经“过时”，但这依然不失为一道有趣的研究题目。

共享权重

在语言模型的输出端重用 Embedding 权重的做法，英文称之为 “Tied Embeddings” 或者 “Coupled Embeddings”，其思想主要是 Embedding 矩阵跟输出端转换到 logits 的投影矩阵大小是相同的（只差个转置），并且由于这个参数矩阵比较大，所以为了避免不必要的浪费，干脆共用同一个权重，如下图所示：

▲ 共享 Embedding 权重的 Transformer 示意图

共享 Embedding 最直接的后果可能是——它会导致预训练的初始损失非常大。这是因为我们通常会使用类似 DeepNorm 的技术来降低训练难度，它们都是将模型的残差分支初始化得接近于零。换言之，模型在初始阶段近似于一个恒等函数，这使得初始模型相当于共享 Embedding 的 2-gram 模型。接下来我们将推导这样的 2-gram 模型损失大的原因，以及分析一些解决方案。