PaperlayerView exact revision on arXiv

Source-linked AI summary

Blockwise Self-Attention for Long Document Understanding

Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, Jie Tang

arXiv:1911.02972v2cs.CLcs.LG

TL;DR

Long-sequence BERT models face memory and time costs, particularly from self-attention. BlockBERT addresses this with sparse block attention and reports lower memory, training, and inference costs while retaining comparable or better accuracy. Its broader evaluation across efficient Transformers and additional long-sequence applications remains future work.

  • Problem

    BERT self-attention has quadratic memory growth with sequence length, making long inputs and resource-efficient training difficult.

  • Method

    BlockBERT sparsifies BERT attention matrices into sparse blocks to reduce memory consumption and computation while supporting short- and long-range context.

  • Results

    BlockBERT achieves comparable or sometimes better prediction accuracy while reducing memory by 18.7-36.1%, training time by 12.0-25.1%, and test time by 27.8%.

  • Takeaways & Limitations

    Sparse block attention provides time and memory savings without significant performance loss.

  • Takeaways & Limitations

    The paper leaves broader benchmarking of efficient Transformers and applications to other long-sequence tasks for future work.

Abstract

from arXiv · show

We present BlockBERT, a lightweight and efficient BERT model for better modeling long-distance dependencies. Our model extends BERT by introducing sparse block structures into the attention matrix to reduce both memory consumption and training/inference time, which also enables attention heads to capture either short- or long-range contextual information. We conduct experiments on language model pre-training and several benchmark question answering datasets with various paragraph lengths. BlockBERT uses 18.7-36.1% less memory and 12.0-25.1% less time to learn the model. During testing, BlockBERT saves 27.8% inference time, while having comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.

1 Introduction

BERT-family models are resource-intensive because self-attention consumes memory quadratically with sequence length, limiting long-document modeling. BlockBERT sparsifies attention into blocks to reduce resource use while preserving or improving performance on long-sequence tasks.

  • 1 Introduction: BERT-family pre-training and fine-tuning require substantial computational resources and memory.BERT pre-training can take days on specialized hardware, while limited GPU memory restricts feasible batch sizes.
  • 1 Introduction: Dot-product self-attention is a major memory bottleneck because its cost grows quadratically with sequence length.This limits input length and complicates downstream processing through truncation or sliding windows.
  • 1 Introduction: Reducing layers, heads, or hidden units degrades performance without resolving the long-sequence problem.Microbatching and gradient checkpointing instead trade training time for lower memory consumption.
  • 1 Introduction: BlockBERT introduces sparse block structures into attention matrices to reduce memory consumption and floating-point operations.The block structure also lets attention heads capture short- or long-range contextual information.
  • 1 Introduction: 18.7-36.1% less memory, 12.0-25.1% less training time, and 27.8% less inference time are reported for BlockBERT.Across benchmark question answering datasets, prediction accuracy is comparable to or sometimes better than the original BERT-family models.

2 Background: Memory Bottleneck in Training BERT

BERT training memory is dominated by activations rather than parameters or optimizer state. Profiling and regression analysis identify a sequence-length-dependent quadratic attention component, motivating targeted memory-reduction strategies.

  • 2 Background: Memory Bottleneck in Training BERT: BERT consists of stacked bidirectional Transformer encoders with multi-head self-attention and position-wise feed-forward layers.The notation uses L for layers, H for hidden units, A for attention heads, N for sequence length, and B for batch size.
  • 2.1 Memory Profiling: Training memory comprises model, optimizer, and activation memory, with activations storing layer outputs for backpropagation.This categorization is used to identify the principal memory bottleneck.
  • 2.1 Memory Profiling: 0.21/1.03/8.49 GB are consumed by model/optimizer/activation memory on one GPU, respectively.Activation memory accounts for 87.6% of total GPU memory and is identified as the bottleneck.
  • 2.2 A Regression Analysis on Activation Memory: Activation memory grows linearly with most hyperparameters but includes a quadratic component in sequence length because of attention layers.With b × N fixed at 4096, regression estimates activation memory as 0.00715 × N + 4.83.
  • 2.3 Techniques for Reducing Traing Memory: Microbatching and gradient checkpointing reduce memory by repeating computation, thereby increasing training time.Low-precision training and knowledge distillation are also discussed as alternative memory-reduction approaches with different trade-offs.
  • 2.3 Techniques for Reducing Traing Memory: Common memory-reduction techniques remain limited in reducing both training time and memory usage simultaneously.The paper therefore focuses on optimizing dot-product attention layers.

3 Model: BlockBERT

BlockBERT sparsifies Transformer attention with permutation-defined block masks, reducing quadratic attention costs while allowing heads to model local and long-distance context. Empirical profiling shows substantial overall memory savings, especially for longer sequences.

  • Blockwise Attention: BlockBERT applies a sparse block mask to the N × N attention matrix, reducing attention memory and FLOPs while retaining efficient dense matrix operations.The input is split into n blocks, and the attention matrix is partitioned into n×n blocks.
  • Computational Reduction: BlockBERT reduces O(N^2) attention memory and FLOPs by a factor of n when the sequence is divided into n blocks.This follows from replacing full-sequence attention with block-structured sparse attention.
  • Implementation Assumption: If N is not divisible by n, the input sequence is padded so it can be divided into n equal-length blocks.The method assumes divisibility between sequence length and block count.
  • Blockwise Multi-head Attention: Each blockwise attention head can use a different permutation-based masking matrix, enabling heads to capture short- or long-range dependencies.The model selects cyclic-shift permutations such as σ, σ^2, through σ^n for different heads.
  • Memory Analysis: 18.7% and 23.8% overall memory savings occur at N = 512 with 2 and 3 blocks, respectively.The estimated O(N^2) activation memory decreases to 1/2 and 1/3 of BERT’s corresponding memory.

4 Experiments

The experiments evaluate BlockBERT’s pre-training efficiency and question-answering performance across sequence lengths and paragraph distributions. They also examine attention-head assignments, block counts, and inference efficiency.

  • Experimental setup: BlockBERT is evaluated against BERT-family baselines, including Google BERT, RoBERTa-1seq/2seq, SparseBERT, and XLNet.The experiments cover pre-training and multiple question-answering datasets with varied paragraph lengths.
  • Pre-training efficiency: 2.2 days are saved at N = 1024 when BlockBERT with n = 2 reduces RoBERTa-1seq training time from 9.7 to 7.5 days.The comparison is reported for pre-training.
  • Fine-tuning tasks: BlockBERT with N = 1024 achieves more comparable or better performance than RoBERTa-1seq on several question-answering tasks.At N = 512, the absolute F1 difference ranges from 0.04 on NaturalQA to 1.18 on NewsQA, averaging 0.55.
  • Fine-tuning tasks: Long-sequence pre-training benefits long-sequence fine-tuning on TriviaQA and SearchQA, but longer pre-training can hurt performance on HotpotQA and NewsQA.The reported effects depend on the match between pre-training and fine-tuning sequence lengths.
  • Inference efficiency: 27.8% test-time savings are observed for BlockBERT with 2 blocks at batch size 8 and sequence length 1024.Three blocks save 30.4%; RoBERTa cannot handle an input of size 16×1024, whereas BlockBERT can.

5 Related Work

Related work addresses efficient Transformer training through memory optimization, attention sparsification, and model-size reduction. The paper emphasizes that language-model perplexity alone may not predict downstream-task quality and calls for broader evaluation.

  • Memory-optimization methods include low-precision training, microbatching, and gradient checkpointing.
  • Attention layer simplification: Attention-sparsification approaches reduce connectivity, but dense masked representations or custom CUDA kernels can limit practical speed and memory gains.
  • Reducing model size for pre-training: Knowledge distillation compresses BERT into smaller students, while ALBERT reduces parameters through parameter sharing.
  • Comparable language-model perplexity is not sufficient evidence of comparable downstream-task performance.
  • SparseBERT can trail XLNet on SQuAD despite similar efficient-Transformer language-model results, motivating comprehensive evaluation of masked-language-model pre-training.

6 Conclusion

The paper presents BlockBERT as a lightweight BERT variant designed to improve efficiency while retaining effectiveness. It reports time and memory savings without significant performance loss and identifies broader benchmarking and long-sequence applications as future work.

  • BlockBERT sparsifies attention matrices into sparse block matrices to reduce time and memory consumption.
  • BlockBERT achieves time and memory savings without significant performance loss.
  • Future work includes benchmarking more efficient Transformers and applying BlockBERT to long-sequence NLP and protein-sequence tasks.

A.1 Notations and Pre-training Hyper-parameters

This appendix section points readers to the paper’s notation and pre-training hyper-parameter tables. The supplied passages identify the notation table but do not provide its entries.

  • Pre-training hyper-parameters are listed in Table 7.
  • Table 6 is the table for BERT notations.

A.2 Profiler Implementation

The profiler separates model, optimizer, and activation memory, using autograd-graph traversal to estimate activation storage. BERT profiling combines a 1,000-step pre-training run with this activation-memory estimate.

  • Model and optimizer memory are profiled by summing tensor element counts multiplied by element size.
  • Activation memory is estimated by traversing PyTorch’s autograd graph and summing required storage.
  • The three memory components together match PyTorch’s memory-profiling tool.
  • BERT is pre-trained for 1,000 steps before model and optimizer memory are computed and activation memory is estimated.

A.3 SparseBERT

SparseBERT adopts Sparse Transformer’s masking pattern and a Fairseq implementation that prioritizes performance comparison over speed.

  • SparseBERT uses sparse masking matrices based on the Sparse Transformer pattern.The patterns are shown in Figure 5.
  • The implementation first computes the N2 attention matrix, then masks it according to the defined sparse pattern.
  • The direct Fairseq implementation targets performance comparison rather than execution speed.Consequently, its training time and memory use are close to those of the dense approach.

A.4 Fine-tuning Settings

Fine-tuning allows inputs up to the pre-training model’s maximum sequence length and uses sliding-window splitting for longer inputs. The settings include grid searches over learning rate and batch size, four training epochs, and dataset paragraph-length distributions reported for context.

  • Fine-tuning Settings: Fine-tuning uses max sequence length=N, matching the pre-training model’s allowed input length.Sequences exceeding this constraint are split with a sliding window of stride 128.
  • Fine-tuning Settings: Figure 5 presents sparse masking matrices for N = 512 and N = 1024 with densities of 44.20% and 34.97%, respectively.
  • Fine-tuning Settings: A sliding window with stride 128 splits inputs that exceed the maximum sequence-length constraint.
  • Fine-tuning Settings: Learning rates are grid-searched over {5e-6, 1e-5, 2e-5, 3e-5, 5e-5}, with batch sizes of 16 or 32.
  • Fine-tuning Settings: Fine-tuning is performed for 4 epochs using code based on HuggingFace and SpanBERT.
  • Fine-tuning Settings: The SQuAD and MrQA paragraph-length distributions are shown in Figure 6.SQuAD 2.0 is omitted because its distribution is very similar to SQuAD 1.1.
Loading 1911.02972v2…