Blockwise Self-Attention for Long Document Understanding

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.Source: §1 · p. 1
• 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.Source: §1 · p. 1
• 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.Source: §1 · p. 2
• 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.Source: §1 · p. 2
• 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.Source: §1 · p. 2

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.Source: §2 · p. 2

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.Source: §2.1 · p. 2
• 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.Source: §2.1 · p. 2

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.Source: §2.2 · p. 3

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.Source: §2.3 · p. 3
• 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.Source: §2.3 · p. 3

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.Source: §3 · pp. 3–4

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.Source: §3 · pp. 3–4

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.Source: §3.1 · p. 4

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.Source: §3.1 · p. 4

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.Source: §3.2 · p. 5

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.Source: §4 · p. 5, §4.2 · pp. 5–6

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.Source: §4.1 · p. 5

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.Source: §4.2 · p. 6
• 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.Source: §4.2 · p. 6

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.Source: §4.2 · p. 7

• Memory-optimization methods include low-precision training, microbatching, and gradient checkpointing.Source: §5.1 · p. 7

Attention layer simplification

• Attention-sparsification approaches reduce connectivity, but dense masked representations or custom CUDA kernels can limit practical speed and memory gains.Source: §5.2 · p. 8

Reducing model size for pre-training

• Knowledge distillation compresses BERT into smaller students, while ALBERT reduces parameters through parameter sharing.Source: §5.2 · p. 8

• Comparable language-model perplexity is not sufficient evidence of comparable downstream-task performance.Source: §5.2 · pp. 8–9
• SparseBERT can trail XLNet on SQuAD despite similar efficient-Transformer language-model results, motivating comprehensive evaluation of masked-language-model pre-training.Source: §5.2 · p. 9

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

• Pre-training hyper-parameters are listed in Table 7.Source: §A.1 · p. 11
• Table 6 is the table for BERT notations.Source: Table 6 · p. 11

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

• SparseBERT uses sparse masking matrices based on the Sparse Transformer pattern.The patterns are shown in Figure 5.Source: §A.3 · p. 11
• The implementation first computes the N2 attention matrix, then masks it according to the defined sparse pattern.Source: §A.3 · p. 11
• 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.Source: §A.3 · p. 11

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.Source: §A.4 · pp. 11–12
• Figure 5 presents sparse masking matrices for N = 512 and N = 1024 with densities of 44.20% and 34.97%, respectively.Source: Figure 5 · p. 12
• A sliding window with stride 128 splits inputs that exceed the maximum sequence-length constraint.Source: §A.4 · pp. 11–12
• Learning rates are grid-searched over {5e-6, 1e-5, 2e-5, 3e-5, 5e-5}, with batch sizes of 16 or 32.Source: §A.4 · p. 12
• Fine-tuning is performed for 4 epochs using code based on HuggingFace and SpanBERT.Source: §A.4 · pp. 11–12
• 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.Source: §A.5 · p. 12, Figure 6 · p. 12