Question

1. 像 Transformer 这样的 parallel-in-time 架构解决了RNN 的一个重要缺陷：RNN 固有的序列计算sequential computation 阻止了输入序列中各元素的并行化。同时 Transformer 也解决了序列较长时梯度消失的问题。Transformer 模型完全依靠自注意力机制。这种机制不仅可以直接并行化，而且由于每个symbol 的 representation 也直接访问到所有其他symbol 的 representation ，这导致了整个序列中有效的全局感受野 global receptive field 。

   然而，值得注意的是，Transformer 具有固定数量的层，放弃了RNN 对学习迭代式转换 iterative transformation或递归式转换recursive transformation的归纳偏置 inductive bias。我们的实验表明：这种归纳偏置可能对一些算法和语言理解任务至关重要，Transformer 不能很好地泛化到训练期间没有遇到的输入长度。

   在论文 《Universal Transformer》 中，作者介绍了Universal Transformer: UT，这是一个parallel-in-time 的递归的自注意力序列模型，它可以作为Transformer 模型的泛化，在广泛的具有挑战性的sequence-to-sequence 任务上产生提升的理论能力和改进的结果。 Universal Transformer结合了像Transformer这样的前馈序列模型的并行能力 parallelizability 和全局感受野、以及RNN的循环归纳偏置，这似乎更适合一系列的算法和自然语言理解的sequence-to-sequence 问题。顾名思义，与标准的 Transformer 相比，在某些假设下，Universal Transformer 可以被证明是图灵完备的。

   在每个recurrent step中，Universal Transformer 使用自注意力机制并行地refine 序列中所有 symbol 的 representation ，然后跟着一个深度可分离卷积depth-wise separable convolution 或 position-wise fully-connected layer （在所有position 和time-step 中共享），如下图所示。作者还添加了一个动态的逐位置停机机制 per-position halting mechanism ，允许模型动态选择每个symbol 所需的refinement step 数量，并首次表明这样的条件计算机制conditional computation mechanism 事实上可以提高几个较小的结构化的算法和语言推理任务的准确性（尽管它在机器翻译任务上的结果略有下降）。

   论文强有力的实验结果表明：Universal Transformer 在广泛的任务中的表现优于 Transformer 和LSTM 。

   

2. 相关工作：

   • 当运行一个fixed 数量的 step 时，Universal Transformer 相当于一个跨所有层共享参数的 multi-layer Transformer 。这在一定程度上类似于Recursive Transformer ，它在不同的深度上共享了self-attention layer 的权重。

     另一种可能更 informative 的描述，Universal Transformer 的方式是：作为一个并行的 RNN block （每个symbol 一个RNN block ，共享参数）来同时演化每个symbol 的 hidden state （通过在每个 step 关注前一个 step 的hidden state 序列来产生）。这样，它与 Neural GPU 和 Neural Turing Machine 等架构有关。因此，Universal Transformer 保留了原始 feed-forward Transformer 模型有吸引力的计算效率，但增加了RNN的recurrent inductive bias。此外，利用动态停机机制，Universal Transformer 可以根据输入数据自动选择 processing step 的数量。

   • Universal Transformer 和其他序列模型之间的联系从架构上看是显而易见的： 如果我们将recurrent step 限制为一个，它将是一个Transformer。

     但是，考虑 Universal Transformer 和 RNN 以及其他网络之间的关系更为有趣，在这些网络中，递归发生在时间维度上。表面上看，这些模型似乎密切相关，因为它们也是递归的。但有一个关键的区别：

     • 像RNN这样的time-recurrent模型不能在递归步骤中访问所有的记忆 memory （只能访问最近的一个记忆）。这使得它们在计算上更类似于自动机automata，因为在递归部分唯一可用的记忆是一个固定尺寸的状态向量fixed-size state vector。
     • 另一方面，Universal Transformer 可以关注整个previous layer，允许它在recurrent step 中访问所有的记忆。

   • 给定足够的内存，Universal Transformer 在计算上是通用的，即它可以模拟任何图灵机Turing machine，从而解决了标准Transformer模型的一个缺陷。

     除了在理论上有吸引力外，我们的结果表明：这种增加的表达性也导致了在几个具有挑战性的序列建模任务上的准确性的提高。这缩小了在机器翻译等大规模任务上与Neural Turing Machine或Neural GPU等模型之间的差距。

     为了说明这一点，我们可以将 Universal Transformer 简化为一个 Neural GPU。忽略decoder 并将 self-attention 模块（即带有残差连接的self-attention）简化为恒等映射，并且假设转移函数是卷积。如果我们将循环step的总数$T$$ T $ 设定为输入长度，我们得到的正是一个Neural GPU。

     Universal Transformer 和 Neural Turing Machine 之间也存在类似的关系。Neural Turing Machine 每个 step 的single read/write 操作可以由 Universal Transformer 的全局的、并行的representation revision来表达。

   • 另一个相关的模型架构是端到端记忆网络Memory Network 。然而，与端到端记忆网络相比，Universal Transformer 使用的 memory 对应于与其输入或输出的individual position 对齐的状态。此外，Universal Transformer 遵循 encoder-decoder 的配置，在大规模 sequence-to-sequence 的任务中实现了有竞争力的性能。

3. Universal Transformer将以下关键属性结合到一个模型中：

   • Weight sharing：我们用一种简单的权重共享形式来扩展Transformer ，从而在inductive bias 和模型表达力之间取得了有效的平衡。

     由于权重共享，所以 Universal Transformer 的模型大小要远远小于标准的 Transformer 。

   • Conditional computation：我们为Universal Transformer 配备了动态停机能力。与固定深度的 Universal Transformer 相比，它显示了更强的结果。

2.1 模型

1. Universal Transformer 是基于常用的encoder-decoder 架构，它的encoder 和decoder都是通过将一个 RNN 分别应用于输入和输出序列的每个位置的 representation 来操作的，如下图所示。然而，与常规 RNN 不同，Universal Transformer 不是在序列中的 position 上递归（沿着宽度方向），而是在每个 position 的 depth 上递归（沿着高度方向）。

   在每个 recurrent time-step，每个位置的 representation 在两个 sub-step 中被并行地修改：

   • 首先，使用自注意力机制在序列中的所有位置上交换信息，从而为每个位置生成一个 vector representation ，该 representation 是由前一个time step的所有其他位置的representation 融合而来。
   • 然后，在每个位置独立地对自注意力机制的输出应用一个转移函数（跨position 和 time 共享）。

   由于递归可以应用任何次数，这意味着Universal Transformer 可以有可变的深度（per-symbol processing step 的数量）。相比之下，Transformer 和 stacked RNN 只有固定的深度。

   

   更完整的架构如下图所示：

   

2.1.1 Encoder

1. 给定一个长度为$m$$ m $ 的输入序列，我们定义${\mathbf{H}}^{(0)}\in {\mathbb{R}}^{m\times d}$$ \mathbf H^{(0)} \in \mathbb R^{m\times d} $ 为 embedding 矩阵，它的第$i$$ i $ 行表示输入序列位置$i$$ i $ 的 symbol 的$d$$ d $ 维 embedding 向量${\overrightarrow{\mathbf{h}}}_i^{(0)}\in {\mathbb{R}}^d,0\le i<m$$ \mathbf{\vec h}_i^{(0)}\in \mathbb R^d, 0\le i\lt m $ 。 然后 Universal Transformer 通过应用Transformer的多头内积自注意力机制 multi-headed dot-product self-attention mechanism ，在 step$t$$ t $ 为所有$m$$ m $ 个位置并行地迭代计算 representation${\mathbf{H}}^{(t)}$$ \mathbf H^{(t)} $ ，然后是一个转移函数transition function 。我们还在这些功能块function block周围添加了残差连接，并应用了 dropout 和layer normalization 。

2. 更具体地说，我们使用缩放的内积注意力：

   $$\text{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\text{softmax}\left(\frac{\mathbf{Q}{\mathbf{K}}^{\mathrm{\top }}}{\sqrt{d}}\right)\mathbf{V}$$

   其中$\mathbf{Q}\in {\mathbb{R}}^{m\times d},\mathbf{K}\in {\mathbb{R}}^{m\times d},\mathbf{V}\in {\mathbb{R}}^{m\times d},\mathbf{Q}{\mathbf{K}}^{\mathrm{\top }}\in {\mathbb{R}}^{m\times m},\text{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})\in {\mathbb{R}}^{m\times d}$$ \mathbf Q\in \mathbb R^{m \times d}, \mathbf K\in \mathbb R^{m \times d}, \mathbf V\in \mathbb R^{m \times d}, \mathbf Q\mathbf K^\top\in \mathbb R^{m\times m}, \text{Attention}(\mathbf Q, \mathbf K, \mathbf V)\in \mathbb R^{m\times d} $ 。

   我们使用有$k$$ k $ 个头的多头自注意力版本：

   $$\begin{array}{c}\text{MultiHeadSelfAttn}({\mathbf{H}}^{(t)})=\text{Concat}({\overrightarrow{\text{head}}}_1,{\overrightarrow{\text{head}}}_2,\cdots ,{\overrightarrow{\text{head}}}_k){\mathbf{W}}^O\\ {\overrightarrow{\text{head}}}_i=\text{Attention}({\mathbf{H}}^{(t)}{\mathbf{W}}_i^Q,{\mathbf{H}}^{(t)}{\mathbf{W}}_i^K,{\mathbf{H}}^{(t)}{\mathbf{W}}_i^V)\end{array}$$

   其中：${\mathbf{W}}^Q\in {\mathbb{R}}^{d\times (d/k)},{\mathbf{W}}^K\in {\mathbb{R}}^{d\times (d/k)},{\mathbf{W}}^V\in {\mathbb{R}}^{d\times (d/k)},{\mathbf{W}}^O\in {\mathbb{R}}^{d\times d}$$ \mathbf W^Q\in \mathbb R^{d \times (d/k)}, \mathbf W^{K}\in \mathbb R^{d \times (d/k)}, \mathbf W^V\in \mathbb R^{d\times (d/k)}, \mathbf W^O\in \mathbb R^{d\times d} $ 。

3. 在 step$t$$ t $ ， Universal Transformer 然后计算所有$m$$ m $ 个输入位置的修正的 representation${\mathbf{H}}^{(t)}\in {\mathbb{R}}^{m\times d}$$ \mathbf H^{(t)}\in \mathbb R^{m\times d} $ ，即：

   $$\begin{array}{c}{\mathbf{H}}^{(t)}=\text{LayerNorm}({\mathbf{A}}^{(t)}+\text{Transition}\left({\mathbf{A}}^{(t)}\right))\\ {\mathbf{A}}^{(t)}=\text{LayerNorm}(({\mathbf{H}}^{(t-1)}+{\mathbf{P}}^{(t)})+\text{MultiHeadSelfAttn}({\mathbf{H}}^{(t-1)}+{\mathbf{P}}^{(t)}))\end{array}$$

   LayerNorm 作用在残差连接之后，而不是之前。

   其中：

   • Transition 为转移函数。根据任务，我们使用两种不同的转移函数之一：要么是可分离的卷积、要么是前馈神经网络（由两个线性投影以及中间的的单个 Relu 激活函数组成，逐位置应用）。

   • ${\mathbf{P}}^{(t)}\in {\mathbb{R}}^{m\times d}$$ \mathbf P^{(t)}\in \mathbb R^{m\times d} $ 为固定的、常量的、二维的（位置维度、时间维度）的 coordinate embedding，通过计算正弦的 position embedding 向量获得：

     $$\begin{array}{c}{\mathbf{p}}_i^{(t)}=(p_{i,1}^{(t)},\cdots ,p_{i,d}^{(t)})\in {\mathbb{R}}^d\\ p_{i,2j}^{(t)}=\sin \left(\frac{i}{{10000}^{2j/d}}\right)+\sin \left(\frac{t}{{10000}^{2j/d}}\right)\\ p_{i,2j+1}^{(t)}=\cos \left(\frac{i}{{10000}^{2j/d}}\right)+\cos \left(\frac{t}{{10000}^{2j/d}}\right)\end{array}$$

     $i$$ i $ 给出了位置维度的信息，$t$$ t $ 给出了时间维度的信息。

   经过$T$$ T $ 个 step （每个 step 并行地更新输入序列的所有位置），Universal Transformer encoder 的 final output 是${\mathbf{H}}^{(T)}\in {\mathbb{R}}^{m\times d}$$ \mathbf H^{(T)}\in \mathbb R^{m \times d} $ 。

2.1.2 Decoder

1. decoder 与 encoder 的基本递归结构相同。然而，在self-attention 函数之后，decoder 额外还使用多头内积注意力函数来关注 encoder 的 final representation${\mathbf{H}}^{(T)}$$ \mathbf H^{(T)} $ ，但query$\mathbf{Q}$$ \mathbf Q $ 是从 decoder representation 投影得到，而 key$\mathbf{K}$$ \mathbf K $ 和 value$\mathbf{V}$$ \mathbf V $ 是从 encoder final representation 投影得到。

   是否可以对 decoder 和 encoder 应用不同的$T$$ T $ ？理论上并没有约束强制要求 decoder 和 encoder 迭代相同的次数。

2. 像 Transformer 模型一样，Universal Transformer 是自回归的。

   • 在生成阶段，Universal Transformer 一次产生一个 symbol 的输出，其中 decoder 仅仅依赖已经产生的输出位置。
   • 在训练阶段，decoder 的目标输出就是输入左移一个位置。

   最后，per-symbol target distribution 是通过应用线性变换$\mathbf{O}\in {\mathbb{R}}^{d\times |\mathcal{V}|}$$ \mathbf O \in \mathbb R^{d\times |\mathcal V|} $ 从 final decoder state 映射到 output vocabulary$\mathcal{V}$$ \mathcal V $ ，然后通过 softmax 进行按行的归一化：

   $$p(y_{\text{pos}}\mid y_{[1:\text{pos}-1]},{\mathbf{H}}_{\text{encoder}}^{(T)})=\text{softmax}\left({\mathbf{H}}_{\text{decoder}}^{(T)}\mathbf{O}\right)\in {\mathbb{R}}^{n\times |\mathcal{V}|}$$

   其中：$n$$ n $ 为输出序列的长度。

   这里我们用${\mathbf{H}}_{\text{encoder}},{\mathbf{H}}_{\text{decoder}}$$ \mathbf H_\text{encoder}, \mathbf H_\text{decoder} $ 来区分 encoder 和 decoder 的最后一层的 representation 。$y_{[1:\text{pos}-1]},{\mathbf{H}}_{\text{encoder}}^{(T)}$$ y_{[1:\text{pos}-1]},\mathbf H^{(T)}_\text{encoder} $ 都是 decoder 的输入。

3. 为了从模型中生成：

   • 首先， encoder 要对输入序列运行一次前向传播。
   • 然后重复运行decoder 从而消耗所有已经生成的symbol ，同时为下一个输出位置的symbol 生成一个额外的分布。然后，我们通常对最高概率的symbol 进行采样从而作为next symbol （用于decoder 的下一轮 RNN 迭代）。

   如果是训练期间，那么 decoder 每一轮迭代都使用 decoder 的目标输出右移一位作为 decoder input 。但是在推断期间，decoder 的目标输出不存在，因此需要用到 decoder 在每一轮的输出。

2.1.3 动态停机

1. 在序列处理系统中，某些 symbol 通常比其他 symbol 更 ambiguous 。因此，将更多的处理资源分配给这些更ambiguous 的 symbol 是合理的。Adaptive Computation Time: ACT 是一种机制，用于动态调节标准 RNN 中处理每个 input symbol 所需的计算步（称为 ponder time ），它基于模型在每个step 预测的标量停机概率 halting probability 。

   受此启发，Universal Transformer 为每个位置（即每个 per-symbol self-attentive RNN）添加了一个动态 ACT 停机机制。一旦某个 symbol 的 recurrent block 停机，它的状态就被简单地复制到下一步，直到所有的block 都停机，或者我们达到最大的 step 数量$T$$ T $ 。然后，encoder 的final output 是以这种方式产生的最后一层 representation 。

2.2 实验

2.2.1 BABI 问答

1. bAbi question answering 数据集：由 20 个不同的任务组成，其目标是在给定一些英文句子的前提下回答一个问题，这些给定的句子编码了潜在的 multiple supporting facts 。

   这么做的目的是：通过要求对每个故事中呈现的语言事实进行某种类型的推理，来衡量各种形式的语言理解 language understanding 。

2. 为了对输入进行编码，与《Tracking the world state with recurrent entity networks》类似：

   • 首先，我们通过应用一个学到的乘性的 positional mask 到每个单词的 embedding ，并 sum 聚合所有单词的 embedding ，从而编码故事中的每个 fact 。

     因为每个 fact 包含多个单词，所以需要对 word embedding 进行聚合从而得到 fact emebdding 。

   • 然后，我们以同样的方式嵌入问题，然后用这些 fact embedding 和 question embedding 馈入 Universal Transformer。

3. 模型可以在每个任务上单独训练（train single）或在所有任务上联合训练（train joint）。下表总结了我们的结果。我们用不同的初始化运行了 10次，并根据验证集上的表现挑选出最佳模型。

   可以看到：Universal Transformer和带有动态停机的Universal Transformer 在所有任务中的平均误差和失败任务数方面都取得了SOTA 的结果。

   括号中的指标代表失败任务数，如 (10/20) 表示在 20 个任务中失败了 10 个。任务失败指的是错误率大于 5% 。

   

4. 为了更好地理解模型的工作，我们分析了这项任务的注意力分布和平均ACT ponder time 。下图为注意力权重在不同 step 上的可视化，左侧不同的颜色条表示不同 head 的注意力权重（一共4个head），即，列代表 head 、行代表句子。

   • 一个 supportive fact 的问题：

     

   • 两个 supportive fact 的问题：

     

     

   • 三个 supportive fact 的问题：

     

     

     

     

     

   结论：

   • 首先，我们观察到注意力分布开始时非常均匀，但在后面的step 中，围绕着回答每个问题所需的正确的supporting fact，注意力分布逐渐变得尖锐，这确实非常类似于人类如何解决这个任务。

   • 其次，通过动态停机我们观察到：需要三个supporting fact 的任务的所有位置的平均 ponder time（即per-symbol recurrent processing chain 的深度（$3.8\pm 2.2$$ 3.8 \pm 2.2 $ ）高于只需要两个supporting fact 的任务（$3.1\pm 1.1$$ 3.1 \pm 1.1 $ ），而这又高于只需要一个supporting fact 的任务（$2.3\pm 0.8$$ 2.3 \pm 0.8 $ ）。这表明，该模型根据回答问题所需的supporting fact 的数量来调整processing step的数量。

   • 最后，我们观察到：在只需要一个supporting fact 的任务中，不同位置的 ponder time 直方图与需要两个/三个supporting fact 的任务相比更加均匀。

     同样，在需要两个supporting fact 的任务与需要三个supporting fact 的任务也是如此。特别是对于需要三个 supporting fact 的任务，许多位置已经在第1步或第 2 步停机了，只有少数位置需要更多的step （如下图所示）。这一点特别有意思，因为在这种情况下，故事的长度确实要高得多，有更多不相关的fact ，而模型似乎成功地学会了以这种方式忽略这些fact 。

   

2.2.2 主谓一致

1. 主谓一致Subject-Verb Aggrement任务：使用语言建模的 training setup 来解决这个任务，即next word prediction 的目标，然后在测试时计算目标动词的 ranking accuracy 。该任务是衡量模型在自然语言句子中捕捉层级结构hierarchical structure的能力的一个代理。

   attractors 用于增加难度，它表示与主语相反数量opposite number 的中间名词intervening noun 的个数（旨在混淆模型）。如：给定句子 The keys __ to the cabinet ，训练期间的目标动词是 are，但是测试期间我们要求 are 的排名比 is 更高。

   

2.2.3 LAMBADA 语言建模

1. LAMBADA 任务：是一项语言建模任务，在给定前面 4 到 5 个句子的条件下预测一个缺失的目标词。

   该数据集是专门设计的，当显示完整的上下文时，人类能够准确地预测目标词，但当只显示出现该词的目标句时，则人类无法预测目标词。因此，它超越了语言建模，并测试了模型在预测目标词时纳入更长上下文的能力。

   该任务在两种情况下被评估：

   • 语言建模（标准设置，更具挑战性）：模型只是在训练数据上进行next-word prediction 的训练，并在测试时对目标词进行评估。

   • 和阅读理解：目标句子（减去最后一个词）被用作 query ，用于从上下文句子中选择目标词（即，检索式的方法）。

     请注意，目标词在81% 的时间里出现在上下文中，这使得这种设置更加简单。然而，在其余19% 的情况下，这个任务是不可能完成的（因为目标词不存在上下文中）。

2. 实验结果：Universal Transformer 在语言建模和阅读理解设置中都取得了SOTA 结果。

   我们 fixed Universal Transformer 结果使用了6个step （即，$T=6$$ T=6 $ ）。然而，带有动态停机的最佳 Universal Transformer 的平均step 数为$8.2\pm 2.1$$ 8.2 \pm 2.1 $ 。

   为了看看动态停机是否仅仅因为采取了更多的step 而做得更好，我们训练了两个 fixed Universal Transformer ，分别用$T=8$$ T=8 $ 和$T=9$$ T=9 $ （见最后两行）。有趣的是，与$T=6$$ T=6 $ 的模型相比，这两个模型取得了更好的结果，但并没有超过动态停机的 Universal Transformer 的表现。

   这使我们相信：动态停机可以作为模型的一个有用的正则器，通过激励一些input symbol 以较小的 step 数量，同时允许其他input symbol 以更多的计算。

   

2.2.4 算术任务

1. 我们在三个算术任务上训练 Universal Transformer ，即 Copy 、Reverse 和（整数）加法，都是在由十进制符号（0 ~ 9）组成的字符串上。在所有的实验中，我们在长度为40 的序列上训练模型，在长度为400的序列上进行评估。 我们使用从randomized offset 开始的位置来训练 Universal Transformer （训练样本需要选择从哪里开始切分，因为输入语料库可以视为一个超长的文本），以进一步鼓励模型学习position-relative 的转换。

   结果如下表所示， Universal Transformer 在所有三个任务上的表现都远远超过了 LSTM 和常规的 Transformer。

   

2.2.5 Learnung To Execute: LTE

1. 作为另一种 sequence-to-sequence learning 问题，我们评估表明模型学习执行计算机程序的能力的任务。这些任务包括程序评估program evaluation 任务（编程、控制和加法），以及记忆memorization 任务（copy、double 和reverse）。

   结果如下表所示， Universal Transformer 在所有任务上的表现都远远超过了 LSTM 和常规的 Transformer。

   

2.2.6 机器翻译

1. 我们使用与 Transformer报告中相同的设置在 WMT2014 英德翻译任务上训练 Universal Transformer ，以评估其在大规模 sequence-to-sequence 任务上的性能。结果总结如下。采用前馈神经网络（而不是可分离卷积）作为转移函数且没有动态停机的 Universal Transformer 比 Transformer 提高了 0.9 BLEU 。

   根据作者在引言部分的介绍，如果采用动态停机，则会导致 Universal Transformer 的效果更差。