[Paper Review] Depth-Adaptive Transformer

Maha Elbayad, Jiatao Gu+2 more|arXiv (Cornell University)|Oct 22, 2019

Neural Networks and ApplicationsComputer Science18 references61 citations

TL;DR

This paper extends Transformer-based sequence-to-sequence models with adaptive decoding depth, allowing predictions at intermediate decoder layers and learning when to exit to balance speed and accuracy. It shows substantial speedups on translation tasks with little to no loss in quality.

ABSTRACT

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.

Motivation & Objective

• Motivate reducing fixed computation in seq2seq models by enabling depth-adaptive decoding.
• Develop mechanisms to predict deployment depth per sequence or per token.
• Investigate training regimes (aligned vs mixed) for stable dynamic computation in Transformers.
• Explore multiple exit classifiers and training targets to optimize speed-accuracy trade-offs.
• Demonstrate effectiveness on IWSLT and WMT translation benchmarks.

Proposed method

• Attach exit classifiers to the output of each decoder block to enable predictions at multiple depths.
• Use aligned or mixed training to handle mismatches between training and inference when using early exits.
• Model exit probabilities q_t(n) for sequence- or token-specific depths and train against exit-based oracles.
• Combine decoding loss with an exit loss via a scalar α to control the speed-accuracy trade-off.
• Experiment with sequence-specific depth ( multinomial and geometric-like exits ) and token-specific depth ( multinomial and geometric-like exits ).
• Evaluate on IWSLT’14 De-En and WMT’14 En-Fr with standard Transformer baselines for comparison.

Experimental results

Research questions

• RQ1 Can Transformer decoders emit outputs at intermediate layers without retraining the entire model?
• RQ2 What mechanisms for depth prediction yield the best speed-accuracy trade-off in neural machine translation?
• RQ3 How do aligned vs mixed training regimes affect performance and stability of dynamic depth models?
• RQ4 Which exit classifier designs ( multinomial vs geometric-like ) and oracle guidance optimize decoding efficiency?
• RQ5 How do depth-adaptive approaches scale from small (IWSLT) to large (WMT) translation tasks?

Key findings

• Adaptive depth can match a six-layer Transformer’s accuracy while using fewer than three decoder blocks on IWSLT De-En.
• Aligned training outperforms mixed training in both fixed-exit and randomly exited scenarios and is computationally efficient.
• Token-specific depth with geometric-like exits and a correctness-based oracle delivers strong speed-accuracy trade-offs, approaching or matching baseline BLEU with substantially fewer blocks.
• On WMT’14 En-Fr, adaptive depth methods achieve BLEU near the best baseline with 40-60% fewer decoding blocks, though gains are more modest than for smaller tasks.
• Confidence-thresholding exits can match full-model accuracy with significant speedups but incur higher per-token classifier overhead on large vocabularies.