ICLR2020 | 深度自适应Transformer

作者 | kaiyuan

整理 | NewBeeNLP

一篇来自ICLR2020，facebook人工智能实验室的论文，关于改进Transformer推理速度的。

• 论文地址：https://arxiv.org/abs/1910.10073

写在前面

大模型在困难任务上表现非常好，而小模型也可以在比较简单的任务上表现出色。但是目前大模型在应用简单任务时的运算量同复杂任务一样，这就导致了不必要的资源浪费，此外还会导致推理速度变慢。对此，提出了一种depth-adaptive transformer，可以在推理阶段针对不同的token自动选择在合适的层进行停止并输出，在不损失模型效果的同时大大提高推理速度。

Anytime Prediction

传统的transformer decoder，所有的token都需要经过N个block，且最后的分类器是唯一的。但是一些简单的任务并不需要推理如此多步，可能在第

t

层模型已经得出结果了，再往上走并没有意义，于是可以选择在第

t

层就输出，这称为「动态计算（Dynamic Computation）」，

forall n, pleft(y_{t+1} | h_{t}^{n}right)=operatorname{softmax}left(W_{n} h_{t}^{n}right)

并且，每一层的分类器

W_{n}

可以选择不同或者相同。接下去就是怎么去实现上面的想法，即对于每个输入token，如何自动确定其在哪一层退出？文中提出了两种策略：

Aligned training

对齐训练会同时优化所有分类器，并假设当前状态之前的所有状态

h_{1}^{n-1}, ldots, h_{t}^{n-1}

均可以被看见，然后计算

N

项loss，如下图(a)

mathrm{LL}_{t}^{n}=log pleft(y_{t} | h_{t-1}^{n}right)

mathrm{LL}^{n}=sum_{t=1}^{|boldsymbol{y}|} mathrm{LL}_{t}^{n}

mathcal{L}_{d e c}(boldsymbol{x}, boldsymbol{y})=-frac{1}{sum_{n} omega_{n}} sum_{n=1}^{N} omega_{n} mathrm{LL}^{n}

其中，

LL^{n}

是第

n

层所有

y_{t}

损失和，

mathcal{L}_{d e c}

是

LL^{n}

的加权平均（论文附录A实验说明平均加权的效果最好）。

Mixed training

但是上述对齐训练是不好实现的，因为不是所有token都在同一层输出，那么当有的token在较早的层输出后，在后续层计算时就无法看见它。解决这个问题一个直观的想法就是将已经输出的单元直接copy到上面一层即可，如下图(b)。对输入序列采样

M

个退出序列

left(n_{1}^{(m)}, ldots n_{|boldsymbol{y}|}^{(m)}right)_{1 leq m leq M}

，并对每个退出序列计算损失，

mathrm{LL}_{t}=log pleft(y_{t} | h_{t-1}^{n}right)

mathrm{LL}=sum_{t=1}^{|boldsymbol{y}|} mathrm{LL}_{t}

mathcal{L}_{d e c}(boldsymbol{x}, boldsymbol{y})=-frac{1}{M} sum_{m=1}^{M} mathrm{LL}

Adaptive Depth Estimation

自适应深度估计。上文还有一个问题待解决：如何采样退出序列？提出了两种策略：

• 「sequence-specific」：对应于aligned training，所有token都在同一层退出；
• 「token-specific」：对应于mixed training，每个token可能在不同的层退出；

q_{t}

表示第

t

步对应的退出层参数分布，其中

q_{t}(n)

表示

y_{t}

在第

t

层退出并得到预测

L_{n}

的概率，

q_{t}^{*}

表示真实分布，使用下式优化：

mathcal{L}_{mathrm{exit}}(boldsymbol{x}, boldsymbol{y})=sum_{t} Hleft(q_{t}^{*}(boldsymbol{x}, boldsymbol{y}), q_{t}(boldsymbol{x})right)

总体的损失函数为：

mathcal{L}(boldsymbol{x}, boldsymbol{y})=mathcal{L}_{operatorname{dec}}(boldsymbol{x}, boldsymbol{y})+alpha mathcal{L}_{mathrm{exit}}(boldsymbol{x}, boldsymbol{y})

Sequence-Specific Depth

在该种模式下，所有token都在同一层退出，所以我们不考虑time-step

t

。通过将encoder输出的平均值

s

送到多项式分类器，我们得到分布

q

：

s=frac{1}{|x|} sum_{t} s_{t}

q(n | boldsymbol{x})=operatorname{softmax}left(W_{h} s+b_{h}right) in mathbb{R}^{N}

接下来怎么得到真实分布

q_{t}^{*}

呢？本文提出了两种策略：

• 「Likelihood-based」：基于每个块之后整个序列的似然性，并且以Dirac delta为中心对退出进行了优化，从而使序列似然性最高，

q^{*}(boldsymbol{x}, boldsymbol{y})=deltaleft(arg max _{n} mathrm{LL}^{n}right)

并且还可以加入正则项鼓励更低的退出层

q^{*}(boldsymbol{x}, boldsymbol{y})=deltaleft(arg max _{n} L^{n}-lambda nright)

• 「Correctness-based」：基于似然的方式无法判断最大的似然是不是正确的序列，因此基于正确性的方式选择正确字符最大分数的最低的层。具体来说，对每个层，我们计算在该层上所有正确预测的字符数，然后选择最大正确数的一层，并同样加上正则项：

C^{n}=#left{tleft|y_{t}=arg max _{y} pleft(y | h_{t-1}^{n}right)right}, quad q^{*}(x, y)=deltaleft(arg max _{n} C^{n}-lambda nright)right.

Token-Specific Depth

基于token的方式对每个token自动学习在哪一层退出，提出两种exit distribution

q_{t}

的计算方式：

• 「Multinomial」：多项分布，

h_{t}^{1}

表示第一层decoder的输出

q_{t}(n | boldsymbol{x}, y_{＜t})=operatorname{softmax}left(W_{h} h_{t}^{1}+b_{h}right)

• 「Poisson Binomial」：泊松二项分布，

chi_{t}^{n}

表示停止概率(halting probability)，其对应有一个阈值

tau_{n}

，当

chi_{t}^{n}

超过这个阈值时即退出否则继续往上走直至结束

forall n in[1 . . N-1], chi_{t}^{n}=operatorname{sigmoid}left(w_{h}^{top} h_{t}^{n}+b_{h}right)

q_{t}left(n | boldsymbol{x}, boldsymbol{y}_{＜t}right)=left{begin{array}{l} {chi_{t}^{n} prod_{n^{prime}＜n}left(1-chi_{t}^{n^{prime}}right), text { if } n＜N} \ {prod_{n^{prime}}left(1-chi_{t}^{n^{prime}}right), text { otherwise }} end{array}right.

接下去也是对真实分布的计算，提出了三种策略：

• 「Likelihood-based」：基于似然的方式在每一步选择似然最大的层退出，

q_{t}^{*}(boldsymbol{x}, boldsymbol{y})=deltaleft(arg max _{n} mathrm{LL}_{t}^{n}-lambda boldsymbol{n}right)

但是显然，这种方法是贪心做法，我们还想要考虑未来的字符，所以本文考虑一个光滑化（smoothed）似然，它考虑附近的几个字符：

kappaleft(t, t^{prime}right)=e^{-frac{left|t-t^{prime}right|^{2}}{sigma}}

widehat{mathrm{LL}_{t}^{n}}=sum_{t^{prime}=1}^{|mathbf{y}|} kappaleft(t, t^{prime}right) mathrm{LL}_{t^{prime}}^{n}

q_{t}^{*}(boldsymbol{x}, boldsymbol{y})=deltaleft(arg max _{n} widehat{mathrm{LL}_{t}^{n}}-lambda boldsymbol{n}right)

• 「Correctness-based」：基于正确性的方式，

C_{t}^{n}=left(y_{t}=arg max _{y} pleft(y | h_{t-1}^{n}right)right), quad widetilde{C}_{t}^{n}=sum_{t^{prime}=1}^{|y|} kappaleft(t, t^{prime}right) C_{t}^{n}

q_{t}^{*}(boldsymbol{x}, boldsymbol{y})=deltaleft(arg max _{n} widetilde{C_{t}^{n}}-lambda nright)

• 「Confidence thresholding」：基于自信度阈值，定义了一个自信度阈值

tau_{n}

，当当前输出

pleft(y_{t+1} | h_{t}^{n}right)

的最大得分超过阈值时，即选择退出。其中

tau

是通过验证集确定的。

上述几种方法的示意图：

一些结论

• Aligned training效果要优于Mixed training；
• 用对齐训练方法得到的模型，在

n>2

的时候，差不多就能超过baseline模型（传统Transformer），而Possion方法大体上比多项分布好；

• 更多更详细的实验分析推荐阅读reference里的作者PPT分享；
• 整体感觉和之前介绍过的Universal Transformer很像，都是采用dynamic computation的思想。但是很多细节还是不同，比如没有大型共享的Transformer层，并且本文主要关注动态停止策略设计等等：

There are a number of differences to universal transformer (UT): UT repeatedly applies the same layer for a number of steps, while as our approach applies different layers at every step. The dynamic computation in UT considers a single type of mechanism to estimate the number of steps (=network depth) while as we consider a number of mechanisms and supervision regimes. Moreover, dynamic computation was not actually used for the machine translation experiments in the UT paper (it was used for other tasks though), the authors used a fixed number of steps for every input instead.

• Open Review[1]DAT作者PPT分享[2]

本文参考资料

[1]

Open Review: https://openreview.net/forum?id=SJg7KhVKPH

[2]

DAT作者PPT分享: http://elbayadm.github.io/assets/talks/anytime/decore_04102019.pdf