Source-linked AI summary
Depth-Adaptive Transformer
Maha Elbayad, Jiatao Gu, Edouard Grave, Michael Auli
TL;DR
Modern sequence-to-sequence models use fixed computation despite differing input difficulty. This paper trains Transformers with intermediate output predictions and adaptive depth, matching tuned baselines while using far fewer decoder layers. The authors also identify practical limits in computing metric-based training oracles.
Problem
Current sequence-to-sequence models apply the same amount of computation to easy and hard inputs, despite large models being mainly needed for difficult examples.
Method
The paper attaches output classifiers at different decoder depths and investigates sequence- or token-level mechanisms for selecting the required computation, using distinct layers at successive steps.
Results
More than three quarters of decoder layers can be removed with no accuracy loss relative to a well-tuned Transformer baseline; on IWSLT14 German-English and WMT’14 English-French, the models match baseline performance at up to 76% less computation.
Takeaways & Limitations
A correctness-based geometric-like classifier provides the best reported speed-accuracy trade-off, while adaptive depth can preserve baseline translation accuracy with substantially less decoding computation.
Takeaways & Limitations
Metric-based test oracles are expensive because each training sentence would need decoding N times, so the authors leave them for future work.
Abstract
from arXiv · showhide
State of the art sequence-to-sequence models for large scale tasks perform a fixed 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.
1 INTRODUCTION
The paper addresses the inefficiency of applying fixed computation to every input by adapting Transformer depth to each sequence or token. It develops training and depth-estimation mechanisms that preserve translation performance while substantially reducing computation.
- Motivation: Large sequence models may need substantial capacity for hard examples, but current models apply the same computation to easy and hard inputs.The paper illustrates this mismatch with short phrases that likely do not require a billion-parameter ensemble.
- Approach: The proposed Transformers adapt the number of decoder layers to each input for a speed-accuracy trade-off at inference time.The approach applies different layers at successive stages and explores training targets for the halting module.
- Approach: Unlike Universal Transformers, the approach applies a different layer at each step and varies the number of steps rather than repeatedly applying one layer with a fixed step count.This is intended to adjust both computation and model capacity.
- Approach: The decoder uses sequence-level or token-level mechanisms, including multinomial and binomial classifiers, correctness-based supervision, and model-score thresholding.These mechanisms estimate how much computation is required during decoding.
- Results: Up to 76% less computation matches the performance of well-tuned baseline models on IWSLT14 German-English and WMT’14 English-French translation.The result is reported across both benchmarks in the introduction.
2 ANYTIME STRUCTURED PREDICTION
The model equips Transformer decoders to emit predictions from intermediate layers, enabling dynamic exits during structured prediction. Training addresses the resulting mismatch in self-attention histories by exposing the decoder to aligned or mixed exit patterns.
- Anytime structured prediction: Anytime structured prediction extends anytime prediction by allowing sequence models to make outputs at different points in the network.This enables intermediate decoder layers to serve as prediction points.
- Transformer architecture: A standard Transformer encoder produces source representations, while the decoder generates target tokens through stacked blocks with self-attention, source attention, and feed-forward sub-blocks.The decoder architecture follows the Transformer sequence-to-sequence model.
- Multiple output classifiers: Each decoder block receives a token representation and contributes a hidden state that can support a prediction at that depth.The decoder processes each newly input target token through its blocks.
- Multiple output classifiers: Output classifiers Cn attached to every decoder block let the model use any of N intermediate exits instead of only the final classifier.Classifiers may be parameterized independently or share weights across blocks.
- Training regimes: Aligned training jointly optimizes all exit classifiers assuming required previous hidden states are available, whereas mixed training samples exit sequences with hidden states from different blocks.Mixed training better reflects test-time misalignment when different tokens exit at different depths.
- Training regimes: When an earlier token exits below the current layer, its last computed hidden state is copied upward so future self-attention can operate.The same copying strategy is used during mixed training and inference, with layer-specific key and value projections applied during aligned inference.
3 ADAPTIVE DEPTH ESTIMATION
The paper estimates adaptive decoder depth using sequence-specific or token-specific exit mechanisms, trained with likelihood- or correctness-based oracle targets. These mechanisms include multinomial, geometric-like, and confidence-thresholding classifiers, with regularization or tuning used to balance speed and accuracy.
- Depth prediction strategies: Sequence-specific depth selects one decoder exit for the entire output sequence, whereas token-specific depth selects a separate exit for each token.The sequence-specific approach conditions on averaged encoder outputs; token-specific approaches use decoder states or per-block halting decisions.
- Exit training: Exit distributions model the probability of emitting a prediction after each decoder block, and inference selects the most probable exit.The exit loss is jointly optimized with the decoding loss, balanced by a hyper-parameter α.
- Depth prediction strategies: Figure 2 presents multinomial sequence-specific, multinomial token-specific, and geometric-like token-specific classifiers for deciding when to emit predictions.The geometric-like classifier makes a binary continue-or-stop decision after every decoder block.
- Oracle targets: The sequence-specific classifiers use likelihood-based or correctness-based oracles, with regularization encouraging lower exits while preserving prediction quality.Likelihood selects the block with the highest sequence likelihood, whereas correctness counts correctly predicted tokens at each block.
- Limitations: BLEU-based oracle targets are feasible but expensive because every training sentence would require decoding N times.The paper leaves these test-metric oracles for future work.
- Token-specific depth: Token-specific classifiers similarly use likelihood or correctness targets, while confidence thresholding exits when the current output score exceeds a tuned threshold.Thresholds are tuned on the validation set to maximize BLEU within computational-cost segments.
4 EXPERIMENTS
Experiments evaluate adaptive-depth Transformers across IWSLT14 German-English and WMT14 English-French translation, comparing training modes, halting mechanisms, speed, and accuracy. Adaptive models match strong baseline accuracy with substantially fewer decoder computations, though gains diminish at larger scale.
- Training regimes: Aligned training outperforms mixed training for both fixed and randomly sampled exits and remains competitive with individual baseline Transformers.Aligned training processes all N exits in one forward/backward pass, whereas mixed training requires multiple passes.
- Training regimes: 1% additional training time is required for an aligned WMT14 En-Fr model with six output classifiers compared with a single-classifier baseline.This comparison keeps the rest of the setup equal.
- IWSLT14 results: 34.73 BLEU at AE = 1.42 matches the N = 6 baseline accuracy with 76% fewer decoding blocks on IWSLT14 De-En.This is achieved by the rightmost Tok-C geometric-like configuration with σ = 0 and λ = 0.1.
- WMT14 results: On WMT14 En-Fr, Tok-LL geometric-like matches BLEU 43.4 from the N = 6 baseline at AE = 2.40, while classifier overhead makes its FLOP trade-off better than confidence thresholding.The adaptive-depth improvements are diminished in this larger-scale setup; the best aligned result of BLEU 43.6 is matched at AE = 3.25.
5 CONCLUSION
The paper extends anytime prediction to structured prediction and finds that adaptive decoder depth can preserve accuracy while substantially reducing computation.
- More than three quarters of decoder layers can be removed without accuracy loss relative to a well tuned Transformer baseline.The conclusion identifies this as the central speed-accuracy result.
APPENDIX A LOSS SCALING
The appendix studies loss weighting and gradient scaling for intermediate decoder classifiers. Equal classifier weighting performs well, while gradient scaling can favor lower layers at the expense of higher layers.
- Equal weighting of classifier losses provides good results on IWSLT De-En.The appendix contrasts uniform weighting with biasing toward selected output classifiers.
- Gradient scaling benefits the lowest decoder layer at the expense of higher layers.The reported comparison indicates that scaling changes the distribution of performance across decoder layers.
- No gradient scaling generally works very well.This is the appendix's overall conclusion about the scaling choice.
APPENDIX B FLOPS APPROXIMATION
The appendix approximates decoder computation by accounting for executed, skipped, exit-prediction, and output-projection costs. It also specifies how adaptive decoding handles skipped blocks and cached source representations.
- Per-token FLOPS are reported for the decoder only because all models use an encoder of the same size.The estimate focuses on decoder computation and counts mainly dot-products and matrix-vector products.
- An executed decoder block incurs the cost of its operations plus source key-value mapping when the block is called for the first time.The first-call cost occurs at t = 1 for the baseline but is input-dependent for depth-adaptive estimation.
- Skipped blocks still compute self-attention keys and values so future time steps can use self-attention normally.This skipped-block cost is denoted FS and equals 4d^2.
- Adaptive decoding adds an exit-prediction cost that depends on the halting mechanism.For confidence thresholding, the final output prediction cost is already included in the exit-prediction cost.
- The FLOPS estimate averages computation per token over source sequences and generated hypotheses.The baseline cost includes the decoder-block costs and final output projection.