Source-linked AI summary

RealFormer: Transformer Likes Residual Attention

Ruining He, Anirudh Ravula, Bhargav Kanagal, Joshua Ainslie

arXiv:2012.11747v3cs.LG

TL;DR

Transformer normalization choices do not transfer straightforwardly from computer vision to NLP, motivating a simple way to propagate attention scores across layers. RealFormer adds residual scores before softmax, and the paper reports improvements across diverse Transformer models and tasks alongside sparser attention.

  • Problem

    Although Pre-LN is usually preferable to Post-LN in computer vision, it does not appear to be so for NLP Transformers.

    Source: §1 · p. 1
  • Method

    RealFormer adds each layer’s residual attention scores to the previous layer’s raw scores before softmax computes attention probabilities.

    Source: §1 · p. 1
  • Results

    RealFormer considerably improves diverse state-of-the-art Transformer models on Masked Language Modeling, Neural Machine Translation, and long-document modeling, and tends to produce sparser attention.

    Source: §5 · p. 9
  • Takeaways & Limitations

    As a drop-in Transformer technique, RealFormer adds no parameters or hyperparameters and usually requires only a few lines of code to implement.

    Source: §1 · p. 2

    The discussion notes that Pre-LN could also be used, but focuses on Post-LN because it tends to outperform Pre-LN.

    Source: §3.2 · p. 3

Abstract

from arXiv · show

Transformer is the backbone of modern NLP models. In this paper, we propose RealFormer, a simple and generic technique to create Residual Attention Layer Transformer networks that significantly outperform the canonical Transformer and its variants (BERT, ETC, etc.) on a wide spectrum of tasks including Masked Language Modeling, GLUE, SQuAD, Neural Machine Translation, WikiHop, HotpotQA, Natural Questions, and OpenKP. We also observe empirically that RealFormer stabilizes training and leads to models with sparser attention. Source code and pre-trained checkpoints for RealFormer can be found at https://github.com/google-research/google-research/tree/master/realformer.

1 Introduction

RealFormer adds a direct path for propagating raw attention scores through Transformer layers. The paper reports gains over Post-LN and Pre-LN baselines, improvements when applied to ADMIN and ETC, and qualitatively sparser, more correlated attention.

1 Introduction

  • Unlike the usual ResNet v1/v2 preference in computer vision, Pre-LN does not appear preferable to Post-LN in the NLP literature.The authors suggest self-attention and Transformer architectures may favor different designs from convolutional networks.Source: §1 · p. 1
  • RealFormer propagates raw attention scores by adding each layer’s residual scores to the previous layer’s scores before softmax computes attention probabilities.The technique is described as a simple skip connection over the backbone Transformer and adds no parameters or hyperparameters.Source: §1 · pp. 1–2
  • RealFormer is presented as a generic, low-cost Transformer improvement that usually requires only a few lines of code to implement.The encoder illustration is only a demonstration; the technique is also described as applicable when decoders are involved.Source: §1 · p. 2, Figure 1 · p. 2
  • RealFormer outperforms Post-LN and Pre-LN Transformers across a wide range of pre-training model sizes and achieves competitive downstream results with half as many pre-training epochs.It is also used as a drop-in replacement for ADMIN and ETC, improving those models on various tasks and leading to new state-of-the-art results.Source: §1 · p. 2
  • RealFormer attention is observed to be sparser and more correlated across layers than baseline attention; the authors suggest this may regularize training and benefit fine-tuning.The authors also suggest that this pattern could stabilize training.Source: §1 · p. 2

2 Related Work

Related work spans self-supervised pre-training, Transformer efficiency, and normalization or initialization schemes. The paper positions RealFormer as orthogonal to efficiency methods and evaluates it alongside ADMIN and ETC.

2 Related Work

  • Following BERT, NLP work explored self-supervision objectives, larger pre-training data and hyperparameters, parameter sharing, and multi-task pre-training, typically using Post-LN Transformers.Examples cited include Yang et al. (2019), Clark et al. (2020), Liu et al. (2019), Lan et al. (2019), Sun et al. (2020), and Raffel et al. (2020).Source: §2 · pp. 2–3
  • Efficiency methods target self-attention’s quadratic sequence-length complexity, while RealFormer is described as orthogonal and applicable to such models, including ETC.The cited approaches include low-rank methods, fixed or learnable attention patterns, and memory-based global and local attention.Source: §2 · p. 3
  • Normalization and initialization studies mostly evaluated NMT, and the paper uses ADMIN to test RealFormer with decoders and on top of that line of work.ADMIN is identified as a recent state-of-the-art model on multiple NMT benchmarks.Source: §2 · p. 3

3 RealFormer

The section reviews the standard Transformer encoder and its normalization designs, then presents RealFormer, which carries attention scores across layers and can be applied to Transformer variants.

3.1 Standard Transformer

  • Each encoder layer contains multi-head attention that aggregates value embeddings for queries and a fully connected, one-hidden-layer FFN applied position-wise; Layer Norm stabilizes both sub-layers.The FFN is applied identically at every sequence position.Source: §3.1 · p. 3
  • Post-LN normalizes sub-layer outputs, whereas Pre-LN normalizes their inputs and creates a direct path for token embeddings.Source: §3.1 · p. 3

3.2 Residual Attention Layer Transformer

  • RealFormer adds each layer’s residual attention scores to the corresponding previous-layer scores, then passes the new scores to the next layer.It uses a Post-LN-style Transformer; the authors note Pre-LN is also possible.Source: §3.2 · pp. 3–4
  • RealFormer can be applied to Transformer variations and uses separate skip paths for distinct attention-module types, such as machine-translation attention modules.The paper describes implementation as requiring only a few lines of code.Source: §3.2 · p. 4
  • For very deep networks, the authors find running-mean attention scores helpful as an alternative to the linearly scaling running sum.They describe the running mean as adding a temperature to the softmax.Source: §3.2 · p. 4

4 Experiments

The experiments test RealFormer as a drop-in Transformer replacement across BERT pre-training and downstream NLP, machine translation with ADMIN, and long-context tasks with ETC. They use minimal or no additional hyperparameter tuning, with results spanning model sizes, training budgets, and attention behavior.

4.1 BERT

  • Experiments replace the Transformer backbone in BERT, ADMIN, and ETC, reusing their training or fine-tuning recipes and applying minimal or no extra hyperparameter tuning.The BERT comparisons cover Post-LN, Pre-LN, and RealFormer across a range of model sizes.Source: §4 · p. 4, §4.1 · p. 4

4.1.1 Pre-training Results

  • Across BERT sizes, RealFormer outperforms Post-LN and Pre-LN in MLM accuracy, with its advantage growing as model size increases.The models were pretrained for 1M steps; the authors also report Post-LN instability and occasional divergence for xLarge.Source: §4.1.1 · p. 5, Table 2 · p. 5

4.1.2 Downstream Results

  • On GLUE, RealFormer achieves the best overall performance and beats both baselines on most tasks; it also considerably outperforms them on SQuAD.The reported downstream evaluations fine-tune BERT-Large models on sentence-level GLUE and token-level SQuAD tasks.Source: §4.1.2 · pp. 5–6, Table 3 · p. 6, Table 4 · p. 6

4.1.3 Research Questions

  • With 500K pre-training steps, RealFormer beats 1M-step Post-LN on GLUE by a significant margin and nearly matches it on SQuAD.This comparison uses fine-tuned BERT-Large checkpoints.Source: §4.1.3 · p. 6, Table 5 · p. 6
  • At a 2e-4 learning rate, RealFormer improves from 73.94% to 74.31% MLM accuracy, Pre-LN from 73.21% to 73.46%, and Post-LN diverges.Across dropout settings, RealFormer also outperforms both baselines; increasing dropout alone does not match the apparent regularization from residual attention.Source: §4.1.3 · pp. 6–7, Table 6 · p. 7
  • RealFormer attention has lower variance across layers and is sparser in top layers; adjacent-layer heads also tend to have lower JSD than Post-LN heads.The entropy analysis uses 8,192 held-out examples, and the lower JSD is especially evident for middle-layer heads.Source: §4.1.3 · p. 7, Figure 4 · p. 8

4.2 ADMIN

  • RealFormer outperforms ADMIN baselines across model depths and sets a reported state-of-the-art En-Fr BLEU score of 43.97 without additional data or augmentation.The comparison uses WMT’14 En-De and En-Fr benchmarks.Source: §4.2 · p. 8, Table 7 · p. 9

4.3 ETC

  • RealFormer improves ETC consistently across four benchmarks and reaches 84.4% WikiHop test accuracy, 2.1% above the previous best result.The paper reports this WikiHop leaderboard position as of June 2021.Source: Table 8 · p. 9, §4.3 · p. 9

5 Conclusions

RealFormer adds residual attention to Transformer-based networks as a simple, generic technique. The paper reports improvements across diverse tasks and comparatively sparser attention.

5 Conclusions

  • RealFormer applies the novel idea of residual attention to improve Transformer-based networks.The technique is described as simple, generic, and cheap.Source: §5 · p. 9
  • RealFormer considerably improves diverse state-of-the-art Transformer-based models on Masked Language Modeling, Neural Machine Translation, and long document modeling.Source: §5 · p. 9
  • RealFormer tends to have comparatively sparser attention within heads and across heads in adjacent layers.Source: §5 · p. 9

A.1 Training Details: BERT

BERT experiments used the official repository and TPU cores. The appendix reports hardware-dependent pre-training overhead and specifies the fine-tuning hardware setup.

A.1 Training Details: BERT

  • Experiments used the official BERT TensorFlow repository, and fine-tuning used 8 TPU v2 cores per model.The fine-tuning configurations cover BERT-Large RealFormer checkpoints pre-trained for 500K and 1M steps.Source: §A.1 · p. 12
  • BERT pre-training used 128 TPU v3 cores for Small, Base, and Large models, and 256 cores for xLarge.These correspond to 64 and 128 chips, respectively.Source: §A.1 · p. 12
  • Adding residual attention edges caused a 10%–15% pre-training performance drop for every reported model size except xLarge.The authors suspect the additions were less optimized than matrix multiplications on TPU v3; they report the overhead is nearly negligible on GPUs.Source: §A.1 · p. 12

A.2 Training Details: ADMIN

ADMIN experiments used the official fairseq-based repository and trained models for 50 epochs on WMT’14 English–German and English–French. Adding residual attention edges did not significantly reduce GPU performance across model sizes.

A.2 Training Details: ADMIN

  • NMT experiments used NVIDIA A100 GPUs and the official ADMIN fairseq repository, with the same scripts for data processing and evaluation.The number of GPUs varied by setup, as detailed in Table 11.Source: §A.2 · p. 12, Table 11 · p. 12
  • ADMIN and RealFormer were trained for 50 epochs on WMT’14 English–German and English–French, averaging the last 25 checkpoints.Learning rates were 5e-4 for English–German and 1.2e-3 for English–French across model sizes.Source: §A.2 · p. 12
  • Adding residual attention edges and running-mean attention scores to ADMIN incurred no significant GPU performance drop across model sizes.Source: §A.2 · p. 12

A.3 Training Details: ETC

The ETC appendix reports pre-training and fine-tuning details for RealFormer variants. It notes a shared RoBERTa initialization and a WikiHop setting that disables RealFormer during fine-tuning.

A.3 Training Details: ETC

  • ETC experiments used TPU v3 cores and the official ETC TensorFlow repository.Source: §A.3 · p. 12
  • ETC-Large with RealFormer benefited significantly from lifting weights from RoBERTa during pre-training.The initialization came from the same RoBERTa checkpoint used for the ETC-Large baseline, which may disadvantage RealFormer because that checkpoint was pre-trained without residual attention.Source: §A.3 · pp. 12–13
  • On WikiHop, development accuracy was 79.21 versus 78.96 when RealFormer was turned off during fine-tuning, so this setup was adopted for the reported results.The pre-trained RealFormer checkpoint was still loaded.Source: §A.3 · p. 13
  • ETC-Large RealFormer fine-tuning covered four benchmarks: WikiHop, HotpotQA, Natural Questions, and OpenKP.The appendix collects benchmark statistics and best-performing configurations in Table 12.Source: §A.3 · p. 13, Table 12 · p. 13

A.4 Entropy Distribution of Pre-trained Baseline Transformer Models

Appendix A.4 directs readers to Figure 5 for violin plots of entropy distributions from the pre-trained baseline Transformers.

A.4 Entropy Distribution of Pre-trained Baseline Transformer Models

  • Figure 5 contains violin plots of entropy distributions for the pre-trained BERT-Base models with Post-LN and Pre-LN Transformers.The passage identifies the figure but does not report its distributional findings.Source: §A.4 · p. 13

A.5 Jensen-Shannon Divergence of Different Pre-trained Transformers

This appendix compares attention patterns across vertically adjacent heads using Jensen-Shannon divergence distributions. It also notes that baseline Transformers have higher entropy in top layers than RealFormer, indicating relatively denser attention.

A.5 Jensen-Shannon Divergence of Different Pre-trained Transformers

  • Each JSD data point compares token-level attention probabilities for a head pair across adjacent layers, with the full distributions shown rather than one scalar per pair.Blue, yellow, and red indicate median JSD below 0.25, from 0.25 to 0.75, and above 0.75, respectively.Source: §A.5 · p. 13
  • Post-LN JSD serves only as a reference because adjacent-layer head correspondence is not meaningful for Post-/Pre-LN, whose layers can permute head order equivalently.The passage characterizes Post-LN JSD as expected to be random.Source: §A.5 · p. 13
  • Top layers 9–11 tend to have larger attention entropies in Post-LN and Pre-LN Transformers than in RealFormer, indicating relatively denser attention.Figure 5 shows entropy distributions for 8,192 held-out examples using pre-trained BERT-Base models.Source: Figure 5 · p. 14
  • Figure 6 shows distributions of Jensen-Shannon divergence between attention probabilities in vertically adjacent heads for RealFormer and Post-LN BERT-Base.The distributions use 8,192 held-out examples.Source: Figure 6 · p. 15
Loading 2012.11747v3…