文章标题：多头注意力模型中的低秩瓶颈
作者/机构：Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank J. Reddi, Sanjiv Kumar

A1 主要贡献

本文旨在识别并解决当前多头注意力模型中一个关键的性能瓶颈。

核心问题：基于Transformer的注意力架构，特别是像BERT和GPT-2这样的大型模型，其性能的提升在很大程度上依赖于使用非常大的词元（token）嵌入维度。这导致模型规模过于庞大，难以在下游任务中部署。本文的分析指出，造成这一问题的一个重要因素是多头注意力机制中的一个固有缺陷。当前普遍采用的做法是，为了在增加注意力头（head）数量的同时保持参数量不变，会将每个头的尺寸（head size）与头的数量成反比地缩放（即 head_size = embedding_size / num_heads）。

研究目标与发现：本文的研究目标是揭示这种缩放策略带来的负面影响。作者通过理论分析和实验证明，当头的尺寸减小到低于输入序列长度时，会产生一个“低秩瓶颈”，从而损害了每个注意力头的表达能力。这解释了为什么在实践中，当头的数量超过某个阈值后，模型性能会下降，如BERT-LARGE模型在超过8个头后性能开始下降（见表1）。为了避免这个瓶颈，现有模型不得不增加整体的嵌入维度，从而导致模型参数量和计算成本的增加。

表 1: 拥有1024嵌入维度的24层Transformer模型BERTLARGE【【8】Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.】的性能在注意力头数量超过8个后开始下降。

创新点/核心贡献：
1. 识别并证明低秩瓶颈：本文分析了多头自注意力层的表示能力，并从理论上证明了头的尺寸对注意力单元施加了低秩瓶颈（定理1）。
2. 提出固定头尺寸方案：作为解决方案，本文提出将注意力单元的头尺寸设置为一个固定值，具体建议为输入序列的长度，使其独立于头的数量和模型的嵌入维度。作者从理论上证明，与标准启发式方法相比，固定头尺寸能严格提升多头注意力层的表达能力（定理2）。这种解耦允许在不牺牲性能的前提下，增加每层的头数量并减小嵌入维度。
3. 实验验证：通过在语言建模（LM1B）、自然语言推理（MNLI）和问答（SQuAD）三个标准NLP任务上的实验，本文证明了采用固定头尺寸的Transformer模型能够以更好的性能扩展性和更小的嵌入维度进行训练。例如，实验表明，一个嵌入维度为512的固定头尺寸模型，其性能可以匹配嵌入维度为1024的BERT-LARGE模型。

A3 背景知识与关键分析

本节介绍了Transformer架构并分析了其核心组件——多头自注意力的表示能力。

2. Transformer架构与分析

模型输入。Transformer网络的输入是一个包含n个词元（token）的序列。通常，每个词元通过一个嵌入层被转换成一个d维的词元嵌入。我们用 $X \in R^{d \times n}$ 表示输入序列中n个词元对应的嵌入矩阵。

2.1 单头注意力

Transformer模块的构成。一个Transformer模块由一个自注意力层和一个前馈网络层组合而成【【20】Vaswani, A., et al. (2017). Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998–6008.】。这两个层都带有残差连接（skip connection）并使用层归一化（Layer Normalization, LN）【【1】Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). Layer normalization. arXiv preprint arXiv:1607.06450.】。具体来说，对于词元嵌入X，点积注意力的计算方式如下：

注意力计算。在此公式中，$W_q \in R^{d_q \times d}$，$W_k \in R^{d_k \times d}$ 和 $W_v \in R^{d_v \times d}$ 分别代表注意力单元中与查询（query）、键（key）和值（value）相关联的投影矩阵【【20】Vaswani, A., et al. (2017). Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998–6008.】。对于单头注意力单元，我们有 $d_q = d_k = d_v = d$。在点积注意力（见公式(1)）中，矩阵P旨在根据输入序列中的其余词元，为给定的词元捕获上下文信息。随后，注意力层的输出形式如下：

表示能力的探讨。其中LN(·)代表层归一化操作。鉴于注意力模块如公式(1)所定义，一个自然的问题是，对于给定的输入序列X，它是否有能力表示任意的上下文P。在接下来的结论中，我们证明了当投影尺寸足够大时，一个注意力单元可以表示任何数据对(X, P)。我们还表明，当d小于n时，模型无法表示任意上下文，从而产生低秩瓶颈。

定理1（表示定理）。如果 $d_q = d_k = d \ge n$，那么对于任意一个满列秩矩阵 $X \in R^{d \times n}$ 和任意一个 $n \times n$ 的正列随机矩阵（positive column stochastic matrix）P，总存在 $d \times d$ 的投影矩阵 $W_q$ 和 $W_k$，使得：

如果 $d_q = d_k = d < n$，则存在X和P，使得对于所有的 $W_q$ 和 $W_k$，公式(3)都不成立。这个结果表明，投影维度 $d_q = d_k = d$ 需要大于序列长度n，注意力单元才能表示任何期望的上下文P。尽管这个结果描述的是单个样本序列的情况，但它揭示了模型架构的一个基本属性：将投影尺寸减小到某个阈值以下会引入瓶颈。

定理1的证明（d ≥ n的情况）。为了证明结果的第一部分，我们给出了 $W_k$ 和 $W_q$ 的一个显式构造，它允许我们使用点积注意力从X生成P。由于X是满列秩的，存在一个左逆 $X^\dagger = (X^T X)^{-1}X^T \in R^{n \times d}$，使得 $X^\dagger X = I_n$。令 $W_k = \tilde{W}_k X^\dagger$ 和 $W_q = \tilde{W}_q X^\dagger$。那么：

既然上述 $W_q$ 和 $W_k$ 的选择已经处理了对X的依赖，我们将根据P选择一个 $\tilde{W}_{kq}$ 来完成构造。下面我们表达查询和键内积上的Softmax操作。注意这里的Softmax是一个列式操作符，为每个查询计算注意力分数。通过使用公式(4)，我们得到：

其中 $D_{\tilde{W}_{kq}}$ 是一个 $n \times n$ 的对角矩阵，使得：

因此，我们可以通过证明总存在一个 $\tilde{W}_{kq}$ 满足以下不动点方程来建立期望的结果：

给定P，为了构造这样一个 $\tilde{W}_{kq}$，我们选择一个任意的正对角矩阵 $D_0$，并设置：

由于P是一个正矩阵，这样的 $\tilde{W}_{kq}$ 总是存在的。接下来，我们验证这个构造确实满足不动点方程（见公式(5)）。注意到：

最后一个等式成立是因为P是一个列随机矩阵。现在，使用公式(6)和(7)：

这完成了证明的第一部分。

定理1的证明（d < n的情况）。考虑d = 1和n = 2的情况。那么 $X \in R^{1 \times 2}$ 且 $W_q, W_k \in R^{1 \times 1}$。令 $X = [1, 0]$。那么：

这个矩阵显然不能用来生成第二列元素不同的P，例如 $P = \begin{pmatrix} 0.5 & 0.75 \\ 0.5 & 0.25 \end{pmatrix}$。

2.2 多头注意力

多头注意力的机制。如2.1节所讨论，一个注意力单元使用上下文P（见公式(1)），基于序列中所有词元的嵌入的加权平均来更新输入词元的嵌入。Vaswani等人【【20】Vaswani, A., et al. (2017). Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998–6008.】提出了多头注意力机制，它通过让多个注意力单元（每个单元被称为一个“头”）在输入的不同低维投影上操作，来增强注意力层的表示能力。这之后是来自不同头的输出的拼接。具体来说，一个有h个头的多头注意力内部的计算形式如下：

多头注意力的输出。多头注意力层的输出则变为：

其中 $W_o \in R^{d \times d}$。对于一个有h个头的模型，查询、键和值的投影矩阵 $\{W_{iq}\}$, $\{W_{ik}\}$ 和 $\{W_{iv}\}$ 是 $d_h \times d$ 的矩阵。因此，每个头将输入投影到一个 $d_h$ 维的子空间来计算上下文，并保持每层的参数数量固定。使用多头注意力相比单头注意力层在经验上取得了更好的性能【【20】Vaswani, A., et al. (2017). Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998–6008.】。

图 1: 在语言建模任务（LM1B）测试集上，使用流行头尺寸启发式方法（$d_p = d/h$）（基线）训练的Transformer与使用固定头尺寸（$d_p = 32$）的Transformer的性能对比。我们训练了嵌入维度从256到512的基线模型。我们训练的固定头尺寸模型，其嵌入维度固定为256，头尺寸为32，头数从4变化到70，同时匹配参数数量。这些图清楚地表明，固定头尺寸使我们能够以更小的嵌入维度（图(b)）和更好的性能扩展性（图(a)）来训练Transformer。注意，对于困惑度（perplexity），值越低越好。

2.3 低秩瓶颈

头尺寸与表达能力的权衡。虽然增加头的数量表面上赋予了模型更强的表达能力，但同时我们也在减小头的尺寸，这可能会降低表达能力。当头的数量h大于 $d/n$ 时，每个头内部的注意力单元会将输入投影到一个小于n的维度上，从而产生低秩瓶颈，并失去表示任意上下文向量的能力（参考定理1）。有趣的是，这与表1中的经验观察一致，即在BERT-LARGE模型【【8】Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.】中，将h增加到8以上会导致性能下降；需要注意的是，在BERT-LARGE的大部分预训练阶段，$d = 1024$，$n = 128$。

现有启发式方法的局限性。由于序列长度n由手头的数据/任务固定，在不引入低秩瓶颈的情况下增加头数量的唯一方法是增加嵌入维度d。这是当前主流的头尺寸启发式方法的一个根本限制，即我们需要增加嵌入维度来支持更多的头。不幸的是，增加嵌入维度会导致训练和存储模型的计算和内存需求更高。此外，由于通常会将基于Transformer的模型学习到的嵌入用于下游任务【【8】Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.】，更大的嵌入维度也增加了所有下游任务的模型大小和所需计算量。

A2 方法细节

3. 固定多头注意力

提出固定头尺寸。在本节中，我们提议固定Transformer的头尺寸，这使我们能够在不需要很大嵌入维度的情况下，享受到多头带来的更高表达能力的优势。关键在于解耦头内部的投影尺寸与模型的嵌入维度之间的依赖关系。现在，投影矩阵将输入投影到固定维度 $d_p$ 的子空间，而与头的数量h无关。这种 $d_p$ 独立于d和h的方法导致了以下注意力机制：

新注意力机制的输出。请注意，这里使用的投影矩阵 $\{V_{iq}\}, \{V_{ik}\}$ 和 $\{V_{iv}\}$ 是 $d_p \times d$ 的矩阵。当 $V_o \in R^{d \times h \cdot d_p}$ 时，这个新的多头注意力层的输出形式如下：

固定头尺寸的优势。这个修改使得每个注意力头更像前馈网络中的一个隐藏单元或卷积网络中的一个滤波器，并允许我们在不担心降低每个头表示能力的情况下改变头的数量。缺点是，与标准的多头注意力不同，每层的参数数量会随着头的数量增加而增加。然而，这一修改使我们能够用更小的嵌入维度训练模型而不产生低秩瓶颈，最终使我们能够减少模型的总参数数量。

3.1 多头注意力 vs. 固定多头注意力

表达能力的比较。给定一个多头（MultiHead）层，只要头尺寸 $d_p \ge d/h$，我们总可以用一个固定多头（FixedMultiHead）层来表示它。虽然这表明当头数量h超过 $d/d_p$ 时，固定多头注意力的单个头与多头注意力中的头一样具有表达能力，但尚不清楚固定多头注意力是否严格更具表达能力。固定多头注意力层能否表示标准多头注意力层无法表示的函数？在本小节中，我们表明，在多头机制中，固定多头注意力层在表达能力方面确实严格优于标准多头注意力层。

函数类的定义。考虑公式(8)中的标准多头注意力单元。

我们将所有参数矩阵的集合表示为W。类似地，考虑由固定头尺寸注意力单元表示的函数：

函数类的包含关系。令V为所有这些参数矩阵的集合。我们定义F和G分别为函数类 $f_W(\cdot)$ 和 $g_V(\cdot)$。如上所述，如果 $d_p \ge d/h$，我们有 $F \subset G$。

定理2（严格更优的表达能力）。下面的定理表明，即使对于G中的简单例子，F中的函数也无法表示它们；这已经表明F是G的一个严格子集。

定理2：令 $n \ge 2$, $d \ge d_p$，且 $h > d/d_p$。考虑一个固定多头注意力层 $g_V(\cdot)$，其参数满足以下条件：

那么，对于任何 $f_W \in F$，都存在一个 $X \in R^{d \times n}$，使得 $f_W(X) \neq g_V(X)$。

定理2的意义。因为 $\|f_W(X) - g_V(X)\|$ 是X的连续函数，所以存在这样的X意味着差值范数的积分（即近似误差）是严格为正的。我们注意到，上述定理中对 $V_{ik}$ 和 $V_{iq}$ 的假设是为了提供一个简单且具有建设性的证明；事实上，多头注意力（F）无法表示如此简单的注意力层，这表明对于更复杂的函数，情况可能更糟。定理2表明，固定多头注意力函数类的表达能力严格优于标准多头注意力函数类。因此，随头数增加而减小头尺寸的启发式方法限制了多头注意力的表达能力，而使用固定的头尺寸将增强注意力层的表达能力。

A4 实验

本节的目标是证明，采用原则性方法设置头尺寸比使用流行的启发式方法能带来更好的性能。我们再次强调，虽然这只是对Transformer架构的一个简单超参数更改，但根据我们的分析将其设置为输入序列长度，可以让我们用更小的嵌入维度训练出更好的模型。

图 2: 在SQuAD和MNLI开发集上，使用流行头尺寸启发式方法的24层Transformer模型BERTLARGE（基线）与固定头尺寸模型的性能比较。我们改变基线模型的嵌入维度，从512到1024。我们训练的固定头尺寸模型，其嵌入维度固定为512，头尺寸为128，头数从8变化到32，同时匹配参数数量。固定头尺寸使我们能够以更小的嵌入维度512训练出性能更好的模型。

图 3: LM1B上的消融研究：（a）我们将所有模型的嵌入维度固定为256，通过增加前馈层的尺寸来改变使用流行头尺寸启发式方法（基线）训练的Transformer的容量。对于固定头尺寸模型，我们将头尺寸固定为32，因此8个头的固定头尺寸模型与8个头的基线模型相同。我们再次注意到，使用标准启发式方法，当头数超过16个时，增加头数会损害性能，而使用固定头尺寸，增加头数会单调地提升性能。（b）我们展示了不同头数下，头尺寸对性能的影响。两张图都清楚地显示了在固定嵌入维度的情况下，增加一种调整Transformer容量的方法所带来的优势。

本节中，我们在三个标准NLP任务上进行了实验，包括语言建模（LM1B）、问答（SQuAD）和句子蕴含（MNLI），以证明：
1. 对于流行的头尺寸启发式方法，将Transformer中的头数增加到某一点以上会损害性能，但对于固定头尺寸的注意力层，增加头数总是有益的。
2. 将头尺寸与嵌入维度解耦，使我们能够用更小的嵌入维度训练模型。
3. 在Transformer中适当地设置头尺寸，可以让我们训练出具有更好性能扩展性的模型。

实验环境

数据集
* 语言建模 (LM1B): 使用One Billion Word Benchmark数据集【【2】Chelba, C., et al. (2013). One billion word benchmark for measuring progress in statistical language modeling. CoRR, abs/1312.3005.】。该数据集包含约3000万训练样本和约30万测试样本。使用拥有32k词汇表的子词分词器，并将输入序列长度限制为256。
* 自然语言推理 (MNLI): Multi-Genre Natural Language Inference数据集【【24】Williams, A., et al. (2018). A broad-coverage challenge corpus for sentence understanding through inference. In NAACL-HLT.】，任务是判断两个句子之间的关系（蕴含、矛盾、中性）。
* 问答 (SQuAD): Stanford Question Answering Dataset【【17】Rajpurkar, P., et al. (2016). Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250.】，任务是从给定段落中预测问题的答案。
* 预训练数据: 对于MNLI和SQuAD任务，模型首先在English Wikipedia和BooksCorpus数据集【【32】Zhu, Y., et al. (2015). Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In ICCV.】上进行预训练。输入使用WordPiece表示法进行分词，词汇表大小为30000。

模型架构与参数
* LM1B: 6层Transformer模型。前馈网络层的宽度固定为1024。
* MNLI/SQuAD: 24层Transformer模型，遵循BERT【【8】Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.】的预训练-微调范式。
* 头尺寸选择: 对于BERT相关实验，选择的头尺寸 $d_p$ 为128，因为大部分预训练是在序列长度为128的数据上完成的。

硬件配置
* LM1B: 在4个TPUv2芯片上训练。
* BERT模型: 预训练第一阶段在16个TPUv3芯片上进行，第二阶段在64个TPUv3芯片上进行。

软件配置与训练细节
* 实现: LM1B实验使用tensor2tensor库【【21】Vaswani, A., et al. (2018). Tensor2tensor for neural machine translation. CoRR, abs/1803.07416.】。MNLI/SQuAD实验使用BERT的开源代码库。
* 优化器: 使用ADAM优化器。
* LM1B训练: 批量大小为4096，训练25万步。学习率为0.1，前1万步进行线性预热，之后学习率随步数的平方根衰减。
* BERT预训练: 批量大小为1024，在序列长度n=128上训练45万步，之后在序列长度512上训练5万步。学习率为1e-4，采用线性预热和衰减策略。
* 正则化: 所有层均使用0.01的权重衰减和概率为0.1的dropout。

实验结果

固定头尺寸 vs. 标准启发式方法
* 实验内容：我们比较了使用固定头尺寸和较小嵌入维度的Transformer与使用标准头尺寸启发式方法但较大嵌入维度的Transformer的性能。
* LM1B结果：在语言建模任务中，嵌入维度为256、头尺寸固定的模型性能优于嵌入维度为512的基线模型（见图1）。
* SQuAD/MNLI结果：在另外两个任务中，嵌入维度为512、头尺寸固定的24层Transformer模型（BERTLARGE变体）性能优于嵌入维度高达1024的基线BERTLARGE模型（见图2）。
* 补充说明：简单地通过减少头数来增加基线模型的头尺寸并不能提升性能，因为减少头数会降低模型的表达能力（见附录图4）。因此，为了获得更好的性能，头尺寸和头数都需要足够大。

消融研究：增加头数的影响
* 实验内容：我们研究了随着头数增加，两种模型的性能变化。
* 结果：如表1和图1a所示，标准模型的性能在头数超过一定数量后下降。而对于固定头尺寸的Transformer，性能随着头数的增加而单调提升（见表2(A)和图3a）。这是因为模型的容量（头尺寸的函数）不再因头数增加而减少。

(A) 增加头数

(B) 增加头尺寸

表 2: SQuAD和MNLI上的消融研究：(A) 具有固定头尺寸128和512嵌入维度的24层Transformer模型，随着头数的增加，准确率有所提高。(B) 具有512嵌入维度和8个头的固定头尺寸模型，随着头尺寸的增加，准确率有所提高。这表明头尺寸确实是自注意力架构中一个重要的容量控制参数。

消融研究：增加头尺寸的影响
* 实验内容：我们比较了具有不同头尺寸的模型的性能。
* 结果：如表2(B)和图3b所示，性能随着头尺寸的增大呈现出明显的提升趋势。这表明，固定头尺寸模型性能的提升确实来自于调整注意力单元中查询、键和值层的头尺寸，并证实了头尺寸是影响注意力模型性能的一个重要因素。

A5 结论

本文研究了多头自注意力模型的表示能力，并证明了多头注意力中较小的头尺寸会导致低秩瓶颈。我们指出，当前模型中使用较大的嵌入维度是这种低秩瓶颈的后果。我们提议使用固定的头尺寸注意力单元，并将头尺寸设置为输入序列的长度，以避免此瓶颈。我们证明了这使我们能够在不增加嵌入维度的情况下增加头的数量。因此，我们能够以更小的嵌入维度和更少的参数训练出性能更好的Transformer。未来，在一个注意力块内部和跨层使用变化的头尺寸进行实验将是很有趣的。这需要进一步理解每一层在计算上下文中的作用，这是未来工作的一个有趣方向。

A6 附录

A. 符号表

B. 证明

定理2的证明。首先，我们将多头（MultiHead）和固定多头（FixedMultiHead）层重写如下。多头层可以重写为：

其中 $W_{io}$ 是 $d \times d/h$ 的矩阵，而 $W_{iv}, W_{ik}, W_{iq}$ 是 $d/h \times d$ 的矩阵。我们将所有参数矩阵的集合表示为W。类似地，重写固定头尺寸注意力层为：

其中 $V_{io} \in R^{d \times d_p}$，且 $V_{iv}, V_{ik}, V_{iq} \in R^{d_p \times d}$。令V为所有这些矩阵的集合。证明的大纲基本上是案例分析：我们将W的可能值分为三类，并在每种情况下证明存在一个X使得 $f_W(X) \neq g_V(X)$。这三种情况是：

案例1: $\sum_{i=1}^h W_{io}W_{iv} \neq \sum_{i=1}^h V_{io}V_{iv}$。
案例2: $\sum_{i=1}^h W_{io}W_{iv} = \sum_{i=1}^h V_{io}V_{iv}$，且存在 $i \in \{1, \dots, h\}$ 使得 $U/\sqrt{d_p} - (W_{ik})^T(W_{iq})/\sqrt{d/h}$ 不是斜对称的。
案例3: $\sum_{i=1}^h W_{io}W_{iv} = \sum_{i=1}^h V_{io}V_{iv}$，且所有的 $U/\sqrt{d_p} - (W_{ik})^T(W_{iq})/\sqrt{d/h}$ 都是斜对称的。

案例1的证明。在第一种情况下，我们可以选择任意的v使得 $(\sum_{i=1}^h W_{io}W_{iv} - \sum_{i=1}^h V_{io}V_{iv})v \neq 0$。选择 $X = v\mathbf{1}^T = [v, v, \dots, v]$。那么，注意到对于任何列随机矩阵P，我们有 $XP = X$。因此，

案例2的证明。在 $\sum_{i=1}^h W_{io}W_{iv} = \sum_{i=1}^h V_{io}V_{iv}$ 的情况下，由于根据假设 $\sum_{i=1}^h V_{io}V_{iv}$ 是满秩的，且每个 $W_iW_i$ 的秩最多为 $d/h$，因此 $W_i \in R^{d \times d/h}$ 中的所有列必须是线性无关的。因此，对于任何 $v \neq 0$，$\{W_iW_iv, i=1,\dots,h\}$ 是一组线性无关的向量。现在考虑任何 $v \in R^d$，并设 $X = ve_1^T$，其中 $e_1 = (1, 0, \dots, 0) \in R^n$。定义 $\phi(t) = \exp(t)/(\exp(t) + n - 1)$。那么，我们有：

类似地，我们可以计算出：

注意到 $f_W(X)$ 和 $g_V(X)$ 的所有列，从第二列到最后一列都是相同的。我们现在比较第一列：

回顾一下，对于任何 $v \neq 0$，$W_{io}W_{iv}v$ 都是线性无关的，所以 $f_W(X)_{:,1} - g_V(X)_{:,1} = 0$ 当且仅当所有的 $\phi(v^T(W_{ik})^T W_{iq}v / \sqrt{d/h}) - \phi(v^T U v / \sqrt{d_p})$ 都为零。然而，由于存在 $i \in \{1, \dots, h\}$ 使得 $U/\sqrt{d_p} - (W_{ik})^T(W_{iq})/\sqrt{d/h}$ 不是斜对称的，我们可以选择一个v满足 $v^T(W_{ik})^T W_{iq}v \neq v^T U v$，从而使得差值不为零，因此 $f_W(X)_{:,1} - g_V(X)_{:,1} \neq 0$。

案例3的证明。现在考虑任何 $X = [v_1, v_2, 0, \dots, 0]$，其中 $v_1$ 和 $v_2$ 将在稍后选择。定义 $\phi_1(t_1, t_2) = \exp(t_1)/(\exp(t_1) + \exp(t_2) + n - 2)$，$\phi_2(t_1, t_2) = \exp(t_2)/(\exp(t_1) + \exp(t_2) + n - 2)$。那么，我们有：

因此，$g_V(X)$ 的第一列可以写成：

类似地，$f_W(X)$ 的第一列是：

由于根据假设 $U/\sqrt{d_p} - (W_{1k})^T(W_{1q})/\sqrt{d/h}$ 是斜对称的，我们有 $v_1^T(U/\sqrt{d_p} - (W_{1k})^T(W_{1q})/\sqrt{d/h})v_1 = 0$ 对所有 $v_1$ 成立。回顾一下，U根据假设是秩为 $d_p$ 的，所以 $U/\sqrt{d_p} - (W_{1k})^T(W_{1q})/\sqrt{d/h}$ 的秩至少为 $d_p - d/h \ge 1$，所以我们可以选择任意一个 $v_1$ 使得 $v_1^T(U/\sqrt{d_p} - (W_{1k})^T(W_{1q})/\sqrt{d/h}) \neq 0$。通过选择合适的 $v_2$ 并进行缩放，可以证明 $f_W(X)_{:,1} - g_V(X)_{:,1} \neq 0$。

C. 实验设置

LM1B实验设置。对于语言建模（LM1B数据集）的实验，我们训练了6层Transformer模型。我们使用的批量大小为4096，训练了25万步。学习率为0.1，前1万步进行线性预热。我们用步数的平方根来衰减学习率。我们训练了基线模型，采用流行的头尺寸启发式方法，嵌入维度从256变化到512。我们将Transformer中的前馈层宽度固定为1024。此外，我们在所有层上使用0.01的权重衰减和概率为0.1的dropout。

BERT实验设置。对于BERT的实验，我们遵循了与【【8】Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.】中相同的实验设置。我们在此呈现关键细节，并请读者参考【【8】Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.】。我们使用1024的批量大小训练45万步，输入序列长度为n=128，然后使用序列长度为512的输入再训练5万步。相比之下，BERT论文使用512的批量大小，在128序列长度的输入上预训练90万步，在512序列长度的输入上预训练10万步。我们使用ADAM优化器，学习率为1e-4，并采用与BERT中相同的线性预热和衰减策略。我们为第一阶段使用5k的预热步数，为第二阶段使用3k的重新预热步数【【28】You, Y., et al. (2019). Reducing bert pre-training time from 3 days to 76 minutes. arXiv preprint arXiv:1904.00962.】。同样，我们在所有层上使用0.01的权重衰减和概率为0.1的dropout。

训练资源。对于语言建模任务，训练在4个TPUv2芯片上进行了几个小时。对于BERT模型，第一阶段的训练在16个TPUv3芯片上进行，第二阶段在64个TPUv3芯片上进行。使用此配置的预训练需要2到3天。我们没有尝试为固定头尺寸架构寻找最优的超参数，而是使用了与训练BERT模型相同的超参数。

D. 附加实验结果

图 4: 在语言建模任务（LM1B）测试集上，使用流行头尺寸启发式方法（基线）训练的Transformer与固定头尺寸（dp）模型的性能对比。与图1不同，我们同时改变基线模型的嵌入维度和头数，以使其头尺寸固定为32。我们训练的固定头尺寸模型，其嵌入维度固定为256，头尺寸为32，头数从4变化到70，同时匹配参数数量。该图再次清楚地表明了固定头尺寸模型的优势。基线模型的主要问题是，将头尺寸固定为32会迫使在嵌入维度较小时，头数也较小。将头数减少到某个阈值以下会损害Transformer的性能。