Incorporating BERT into Neural Machine Translation

In plain words

A translation model reads a sentence in one language and writes it in another. A large pre-trained language model such as BERT already knows a lot about language, but the usual way to use it is to start the translator from BERT's weights, which helps little here. This paper takes a different route. BERT is run over the source sentence once, and it stays frozen. Its output is then made available to every layer of the translator, so each layer can look at BERT's view of the sentence through an extra attention step, next to its normal attention to the words.

To keep the translator from leaning on BERT alone, the training randomly switches off one of the two attention paths in some layers. The authors report that this helps on several translation settings, including sentence and document level, and on translation without parallel data. The cost is that every translation also runs BERT, so inference is slower and the model needs room for a second network.

TL;DR

• Frozen BERT output is fed into every encoder and decoder layer through an extra attention branch, averaged with the usual attention §4.1 eq. 1 eq. 2.
• A drop-net rule randomly uses one branch only during training, which the authors report reduces overfitting §4.2 fig. 2.
• Author-reported gains over Transformer baselines on IWSLT, WMT, document-level, semi-supervised and unsupervised translation §5 §7 tab. 2 tab. 3 tab. 5 tab. 7.
• Costs: about 39 to 49 percent more inference time, and a second network stored with the model §4.3 tab. 13.

Problem

Pre-trained encoders such as BERT are usually used by fine-tuning them on the downstream task, or by feeding their output as input embeddings. The authors first try both for NMT. Initialising the encoder from BERT gave a lower BLEU than a plain Transformer on IWSLT'14 En-De, and using BERT output as input embeddings gave a small gain §3 tab. 1. Training BERT from scratch for NMT is too costly for many researchers, so the work assumes a released pre-trained BERT is used as is §1. BERT and the NMT model may also use different tokenisations, so their sequence lengths can differ §4.1.

Method

The NMT model stays a standard encoder-decoder Transformer. A frozen BERT encodes the source sentence once. Its last-layer output is a fixed sequence of representations, one per BERT wordpiece §4.1, Step 1. Every encoder layer then has a second attention module that queries this sequence, and every decoder layer has one that queries it too §4.1, Steps 2 and 3.

Encoder layer (Step 2). For layer l, the output of self-attention over the previous layer's states is averaged with the output of BERT attention over the BERT output. The same feed-forward block as the Transformer follows eq. 1 eq. 6 eq. 7. The two attention modules have separate parameters eq. 1.

Decoder layer (Step 3). Self-attention runs first. The result is averaged with the output of BERT-decoder attention and the output of encoder-decoder attention, then passed through the feed-forward block eq. 2. The text says the model uses the encoder's final output for the encoder-decoder branch eq. 2.

Drop-net (§4.2). During training, each layer draws one uniform number per iteration. With probability p/2 only one branch is kept, and with probability p/2 only the other; otherwise both are kept, each halved eq. 3 eq. 4. Inference uses the expected value, which equals the averaged form eq. 3. The paper fixes p = 1.0 for all main experiments unless stated §5.1 §6.

The training recipe (§5.1). First train a plain NMT model to convergence. Copy its encoder and decoder weights into the BERT-fused model. The BERT-encoder and BERT-decoder attention modules are new and start randomly initialised. Keep BERT frozen. Train the rest with Adam and an inverse-square-root learning rate schedule §5.1 app. A.2.

# BERT-fused encoder layer  [eq. 1]
# Input:  x (B, L, d) NMT states; bert_out (B, L_b, d_bert) frozen BERT output
# Output: fused states (B, L, d)

def bert_fused_encoder_layer(x, bert_out):
    # Attention over the layer's own states (self-attention)  [eq. 1]
    self_out = self_attention(query=x, key=x, value=x)
    # Attention over the frozen BERT output  [eq. 1]
    bert_attn_out = bert_attention(query=x, key=bert_out, value=bert_out)
    # Average the two branches  [eq. 1]
    fused = 0.5 * (self_out + bert_attn_out)
    # Feed-forward block of the Transformer  [eq. 7] [§4.1]
    return feed_forward(fused)

# BERT-fused decoder layer  [eq. 2]
# Input:  s (B, T, d) decoder states; bert_out (B, L_b, d_bert); enc (B, L, d)
# Output: (B, T, d)

def bert_fused_decoder_layer(s, bert_out, enc):
    # Causal self-attention over the decoder states  [eq. 2]
    s_hat = self_attention(query=s, key=s, value=s, causal=True)
    # Average BERT-decoder attention with encoder-decoder attention  [eq. 2]
    s_tilde = 0.5 * (
        bert_attention(query=s_hat, key=bert_out, value=bert_out)
        + encoder_attention(query=s_hat, key=enc, value=enc)
    )
    # Feed-forward block  [eq. 7] [§4.1]
    return feed_forward(s_tilde)

# Drop-net branch weights for one layer  [eq. 3]
# Input:  p (drop-net rate in [0, 1]); training flag
# Output: (w_self, w_bert) weights for the two branches

def drop_net_weights(p, training):
    # Inference uses the expectation, the averaged form  [eq. 3]
    if not training:
        return 0.5, 0.5
    # One uniform draw per layer and iteration  [§4.2]
    draw = uniform(0, 1)
    if draw < p / 2:
        # Keep the self-attention branch only  [eq. 3]
        return 1.0, 0.0
    if draw > 1 - p / 2:
        # Keep the BERT branch only  [eq. 3]
        return 0.0, 1.0
    # Keep both, each halved  [eq. 3]
    return 0.5, 0.5

Key equations, as stated in the paper:

\tilde h^l_i = \tfrac{1}{2}\big(\mathrm{attn}_S(h^{l-1}_i, H^{l-1}_E, H^{l-1}_E) + \mathrm{attn}_B(h^{l-1}_i, H_B, H_B)\big) eq. 1

\tilde s^l_t = \tfrac{1}{2}\big(\mathrm{attn}_B(\hat s^l_t, H_B, H_B) + \mathrm{attn}_E(\hat s^l_t, H^L_E, H^L_E)\big) eq. 2

Attention is standard multi-head Transformer attention with learned query, key and value projections app. E, eq. 6. The feed-forward block is W2 max(W1 x + b1, 0) + b2 with layer normalisation as in the Transformer app. E, eq. 7.

Hyperparameters

Values in the paper are cited to the paper. Where the paper says "not stated", the code's training command is given in the Notes column.

Results

Author-reported. Setup: each row is a single run of each system unless stated; the baseline is a Transformer trained with the same settings §5.1. Numbers are BLEU.

Table 2, IWSLT'14 and IWSLT'17, Transformer vs BERT-fused tab. 2 §5.2:

Table 3, WMT'14 tab. 3 §5.2: compared with published and reproduced systems.

Table 4, document-level translation, IWSLT'14 tab. 4 §5.3:

The document-level BERT-fused model is reported as significantly better than the sentence-level baseline at p < 0.01 §5.3.

Table 5, WMT'16 Romanian→English, semi-supervised tab. 5 §5.4:

Table 6, ablation on IWSLT'14 En→De tab. 6 §6:

Table 7, unsupervised NMT tab. 7 §7:

Table 1, preliminary exploration, IWSLT'14 En→De tab. 1 §3:

Appendix B.2 ensembles and deeper models on IWSLT'14 En→De tab. 8 app. B.2:

Appendix B.2 Part III, ensembles on IWSLT'14 De→En tab. 9 app. B.2:

Appendix B.3, other feeding strategies tab. 10 app. B.3:

Appendix B.4, prior De→En results cited by the paper tab. 11 app. B.4:

Appendix B.5, back-translation with Wikipedia sentences on IWSLT'14 En→De tab. 12 app. B.5:

Appendix C, inference time in seconds tab. 13 §4.3:

Figure 2 fig. 2 plots training loss, validation loss and validation BLEU for each drop-net rate. The text says larger rates give larger training loss and smaller validation loss and better validation BLEU. The numbers are not given in the text.

Paper vs code

The official code is at the link in the abstract: bert-nmt/bert-nmt at commit fcb616d28091ac23c9c16f30e6870fe90b8576d6. It is a fork of fairseq and includes a bert/ folder. The fairseq files are BSD-3-Clause, and the bert/ files are Apache-2.0 LICENSE modeling.py. The table above uses the training flags in the README. Discrepancies:

• Combination weights, default path. Eq. 1 and eq. 2 average the two branches with ½ each eq. 1 eq. 2. The code's default weights are 1 and 1 translation.py, so with drop-net off the output is the sum of the two branches. With --encoder-bert-dropout on, the code uses the paper's ½ at inference transformer.py, which matches eq. 3. The paper's reported numbers come from drop-net with p = 1.0 §5.1, so the ½ path is the one that matches them.
• Residual and dropout. Eq. 1 and eq. 2 have no residual term and no dropout. The code adds a residual connection, applies dropout to each branch before combining, and then layer-normalises transformer.py. Figure 1 draws residual connections fig. 1, so the residual is consistent with the figure but not with the equation as written.
• Layer norm placement. The paper says layer normalisation follows the Transformer app. E. The BERT-fused code path defaults to post-norm (encoder_normalize_before False) transformer.py.
• Decoder drop-net labelling. For the decoder, eq. 4 pairs the low-draw event with BERT-only. The code pairs it with encoder-attention-only transformer.py. Both events have the same probability, so the expected output is the same.
• Drop-net rate. The paper states a rate p = 1.0 §5.1. The code's flag --encoder-bert-dropout-ratio is the per-side probability: a draw below the flag keeps one branch, and a draw above one minus the flag keeps the other transformer.py. The README allows [0, 0.5] and its example uses 0.5 README.md README.md. This implies the flag should be p/2 = 0.5 when p = 1.0. That reading is inferred from the code, not stated in the paper. The code default is 0.25 translation.py.
• Architecture name. The paper names the config transformer_iwslt_de_en §5.1. The BERT attention branch exists only in the transformers2 model family transformer.py, and the README trains with transformer_s2_iwslt_de_en README.md. The plain transformer encoder layer has no BERT branch transformer.py.
• Extra options not in the paper. bert_gates (all on by default, so BERT is used in every layer, as in the paper) transformer.py. --encoder-bert-mixup replaces the drop-net draw with a continuous mix translation.py transformer.py. --mask-cls-sep (off by default) removes the [CLS] and [SEP] positions from the BERT attention mask translation.py fairseq_model.py. The paper does not say whether any of these were used.
• BERT trainability. The paper says BERT is fixed §5.1. By default the code sets BERT parameters to not require gradients train.py. With --finetune-bert, BERT trains and only the pooler is frozen train.py.
• Doc-level input. The paper says the BERT input for document-level translation is the previous sentence and the current sentence, each wrapped in special tokens §5.3. The code wraps each BERT line in [CLS] and [SEP] binarizer.py. How the previous sentence is inserted was not traced.
• Decoder notation. Eq. 2 uses the decoder states up to and including position t+1 eq. 2, while the text defines S before t as the states up to t-1 §4.1, Step 3. The code applies causal self-attention, which is what the README decoding path implies. This reading is inferred.

Pitfalls

• Default flags do not match eq. 1 and eq. 2. Without --encoder-bert-dropout, the code sums the branches (weights 1 and 1), while the equations average them (½ each) translation.py. Use the drop-net path to reproduce the paper's numbers.
• Drop-net flag value. The flag is the per-side probability (p/2), not p. Its default of 0.25 is half the p = 1.0 that the paper uses translation.py.
• Extra network and runtime. BERT runs on every source sentence, in training and in decoding. The authors measure 38.6 to 49.3 percent more inference time tab. 13; Section 4.3 says "about 45%" §4.3. The paper acknowledges extra storage but does not quantify it §4.3.
• BERT tokenisation. The BERT input is a separate, BERT-tokenised file made at preprocessing (*.bert.*), with the same --bert-model-name at training and generation README.md. Interactive decoding must re-tokenise the raw input with BERT's tokenizer README.md. Use the same BERT model for preprocessing, training and generation, as the README asks.
• Case mismatch. The paper lowercases En-De data §5.1. For De→En it uses bert-base-german-cased app. D. Whether the BERT input is lowercased is not stated.
• Random initialisation of the new attention modules. The paper says the BERT attention modules start random §5.1. The init scheme is not stated. The code creates them with the same multi-head attention class as the Transformer transformer.py; the initialiser was not traced.
• Baseline inconsistency in the paper. The De→En Transformer is 34.64 in Table 2 and in the document-level table tab. 2 tab. 4, and 34.67 in appendix Table 9 tab. 9. Pick one and cite it.
• Drop-net sweep has no table. The sweep over p is given only as curves fig. 2. The BLEU at each rate cannot be read from the text.
• Training recipe is only partly stated. Betas, weight decay, warm-up and label smoothing are stated for the Section 3 runs app. A.2, not for the main runs. The main-run values in the Hyperparameters table come from the README README.md. The number of updates is also only in the README.
• Stacked-decoder variant. Appendix B.2 Part I describes a different decoder ordering that gives 29.35 BLEU eq. 5 app. B.2. It is not the main method and is not in the default code path.
• Author-reported only. Baselines are the authors' runs or cited numbers. Significance is reported only for the document-level comparison §5.3. The "state of the art" claims are against the named systems, not a complete survey §5.2 §7.

Lineage

• Builds on: the Transformer, which supplies the attention and feed-forward blocks §2 app. E; BERT, used here as a frozen feature provider §2 §4.1; feature use of ELMo-style context embeddings, which the paper compares against §4.3.
• Compared against: Transformer baselines; DynamicConv and the Evolved Transformer on WMT tab. 3; XLM and MASS on unsupervised translation tab. 7; back-translation §5.4 app. B.5; the document-level model of Miculicich et al. tab. 4.
• Superseded by: not known from content in this wiki.