Source-linked AI summary
Accelerating Neural Transformer via an Average Attention Network
Biao Zhang, Deyi Xiong, Jinsong Su
TL;DR
Transformer training is parallelizable, but autoregressive generation and decoder self-attention slow decoding. The paper replaces decoder self-attention with an average attention network using cumulative averaging, gating, masking, and dynamic programming. Across six WMT17 language pairs, it reports over 4 times faster decoding with almost no loss in translation quality or training speed.
Problem
Transformer decoding remains sequential under autoregressive generation, while decoder self-attention takes increasing time as target words are generated.
Source: §1 · pp. 1–2Method
AAN replaces decoder self-attention with cumulative averaging and a gating layer, using masking for parallel training and dynamic programming for decoding.
Source: §1 · p. 2Results
On six WMT17 language pairs, AAN decodes over 4 times faster with almost no loss in translation quality or training speed.
Source: §1 · p. 2, §5.2.1 · p. 9Takeaways & Limitations
The authors report that the decoding-speed benefit is language-invariant and can be adapted to other translation tasks.
Source: §5.2.1 · p. 9
Abstract
from arXiv · showhide
With parallelizable attention networks, the neural Transformer is very fast to train. However, due to the auto-regressive architecture and self-attention in the decoder, the decoding procedure becomes slow. To alleviate this issue, we propose an average attention network as an alternative to the self-attention network in the decoder of the neural Transformer. The average attention network consists of two layers, with an average layer that models dependencies on previous positions and a gating layer that is stacked over the average layer to enhance the expressiveness of the proposed attention network. We apply this network on the decoder part of the neural Transformer to replace the original target-side self-attention model. With masking tricks and dynamic programming, our model enables the neural Transformer to decode sentences over four times faster than its original version with almost no loss in training time and translation performance. We conduct a series of experiments on WMT17 translation tasks, where on 6 different language pairs, we obtain robust and consistent speed-ups in decoding.
1 Introduction
Transformer training is parallelizable, but autoregressive decoding and decoder self-attention make generation sequential and increasingly costly. The paper proposes AAN to replace decoder self-attention with cumulative averaging and gating, using masking and dynamic programming to preserve parallel training and accelerate decoding.
1 Introduction
- Autoregressive decoding is sequential, and decoder self-attention takes more time as target words accumulate because it models dependencies on previous words.Source: §1 · pp. 1–2
- AAN combines a cumulative-average layer over previous positions with a feed-forward gating layer that improves expressiveness.Source: §1 · p. 2
- Masking enables parallel AAN training, while dynamic programming lets decoding use only the previous hidden state.Source: §1 · p. 2
- By treating each previous word as an equal contributor, AAN maintains connection signals from prior inputs at any input length, which the authors argue supports long-range dependencies.Source: §1 · p. 2
- On six WMT17 language pairs, AAN decodes over 4 times faster with almost no loss in translation quality or training speed.Source: §1 · p. 2
2 Related Work
Related work contrasts the parallelization and dependency modeling of Transformer with the sequential RNN and local dependencies of shallow CNNs. The paper positions AAN as a simpler attention design for faster decoding that aims to preserve translation performance.
2 Related Work
- RNNs model long-range dependencies but cannot parallelize; CNNs enable alternatives to recurrence, while shallow CNNs capture only local dependencies and Transformer offers parallelization with short dependency paths.Source: §2 · p. 2
- Prior attention studies explored accuracy and capability, whereas this work focuses on simplifying self-attention to accelerate decoding.Source: §2 · pp. 2–3
- Unlike Gu et al. (2018), which speeds decoding by generating words without prior predictions at the cost of translation quality, this work aims to preserve translation performance.Source: §2 · p. 3
3 The Average Attention Network
The Average Attention Network (AAN) replaces adaptive self-attention weights with cumulative averages of previous inputs and a gating layer. Masking enables parallel training, while dynamic programming lets decoding use only the previous state.
3 The Average Attention Network
- The cumulative-average layer replaces adaptive attention weights with fixed average weights, connecting each representation to previous inputs with invariant strength for long-range dependencies.The average over previous positions produces context-sensitive representations.Source: §3 · p. 3
- A feed-forward gating layer combines current input with averaged context, controlling how much past information is preserved and how much new input is captured.The input and forget gates regulate these contributions.Source: §3 · p. 3
3.1 Parallelization in Training
- A masking trick enables parallel computation of the cumulative averages, allowing all masked AAN components to be fully parallelized during training.Figure 3 visualizes the mask-based cumulative-average operation on input embeddings.Source: §3.1 · pp. 3–4, Figure 3 · p. 4
3.2 Model Analysis
- Masked AAN and self-attention both have O(n^2·d + n·d^2) complexity and O(1) sequential operations; AAN has O(1) maximum path length.Original AAN requires O(n) sequential operations.Source: §3.2 · p. 4
3.3 Decoding Acceleration
- Dynamic programming computes each decoding representation from only the previous state, rather than requiring all previous states as self-attention does.This is possible because the average calculation can be decomposed into two steps.Source: §3.3 · p. 4
4 Neural Transformer with AAN
The Transformer’s decoder replaces target-side self-attention with AAN, while retaining an encoder-decoder architecture. Decoder layers use AAN to model dependencies on previous target words, then incorporate source semantics before predicting the next token.
4 Neural Transformer with AAN
- AAN replaces the decoder’s self-attention network in the Transformer architecture.The overall model retains an encoder-decoder framework.Source: §4 · p. 4
- The encoder comprises six identical layers, and MHAtt denotes the multi-head attention mechanism used in the Transformer.Source: §4 · p. 5
- Each of the six decoder layers uses AAN as its first sub-layer to capture dependencies on previously predicted target words.Two further sub-layers seek translation-relevant source semantics.Source: §4 · p. 5
- The decoder’s top-layer representation sN is used to compute next-token probabilities through a linear transformation and softmax.Source: §4 · p. 5
- Frequency-based positional encodings are added to the input layer, and the model is trained on bilingual data with a maximum-likelihood objective.Source: §4 · p. 5
5.1 WMT14 English-German Translation
On WMT14 English–German, the AAN matches the Transformer’s translation quality and convergence while substantially reducing decoding time. Analyses also report similar training speed and promising translation quality on long source sentences.
5.1.1 Model Settings
- The WMT14 English–German evaluation uses a 4.5M-pair training corpus and case-sensitive BLEU; models use 512-dimensional inputs and outputs, 8 attention heads, and a 32K vocabulary.Training was capped at 100K steps, and evaluation averaged the last five saved models.Source: §5.1 · p. 5, §5.1.1 · pp. 5–6
5.1.2 Translation Performance
- The Transformer scores 26.37 BLEU and AAN scores 26.31, with no significant difference reported.Removing the FFN lowers AAN by 0.26 BLEU points, while removing the gating layer lowers it by 0.40 points.Source: §5.1.2 · p. 6, Table 2 · p. 6
5.1.3 Analysis on Convergency
- AAN and the Transformer converge similarly and require almost the same number of training steps for high-quality translation.The convergence plot shows training loss against training steps scaled by 10^2.Source: §5.1.3 · p. 6, Figure 5 · p. 6
5.1.4 Analysis on Speed
- AAN processes about 0.2464 training steps per second, versus 0.2474 for the Transformer.The authors describe these training speeds as similar.Source: §5.1.4 · p. 6, Table 3 · p. 6
- Across beam sizes 4–20, AAN decoding takes only one quarter of the Transformer’s time.Table 3 reports training and decoding time, averaged over newstest2014.Source: Table 3 · p. 6, §5.1.4 · p. 7
5.1.5 Effects on Sentence Length
- On the longest source sentences, AAN produces relatively shorter translations but significantly better translation quality.Figure 6 compares BLEU and average translation length by source-sentence length; Figure 7 plots translation time by source length.Source: §5.1.5 · p. 7, Figure 6 · p. 7, Figure 7 · p. 7
5.2 WMT17 Translation Tasks
The WMT17 evaluation spans six language pairs with training corpora ranging from 0.21M to 52M sentence pairs. Across the 12 translation directions, AAN delivers consistently faster decoding and competitive translation results relative to the Transformer.
5.2 WMT17 Translation Tasks
- The six WMT17 language pairs use training corpora ranging from 0.21M to 52M sentence pairs, with newstest2017 as the test set.The evaluation covers both translation directions for each pair.Source: §5.2 · p. 8
5.2.1 Translation Results
- AAN decodes more than 4 times faster than the Transformer across WMT17 language pairs, while remaining competitive in translation quality across all 12 directions.Table 5 reports average seconds per source sentence; the authors describe the speed benefit as language-invariant.Source: §5.2.1 · p. 9, Table 5 · p. 9
6 Conclusion and Future Work
The paper concludes that AAN alleviates the Transformer’s decoding bottleneck, reporting over fourfold decoder speedups in experiments on WMT14 and six WMT17 language pairs. Future work will extend evaluation to other sequence-to-sequence tasks and seek stronger modeling ability.
6 Conclusion and Future Work
- AAN combines cumulative averaging and a feed-forward gating layer, with masking and dynamic programming, and speeds up the Transformer’s decoder by over 4 times in the reported experiments.The experiments cover one WMT14 and six WMT17 language pairs.Source: §6 · p. 9
- Future work will apply AAN to other sequence-to-sequence tasks and aim to improve its modeling ability to consistently outperform the Transformer.Source: §6 · p. 9