ABC: Attention with Bounded-memory Control

• Transformer attention uses a memory that grows with sequence length, making its overall computational complexity quadratic and especially costly for long sequences.Source: §1 · p. 1
• ABC unifies efficient attention approaches as bounded memories organized by different control strategies, connecting models that might otherwise seem distinct.Source: §1 · p. 1
• ABCMLP replaces largely context-agnostic memory organization with learned, contextualized control that decides how to store each token.Source: §1 · p. 1
• ABC models have linear complexity in sequence length and recurrent computation graphs for causal attention, supporting uses such as language modeling and text generation.Source: §1 · p. 1
• Linformer (Wang et al., 2020b), previously thought inapplicable to causal attention or autoregressive decoding (Tay et al., 2020), is applicable under the ABC view and performs competitively in machine translation.Source: §1 · p. 2
• Across language modeling, machine translation, and masked language model finetuning, ABCMLP outperforms previous ABC approaches and improves inference speed and memory versus strong Transformer baselines with no or negligible accuracy loss.Source: §1 · p. 2

• Attention stores key-value pairs in memory slots, and each query reads a softmax-normalized combination to produce a d-dimensional vector.Source: §2 · p. 2
• With N queries, parallelizable attention takes linear time and quadratic space, while sequential decoding takes quadratic time and linear space.Source: §2 · p. 2
• The outer-product representation expresses key and value memory as sums of slot-wise products; standard-basis control vectors place each token’s key and value in its corresponding slot.Source: §2 · p. 2

3 Attention with Bounded Memory

• ABC uses a constant-size key/value memory, so reading costs O(n) and overall complexity is O(Nn), linear in sequence length.The memory size n is set a priori.Source: §3 · p. 2

3.1 Linformer

• Linformer is an ABC instance that down-projects keys and values into memory and admits recurrent causal computation, contrary to prior claims of inapplicability.The paper reports competitive machine-translation performance for Linformer.Source: §3.1 · p. 3

3.2 Clustering-Based Attention

• Clustering-based attention groups keys and values into shared clusters and attends to their centroids as a constant-size memory.Its linear sequence-length complexity holds aside from clustering and centroid-computation overhead.Source: §3.2 · p. 3

3.3 Sliding-Window Attention

• Sliding-window attention can use a FIFO memory: a shift matrix removes the oldest entry, then the newest token fills the freed slot.Dilated and stride convolution patterns can also be recovered; learning the shift matrix is beyond this work’s scope.Source: §3.3 · pp. 3–4
• ABC also encompasses variants of Rae et al. (2020) and sparse attention patterns, whose control strategies are context-agnostic.The authors propose learning a contextualized strategy from data in the next section.Source: §3.3 · p. 4

4 Learned Memory Control

• ABCMLP uses an MLP on each token representation x_i to produce control vectors that determine which memory slots receive the token and how much.The resulting memory is read using context-dependent real queries that retrieve information for the prediction task.Source: §4 · pp. 4–5
• Normalizing the nonnegative control weights offsets sequence-length-related growth in memory scales; the resulting computation has linear sequence-length complexity and a recurrent graph.The normalization is over the full sequence for encoder self-attention or cross attention, and over the prefix for causal attention.Source: §4 · p. 4
• Because φ_i depends only on current input x_i, rather than prior memory, ABCMLP supports parallel control computation and causal attention.Causal normalization uses the sum over the prefix, so it does not require future tokens.Source: §4 · p. 4
• ABCMLP’s memory construction can be viewed as context-agnostic pseudo-query attention followed by retrieval with context-dependent real queries.For one memory slot, a learned vector wφ acts as the pseudo-query; the construction retains linear complexity because wφ is context-independent.Source: §4 · pp. 4–5
• ABCMLP resembles hierarchical attention, but attends to the same sequence in parallel rather than summarizing non-overlapping segments in a cascade.Its learned pseudo-queries also relate to inducing points in set attention (Lee et al., 2019), whose intervening feedforward network precludes ABC’s outer-product memory computation and efficient recurrences.Source: §4 · p. 5
• Unlike kernel-linearized attention methods that approximate, heuristically design, or learn feature maps, ABCMLP keeps softmax over a smaller constant-sized context.The paper reports that ABCMLP can use a much smaller memory while retaining accuracy.Source: §4 · p. 5

5.1 Language Modeling

• The WikiText-103 comparison uses 16- and 32-layer settings, with evaluation contexts of 0 or 480 tokens in the former and 256 tokens in the latter.The 16-layer models have about 242M parameters and train on 512-token segments; the 32-layer models have about 484M parameters.Source: §5.1 · p. 6
• With n = 64, ABCMLP beats Linformer and ABCRD by more than 2.9 test perplexity, widening to more than 3.6 with 480-token context; ABCMLP-64 is 0.5 perplexity behind BASE.ABCMLP-32 also beats kernel-based methods by more than 0.9 test perplexity in the 32-layer setting.Source: §5.1 · p. 6, Table 3 · p. 6

5.2 Machine Translation

• On WMT14 EN-DE, ABCMLP is on par with BASE and beats other ABC variants by more than 1.1 BLEU, including with smaller memories.Linformer beats ABCRD by more than 0.5 BLEU in this setting; ABCMLP beats other ABC models by more than 0.4 BLEU even with smaller memories.Source: §5.2 · p. 7, Table 4 · p. 7
• On IWSLT14 ES-EN, ABCMLP underperforms BASE by 0.2 BLEU, while Linformer fails to converge across multiple random seeds.The authors suggest ABCMLP’s result indicates effectiveness with longer sequences and attribute Linformer’s convergence issue to its purely position-based strategy for varying-length decoding.Source: §5.2 · p. 7, Table 4 · p. 7

5.3 Masked Language Model Finetuning

• The classification study warm-starts ABC variants from pretrained RoBERTa-base, continues MLM pretraining on assembled text data, then finetunes and evaluates on GLUE.The authors lacked resources to pretrain ABC from scratch; direct downstream finetuning after swapping attention, without continued MLM pretraining, failed for all models.Source: §5.3 · p. 7
• For GLUE text classification after continued RoBERTa-base pretraining, ABCMLP beats Linformer and performs on par with or better than BASE.ABCRD fails to converge in continued pretraining, even with multiple seeds.Source: §5.3 · p. 8, Table 5 · p. 8

• ABCMLP decodes about 1.25 times faster than RFA, while all linear-time models maintain consistent speed across sequence lengths.The comparison uses RFA memory sizes of 256/128 and ABCMLP sizes of 32/8; the models are selected for efficiency without accuracy drop.Source: §6 · p. 8, Figure 1 · p. 8
• T2R decodes about 20% faster than ABCMLP despite similar memory sizes, because it omits the softmax computation used by ABCMLP.Source: §6 · p. 8
• ABCMLP, RFA, and T2R have similar memory-overhead curves, and ABCMLP achieves significant memory savings.Source: §6 · p. 8, Figure 1 · p. 8
• For text encoding, both ABCMLP and Linformer improve inference speed and memory over BASE, but Linformer is faster because its linear projection is cheaper than ABCMLP’s MLP.The comparison uses 512-token inputs with batch size 16.Source: §6 · p. 8
• Across all five experiments, ABCMLP outperforms other ABC models in accuracy, while Linformer sometimes fails to converge or gives sub-optimal performance.ABCMLP slightly underperforms Linformer in text-encoding inference speed.Source: §6 · pp. 8–9
• On WMT14 EN-DE, reducing causal memory from 64 to 8 costs 0.3 BLEU when cross-attention memory is large enough; an 8-8 configuration is 1.1 BLEU behind the best configuration.Source: §6 · p. 9

• ABC provides a unified perspective on several recently proposed models by showing that they vary in how they organize bounded memory.Source: §7 · p. 9
• ABCMLP is a particular ABC instance that learns contextualized memory control and outperforms previous ABC models on language modeling, machine translation, and masked language model finetuning.Source: §7 · p. 9
• Compared with strong transformer baselines, ABCMLP achieves substantial efficiency improvements with no or negligible accuracy loss.Source: §7 · p. 9

A.1 Sparse Local-to-global Attention

• Local-to-global attention is an ABC instance when each global token maps to its own memory slot and other tokens map to zero vectors.Source: §A.1 · p. 12

A.2 Random Memory Control

• ABCRD randomly assigns every token to one of n memory slots, serving as a baseline for comparing random with learned bounded-memory control.Unlike random sparse attention, it attends to all tokens but squashes them into an n-slot memory.Source: §A.2 · p. 12

A.3 Compressive Transformer with Mean Pooling

• The mean-pooling compressive transformer is an ABC instance: it groups keys into n means with compression ratio c = N/n, assuming N mod n = 0 or padding.Source: §A.3 · p. 12

A.4 Dilated Convolution Attention Patterns

• Dilated attention, which attends to every other token, can be simulated with two queues, but this may add enough overhead to make the implementation more expensive than the original pattern.The appendix characterizes this implementation as having more conceptual than practical value.Source: §A.4 · p. 12

A.5 Shared Workspace and Linear Unified Nested Attention

• SW and LUNA derive control vectors from previous-layer memory, which exposes causal control to future information; ABCMLP instead conditions control on the current layer’s input.Source: §A.5 · pp. 12–13

B.1 Normalization in Causal Attention

• Normalizing eK and eV instead of the φi vectors admits a linear-complexity computation graph for causal ABCMLP.The appendix states that M/z divides each row of M by the corresponding dimension of z.Source: §B.1 · p. 13

B.2 Higher-Dimensional Case of Example 1

• With n memory slots, n separately parameterized pseudo-queries summarize the context in parallel, producing n-by-d memory representations.The appendix describes eK this way and says eV is similar.Source: §B.2 · p. 13

• The language-modeling experiments closely build on Baevski and Auli (2019).Source: §C.1 · p. 13
• The experimental setup also closely builds on Kasai et al. (2021b).Source: §C.1 · p. 13
• All models are trained on 4 A100 GPUs, and the hyperparameters are summarized in Table 10.Source: §C.1 · p. 13

C.2 Machine Translation

• IWSLT14 contains 1,421 training, 8 development, and 42 test documents, averaging 126.7 sentences per document; examples use length-4 sliding windows with stride one.Each sentence averages 21.7 Spanish and 22.5 English BPE subwords, and inference uses predicted target-side context.Source: §C.2 · p. 13
• Training averages checkpoints from the last five epochs and uses beam search of size 5 with length penalty 0.6 at inference.Models are trained on four RTX 2080 Ti GPUs.Source: §C.2 · p. 13
• Clustering- and sliding-window ABC underperform ABCMLP with the same memory sizes by more than 0.5 BLEU on WMT14 EN-DE; ReLU and sigmoid variants also underperform the exp variant.MLP-ReLU with 32/8 memory sizes fails to converge; MLP-exp-all applies ABC in both encoder and decoder, unlike the decoder-only alternatives.Source: §C.2 · p. 13, Table 8 · p. 14
• MLP-exp-all, which applies ABC in both encoder and decoder, underperforms ABCMLP by only 0.3 BLEU.The other listed variants apply ABC only in the decoder.Source: §C.2 · p. 13, Table 8 · p. 14
• The attention-overhead comparison uses a synthetic sequence-to-sequence generation task with varying sequence lengths, batch size 16, and greedy decoding.Figure 1b compares ABCMLP with 32/8 memory sizes against softmax attention; the models match those in §5.2.Source: §C.2 · p. 14

C.3 Masked Language Model Finetuning

• Continued pretraining concatenates BookCorpus, English Wikipedia, OpenWebText, and RealNews, replacing RoBERTa’s CC-News and omitting Stories because public access was broken.The authors lacked access to RoBERTa’s original pretraining data and split the shuffled training data into 20 shards because their machine could not load it all at once.Source: §C.3 · p. 14
• Downstream finetuning uses the hyperparameters of Liu et al. (2019), and Table 13 describes the GLUE tasks and their statistics.The cited table caption names MNLI, QNLI, QQP, and SST-2.Source: §C.3 · p. 14, Table 13 · p. 15
• The appendix identifies tables for dataset statistics and for language-modeling, machine-translation, and continued-pretraining hyperparameters.The supplied passages give the tables’ scope but not their parameter values.Source: Table 9 · p. 15, Table 10 · p. 15, Table 11 · p. 15