Cross-Attention is All You Need: Adapting Pretrained Transformers for Machine Translation

In plain words

A translation model trained for one pair of languages can be reused for a new pair, but the usual way is to update every weight again. This paper asks how much of that work is needed. It freezes most of the model and updates only the attention layers that link the reading half to the writing half, plus the word-piece tables for the new language. On the tested pairs, the result is close to retraining everything, while far fewer values have to be saved for each new pair.

The authors also noticed that the new word tables end up lined up with the original model's tables, so the new language is placed in the same space as the old one. This helps in two ways the paper tests: the model keeps working on its original language pair, and it can translate between two new languages without any direct training data for that pair.

TL;DR

• Starting from a French-English parent translation model, the authors update only the cross-attention layers and new embeddings for each child pair §2.1 §3.1.
• Across six child pairs, this trails full fine-tuning by 0.1 to 2.0 BLEU, and it stores about 17% of parameters per pair, averaged over the pairs tab. 2.
• Cross-attention updates need the pretrained cross-attention values: random cross-attention is worse, and in two pairs falls below training from scratch §4.1 tab. 2.
• Learned embeddings align with the parent's, which the authors use against forgetting and for zero-shot De–Es, Ro–Es and Ro–De translation §4.2 §5.1 §5.2.
• Results are author-reported; the paper gives no seeds or variance, and the parent and mBART recipes are not fully stated in the paper §3.2 §4.

Problem

Fine-tuning a pretrained translation model on a new language pair updates all its parameters. The paper notes that prior work found cross-attention more important than self-attention when pruned or hard-coded, but studied this when training from scratch, not when transferring §1 §2. Full fine-tuning also means a complete new model must be stored for each pair §1.

Scope: the child pair must share one language with the parent, so either the source or the target side changes §2.1. Target-side transfer is the harder case in the results §4.1.

Method

The parent is a French-English Transformer. For each child pair, the new-language embeddings (token, positional and layer-norm parameters) are initialised randomly, and every other parameter starts from the parent §3.1. The shared-language embeddings stay frozen §3.1. The paper groups parameters as source embeddings, target embeddings, encoder body, decoder body and cross-attention (with its layer norm) §2.1.

Five runs are compared per child pair §4 tab. 2:

1. scratch: trained from random initialisation on the child data §3.1.
2. {src,tgt}: only the new-side embeddings are trained §3.1.
3. {src,tgt}+body: the new embeddings and all of the encoder, decoder and cross-attention are trained §3.1.
4. {src,tgt}+xattn: the new embeddings and cross-attention are trained; encoder and decoder stay frozen §3.1 fig. 1.
5. {src,tgt}+randxattn: as variant 4, but cross-attention starts random §3.1.

An mBART-based ablation uses a denoising-pretrained parent whose cross-attention has no translation training §3.1 §4.1 fig. 2.

Analysis methods:

• Bilingual lexicon induction: nearest-neighbour retrieval of child-side subwords against the parent embeddings, scored by exact matches in MUSE dictionaries §4.2.
• Forgetting: the fine-tuned child embeddings are plugged back into the child model, and the model is evaluated on the parent's Fr–En test set §5.1 fig. 4.
• Zero-shot: the encoder embeddings of a De–En model are swapped into a Fr–Es model, giving a De–Es model with no De–Es training §5.2.

Child-model construction follows the training-setup code in fairseq_cli/train.py lines 71 to 84 train.py.

# Build the child model from the parent  [§3.1]
# Input:  parent weights, new source vocabulary
# Output: child model

def build_child_model(parent_weights, new_src_vocab):
    # Architecture with the new source vocabulary  [§3.1]
    child = new_model(new_src_vocab)
    for name in child.parameter_names:
        # New source embeddings keep random init  [§3.1]
        if name in SOURCE_EMBEDDING_NAMES:
            continue
        # Everything else, cross-attention included, is copied  [§2.1]
        child[name] = parent_weights[name]
    return child

Trainable-parameter selection is at fairseq_cli/train.py lines 97 to 116 train.py.

# Choose the trainable parameters  [§2.1] [§3.1]
# Input:  child model
# Output: sets requires_grad on every parameter

def mark_trainable(child):
    for name in child.parameter_names:
        # Cross-attention and its layer norm ("encoder_attn")  [§2.1]
        is_cross_attention = "encoder_attn" in name
        # New source-side embeddings, positions and layer norm  [§3.1]
        is_new_src_embedding = name in SOURCE_EMBEDDING_NAMES
        # Encoder, decoder and shared-side embeddings stay frozen  [§3.1]
        child.requires_grad[name] = (
            is_cross_attention or is_new_src_embedding
        )

The optimiser and loss settings are in the released script src+xattn.sh lines 16 to 37 …_src+xattn.sh.

# One training step  [§3.1]
# Input:  batch of (source tokens, target tokens)
# Output: updated trainable parameters

def train_step(batch):
    # Forward pass through the whole child model  [§3.1]
    logits = child(batch.source_tokens, batch.target_input)
    # Label-smoothed cross-entropy, smoothing 0.1  [code: src+xattn.sh#L16]
    loss = label_smoothed_cross_entropy(logits, batch.target_output, 0.1)
    # Adam (0.9, 0.98), eps 1e-6, no weight decay  [code: src+xattn.sh#L18]
    # Peak LR 3e-5, polynomial decay, 2500 warmup  [code: src+xattn.sh#L21]
    optimizer.step(loss, parameters=trainable_parameters(child))

Hyperparameters

The paper states the architecture and the vocabulary sizes, but none of the optimiser, learning-rate, batch, dropout, stopping or decoding settings. Those values below come from the released scripts, not from the paper. Script citations give the directory scripts/training/, a …_ prefix that stands for the common stem reinit-src-embeddings-and-finetune-parent-model-on-translation, and the suffix (src+xattn.sh is the cross-attention run, src+body.sh the full body run, and src.sh the embedding-only run).

Child data and vocabulary from tab. 1:

Results

Author-reported. Setup: French-English parent, transferred to six child pairs, BLEU on each pair's test set, Transformer base §3.2 §4. The percentage in parentheses is the share of parameters that must be stored for each setting; it is averaged over the six pairs and includes each pair's new vocabulary §4.1. Differences in parentheses are relative to {src,tgt}+body tab. 2.

Other reported numbers, all author-reported:

• Storage for the six new pairs: {src,tgt}+xattn stores 124,430,336 parameters, against 313,583,616 for {src,tgt}+body §4.1.
• Excluding embeddings, cross-attention needs about 9% of a full model's storage and the full body about 67% §4.1 tab. 2.
• Bilingual lexicon induction (German to French, 2,025 of 8,000 German subwords in MUSE): 55% exact-match accuracy for src+xattn against 19.7% for src+body; on a manual sample of 40 words, about 80% against about 30% §4.2 app. A, tab. 4. Other pairs are in fig. 3 and are not given in the text.
• Forgetting: the parent scores 35.0 BLEU on Fr–En. Child models that keep their cross-attention retain more of it than full-body models, and tgt+body retains 0.2 BLEU §5.1 fig. 4.
• Zero-shot (De–Es on newstest2013, Ro–Es and Ro–De on TED test sets): 9.2, 14.7 and 9.8 BLEU, against supervised from-scratch baselines of 18.3, 18.6 and 13.4 §5.2 tab. 3. Two De–Es models trained on 50K and 100K sentences score 7.2 and 12.0 §5.2.
• mBART as parent: embed-only fine-tuning fails. Source-side fine-tuning beats scratch on Ro–En, De–En and Ha–En, and embed beats scratch only on Ja–En, by 0.1 BLEU §4.1 fig. 2. The chart values are not in the text.

Paper vs code

The official repository is MGheini/xattn-transfer-for-mt (MIT, owner Mozhdeh Gheini, first author), linked from the paper's abstract §1 footnote. Checked at bf86b02. It is a fairseq fork with shell scripts and does not include the parent model or the data.

• Hyperparameters: the paper states none of the optimiser, learning-rate, batch, dropout or stopping settings; the Hyperparameters table gives the script values [code: scripts/training/…_src+xattn.sh#L14-L51].
• Cross-attention set: the paper defines cross-attention with its layer norm §2.1. The code selects every parameter whose name contains encoder_attn [code: fairseq_cli/train.py#L109-L116], which covers both the attention and its layer norm. No discrepancy.
• Shared-side freezing: the paper freezes the common-language embeddings §3.1. The code freezes the decoder embedding, its positions, its layer norm and the output projection [code: fairseq_cli/train.py#L97-L105]. No discrepancy.
• Architecture: the paper says Transformer base §3.2. The scripts also use pre-norm, layer norm on embeddings and learned positions [code: ...src+xattn.sh#L41-L46]. Fairseq's base defaults are not these [code: fairseq-modified/fairseq/models/transformer.py#L934-L935].
• Dictionary induction similarity: the paper's footnote and text say nearest neighbour by cosine similarity §4.2. The code takes a plain dot product of the stored vectors [code: scripts/dict_induction/induce_dict.py#L51], and the extraction script writes them without normalising [code: scripts/dict_induction/extract_embeddings.py#L23-L25]. The released induction therefore does not compute cosine.
• Match rule: the paper says exact match against the gold dictionary §4.2. The code lowercases both the induced and the gold words and drops the first character of each (the word-start marker) before comparing [code: scripts/dict_induction/induce_dict.py#L58].
• Vocabulary rows: extraction skips the first four rows, the special tokens [code: scripts/dict_induction/extract_embeddings.py#L23]. The stitching script drops the last row of the source matrix on the default path [code: scripts/model_stitiching/plug_in_src_embeddings.py#L27]. The paper says neither.
• Recipe differences between runs: the embedding-only scripts use peak LR 3e-4, inverse square root and warmup 4000 [code: scripts/training/…_src.sh#L21, #L30, #L37]. The cross-attention and body scripts use 3e-5, polynomial decay and warmup 2500 [code: ...src+xattn.sh#L21, #L30, #L37]. The paper does not mention this.
• Stopping: patience 25 is set, but the metric is validation loss, the fairseq default [code: fairseq-modified/fairseq/options.py#L520], not BLEU. Early stopping is not described in the paper.

Pitfalls

• No training recipe in the paper. Learning rate, batch size, schedule, dropout, label smoothing and stopping are absent from the text. Anything a reimplementer needs comes from the scripts, not the paper, so the card's Hyperparameters table is the reference.
• Confounded baselines. The embedding-only scripts use a peak learning rate of 3e-4 with inverse square root and a 4000-step warmup, while the cross-attention and body scripts use 3e-5 with polynomial decay [code: ...src.sh#L21, #L30; ...src+xattn.sh#L21, #L30]. Row differences in Table 2 therefore mix the training recipe with the choice of parameters.
• mBART comparison. The mBART scripts use label smoothing 0.2, dropout 0.3, --max-tokens 512 and a different update frequency from the translation-parent runs [code: scripts/training/finetune-mbart-on-transaltion_embed.sh#L18, #L40, #L27, #L31]. The paper's mBART comparison is not on one recipe, and its chart values are not in the text.
• Batch size depends on hardware. The scripts set --max-tokens 2048 and --update-freq 8 [code: …_src+xattn.sh#L25, #L29]. In fairseq these are a per-device token budget and a gradient-accumulation count (fairseq behaviour, not checked in this repo), so the effective batch depends on the GPU count, which the paper does not state.
• Stopping on loss, not BLEU. Checkpoint selection uses validation loss (fairseq default), and evaluation loads checkpoint_best.pt [code: scripts/evaluation/decode_and_score_valid_and_test.sh]. A reimplementation that stops on BLEU will select different checkpoints.
• Single seed. Each script sets one seed (1 for translation-parent runs, 222 for mBART) [code: ...src+xattn.sh#L14]. The paper reports no variance across seeds, so differences of a few tenths of a BLEU point in Table 2 are not established as meaningful.
• Vocabulary and storage numbers. Table 2's percentages average the six pairs and include the new vocabulary's embeddings, so they move with vocabulary size. The 9% and 67% figures in §4.1 are obtained by subtracting the {src,tgt} row from the cross-attention and body rows §4.1.
• Target-side transfer shows larger gaps. The drop from full-body fine-tuning is bigger for Fr–Es (-0.8) and Fr–De (-1.8) than for the source-side pairs, and the {src,tgt} gains are larger there too tab. 2. The authors attribute this to two of three cross-attention matrices (key and value) seeing the new language when the target changes §4.1.
• Random initialisation of new embeddings. The paper says new embeddings are random §3.1 but does not give the scheme. The code keeps whatever fairseq's model build produces for the new vocabulary; this was not checked beyond that.
• Zero-shot is not directly comparable. The zero-shot pairs are evaluated on TED talks and newstest2013, and the supervised baselines use different corpora from Table 2 §5.2 tab. 3. Decoding settings for the stitched models are not stated in the paper.
• Data and parent are not in the repository. The README places data on the main branch; this checkout is only-code, and the Fr–En parent is not provided [README]. Reproduction needs the parent recipe from elsewhere.

Lineage

• Related, not compared in experiments: adapter modules (Houlsby et al., 2019) §6 and module-freezing studies (Artetxe et al., 2020, and a cited work on frozen Transformers) §6. The paper contrasts its own approach with adapters, which add new parameters; this one reuses the model's own cross-attention §6.
• Prior work on cross-attention importance when training from scratch: pruning (Voita et al., 2019, and a NeurIPS 2019 pruning study) and hard-coding (You et al., 2020) §1 §6.
• Parent transfer baseline: fine-tuning a pretrained MT model on a new pair, with the vocabulary handling of Zoph et al. (2016) §2.1 §3.1.
• No later work is linked from this wiki.