[Paper Review] Training Deeper Neural Machine Translation Models with Transparent Attention

Ankur Bapna, Mia Xu Chen+3 more|arXiv (Cornell University)|Aug 22, 2018

Natural Language Processing Techniques26 references16 citations

TL;DR

This paper proposes Transparent Attention, a method that enhances gradient flow in deep neural machine translation models by allowing the decoder to attend to weighted combinations of all encoder layer outputs, not just the top layer. The technique enables successful training of deeper Transformer and Bi-RNN encoders (up to 20 layers), achieving consistent 0.7–1.1 BLEU gains on WMT’14 En→De and WMT’15 Cs→En benchmarks, with improved optimization stability and performance over deeper models.

ABSTRACT

While current state-of-the-art NMT models, such as RNN seq2seq and Transformers, possess a large number of parameters, they are still shallow in comparison to convolutional models used for both text and vision applications. In this work we attempt to train significantly (2-3x) deeper Transformer and Bi-RNN encoders for machine translation. We propose a simple modification to the attention mechanism that eases the optimization of deeper models, and results in consistent gains of 0.7-1.1 BLEU on the benchmark WMT'14 English-German and WMT'15 Czech-English tasks for both architectures.

Motivation & Objective

• To address the challenge of training extremely deep neural machine translation models, particularly beyond 12 layers, where standard attention mechanisms fail due to poor gradient flow.
• To improve optimization stability in deep Transformer and Bi-RNN encoders by modifying the attention mechanism to support direct gradient paths across all encoder layers.
• To investigate whether deeper architectures can yield better translation performance than wider ones, despite increased training difficulty.
• To validate that enhanced gradient flow leads to consistent performance gains on standard NMT benchmarks.

Proposed method

• Introduce a learnable weight matrix W of size (N+1)×M, where N is the number of encoder layers and M is the number of decoder attention modules, to assign attention weights across all encoder layers.
• Apply softmax normalization to the weights W to produce attention distributions s_{i,j} = exp(W_{i,j}) / Σ_k exp(W_{k,j}) for each attention head j and layer i.
• Compute a weighted combination z^j_t = Σ_i s_{i,j} h^i_t of encoder hidden states h^i_t across all layers i, which the decoder attends to instead of just the top layer.
• Apply a projection to the weighted combination in RNMT+ before attention to maintain compatibility with the original architecture.
• Use dropout on the weight matrix W during training to stabilize learning and prevent overfitting to attention weights.
• Train models using Adam optimizer with synchronized gradient updates across eight P100 GPUs, using standard WMT’14 and WMT’15 datasets with subword tokenization and shared 32k vocabularies.

Experimental results

Research questions

• RQ1 Can deeper Transformer and Bi-RNN encoders be successfully trained for neural machine translation when standard attention mechanisms fail beyond 12 layers?
• RQ2 Does introducing a weighted, learnable attention mechanism across all encoder layers improve gradient flow and optimization stability in deep models?
• RQ3 To what extent do deeper models with Transparent Attention outperform standard wide models in terms of translation quality and parameter efficiency?
• RQ4 How does the distribution of attention weights evolve during training, and does it reflect an ensembling-like behavior across encoder layers?

Key findings

• The vanilla Transformer model fails to train with 12 or more encoder layers, but Transparent Attention enables successful training of up to 20-layer encoders.
• Transparent Attention achieves consistent improvements of 0.7–1.1 BLEU points on both WMT’14 En→De and WMT’15 Cs→En translation tasks across all depth settings.
• The model with 20-layer encoder and Transparent Attention outperforms standard Transformer Big despite having less than half the model capacity.
• The attention weights s_{i,j} evolve such that lower layers (e.g., embeddings and early layers) receive higher attention early in training, while deeper layers (e.g., 16 and 20) gain prominence as training progresses, suggesting dynamic feature integration.
• The gradient norm ratio (r_t) plots show that Transparent Attention stabilizes training dynamics, resembling those of well-trained RNMT+ models and reducing vanishing gradient effects.
• The performance gains are partially attributed to an implicit ensembling effect, where multiple encoder layer representations are dynamically combined via attention, similar to contextualized embeddings in models like BERT.