D EPTH -A DAPTIVE T RANSFORMER
A BSTRACT
State of the art sequence-to-sequence models for large scale tasks perform a ﬁxed number of computations for each input sequence regardless of whether it is easy or hard to process.
In this paper, we train Transformer models which can make output predictions at different stages of the network and we investigate different ways to predict how much computation is required for a particular sequence.
Unlike dynamic computation in Universal Transformers, which applies the same set of layers iteratively, we apply different layers at every step to adjust both the amount of computation as well as the model capacity.
On IWSLT German-English translation our approach matches the accuracy of a well tuned baseline Transformer while using less than a quarter of the decoder layers.
I NTRODUCTION
The size of modern neural sequence models (Gehring et al., 2017; Vaswani et al., 2017; Devlin et al., 2019) can amount to billions of parameters (Radford et al., 2019).
For example, the winning entry of the WMT’19 news machine translation task in English-German used an ensemble totaling two billion parameters (Ng et al., 2019).
While large models are required to do better on hard examples, small models are likely to perform as well on easy ones, e.g ., the aforementioned ensemble is probably not required to translate a short phrase such as "Thank you".
However, current models apply the same amount of computation regardless of whether the input is easy or hard.
In this paper, we propose Transformers which adapt the number of layers to each input in order to achieve a good speed-accuracy trade off at inference time.
We extend Graves (2016; ACT) who introduced dynamic computation to recurrent neural networks in several ways: we apply different layers at each stage, we investigate a range of designs and training targets for the halting module and we explicitly supervise through simple oracles to achieve good performance on large-scale tasks.
Universal Transformers (UT) rely on ACT for dynamic computation and repeatedly apply the same layer (Dehghani et al., 2018).
Our work considers a variety of mechanisms to estimate the network depth and applies a different layer at each step.
Moreover, Dehghani et al.
ﬁx the number of steps for large-scale machine translation whereas we vary the number of steps to demonstrate substantial improvements in speed at no loss in accuracy.
UT uses a layer which contains as many weights as an entire standard Transformer and this layer is applied several times which impacts speed.
Our approach does not increase the size of individual layers.
We also extend the resource efﬁcient object classiﬁcation work of Huang et al.
(2017) and Bolukbasi et al.
(2017) to structured prediction where dynamic computation decisions impact future computation.
Related work from computer vision includes Teerapittayanon et al.
; Figurnov et al.
(2017) and Wang et al.
(2018) who explored the idea of dynamic routing either by exiting early or by skipping layers.
We encode the input sequence using a standard Transformer encoder to generate the output sequence with a varying amount of computation in the decoder network.
Dynamic computation poses a challenge for self-attention because omitted layers in prior time-steps may be required in the future.
We experiment with two approaches to address this and show that a simple approach works well (§2).
Next, we investigate different mechanisms to control the amount of computation in the decoder network, either for the entire sequence or on a per-token basis.
This includes multinomial and binomial classiﬁers supervised by the model likelihood or whether the argmax is already correct as well as simply thresholding the model score (§3).
Experiments on IWSLT14 German-English
[Figure]
translation (Cettolo et al., 2014) as well as WMT’14 English-French translation show that we can match the performance of well tuned baseline models at up to 76% less computation (§4).
A NYTIME STRUCTURED PREDICTION
We ﬁrst present a model that can make predictions at different layers.
This is known as anytime prediction for computer vision models (Huang et al., 2017) and we extend it to structured prediction.
2. T RANSFORMER WITH MULTIPLE OUTPUT CLASSIFIERS
We base our approach on the Transformer sequence-to-sequence model (Vaswani et al., 2017).
Both encoder and decoder networks contain N stacked blocks where each has several sub-blocks surrounded by residual skip-connections.
The ﬁrst sub-block is a multi-head dot-product self-attention and the second a position-wise fully connected feed-forward network.
For the decoder, there is an additional sub-block after the self-attention to add source context via another multi-head attention.
Given a pair of source-target sequences ( x , y ) , x is processed with the encoder to give representations s =
( s 1 , . . .
, s | x | ) .
Next, the decoder generates y step-by-step.
For every new token y t input to the decoder at time t , the N decoder blocks process it to yield hidden states ( h nt ) 1 ≤ n ≤ N : h 0 t = embed( y t ) , h nt = [Equation] n ( h n
[Equation]
where block n is the mapping associated with the n th block and embed is a lookup table.
The output distribution for predicting the next token is computed by feeding the activations of the last decoder layer h
Nt into a softmax normalized output classiﬁer W :
[Equation]
Standard Transformers have a single output classiﬁer attached to the top of the decoder network.
However, for dynamic computation we need to be able to make predictions at different stages of the network.
To achieve this, we attach output classiﬁers C n parameterized by W n to the output h
nt of each of the N decoder blocks:
[Equation]
The classiﬁers can be parameterized independently or we can share the weights across the N blocks.
2. T RAINING MULTIPLE OUTPUT CLASSIFIERS
Dynamic computation enables the model to use any of the N exit classiﬁers instead of just the ﬁnal one.
Some of our models can choose a different output classiﬁer at each time-step which results in an exponential number of possible output classiﬁer combinations in the sequence length.
We consider two possible ways to train the decoder network (Figure 1).
Aligned training optimizes all classiﬁers simultaneously and assumes all previous hidden states required by the self-attention are available.
However, at test time this is often not the case when we choose a different exit for every token which leads to misaligned states.
Instead, mixed training samples several sequences of exits for a given sentence and exposes the model to hidden states from different layers.
Generally, for a given output sequence y , we have a sequence of chosen exits ( n 1 , . . .
, n | y | ) and we denote the block at which we exit at time t as n t .
2.2. A LIGNED TRAINING
Aligned training assumes all hidden states h n − 1, . .
, h n − 1 t are available in order to compute selfattention and it optimizes N loss terms, one for each exit (Figure 1a):
[Equation]
The compound loss L dec ( x , y ) is a weighted average of N terms w.r.t.
to ( ω 1 , . . . ω N ) .
We found that uniform weights achieve better BLEU compared to other weighing schemes ( c.f . Appendix A).
At inference time, not all time-steps will have hidden states for the current layer since the model exited early.
In this case, we simply copy the last computed state to all upper layers, similar to mixed training (§2.2.2).
However, we do apply layer-speciﬁc key and value projections to the copied state.
2.2. M IXED TRAINING
Aligned training assumes that all hidden states of the previous time-steps are available but this assumption is unrealistic since an early exit may have been chosen previously.
This creates a mismatch between training and testing.
Mixed training reduces the mismatch by training the model to use hidden states from different blocks of previous time-steps for self-attention.
We sample M different exit sequences ( n ( m ) 1 , . . .
n ( m ) |
[Equation]
and evaluate the following loss:
[Equation]
When n t < N , we copy the last evaluated hidden state h nt to the subsequent layers so that the self-attention of future time steps can function as usual (see Figure 1b).
A DAPTIVE DEPTH ESTIMATION
We present a variety of mechanisms to predict the decoder block at which the model will stop and output the next token, or when it should exit to achieve a good speed-accuracy trade-off.
We consider two approaches: sequence-speciﬁc depth decodes all output tokens using the same block (§3.1) while token-speciﬁc depth determines a separate exit for each individual token (§3.2).
We model the distribution of exiting at time-step t with a parametric distribution q t where q t ( n ) is the probability of computing block 1 , . .
, block n and then emitting a prediction with C n .
The parameters of q t are optimized to match an oracle distribution q ∗ t with cross-entropy:
[Equation]
The exit loss ( L exit ) is back-propagated to the encoder-decoder parameters.
We simultaneously optimize the decoding loss (Eq.
and the exit loss (Eq.
(6)) balanced by a hyper-parameter α to ensure that the model maintains good generation accuracy.
The ﬁnal loss takes the form:
[Equation]
In the following we describe for each approach how the exit distribution q t is modeled (illustrated in Figure 2) and how the oracle distribution q ∗ t is inferred.
State
Copied state 1 Halting decision C n Classiﬁer
Copy
D e c ode r dep t h
Decoding step s
C C
C(a [Figure] Sequence-speciﬁc depth
D e c ode r dep t h
Decoding step
C C
C(b) Token-speciﬁc - Multinomial
D e c ode r dep t h
Decoding step
C
S
C
S
C C
C
C(c [Figure]
It then outputs a token at this depth with classiﬁer C n .
The token-speciﬁc multinomial classiﬁer determines the exit after the ﬁrst block and proceeds up to the predicted depth before outputting the next token (2b).
The token geometric-like classiﬁer (2c) makes a binary decision after every block to dictate whether to continue (C) to the next block or to stop (S) and emit an output distribution.
3. S EQUENCE - SPECIFIC DEPTH :
For sequence-speciﬁc depth, the exit distribution q and the oracle distribution q ∗ are independent of the time-step so we drop subscript t .
We condition the exit on the source sequence by feeding the average s of the encoder outputs to a multinomial classiﬁer:
[Equation]
where W h and b h are the weights and biases of the halting mechanism.
We consider two oracles to determine which of the N blocks should be chosen.
The ﬁrst is based on the sequence likelihood and the second looks at an aggregate of the correctly predicted tokens at each block.
Likelihood-based: This oracle is based on the likelihood of the entire sequence after each block
and we optimize it with the Dirac delta centered around the exit with the highest sequence likelihood.
[Equation]
We add a regularization term to encourage lower exits that achieve good likelihood:
[Equation]
Correctness-based:
Likelihood ignores whether the model already assigns the highest score to the correct target.
Instead, this oracle chooses the lowest block that assigns the largest score to the correct prediction.
For each block, we count the number of correctly predicted tokens over the sequence and choose the block with the most number of correct tokens.
A regularization term controls the trade-off between speed and accuracy.
[Equation]
Oracles based on test metrics such as BLEU are feasible but expensive to compute since we would need to decode every training sentence N times.
We leave this for future work.
3. T OKEN - SPECIFIC DEPTH :
The token-speciﬁc approach can choose a different exit at every time-step.
We consider two options for the exit distribution q t at time-step t: a multinomial with a classiﬁer conditioned on the ﬁrst decoder hidden state h 1 t and a geometric-like where an exit probability χ nt is estimated after each block based on the activations of the current block h nt .
Multinomial q t :
[Equation]
The most probable exit arg max q t ( n | x , y <t ) is selected at inference.
Geometric-like q t :
[Equation]
where, d is the dimension of the decoder states, W h ∈ R N × d and w h ∈ R
d are the weights of the halting mechanisms, and b h their biases.
During inference the decoder exits when the halting signal χ nt exceeds a threshold τ n which we tune on the valid set to achieve a better accuracy-speed trade-off.
If thresholds ( τ n ) 1 ≤ n<N have not been exceeded, then we default to exiting at block N .
The two classiﬁers are trained to minimize the cross-entropy with respect to either one the following oracle distributions:
Likelihood-based:
At each time-step t , we choose the block whose exit classiﬁer has the highest likelihood plus a regularization term weighted by λ to encourage lower exits. [Equation]
[Equation]
This oracle ignores the impact of the current decision on the future time-steps and we therefore consider smoothing the likelihoods with an RBF kernel.
[Equation]
where we control the size of the surrounding context with σ the kernel width.
We refer to this oracle as LL( σ, λ ) including the case where we only look at the likelihood of the current token with σ → 0 .
Correctness-based: Similar to the likelihood-based oracle we can look at the correctness of the prediction at time-step t as well as surrounding positions.
We deﬁne the target q ∗ t as follows:
[Equation]
Conﬁdence thresholding Finally, we consider thresholding the model predictions (§2), i.e., exit when the maximum score of the current output classiﬁer p ( y t +1 | h nt ) exceeds a hyper-parameter threshold τ n .
This does not require training and the thresholds τ = ( τ 1 , . . .
, τ N − 1 ) are simply tuned on the valid set to maximize BLEU.
Concretely, for 10k iterations, we sample a sequence of thresholds τ ∼ U (0 , 1) N − 1 , decode the valid set with the sampled thresholds and then evaluate the
BLEU score and computational cost achieved with this choice of τ .
After 10k evaluations we pick the best performing thresholds, that is τ with the highest BLEU in each cost segment.
E XPERIMENTS 4. E XPERIMENTAL SETUP
We evaluate on several benchmarks and measure tokenized BLEU (Papineni et al., 2002):
IWSLT’14 German to English (De-En).
We use the setup of Edunov et al.
and train on 160K sentence pairs.
We use N = 6 blocks, a feed-forward network (ffn) of intermediate-dimension
[Table]
WMT’14 English to French (En-Fr).
We also experiment on the much larger WMT’14 EnglishFrench task comprising 35.5m training sentence pairs.
We develop on 26k held out pairs and test on newstest14.
The vocabulary consists of 44k joint BPE types (Sennrich et al., 2016).
We use a Transformer big architecture and tie the embeddings of the encoder, the decoder and the output classiﬁers ( ( W n ) 1 ≤ n ≤ 6 ; §2.1).
We average the last ten checkpoints and use a beam of width 4.
Models are implemented in fairseq (Ott et al., 2019) and are trained with Adam (Kingma & Ba, 2015).
We train for 50k updates on 128 GPUs with a batch size of 460k tokens for WMT’14 En-Fr and on 2 GPUs with 8k tokens per batch for IWSLT’14 De-En.
To stabilize training, we re-normalize the gradients if the norm exceeds g clip = 3 .
For models with adaptive exits, we ﬁrst train without exit prediction ( α = 0 in Eq.
(7)) using the aligned mode ( c.f . §2.2.1) for 50k updates and then continue training with α (cid:54)
=
until convergence.
The exit prediction classiﬁers are parameterized by a single linear layer (Eq.
with the same input dimension as the embedding dimension, e.g ., 1024 for a big Transformer; the output dimension is N for a multinomial classiﬁer or one for geometric-like.
We exit when χ t,n > 0 .
5 for geometric-like classiﬁers.
4. T RAINING MULTIPLE OUTPUT CLASSIFIERS
We ﬁrst compare the two training regimes for our model (§2.2).
Aligned training performs selfattention on aligned states (§2.2.1) and mixed training exposes self-attention to hidden states from different blocks (§2.2.2).
We compare the two training modes when choosing either a uniformly sampled exit or a ﬁxed exit n = 1 , . . .
, 6 at inference time for every time-step.
The sampled exit experiment tests the robustness to mixed hidden states and the ﬁxed exit setup simulates an ideal setting where all previous states are available.
As baselines we show six separate standard Transformers with N ∈ [1 .. 6] decoder blocks.
All models are trained with an equal number of updates and mixed training with M =6 paths is most comparable to aligned training since the number of losses per sample is identical.
Table 1 shows that aligned training outperforms mixed training both for ﬁxed exits as well as for randomly sampled exits.
The latter is surprising since aligned training never exposes the self-attention mechanism to hidden states from other blocks.
We suspect that this is due to the residual connections which copy features from lower blocks to subsequent layers and which are ubiquitous in Transformer models (§2).
Aligned training also performs very competitively to the individual baseline models.
Aligned training is conceptually simple and fast.
We can process a training example with N exits in a single forward/backward pass while M passes are needed for mixed training.
In the remaining paper, we use the aligned mode to train our models.
Appendix A reports experiments with weighing the various output classiﬁers differently but we found that a uniform weighting scheme worked well.
On our largest setup, WMT’14 English-French, the training time of an aligned model with six output classiﬁers increases only marginally by about 1% compared to a baseline with a single output classiﬁer keeping everything else equal.
[Figure] [Figure]
4. A DAPTIVE DEPTH ESTIMATION
Next, we train models with aligned states and compare adaptive depth classiﬁers in terms of BLEU as well as computational effort.
We measure the latter as the average exit per output token ( AE ).
As baselines we use again six separate standard Transformers with N ∈ [1 .. 6] with a single output classiﬁer.
We also measure the performance of the aligned mode trained model for ﬁxed exits n ∈ [1 .. 6] .
For the adaptive depth token-speciﬁc models (Tok), we train four combinations: likelihoodbased oracle (LL) + geometric-like, likelihood-based oracle (LL) + multinomial, correctness based oracle (C) + geometric-like and correctness-based oracle (C) + multinomial.
Sequence-speciﬁc models (Seq) are trained with the correctness oracle (C) and the likelihood oracle (LL) with different values for the regularization weight λ .
All parameters are tuned on the valid set and we report results on the test set for a range of average exits.
Figure 3 shows that the aligned model (blue line) can match the accuracy of a standard 6-block
Transformer (black line) at half the number of layers ( n = 3 ) by always exiting at the third block.
The aligned model outperforms the baseline for n = 2 , . . .
For token speciﬁc halting mechanisms (Figure 3a) the geometric-like classiﬁers achieves a better speed-accuracy trade-off than the multinomial classiﬁers (ﬁlled vs. empty triangles).
For geometriclike classiﬁers, the correctness oracle outperforms the likelihood oracle (Tok-C geometric-like vs.
Tok-LL geometric-like)
but the trend is less clear for multinomial classiﬁers.
At the sequence-level, likelihood is the better oracle (Figure 3b).
The rightmost Tok-C geometric-like point ( σ = 0 , λ = 0 . 1 ) achieves 34.73 BLEU at AE = 1 .which corresponds to similar accuracy as the N = 6 baseline at 76% fewer decoding blocks.
[Figure]
The best accuracy of the aligned model is 34.95 BLEU at exit 5 and the best comparable Tok-C geometric-like conﬁguration achieves 34.99 BLEU at AE = 1 . 97 , or 61% fewer decoding blocks.
When ﬁxing the budget to two decoder blocks, Tok-C geometric-like with AE = 1 .
97 achieves
BLEU 35, a 0.64 BLEU improvement over the baseline ( N = 2 ) and aligned which both achieve
BLEU 34.35.
Conﬁdence thresholding (Figure 3c) performs very well but cannot outperform Tok-C geometriclike.
Ablation of hyper-parameters In this section, we look at the effect of the two main hyperparameters on IWSLT’14 De-En: λ the regularization scale ( c.f . Eq.
(9)), and the RBF kernel width σ used to smooth the scores ( c.f . Eq.
We train Tok-LL Geometric-like models and evaluate them with their default thresholds (exit if χ nt > 0 . 5 ).
Figure 4a shows that higher values of λ lead to lower exits.
Figure 4b shows the effect of σ for two values of λ .
In both curves, we see that wider kernels favor higher exits.
S CALING THE ADAPTIVE - DEPTH MODELS
Finally, we take the best performing models form the IWSLT benchmark and test them on the large
WMT’14 English-French benchmark.
Results on the test set (Figure 5a) show that adaptive depth still shows improvements but that they are diminished in this very large-scale setup.
Conﬁdence thresholding works very well and sequence-speciﬁc depth approaches improve only marginally over the baseline.
Tok-LL geometric-like can match the best baseline result of BLEU 43.4 ( N = 6 ) by using only AE
= 2 . 40 which corresponds to 40% of the decoder blocks; the best aligned result of BLEU 43.6 can be matched with AE = 3 . 25 .
In this setup, Tok-LL geometric-like slightly outperforms the Tok-C counterpart.
Conﬁdence thresholding matches the accuracy of the N =6 baseline with AE 2.5 or 59% fewer decoding blocks.
However, conﬁdence thresholding requires computing the output classiﬁer at each block to determine whether to halt or continue.
This is a large overhead since output classiﬁers predict 44k types for this benchmark (§4.1).
To better account for this, we measure the average number of FLOPs per output token (details in Appendix B).
Figure 5b shows that the Tok-LL geometric-like approach provides a better trade-off when the overhead of the output classiﬁers is considered.
4. Q UALITATIVE RESULTS
The exit distribution for a given sample can give insights into what a Depth-Adaptive Transformer decoder considers to be a difﬁcult task.
In this section, for each hypothesis (cid:101) y , we will look at the sequence of selected exits ( n 1 , . . . , n | (cid:101) y | ) and the probability scores ( p 1 , . . .
p | (cid:101) y | ) with p t = p ( (cid:101) y t | h n t t − 1 ) i.e . the conﬁdence of the model in the sampled token at the selected exit.
Figures 6 and 7 show hypotheses from the WMT’14 En-Fr and IWSLT’14 De-En test sets, respectively. For each hypothesis we state the exits and the probability scores. In Figure 6a, predicting
C h i @ @
r ac , l e P r
e m i e r m i n i s t r e , é t a it p
r é s e n t .
< / s >[-
-]
E x it 0 .0 .0 .0 .
[Figure]
M a i s l e
s p a ss a g e r s n e d e v r a
i e n t p a s s ’ a tt e nd r e à d e s c h a ng e m e n t
s
i
mm é d i a t s .
< / s >[-
 E x it-] 0 .0 .0 .0 .
S c
o r e
[Figure]
[Figure]
See Figure 6 for more details.
‘présent’ (meaning ‘present’) is hard.
A straightforward translation is ‘était là’ but the model chooses ‘present’ which is also appropriate.
In Figure 6b, the model uses more computation to predict the deﬁnite article ‘les’ since the source has omitted the article for ‘passengers’.
A clear trend in both benchmarks is that the model requires less computation near the end of decoding to generate the end of sequence marker </ s > and the preceding full-stop when relevant.
In Figure 8, we show the distribution of the exits at the beginning and near the end of test set hypotheses. We consider the beginning of a sequence to be the ﬁrst 10% of tokens and the end as the last 10% of tokens. The exit distributions are shown for three models on WMT’14 En-Fr: Model 1 has an average exit of AE = 2 .
53 , Model 2 exits at AE
= 3 .
79 on average and Model 3 with AE
= 4 .
Within the same models, deep exits late are used at the beginning of the sequence and early exits are selected near the end.
For heavily regularized models such as Model 1 with AE = 2 .
53 , the disparity between beginning and end is less severe as the model exits early most of the time.
Model 2 and
Model 3 are less regularized (higher AE) and tend to use late exits at the beginning of the sequence and early exits near the end.
On the other hand, the more regularized Model 1 with AE = 2 . 53 exits
[Figure]
F r e q u e n c y
Beginning
End
[Figure]
Model
[Figure]
Model
[Figure]
early most of the time.
There is also a correlation between the model probability and the amount of computation, particularly in models with low AE .
Figure 9 shows the joint histogram of the scores and the selected exit.
For both Model 1 and Model 2 , low exits ( n ≤ 2 ) are used in the high conﬁdence range [0 . 8 − 1] and high exits ( n ≥ 4 ) are used in the low-conﬁdence range [0 − 0 .
Model 3 has a high average exit ( AE = 4 . 68 ) so most tokens exit late, however, in low conﬁdence ranges the model does not exit earlier than n = 5 .
C ONCLUSION
We extended anytime prediction to the structured prediction setting and introduced simple but effective methods to equip sequence models to make predictions at different points in the network.
We compared a number of different mechanisms to predict the required network depth and ﬁnd that a simple correctness based geometric-like classiﬁer obtains the best trade-off between speed and accuracy.
Results show that the number of decoder layers can be reduced by more than three quarters at no loss in accuracy compared to a well tuned Transformer baseline.
A CKNOWLEDGMENTS
We thank Laurens van der Maaten for fruitful comments and suggestions.
A PPENDIX A L OSS SCALING
In this section we experiment with different weights for scaling the output classiﬁer losses.
Instead of uniform weighting, we bias towards speciﬁc output classiﬁers by assigning higher weights to their losses.
Table 2 shows that weighing the classiﬁers equally provides good results.
[Table]
Gradient scaling Adding intermediate supervision at different levels of the decoder results in richer gradients for lower blocks compared to upper blocks.
This is because earlier layers affect more loss terms in the compound loss of Eq.
To balance the gradients of each block in the decoder, we scale up the gradients of each loss term ( − LL n ) when it is updating the parameters of its associated block ( block n with parameters θ n ) and revert it back to its normal scale before back-propagating it to the previous blocks.
Figure 10 and Algorithm 1 illustrate this gradient scaling procedure.
The θ n are updated with γ n -ampliﬁed gradients from the block’s supervision and ( N − n ) gradients from the subsequent blocks.
We choose γ n = γ ( N − n ) to control the ratio γ :1 as the ratio of the block supervision to the subsequent blocks’ supervisions.
Table 3 shows that gradient scaling can beneﬁt the lowest layer at the expense of higher layers.
However, no scaling generally works very well.
[Table]
[Figure] 1: for n ∈ 1 ..N do 2: h nt = block n ( h n − 1 t ) 3: p ( y t +1 | h nt )
=
softmax( W n h nt ) 4: p ( y t +1 | h nt ) =
SCALE _ GRADIENT ( p ( y t +1 | h nt ) , γ n ) 5: if n < N then h nt =
SCALE _ GRADIENT
( h nt , γ n +1 ) 6: end for 7: function SCALE _ GRADIENT (Tensor x , scale γ ) 8: return γx + (1 − γ )
STOP _ GRADIENT ( x ) 9: (cid:46) STOP _ GRADIENT in PyTorch with x .detach().
10: end function
Uniform n =n =n =n =n =n =
[Table]
Average
Baseline34.
[Table]
[Table] [Table] [Table]
[Table] [Table]
∅
[Table] [Table]
[Table] [Table] [Table] =
[Table] [Table]
γ = 0 .35.
[Table]
[Table]
γ = 0 .34.
[Table] [Table] [Table] [Table]
[Table]
[Table]
γ = 0 .34.
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
γ = 1 .35.
[Table] [Table]
[Table] [Table] [Table]
[Table]
Uniform n =n =
[Table]
A PPENDIX B FLOPS APPROXIMATION
This section details the computation of the FLOPS we report.
The per token FLOPS are for the decoder network only since we use an encoder of the same size for all models.
We breakdown the FLOPS of every operation in Algorithm 2 (blue front of the algorithmic statement).
We omit non-linearities, normalizations and residual connections.
The main operations we account for are dot-products and by extension matrix-vector products since those represent the vast majority of
FLOPS (we assume batch size one to simplify the calculation).
[Table]
With this breakdown, the total computational cost at time-step t of a decoder block that we actually go through, denoted with FC, is:
[Equation]
where the cost of mapping the source’ keys and values is incurred the ﬁrst time the block is called (ﬂagged with FirstCall).
This occurs at t = 1 for the baseline model but it is input-dependent with depth adaptive estimation and may never occur if all tokens exit early.
If skipped, a block still has to compute the keys and value of the self-attention block so the selfattention of future time-steps can function.
We will denote this cost with FS and we have FS = 4 d 2 d .
Depending on the halting mechanism, an exit prediction cost, denoted wit FP, is added:
Sequence speciﬁc depth:
[Equation]
Token-speciﬁc Multinomial
[Equation]
Token-speciﬁc Geometric-like:
[Equation]
Conﬁdence thresholding:
[Equation]
For a set of source sequences { x ( i ) } i ∈I and generated hypotheses { y ( i ) } i ∈I , the average ﬂops per token is:
[Equation]
In the case of conﬁdence thresholding the ﬁnal output prediction cost ( 2 V d d ) is already accounted for in the exit prediction cost FP .
do 4: ﬂag source’ keys be evaluated.
5: end for 6: while y t (cid:54) do 7:
Embed the last token y
8: for n
N do 9: (cid:46) Self-attention.
10: - Map the input into a key ( value v
FLOPS=
4 d 2 d 11: - Map the input into a query q .
FLOPS=
2 d 2 d 12: - Score the memory keys with q to get the attention weights α .
FLOPS= 4 td d 13: - Map the attention output.
FLOPS=
2 d 2 d 14: (cid:46) Encoder-Decoder interaction.
15: if FirstCall[ n ] then 16:
Map the source states into keys and values for the nth block.
FLOPS=
| x | d e d d 17:
FirstCall[ n ] =
False 18: end if 19: - Map the input into a query q .
FLOPS= 2 d 2 d 20: - Score the memory keys with q to get the attention weights α .
FLOPS=
| x | d d 21: - Map the attention output.
FLOPS=
2 d 2 d 22:
Feed-forward network.
FLOPS= 4 d d d f 23:
Estimate the halting probability χ t,n .
FLOPS= 2 d d 24: if χ t,n > 0 .
5 then 25:
Exit 26: end if 27: end for 28: if n < N then 29:
Skipped do 31:
32: end
33: end
Project the ﬁnal state and sample output
36: end while