UNIVERSAL TRANSFORMERS

代码链接：

摘要

循环神经网络（RNN）通过在每个新数据点更新其状态来顺序处理数据，长期以来一直是序列建模任务的事实标准。然而，其固有的顺序计算特性导致训练速度缓慢。前馈和卷积架构最近在某些序列建模任务（例如机器翻译）中表现出更优异的性能，并且它们还具有并发处理序列中所有输入的优势，从而易于并行化并加快训练速度。尽管如此，流行的前馈序列模型（例如 Transformer）在许多循环模型能够轻松处理的简单任务中泛化能力不足，例如字符串复制，甚至当字符串或公式长度超过训练时观察到的长度时，简单的逻辑推理也难以胜任。我们提出了 Universal Transformer (UT)，这是一种时间并行自注意力循环序列模型，可以看作是 Transformer 模型的泛化，并解决了上述问题。Universal Transformer 结合了 Transformer 等前馈序列模型的并行化能力和全局感受野，以及循环神经网络（RNN）的循环归纳偏置。我们还添加了动态的逐位置停机机制，并发现它提高了多个任务的准确率。与标准 Transformer 不同，在某些假设下，UT 可以被证明是图灵完备的。我们的实验表明，UT 在广泛的算法和语言理解任务中都优于标准 Transformer，包括具有挑战性的 LAMBADA 语言建模任务（UT 在该任务中取得了新的最佳结果）以及机器翻译任务（UT 在 WMT14 英德数据集上比 Transformer 提高了0.9的BLEU值）。

1.INTRODUCTION

近年来，卷积神经网络和全注意力前馈架构（例如 Transformer）已成为循环神经网络（RNN）在序列建模任务（尤其是机器翻译）中的可行替代方案。这些时间并行架构解决了 RNN 的一个重大缺陷，即其固有的顺序计算方式阻碍了输入序列元素间的并行化，同时还解决了序列长度增加时梯度消失的问题。Transformer 模型尤其依赖于自注意力机制来计算输入和输出中符号的一系列上下文信息向量空间表示，然后利用这些表示来预测后续符号的分布，从而逐个符号地预测输出序列。这种机制不仅易于并行化，而且由于每个符号的表示都直接受到所有其他符号表示的影响，因此在整个序列上形成了一个有效的全局感受野。这与卷积神经网络等架构形成鲜明对比，后者通常只有有限的感受野。

值得注意的是，Transformer 模型由于其固定的层堆叠结构，放弃了循环神经网络（RNN）学习迭代或递归变换的归纳偏置。我们的实验表明，这种归纳偏置对于多种复杂度的算法和语言理解任务至关重要：与神经图灵机、神经 GPU 或堆叠 RNN 等模型相比，Transformer 模型对训练过程中未遇到的输入长度泛化能力较差。

本文提出了一种 Universal Transformer (UT) 模型，它是一种时间并行循环自注意力序列模型，可以看作是 Transformer 模型的推广，在理论能力和各种具有挑战性的序列到序列任务上均取得了更优异的成果。UT 模型结合了 Transformer 等前馈序列模型的并行性和全局感受野，以及 RNN 的循环归纳偏置，这使其更适合于一系列算法和自然语言理解相关的序列到序列问题。顾名思义，与标准 Transformer 模型不同，在某些假设条件下，UT 模型可以被证明是图灵完备的（或“计算通用”，如第4节所示）。

在每个循环步骤中，通用 Transformer 使用自注意力机制并行迭代地细化序列中所有符号的表示，随后进行一次变换（所有位置和时间步共享），该变换由深度可分离卷积或位置全连接层组成（见图1）。我们还添加了动态的逐位置停止机制，允许模型动态地为每个符号选择所需的细化步骤数，并首次证明这种条件计算机制确实可以提高几个较小的结构化算法和语言推理任务的准确率（尽管在机器翻译任务上略有下降）。

我们强大的实验结果表明，UT 在广泛的任务中均优于 Transformer 和 LSTM。新增的循环机制显著提升了机器翻译的性能，UT 在该任务中优于标准 Transformer。在多个算法任务和 bAbI 语言理解任务的实验中，UT 也始终显著优于 LSTM 和标准 Transformer。此外，在极具挑战性的 LAMBADA 文本理解数据集上，采用动态停机机制的 UT 达到了新的最佳水平。

2.MODEL DESCRIPTION

2.1 THE UNIVERSAL TRANSFORMER

Universal Transformer（UT；见图 2）基于大多数神经序列到序列模型中常用的编码器-解码器架构。UT 的编码器和解码器分别通过将循环神经网络应用于输入和输出序列中每个位置的表示来实现。然而，与大多数循环神经网络在序列数据上的应用不同，UT 并非对序列中的位置进行循环，而是对每个位置的向量表示进行连续的迭代（即“深度”）。换句话说，UT 的计算量不受序列中符号数量的限制，而仅受每个符号表示的迭代次数的限制。

在每个循环时间步中，每个位置的表示都会同时（并行）进行修正，分为两个子步骤：首先，使用自注意力机制在序列中的所有位置之间交换信息，从而为每个位置生成一个向量表示，该表示包含了前一个时间步所有其他位置的表示信息。然后，将一个转移函数（位置和时间共享）独立地应用于自注意力机制的输出。由于循环转移函数可以应用任意次数，这意味着 UT 可以具有可变的深度（每个符号的处理步骤数）。关键在于，这与大多数流行的神经序列模型（包括 Transformer 或深度 RNN）形成鲜明对比，后者由于使用了固定的层堆叠而具有恒定的深度。接下来，我们将更详细地描述编码器和解码器。

ENCODER。给定一个长度为 $m$ 的输入序列，我们首先创建一个矩阵，该矩阵的行被初始化为序列中每个位置符号的 $d$ 维嵌入 $H^0 ∈\mathbb R^{m×d}$ 。然后，UT 迭代地并行计算步骤 $t$ 中所有 $m$ 个位置的表示 $H_t$ ，方法是应用 Vaswani et al. (2017) 提出的多头点积自注意力机制，并随后应用一个循环转移函数。我们还在每个功能块周围添加残差连接，并应用 dropout 和层归一化（参见图 2 的简化图，以及附录 A 中的图 4 的完整模型）。

更具体地说，我们使用缩放点积注意力机制，它将查询 $Q$ 、键 $K$ 和值 $V$ 组合在一起，如下所示。

ATTENTION(Q,K,V)=SOFTMAX(\frac{QK^T}{\sqrt{d}})V,\tag{1}

其中 $d$ 是 $Q$ 、 $K$ 和 $V$ 的列数。我们使用 (Vaswani et al., 2017) 中引入的具有 $k$ 个头的多头版本。

\text{MULTI-HEAD-SELF-ATTENTION}(H^t)=CONCAT(head_1,...,head_k)W^O\tag{2}

where~head_i=ATTENTION(H^tW^Q_i,H^tW^k_i,H^tW^V_i)\tag{3}

我们使用学习到的参数矩阵 $W^Q ∈\mathbb R^{d×d/k}$ 、 $W^K ∈\mathbb R^{d×d/k}$ 、 $W^V ∈\mathbb R^{d×d/k}$ 和 $W^O ∈\mathbb R^{d×d}$ ，通过仿射投影将状态 $H_t$ 映射到查询、键和值。

在步骤 $t$ ，UT 计算所有 $m$ 个输入位置的修正表示 $H^t∈\mathbb R^{m×d}$ ，如下所示。

H^t=LAYERNORM(A^t+TRANSITION(A^t))\tag{4}

where~A^t=LAYERNORN((H^{t-1}+P^t)+\text{MULTIHEADSELFATTENTION}(H^{t-1}+P^t)),\tag{5}

其中 $LAYERNORM()$ 在 Ba et al. (2016) 中定义， $TRANSITION()$ 和 $P^t$ 将在下文讨论。

根据任务的不同，我们使用两种不同的转移函数之一：可分离卷积或全连接神经网络，该神经网络由两个仿射变换之间的单个修正线性激活函数组成，按位置应用，即分别应用于 $A^t$ 的每一行。

上述 $P^t ∈ \mathbb R^{m×d}$ 是固定的、恒定的二维（位置，时间）坐标嵌入，是通过计算(Vaswani et al., 2017) 中定义的正弦位置嵌入向量得到的，其中位置 $1≤i≤m$ 和时间步长 $1≤t≤T$ 分别针对每个向量维度 $1≤j ≤d$ 进行计算，然后求和：

P^t_{i,2j}=sin(i/10000^{2j/d})+sin(t/10000^{2j/d})\tag{6}

P^t_{i,2j+1}=cos(i/10000^{2j/d})+cos(t/10000^{2j/d}).\tag{7}

经过 $T$ 步（每步并行更新输入序列的所有位置）后，Universal Transformer 编码器的最终输出是一个 $d$ 维向量表示矩阵 $H^T ∈\mathbb R^{m×d}$ ，代表输入序列的 $m$ 个符号。

DECODER。解码器与编码器共享相同的基本循环结构。然而，在自注意力函数之后，解码器还会使用公式 2 中的多头点积注意力函数来关注输入序列中每个位置的最终编码器表示 $H^T$ ，但查询 $Q$ 来自对解码器表示的投影，键和值（K 和 V）来自对编码器表示的投影（此过程类似于标准注意力）。

与 Transformer 模型类似，UT 模型也是自回归的。它使用 teacher-forcing 训练法进行训练，在生成阶段，它一次生成一个符号，解码器则消耗之前生成的输出位置。在训练过程中，解码器的输入是目标输出，并向右平移一个位置。解码器的自注意力分布被进一步掩蔽，使得模型只能关注任何预测符号左侧的位置。最后，通过将最终解码器状态到输出词表大小 $V$ 的仿射变换 $O ∈ \mathbb R^{d×V}$ ，然后应用 softmax 函数，得到每个符号的目标分布，该 softmax 函数产生一个 (m×V) 维的输出矩阵，并对其行进行归一化。

p(y_{pos}|y_{[1:pos-1]},H^T)=SOFTMAX(OH^T)\tag{8}

为了从模型中生成结果，编码器首先针对给定的输入序列运行一次。然后，解码器重复运行，消耗所有已生成的符号，并在每次迭代中为下一个输出位置的符号生成一个基于词表的额外分布。最后，我们通常会采样或选择概率最高的符号作为下一个符号。

2.2 DYNAMIC HALTING

在序列处理系统中，某些符号（例如某些单词或音素）通常比其他符号更具歧义性。因此，为这些更具歧义性的符号分配更多处理资源是合理的。自适应计算时间（ACT）是一种机制，它基于模型在每一步预测的标量停机概率，动态调节标准循环神经网络中处理每个输入符号所需的计算步骤数（称为“思考时间”）。

受 Universal Transformers 的解释启发，即并行地将自注意力 RNN 应用于序列中的所有位置，我们也为每个位置（即每个符号的自注意力 RNN；详见附录C）添加了一个动态的 ACT 停止机制。一旦每个符号的循环块停止，其状态就会被复制到下一步，直到所有块都停止，或者达到最大步数。编码器的最终输出就是以这种方式生成的最后一层表示。

3.EXPERIMENTS AND ANALYSIS

我们对 Universal Transformer 在多种算法和语言理解任务以及机器翻译任务上进行了评估。附录 D 中将更详细地描述这些任务和数据集。