作者/机构: Mostafa Dehghani (University of Amsterdam), Stephan Gouws (DeepMind), Oriol Vinyals (DeepMind), Jakob Uszkoreit (Google Brain), Łukasz Kaiser (Google Brain)
卷积和全注意力前馈架构,如 Transformer【31,Attention is all you need,2017,CoRR】,已成为循环神经网络 (RNN) 在一系列序列建模任务(特别是机器翻译)中的可行替代方案。这些“时间上并行”的架构解决了 RNN 的一个显著缺点,即其固有的顺序计算阻碍了输入序列元素间的并行化,同时仍能解决序列变长时的梯度消失问题。Transformer 模型完全依赖自注意力机制来计算其输入和输出中符号的一系列上下文感知向量空间表示。这种机制不仅易于并行化,而且由于每个符号的表示都直接受到所有其他符号表示的影响,从而在整个序列上实现了有效的全局感受野。
然而,具有固定层堆叠的 Transformer 放弃了 RNN 对学习迭代或递归变换的归纳偏置。实验表明,这种归纳偏置对于几种不同复杂度的算法和语言理解任务可能至关重要。与神经图灵机(Neural Turing Machine)【13,Neural turing machines,2014,CoRR】、神经 GPU(Neural GPU)【18,Neural GPUs learn algorithms,2016,ICLR】或堆栈 RNN(Stack RNNs)【17,Inferring algorithmic patterns with stack-augmented recurrent nets,2015,NIPS】等模型相比,Transformer 在处理训练期间未见过的输入长度时泛化能力不佳。
本文介绍通用 Transformer (Universal Transformer, UT),这是一种时间上并行的循环自注意力序列模型,可以看作是 Transformer 模型的泛化,它解决了上述问题,并带来了更强的理论能力和在广泛挑战性序列到序列任务上的改进结果。UT 结合了 Transformer 等前馈序列模型的并行性和全局感受野,以及 RNN 的循环归纳偏置,后者似乎更适合一系列算法和自然语言理解的序列到序列问题。
通用 Transformer 在每个循环步骤中,通过使用自注意力机制并行地迭代优化序列中所有符号的表示,然后应用一个在所有位置和时间步上共享的变换函数(深度可分离卷积或逐位置全连接层)。我们还增加了一个动态的逐位置暂停机制【12,Adaptive computation time for recurrent neural networks,2016,arXiv】,允许模型为每个符号动态选择所需的优化步数。实验首次证明,这种条件计算机制可以在几个较小的、结构化的算法和语言推理任务上提高准确性。与标准 Transformer 不同,在某些假设下,UT 可以被证明是图灵完备的。
我们的实验结果表明,UT 在广泛的任务中优于 Transformer 和 LSTM。在机器翻译中,UT 优于标准 Transformer;在几个算法任务和 bAbI 语言理解任务中,UT 始终显著优于 LSTM 和标准 Transformer;在具有挑战性的 LAMBADA 文本理解数据集上,带有动态暂停的 UT 达到了新的业界最佳水平。
UT 核心架构。通用 Transformer (UT; 见图 2) 基于大多数神经序列到序列模型中常用的流行编码器-解码器架构【29,Sequence to sequence learning with neural networks,2014,NIPS;4,Learning phrase representations using RNN encoder-decoder for statistical machine translation,2014,CoRR;31,Attention is all you need,2017,CoRR】。UT 的编码器和解码器都通过将一个循环神经网络分别应用于输入和输出序列中每个位置的表示来运作。然而,与大多数将循环神经网络应用于序列数据的应用不同,UT 不是在序列的位置上进行循环,而是在每个位置的向量表示的连续修正上进行循环(即在“深度”上循环)。换句话说,UT 的计算不受序列中符号数量的限制,而只受每个符号表示的修正次数限制。
并行循环更新机制。在每个循环时间步中,每个位置的表示都通过两个子步骤并行地进行修正:首先,使用自注意力机制在序列的所有位置之间交换信息,从而为每个位置生成一个由前一时间步所有其他位置的表示所通知的向量表示。然后,通过对自注意力机制的输出应用一个转换函数(在位置和时间上共享),在每个位置上独立进行。由于循环转换函数可以应用任意次数,这意味着 UT 可以具有可变的深度(每个符号的处理步数)。至关重要的是,这与大多数流行的神经序列模型(包括 Transformer 【31,Attention is all you need,2017,CoRR】或深度 RNN)形成对比,后者由于应用了固定的层堆叠而具有恒定的深度。我们现在更详细地描述编码器和解码器。
编码器。对于长度为 m 的输入序列,我们从一个矩阵开始,其行被初始化为序列中每个位置符号的 d 维嵌入 $H^0 \in \mathbb{R}^{m \times d}$。然后,UT 通过应用【31,Attention is all you need,2017,CoRR】中的多头点积自注意力机制,接着是一个循环转换函数,迭代地计算所有 m 个位置在步骤 t 的表示 $H^t$。我们还在每个功能块周围添加了残差连接,并应用了 dropout 和层归一化【27,Dropout: a simple way to prevent neural networks from overfitting,2014,Journal of Machine Learning Research;2,Layer normalization,2016,arXiv】(简化图见图 2,完整模型见附录 A 的图 4)。
注意力机制。我们使用缩放点积注意力,它结合了查询 Q、键 K 和值 V,如下所示:
$$\text{ATTENTION}(Q,K,V)=\text{SOFTMAX}\left(\frac{QK^{T}}{\sqrt{d}}\right)V,$$其中 d 是 Q、K 和 V 的列数。我们使用【31,Attention is all you need,2017,CoRR】中介绍的 k 个头的多头版本,
$$ \text{MULTIHEADSELFATTENTION}(H^t) = \text{CONCAT}(\text{head}_1, ..., \text{head}_{\text{k}})W^O $$ $$\text{where head}_{\text{i}} = \text{ATTENTION}(H^t W_i^Q, H^t W_i^K, H^t W_i^V)$$并且我们使用学习到的参数矩阵 $W^Q \in \mathbb{R}^{d \times d/k}$、$W^K \in \mathbb{R}^{d \times d/k}$、$W^V \in \mathbb{R}^{d \times d/k}$ 和 $W^O \in \mathbb{R}^{d \times d}$ 通过仿射投影将状态 $H^t$ 映射到查询、键和值。
状态更新。在步骤 t,UT 随后为所有 m 个输入位置计算修正后的表示 $H^t \in \mathbb{R}^{m \times d}$,如下所示:
$$H^t = \text{LAYERNORM}(A^t + \text{TRANSITION}(A^t))$$$\text{where } A^{t} = \text{LAYERNORM}((H^{t-1} + P^{t}) + \text{MULTIHEADSELFATTENTION}(H^{t-1} + P^{t}))$
其中 LAYER NORM() 在【2,Layer normalization,2016,arXiv】中定义,TRANSITION() 和 $P^t$ 将在下面讨论。
转换函数。根据任务的不同,我们使用两种不同的转换函数之一:要么是可分离卷积【5,Xception: Deep learning with depthwise separable convolutions,2016,arXiv】,要么是一个全连接神经网络,该网络由两个仿射变换之间的一个 ReLU 激活函数组成,并逐位置应用,即独立地应用于 $A^t$ 的每一行。
位置和时间嵌入。上述的 $P^t \in \mathbb{R}^{m \times d}$ 是固定的、常数的二维(位置,时间)坐标嵌入,通过为位置 $1 \le i \le m$ 和时间步 $1 \le t \le T$ 分别计算【31,Attention is all you need,2017,CoRR】中定义的正弦位置嵌入向量,并对每个向量维度 $1 \le j \le d$ 进行求和得到:
$$P_{i, 2 j}^t=\sin \left(i / 10000^{2 j / d}\right)+\sin \left(t / 10000^{2 j / d}\right)$$ $$P_{i, 2j+1}^t = \cos(i/10000^{2j/d}) + \cos(t/10000^{2j/d}).$$编码器最终输出。经过 T 个步骤(每个步骤并行更新输入序列的所有位置)后,通用 Transformer 编码器的最终输出是输入序列 m 个符号的 d 维向量表示矩阵 $H^T \in \mathbb{R}^{m \times d}$。
解码器。解码器与编码器共享相同的基本循环结构。然而,在自注意力函数之后,解码器还额外使用公式 2 中的多头点积注意力函数关注输入序列中每个位置的最终编码器表示 $H^T$,但其查询 Q 是通过投射解码器表示得到的,而键和值(K 和 V)是通过投射编码器表示得到的(这个过程类似于标准的注意力机制【3,Neural machine translation by jointly learning to align and translate,2014,CoRR】)。
自回归解码过程。与 Transformer 模型一样,UT 是自回归的【11,Generating sequences with recurrent neural networks,2013,CoRR】。在训练时使用教师强制(teacher-forcing),在生成时它一次产生一个符号,解码器消费先前产生的输出位置。在训练期间,解码器的输入是目标输出,向右移动一个位置。解码器的自注意力分布被进一步掩码,以便模型只能关注任何预测符号左侧的位置。最后,通过将最终解码器状态应用一个到输出词汇大小 V 的仿射变换 $O \in \mathbb{R}^{d \times V}$,然后是一个 softmax,得到每个符号的目标分布,这产生一个在行上归一化的 $(m \times V)$ 维输出矩阵:
$$p(y_{pos}|y_{[1:pos-1]},H^{T})=\mathrm{SOFTMAX}(OH^{T})^{1}$$生成过程。要从模型生成,编码器首先对条件输入序列运行一次。然后解码器重复运行,消费所有已生成的符号,同时在每次迭代中为下一个输出位置的符号生成一个额外的词汇表分布。我们然后通常采样或选择概率最高的符号作为下一个符号。
动态暂停的动机。在序列处理系统中,某些符号(例如一些单词或音素)通常比其他符号更具歧义性。因此,为这些更具歧义性的符号分配更多的处理资源是合理的。自适应计算时间(Adaptive Computation Time, ACT)【12,Adaptive computation time for recurrent neural networks,2016,arXiv】是一种机制,用于根据模型在每一步预测的标量暂停概率,动态调节处理每个输入符号所需的计算步骤数(称为“思考时间”),这在标准的循环神经网络中应用。
将 ACT 应用于 UT。受到将通用 Transformer 解释为并行应用于序列中所有位置的自注意力 RNN 的启发,我们还为每个位置添加了一个动态的 ACT 暂停机制(即,为每个符号的自注意力 RNN 添加;更多细节见附录 C)。一旦每个符号的循环块暂停,其状态就会被简单地复制到下一步,直到所有块都暂停,或者我们达到最大步数。编码器的最终输出就是以这种方式生成的最后一层表示。
与 Transformer 的关系。当运行固定步数时,通用 Transformer 等同于一个在所有层之间参数共享的多层 Transformer。这部分类似于 Recursive Transformer,后者在其自注意力层的深度上共享权重【14,Hyperbolic attention networks,2018,arXiv】。然而,由于每个符号的循环转换函数可以应用任意次数,另一种可能更具信息量的方式来描述 UT 是将其视为一个并行 RNN 块(每个符号一个,参数共享),这些 RNN 并发地演化每个符号的隐藏状态,这些状态在每一步通过关注前一步的隐藏状态序列来生成。
与 RNN、Neural GPU 和 NTM 的关系。通过这种方式,它与诸如神经 GPU(Neural GPU)【18,Neural GPUs learn algorithms,2016,ICLR】和神经图灵机(Neural Turing Machine)【13,Neural turing machines,2014,CoRR】等架构相关。UT 因此保留了原始前馈 Transformer 模型吸引人的计算效率,但增加了 RNN 的循环归纳偏置。此外,使用动态暂停机制,UT 可以根据输入数据选择处理步骤的数量。通用 Transformer 与其他序列模型之间的联系从其架构中显而易见:如果我们将循环步骤限制为一步,它就变成了一个 Transformer。但更有趣的是考虑通用 Transformer 与 RNN 和其他在时间维度上进行循环的网络之间的关系。表面上看,这些模型可能因为也是循环的而显得密切相关。但有一个关键的区别:像 RNN 这样的时间循环模型在循环步骤中无法访问内存。这使得它们在计算上更类似于自动机,因为在循环部分可用的唯一内存是一个固定大小的状态向量。而 UT 则可以关注整个前一层,使其能够在循环步骤中访问内存。
图灵完备性。在有足够内存的情况下,通用 Transformer 是计算上通用的——即它属于可以用来模拟任何图灵机的模型类别,从而解决了标准 Transformer 模型的一个缺点。除了理论上的吸引力,我们的结果表明,这种增加的表达能力还导致在几个具有挑战性的序列建模任务上准确性的提高。这弥合了在诸如机器翻译等大规模任务上具有竞争力的实用序列模型与诸如神经图灵机或神经 GPU 等计算上通用的模型【13,Neural turing machines,2014,CoRR;18,Neural GPUs learn algorithms,2016,ICLR】之间的差距,后者可以使用梯度下降进行训练以执行算法任务。
证明思路。为了证明这一点,我们可以将一个神经 GPU 归约到一个通用 Transformer。忽略解码器并将自注意力模块(即带有残差连接的自注意力)参数化为恒等函数,我们假设转换函数是一个卷积。如果我们现在将总循环步数 T 设置为等于输入长度,我们就精确地得到了一个神经 GPU。请注意,最后一步是通用 Transformer 与香草 Transformer 的关键区别所在,后者的深度不能随输入大小动态扩展。通用 Transformer 和神经图灵机之间也存在类似的关系,后者的每步单个读/写操作可以由通用 Transformer 的全局、并行的表示修正来表达。然而,与这些仅在算法任务上表现良好的模型相比,通用 Transformer 在诸如 LAMBADA 和机器翻译等现实的自然语言任务上也取得了有竞争力的结果。
与端到端记忆网络的关系。另一个相关的模型架构是端到端记忆网络(end-to-end Memory Networks)【28,End-to-end memory networks,2015,NIPS】。然而,与端到端记忆网络相比,通用 Transformer 使用的内存对应于与其输入或输出的单个位置对齐的状态。此外,通用 Transformer 遵循编码器-解码器配置,并在大规模序列到序列任务中取得了有竞争力的性能。
数据集:
模型:
硬件配置:
软件配置:
tensor2tensor 库【32,Tensor2tensor for neural machine translation,2018,CoRR】中提供。本文介绍了通用 Transformer (Universal Transformer),它是 Transformer 模型的一种泛化,扩展了其理论能力,并在广泛的挑战性序列建模任务(如语言理解和各种算法任务)上取得了业界最佳结果,从而解决了标准 Transformer 的一个关键缺陷。
通用 Transformer 将以下关键特性结合在一个模型中:
- 权重共享: 借鉴 CNN 和 RNN 中权重共享的直觉,我们为 Transformer 扩展了一种简单的权重共享形式,它在归纳偏置和模型表达能力之间取得了有效的平衡,这在小型和大型实验中都得到了广泛证明。
- 条件计算: 为了构建一个计算上通用的机器,我们为通用 Transformer 配备了通过最近引入的机制来暂停或继续计算的能力,这与固定深度的通用 Transformer 相比显示出更强的结果。
我们对时间并行序列模型的最新发展感到兴奋。通过在处理深度上增加计算能力和循环性,我们希望在本文提出的基础通用 Transformer 之上的进一步改进将帮助我们构建更强大、数据效率更高、并能超越当前最先进水平进行泛化的学习算法。
计算能力的关键差异。在计算能力方面,Transformer 和通用 Transformer 之间的关键区别在于顺序计算步骤的数量(即深度)。虽然标准 Transformer 执行的总操作数与输入大小成比例,但顺序操作的数量是恒定的,与输入大小无关,仅由层数决定。假设精度有限,这一特性意味着标准 Transformer 不可能是计算上通用的。然而,当选择一个步数作为输入长度的函数时,通用 Transformer 不受此限制。请注意,这与是否采用自适应计算时间无关,但确实假设步数是非恒定的,即使可能是确定性的。在训练后动态改变步数是由通用 Transformer 中跨顺序计算步骤共享权重实现的。
一个直观的例子。一个直观的例子是那些需要顺序处理每个输入元素的函数。在这种情况下,对于任何给定的深度选择 T,都可以构造一个长度为 N > T 的输入序列,该序列无法被标准 Transformer 正确处理。然而,通过适当的、依赖于输入长度的顺序步骤选择,通用 Transformer、RNN 或神经 GPU 可以执行这样的函数。
实现逻辑。我们基于 ACT 【12,Adaptive computation time for recurrent neural networks,2016,arXiv】在 TensorFlow 中实现了动态暂停。在带有动态暂停的 UT 的每一步中,我们都会获得暂停概率、余数、到目前为止的更新次数和前一个状态(都初始化为零),以及一个介于 0 和 1 之间的标量阈值(一个超参数)。然后,我们为每个位置计算新状态,并根据每个位置的状态计算新的逐位置暂停概率。UT 随后决定对某些超过阈值的位置暂停,并更新其他位置的状态,直到模型对所有位置都暂停或达到预定义的最多步数。
# 当此谓词为 FALSE 时,While-loop 停止
# 即 all( (probability < threshold) & (counter < max_steps) ) 都为 false
def should_continue(u0, u1, halting_probability, u2, n_updates, u3):
return tf.reduce_any(
tf.logical_and(
tf.less(halting_probability, threshold),
tf.less(n_updates, max_steps)))
# Do while loop iterations until predicate above is false
(remainder, n_updates, new_state) = tf.while_loop(
should_continue, ut_with_dynamic_halting, (state,
step, halting_probability, remainders, n_updates, previous_state))
每步计算细节。以下代码展示了每一步的计算过程:
def ut_with_dynamic_halting(state, step, halting_probability,
remainders, n_updates, previous_state):
# 根据状态计算概率
p = common_layers.dense(state, 1, activation=tf.nn.sigmoid,
use_bias=True)
# 尚未暂停的输入的掩码
still_running = tf.cast(
tf.less(halting_probability, 1.0), tf.float32)
# 在此步骤暂停的输入的掩码
new_halted = tf.cast(
tf.greater(halting_probability + p * still_running, threshold),
tf.float32) * still_running
# 尚未暂停且在此步骤也未暂停的输入的掩码
still_running = tf.cast(
tf.less_equal(halting_probability + p * still_running,
threshold), tf.float32) * still_running
# 将此步骤的暂停概率添加到尚未暂停的输入的
# 暂停概率中
halting_probability += p * still_running
# 计算在此步骤暂停的输入的余数
remainders += new_halted * (1 - halting_probability)
# 将余数添加到在此步骤暂停的输入中
halting_probability += new_halted * remainders
# 对所有仍在运行的输入增加 n_updates
n_updates += still_running + new_halted
# 计算要应用于新状态和输出的权重:
# 当输入已暂停时为 0,
# 当输入尚未暂停时为 p,
# 当它在此步骤暂停时为余数。
update_weights = tf.expand_dims(p * still_running +
new_halted * remainders, -1)
# 对状态应用转换
transformed_state = transition_function(self_attention(state))
# 对未暂停的输入插值转换后的状态和前一个状态
new_state = ((transformed_state * update_weights) +
(previous_state * (1 - update_weights)))
step += 1
return (transformed_state, step, halting_probability,
remainders, n_updates, new_state)
这里,我们提供有关 bAbI、主谓一致性、LAMBADA 语言建模和学习执行 (LTE) 任务的一些额外细节。
任务描述。bAbI 问答数据集【33,Towards ai-complete question answering: A set of prerequisite toy tasks,2015,arXiv】由 20 个不同的合成任务组成。其目的是让每个任务测试语言理解和推理的一个独特方面,包括:从故事中的支持事实进行推理、回答真/假类型的问题、计数、理解否定和不确定知识、理解共指、时间推理、位置和大小推理、寻路以及理解动机的能力。
数据集版本。数据集有两个版本,一个有 1k 个训练样本,另一个有 10k 个。对于一个模型来说,数据高效地仅使用 1k 个训练样本就取得好结果是很重要的。此外,最初的想法是,单个模型应该在所有任务上进行评估(而不是针对每个任务进行调整),这就是表1 和附录 E 中表格里的 “train joint” 设置。
任务描述。主谓一致性是预测英语句子中主语和动词之间数的一致性的任务。成功完成此任务是模型能够学习近似句法结构的有力指标,因此被 Linzen 等人【22,Assessing the ability of lstms to learn syntax-sensitive dependencies,2016,Transactions of the Association for Computational Linguistics】提议作为评估不同模型捕捉自然语言中层次结构能力的代理任务。
实验设置。Linzen 等人【22,Assessing the ability of lstms to learn syntax-sensitive dependencies,2016,Transactions of the Association for Computational Linguistics】提出了两种实验设置来训练模型完成此任务:1) 使用语言建模目标进行训练,即下一个词预测;2) 作为二元分类,即给定句子预测动词的数。在本文中,我们使用语言建模目标,这意味着我们为模型提供隐式监督,并根据正确形式的动词与不正确形式的动词的排名准确性进行评估。
难度级别。在这个任务中,为了设置不同的难度级别,使用了“一致性吸引子”,即一个或多个与主语数相反的介入名词,目的是迷惑模型。在这种情况下,模型需要正确识别与给定动词相对应的句法主语的中心词,并忽略介入的吸引子,以便预测该动词的正确形式。以下是此任务的一些示例,其中主语和相应的动词用粗体表示,一致性吸引子用下划线表示:
- 无吸引子: The boy smiles.
- 一个吸引子: The number of <u>men</u> is not clear.
- 两个吸引子: The ratio of <u>men</u> to <u>women</u> is not clear.
- 三个吸引子: The ratio of <u>men</u> to <u>women</u> and <u>children</u> is not clear.
任务描述。LAMBADA 任务【23,The lambada dataset: Word prediction requiring a broad discourse context,2016,ACL】是一个广泛上下文的语言建模任务。在这个任务中,给定一个叙事段落,目标是预测段落中最后一个句子(目标句子)的最后一个词(目标词)。这些段落是经过专门挑选的,使得人类受试者在接触到长段落时能够轻松猜出最后一个词,但如果他们只看到目标词之前的目标句子则不能。以下是数据集中的一个样本:
- 上下文: “是的,我以为我要失去这个孩子了。”“我也很害怕,”他说道,真诚充满他的眼睛。“你是吗?”“是的,当然。你为什么还要问?”“这个孩子并不完全是计划中的。”
- 目标句子: “你真的认为我会希望你...吗?”
- 目标词: 流产
评估。LAMBADA 任务在于给定整个段落(即上下文加上目标句子)来预测目标词。还提供了一个“控制集”,该控制集是通过随机抽样与用于构建 LAMBADA 的段落具有相同形状和大小的段落而构建的,但未经过任何过滤。控制集用于在测试 LAMBADA 任务之前评估模型在标准语言建模上的表现,从而确保后者表现不佳不能简单地归因于语言建模能力差。该任务在两种设置下进行评估:作为语言建模(标准设置)和作为阅读理解。在前者(更具挑战性)的情况下,模型仅在训练数据上进行下一个词预测的训练,并在测试时对目标词进行评估(即模型被训练来预测所有词,而不仅仅是具有挑战性的目标词)。在本文中,我们报告了通用 Transformer 在两种设置下的结果。
任务描述。LTE 是一组任务,用于指示模型学习执行计算机程序的能力,由 Zaremba & Sutskever【35,Learning to execute,2015,CoRR】提出。这些任务包括两个子集:1) 程序评估任务(程序、控制和加法),旨在评估模型理解数值运算、if 语句、变量赋值、操作的组合性等的能力,以及 2) 记忆任务(复制、加倍和反转)。
难度参数。程序评估任务的难度由其长度和嵌套来参数化。长度参数是程序中出现的整数的位数(因此整数从 [1, 长度] 中均匀选择),嵌套参数是允许我们将操作相互组合的次数。较高的嵌套值会产生具有更深解析树的程序。例如,这是一个用 length = 4 和 nesting = 3 生成的程序:
- 输入: j=8584 for x in range(8): j+=920 b=(1500+j) print((b+7567))
- 目标: 25011
我们为几个示例展示了 bAbI 任务上注意力分布的可视化。注意力权重的可视化是基于故事中所有事实和问题的不同头在不同时间步上的情况。左侧的不同颜色条表示基于不同头(总共 4 个头)的注意力权重。