Overview

  • Attention is a query-conditioned way to retrieve and combine information from a sequence rather than relying on a single fixed summary. This chapter connects the classical encoder-decoder formulation to self-attention and shows why longer contexts introduce distinct computational and memory constraints. The final roadmap separates methods that store less, read less, or execute the same operation more efficiently.

What attention changes

  • Attention replaces a fixed sequence summary with query-conditioned retrieval, allowing each output step to combine the source positions relevant to that prediction. Classical encoder-decoder attention provides the simplest derivation of this principle.

  • Attention is a mechanism that allows a neural network to dynamically assign different levels of importance to different parts of its input when producing an output. Instead of compressing all source information into a single fixed-size representation, attention constructs an input-dependent weighted combination of representations, allowing the model to focus on the information most relevant to the current prediction. The mechanism was introduced to neural machine translation in Neural Machine Translation by Jointly Learning to Align and Translate by Bahdanau et al. (2014), where the decoder learns a soft alignment over encoder states while generating each target token.

  • Early encoder-decoder models compress the entire source sequence into one fixed-dimensional vector, which the decoder must use for every prediction. As inputs grow longer or more information-dense, that fixed context becomes restrictive. Bahdanau attention instead lets the decoder access all encoder states and build a new context vector for each output token.

  • For source annotations \(h_1,\ldots,h_T\), the context vector used at target step \(i\) is a weighted sum:

\[c_i = \sum_{j=1}^{T} \alpha_{ij} h_j\]
  • The normalized attention weights are obtained from learned compatibility scores:
\[\alpha_{ij} = \frac{\exp(e_{ij})}{\sum_{k=1}^{T}\exp(e_{ik})}\]

  • In additive attention, the compatibility score can be written as:

    \[e_{ij} = v_a^\top \tanh(W_a s_{i-1} + U_a h_j)\]
    • Here, \(s_{i-1}\) is the decoder state before generating the next target token, \(h_j\) is the encoder representation of source position \(j\), and the learned parameters determine how compatible the current decoder state is with each source position. The resulting distribution \(\alpha_{ij}\) acts as a soft alignment between the output being generated and the source sequence. This alignment is learned jointly with the translation objective rather than being supplied as explicit supervision, as demonstrated in Neural Machine Translation by Jointly Learning to Align and Translate by Bahdanau et al. (2014).

  • Attention replaces uniform treatment of input representations with query-conditioned aggregation. This retrieval principle extends beyond translation to language, vision, speech, multimodal models, and long-context reasoning.

From alignment to self-attention

  • Self-attention extends alignment from the encoder-decoder interface to token-to-token communication within the same sequence. The Transformer makes scaled dot-product attention and parallel heads its central sequence-mixing operations.

  • The Transformer generalized this idea by making attention the primary mechanism for communication between token representations. Attention Is All You Need by Vaswani et al. (2017) removed recurrence from the core sequence model and introduced scaled dot-product self-attention, where every token can construct a query, key, and value representation and selectively aggregate information from other tokens.

  • Given query, key, and value matrices \(Q\), \(K\), and \(V\), scaled dot-product attention is:

\[\operatorname{Attention}(Q,K,V) = \operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\]

  • The matrix \(QK^\top\) contains pairwise query-key similarity scores. Division by \(\sqrt{d_k}\) prevents dot products from growing excessively with key dimensionality, while the softmax converts the scores into normalized attention weights. Multiplying these weights by \(V\) produces a context-dependent representation for every query position. In self-attention, the queries, keys, and values are derived from the same sequence; in cross-attention, the queries and key-value representations originate from different sequences or modalities.

  • Autoregressive language models additionally apply a causal mask so that token position \(i\) can attend only to positions at or before \(i\). If \(M\) denotes the causal mask, attention can be written as:

\[\operatorname{Attention}(Q,K,V) = \operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V\]

  • Entries corresponding to disallowed future positions receive effectively negative-infinite logits before the softmax. This preserves autoregressive factorization while still allowing every generated token to use the complete permitted history.

  • Multi-head attention extends this mechanism by performing several attention operations in parallel using different learned projections. Each head can represent a different compatibility structure, and their outputs are concatenated and projected back into the model dimension. Attention Is All You Need by Vaswani et al. (2017) showed that this architecture could replace recurrent and convolutional sequence processing while improving parallelism and modeling long-range dependencies.

Scaling bottlenecks

  • Dense attention introduces quadratic interactions during full-sequence processing, while autoregressive decoding repeatedly accesses a growing KV cache. These are distinct constraints and motivate different architectural and systems responses.

  • Long contexts expose two different bottlenecks: dense attention requires quadratic interactions during training and prefill, while autoregressive generation must retain and repeatedly read historical KV state. The former emphasizes computation and intermediate memory; the latter emphasizes cache capacity and bandwidth.

  • Modern attention mechanisms can therefore be understood through two complementary questions: how much history should the model store, and how much of that history should each query read. Many mechanisms reduce one of these costs while attempting to preserve the quality benefits of dense multi-head attention.

Architectural efficiency mechanisms

  • Efficiency techniques modify different resources: the number of KV representations stored, the history consulted by each query, the form of the attention operator, or the movement of data through hardware. These approaches can often be composed.

  • Multi-Query Attention reduces storage by sharing a single set of keys and values across all query heads. Fast Transformer Decoding: One Write-Head is All You Need by Shazeer (2019) introduced this design to reduce the memory-bandwidth cost of autoregressive decoding while retaining multiple query heads.

  • Grouped-Query Attention provides an intermediate point between Multi-Head Attention and Multi-Query Attention. Instead of maintaining independent key-value heads for every query head or sharing one key-value head globally, query heads are partitioned into groups that share key-value heads. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints by GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints by Ainslie et al. (2023) (2023) showed that this design can approach multi-head attention quality while retaining much of the inference efficiency of multi-query attention.

  • Sliding-window attention bounds each query to a recent neighborhood, producing approximately linear sequence cost for a fixed window. It suits local dependencies but requires other mechanisms when direct long-range retrieval matters.

  • Multi-Head Latent Attention reduces KV-cache storage by compressing key-value information into a lower-dimensional latent representation. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model by DeepSeek-AI (2024) introduced Multi-Head Latent Attention as a core architectural component, using low-rank joint compression of keys and values to reduce inference memory while preserving expressive multi-head query behavior.

  • Linear-attention families attack the dense attention matrix more directly. Instead of explicitly computing the full \(QK^\top\) matrix, they restructure attention so that a compact state or kernelized summary of prior keys and values can be maintained. This can reduce sequence-length complexity from quadratic toward linear, although the resulting operation is generally not identical to softmax attention and therefore introduces different modeling and optimization tradeoffs.

  • Sparse-attention methods retain selective query-key interactions rather than computing all possible pairs. The central challenge is deciding which tokens should interact. Fixed patterns such as windows are simple and hardware-friendly, while learned sparse mechanisms attempt to identify important positions dynamically. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention by Yuan et al. (2025) combines compressed coarse-grained context, selected token blocks, and local windows so that long-context attention can allocate computation to different forms of relevant history.

  • Efficient attention kernels address a different dimension of the problem. They preserve the mathematical result of standard attention while changing how the computation is executed on hardware. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness by Dao et al. (2022) tiles attention computation to reduce reads and writes between high-bandwidth memory and on-chip SRAM, avoiding materialization of the complete attention matrix while remaining exact for standard softmax attention.

Design axes and chapter roadmap

  • The following chapters proceed from the original alignment mechanisms through long-context connectivity, compressed KV state, recurrent attention, exact GPU kernels, and streaming. Distinguishing their optimization targets makes the subsequent comparisons more useful.

  • These mechanisms can be organized into a useful evolutionary view:

\[\text{Additive Attention} \rightarrow \text{Dot-Product Attention} \rightarrow \text{Multi-Head Attention}\] \[\rightarrow \begin{cases} \text{MQA / GQA / MLA} & \text{store less KV state}\\ \text{Local / Sliding-Window / Sparse Attention} & \text{read fewer positions}\\ \text{Linear Attention} & \text{summarize history into compact state}\\ \text{FlashAttention} & \text{move attention data more efficiently} \end{cases}\]

  • These methods address different resources: MQA, GQA, and MLA reduce KV-cache costs; local and sparse attention limit token interactions; linear attention changes the operator and historical state; FlashAttention improves execution without approximating dense softmax.

  • Attention has evolved from decoder-side alignment into a design space for learned information routing. Modern implementations must preserve useful retrieval while balancing computation, stored state, bandwidth, and accelerator constraints.

Classical Encoder-Decoder Attention

  • Classical attention was developed to make sequence-to-sequence decoding depend on the relevant source positions at each output step. This chapter first examines the fixed-vector bottleneck, then derives additive and multiplicative alignment and the global/local variants that shaped later attention systems. The equations distinguish alignment scores, normalized weights, and the resulting context vector.

The fixed-context bottleneck

  • Classical sequence-to-sequence translation compressed the source into a single vector. The earliest attention mechanism removed this constraint by allowing the decoder to select different source positions at each output step.

Origins of Attention

  • Neural attention emerged from sequence-to-sequence modeling, where an encoder maps an input sequence into a representation and a decoder generates an output sequence conditioned on that representation. Early encoder-decoder architectures compressed the entire input into a single fixed-dimensional vector, creating an information bottleneck that became increasingly problematic as sequence length increased. Neural Machine Translation by Jointly Learning to Align and Translate by Bahdanau et al. (2014) explicitly identified this fixed-length representation as a bottleneck and introduced a differentiable mechanism that lets the decoder dynamically select relevant encoder states while generating each output token.

  • The shift was from one static source vector to dynamic decoder-side selection over a sequence of encoder states. Different output steps can therefore retrieve different source information instead of repeatedly consulting the same compressed representation.

  • The following figure (source) illustrates the intuition behind introducing attention into sequence-to-sequence models, allowing the decoder to retrieve relevant source information dynamically instead of relying entirely on a fixed-length context vector.

  • This mechanism became the foundation for subsequent attention architectures. Effective Approaches to Attention-based Neural Machine Translation by Luong et al. (2015) systematized several useful variants, including global attention over the entire source sequence and local attention over a restricted subset of source positions. The conceptual progression ultimately led to Attention Is All You Need by Vaswani et al. (2017), which removed recurrence entirely and made attention the primary mechanism for communication among sequence positions.

Classic Sequence-to-Sequence Model

  • Before attention, neural machine translation commonly followed the encoder-decoder paradigm. Given an input sequence \(X = (x_1,x_2,\ldots,x_T)\) the encoder processes the tokens sequentially and produces hidden states:

    \[h_t = f(x_t,h_{t-1})\]
    • where \(h_t\) summarizes the input observed through position \(t\).

  • A conventional encoder-decoder then reduces the encoder’s hidden-state sequence to a single context vector:

\[c = q(h_1,h_2,\ldots,h_T)\]
  • In a simple recurrent implementation, this may effectively correspond to the final encoder state:
\[c = h_T\]

  • The decoder generates an output sequence \(Y=(y_1,\ldots,y_{T'})\) according to

    \[p(Y\mid X) = \prod_{i=1}^{T'} p(y_i\mid y_{<i},c)\]

    • with a recurrent decoder state such as:

      \[s_i=f(s_{i-1},y_{i-1},c)\]

  • The architecture therefore forces every source-side detail needed throughout decoding through \(c\). The representation has fixed dimensionality regardless of whether the source contains five tokens or hundreds. Neural Machine Translation by Jointly Learning to Align and Translate by Bahdanau et al. (2014) observed that this requirement makes long sequences particularly difficult and cited earlier evidence that basic encoder-decoder performance deteriorates rapidly as input length grows.

  • The following figure (source) shows the classical recurrent neural machine translation architecture, where the encoder processes the complete source sequence before the decoder emits the target sequence.

  • The following figure (source) provides an alternative visualization of the classical encoder-decoder architecture, in which the encoder compresses the input sequence into a fixed-dimensional representation before decoding begins.

The Fixed-Context Bottleneck

  • A fixed context vector must retain information through encoding and serve every subsequent decoding step, even though those steps need different source details.

  • Longer source sequences make the fixed-vector bottleneck more severe because early details must survive recurrent updates and remain accessible throughout decoding.

  • The following figure (source) illustrates how translation performance varies with sentence length, motivating the introduction of attention to address the fixed-context bottleneck.

  • Attention changes the problem from:

    \[X \rightarrow c \rightarrow Y\]

    • to:

      \[X \rightarrow (h_1,h_2,\ldots,h_T) \rightarrow c_i \rightarrow y_i\]
      • where the context \(c_i\) is recomputed for each output position \(i\).

  • The important distinction is therefore not merely that attention supplies the decoder with more information. It provides selective, time-dependent access to the encoder’s representations. Neural Machine Translation by Jointly Learning to Align and Translate by Bahdanau et al. (2014) emphasizes that the model no longer has to squash the entire source sentence into one fixed-length representation and instead adaptively chooses encoder representations during decoding.

  • The following figure (source) illustrates the conventional sequence-to-sequence architecture, in which the decoder relies on a single fixed-context representation produced by the encoder.

  • The following figure (source) illustrates the corresponding attention-based architecture, in which the decoder accesses encoder hidden states dynamically to construct a different context vector at each decoding step.

Alignment and context construction

  • Attention changes the decoder input from a fixed summary to an alignment-weighted combination of encoder states. The context-vector and alignment formulations explain how those weights are learned.

Sequence-to-Sequence with Attention

  • An attention-based encoder produces a collection of source representations rather than only one final representation:
\[H=(h_1,h_2,\ldots,h_T)\]
  • Bahdanau attention uses a bidirectional recurrent encoder, allowing the annotation associated with each source position to summarize both preceding and following source context:
\[h_j= \begin{bmatrix} \overrightarrow{h_j}\\ \overleftarrow{h_j} \end{bmatrix}\]

  • At decoder step \(i\), the model computes how well each source representation \(h_j\) matches the decoder’s current needs. These compatibility values are called alignment scores:

    \[e_{ij}=a(s_{i-1},h_j)\]
    • where \(s_{i-1}\) is the previous decoder hidden state and \(a(\cdot)\) is a learned alignment function.

  • Bahdanau et al. parameterize this alignment model as a small feed-forward neural network:

\[e_{ij} = v_a^\top \tanh \left( W_a s_{i-1} + U_a h_j \right)\]

  • The parameters \(v_a\), \(W_a\), and \(U_a\) are learned jointly with the rest of the translation model.

  • The scores are normalized across source positions with a softmax:

    \[\alpha_{ij} = \frac{\exp(e_{ij})} {\sum_{k=1}^{T}\exp(e_{ik})}\]
    • so that:
    \[\sum_{j=1}^{T}\alpha_{ij}=1\]

  • The values \(\alpha_{ij}\) can therefore be interpreted as the relative amount of attention assigned to source position \(j\) while generating target position \(i\).

Context construction and gradient flow

  • Attention weights determine the context vector used at each decoder step. Because the alignment calculation is differentiable, the translation objective can train both the alignment model and the encoder-decoder jointly.

  • The decoder-specific context vector is the weighted sum

    \[c_i = \sum_{j=1}^{T} \alpha_{ij}h_j\]
    • and the decoder then predicts the next token using its recurrent state, the previous output, and this dynamically constructed context:
    \[p(y_i\mid y_{<i},X) = g(y_{i-1},s_i,c_i)\]

  • Because all operations are differentiable, explicit word-alignment labels are unnecessary. Translation loss supplies the learning signal, and gradients propagate through the context vector, attention weights, alignment network, and encoder. This is why the mechanism is commonly called soft attention.

  • The following figure (source) shows the attention-based encoder-decoder architecture, in which each decoder output receives a dynamically weighted context derived from the sequence of encoder annotations.

Understanding the Context Vector

  • The context vector \(c_i\) is not a persistent memory location. It is a temporary readout from the encoder’s representation at decoder step \(i\). Different decoder positions generally produce different attention distributions and therefore different context vectors:
\[c_1 \neq c_2 \neq \cdots \neq c_{T'}\]

  • Conceptually, the encoder states behave like a memory \(H=\{h_1,\ldots,h_T\}\), the decoder state behaves like a query, the alignment model computes compatibility with each memory location, and the weighted sum retrieves the information most relevant to the current prediction.

  • This interpretation foreshadows the query-key-value abstraction later formalized by Transformer attention. Classical encoder-decoder attention does not explicitly construct separate key and value projections, but it already contains the essential operation:

\[\text{query} \rightarrow \text{compatibility scores} \rightarrow \text{normalized weights} \rightarrow \text{weighted aggregation}\]

Alignment as Learned Soft Selection

  • The complete Bahdanau attention operation can be summarized as:
\[e_{ij} = v_a^\top \tanh(W_as_{i-1}+U_ah_j)\] \[\alpha_{ij} = \operatorname{softmax}_j(e_{ij})\] \[c_i = \sum_j \alpha_{ij}h_j\]

  • The alignment network is therefore responsible for deciding where to look, while the weighted sum determines what information is retrieved. Because the attention distribution is continuous rather than discrete, the model can combine information from several source positions simultaneously.

  • The resulting attention matrix \(A=[\alpha_{ij}]\) provides a two-dimensional representation of source-target alignment. Rows correspond to target positions and columns to source positions. Translation models frequently learn approximately diagonal or reordered alignment patterns corresponding to linguistic relationships between source and target words. Bahdanau et al. report that these learned soft alignments qualitatively agree with intuitive word alignments.

Global and local alignment variants

  • Luong attention explores alternatives for scoring and selecting source positions. The global and local variants highlight how the choice of accessible positions changes the retrieval behavior.

Luong Attention

  • Effective Approaches to Attention-based Neural Machine Translation by Luong et al. (2015) simplified and generalized encoder-decoder attention by separating two design decisions: which source positions should be considered, and which scoring function should measure compatibility between source and target states. The paper distinguishes global attention, which considers all encoder states, from local attention, which restricts computation to a neighborhood.
Global Attention

  • In global attention, every encoder state is available at every decoder step. Given decoder state \(h_t\) and source state \(\bar h_s\), the model computes:

    \[\alpha_t(s) = \frac {\exp(\operatorname{score}(h_t,\bar h_s))} {\sum_{s'}\exp(\operatorname{score}(h_t,\bar h_{s'}))}\]

    • and forms:

      \[c_t = \sum_s \alpha_t(s)\bar h_s\]

  • Luong et al. evaluate several alignment functions. Dot-product attention uses

\[\operatorname{score}(h_t,\bar h_s) = h_t^\top\bar h_s\]
  • General attention introduces a learned matrix:
\[\operatorname{score}(h_t,\bar h_s) = h_t^\top W_a\bar h_s\]
  • Concat attention instead applies a learned nonlinear compatibility model:
\[\operatorname{score}(h_t,\bar h_s) = v_a^\top \tanh \left( W_a[h_t;\bar h_s] \right)\]

  • These formulations make an important distinction explicit: attention does not prescribe one particular similarity function. It describes a general retrieval operation in which compatibility scores are computed, normalized, and used to aggregate representations.

  • The following figure (source) shows the global attention architecture, where the current decoder state is compared against all source-side hidden states to construct the attentional hidden state.

Local Attention

  • Global attention scores every source position. Luong local attention instead restricts retrieval to a source-side window, introducing a tradeoff between alignment flexibility and interaction cost.

  • In the predictive variant, the center position is learned from the decoder state:

    \[p_t = S \cdot \operatorname{sigmoid} \left( v_p^\top \tanh(W_ph_t) \right)\]
    • where \(S\) is the source sequence length. Attention is then restricted to positions around \(p_t\).

  • Luong attention additionally applies a Gaussian prior centered at the predicted position:

    \[\alpha_t(s) \propto \operatorname{align}(h_t,\bar h_s) \exp \left( -\frac{(s-p_t)^2}{2\sigma^2} \right)\]
    • typically with \(\sigma=\frac{D}{2}\) for an attention window of width determined by \(D\).

  • This introduces a principle that reappears much later in efficient Transformer architectures: global pairwise interaction is not always necessary. Restricting attention to a carefully chosen local neighborhood can substantially reduce computation while retaining useful contextual interactions.

Scoring functions and the move to self-attention

  • Additive and multiplicative scoring differ in how they compute compatibility, not in their use of weighted context. That distinction sets up the scaled dot-product formulation of Transformer self-attention.

Additive vs. Multiplicative Attention

  • Bahdanau attention is commonly called additive attention because the transformed query and key representations are added before being passed through a nonlinear scoring network:
\[e_{ij} = v_a^\top \tanh (W_q q_i+W_k k_j)\]

  • Luong’s dot-product formulation is a form of multiplicative attention:

    \[e_{ij} = q_i^\top k_j\]

    • while the general form inserts a learned transformation:

      \[e_{ij} = q_i^\top W_a k_j\]

  • The following figure (source) compares alternative attention alignment functions, illustrating the differences between dot-product, general multiplicative, and additive scoring.

  • Additive attention provides a flexible learned compatibility function but requires a small neural network for every query-key pair. Dot-product attention maps naturally to highly optimized matrix multiplication:

    \[E=QK^\top\]
    • which becomes especially important when attention is computed simultaneously across every position in a sequence.

From Encoder-Decoder Attention to Self-Attention

  • Classical attention primarily connects two different sequences: decoder states query representations generated by an encoder. The next conceptual step is self-attention, where the queries, keys, and values are all derived from the same sequence.

  • Attention Is All You Need by Vaswani et al. (2017) made this operation the central computational primitive of the Transformer, eliminating recurrence and convolution from the architecture. Instead of recurrently propagating information across positions, every token can directly retrieve information from other relevant positions through attention.

  • This transition changes attention from an auxiliary alignment mechanism into the sequence model itself. The encoder-decoder lineage can therefore be viewed as:

\[\text{Fixed Context} \rightarrow \text{Additive Attention} \rightarrow \text{Global/Local Multiplicative Attention} \rightarrow \text{Self-Attention} \rightarrow \text{Multi-Head Attention}\]

Self-Attention and the Transformer

  • Self-attention generalizes decoder-side source retrieval into token-to-token information exchange within a sequence. This chapter develops scaled dot-product and multi-head attention, explains causal masking and cross-attention, and distinguishes the architectural capability of dense retrieval from its computational cost. Each formulation builds on the same query-key-value abstraction.

Token-to-token attention foundations

  • Self-attention generalizes encoder-decoder alignment by letting positions within the same sequence exchange information. Queries, keys, values, and scaling define the core calculation examined below.

From Recurrent Attention to Self-Attention

  • Cross-attention retrieves information from a separate sequence: decoder states supply queries and encoder states supply keys and values. Self-attention uses the same mechanism within one sequence, letting each position retrieve from other permitted positions.

  • Attention Is All You Need by Vaswani et al. (2017) made self-attention the primary sequence-processing operation of the Transformer, replacing the recurrent and convolutional sequence transformations used by preceding architectures. This removes the inherently sequential recurrence across token positions and permits representations for all positions in a layer to be computed in parallel during training.

  • For an input representation matrix \(X\in\mathbb{R}^{n\times d_{\text{model}}}\), the query, key, and value matrices are learned linear projections:

    \[Q=XW^Q\] \[K=XW^K\] \[V=XW^V\]

    • where:

      \[W^Q\in\mathbb{R}^{d_{\text{model}}\times d_k}\] \[W^K\in\mathbb{R}^{d_{\text{model}}\times d_k}\] \[W^V\in\mathbb{R}^{d_{\text{model}}\times d_v}\]

  • The separation between queries, keys, and values is important. Even though they originate from the same token representations in self-attention, the learned projections allow the model to represent the matching problem separately from the information ultimately communicated.

Queries, Keys, and Values

  • Attention can be interpreted as differentiable content-addressable retrieval. For a token at position \(i\), the query \(q_i\) is compared with the keys \(k_j\) of candidate source positions. These comparisons determine how much of each corresponding value \(v_j\) should contribute to the output:
\[o_i = \sum_{j=1}^{n} \alpha_{ij}v_j\]

  • The attention coefficients are normalized compatibility scores:

    \[\alpha_{ij} = \frac {\exp(e_{ij})} {\sum_{k=1}^{n}\exp(e_{ik})}\]
    • where \(e_{ij}\) measures the compatibility between \(q_i\) and \(k_j\).

  • This gives the attention operation three conceptually distinct stages:

\[\text{match queries against keys} \rightarrow \text{normalize scores} \rightarrow \text{aggregate values}\]
  • Queries and keys determine where information is retrieved; values determine the information transferred. This separation later enables head sharing and KV compression.

Scaled Dot-Product Attention

  • The Transformer replaces the neural alignment function of additive attention with scaled dot-product attention. For one query and key, \(e_{ij} = \frac{q_i^\top k_j}{\sqrt{d_k}}\) and for all positions simultaneously:
\[S = \frac{QK^\top}{\sqrt{d_k}}\]
  • The attention output is:
\[\operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V\]
  • Attention Is All You Need by Vaswani et al. (2017) introduced this scaled formulation because ordinary dot products grow in magnitude as the key dimension increases, potentially pushing the softmax into saturated regions with very small gradients.

Why Divide by the Square Root of the Key Dimension?

  • Suppose the components of \(q\) and \(k\) are independent random variables with zero mean and unit variance. Their dot product is
\[q^\top k = \sum_{r=1}^{d_k}q_rk_r\]

  • with variance approximately:

    \[\operatorname{Var}(q^\top k)=d_k\]

    • and therefore standard deviation:

      \[\operatorname{Std}(q^\top k)=\sqrt{d_k}\]

  • As \(d_k\) increases, unscaled logits therefore become larger in magnitude. Softmax can then become highly peaked:

\[\operatorname{softmax}(z)_i = \frac{\exp(z_i)} {\sum_j\exp(z_j)}\]

  • making gradients through low-probability positions extremely small. Scaling by \(1/\sqrt{d_k}\) approximately normalizes the variance of the logits:

    \[\operatorname{Var} \left( \frac{q^\top k}{\sqrt{d_k}} \right) \approx 1\]
    • which keeps the score distribution in a more numerically useful range.

Masks and attention structure

  • The attention matrix represents pairwise communication; masks restrict it according to padding or generation constraints. Multi-head attention then learns several distinct interaction structures in parallel.

Attention Matrix

  • For a sequence containing \(n\) tokens, \(QK^\top \in \mathbb{R}^{n\times n}\) so self-attention explicitly forms pairwise relationships between sequence positions. Entry \(S_{ij}\) measures how strongly token \(i\) matches token \(j\) before normalization.

  • After applying the row-wise softmax:

\[A = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)\]

  • each row forms an attention distribution:

    \[\sum_{j=1}^{n}A_{ij}=1\]
    • and the layer output becomes \(O=AV\).

  • Unlike recurrent architectures, a token can therefore exchange information directly with a distant token within a single self-attention layer. The cost of this global connectivity is that the attention-score matrix contains \(n^2\) entries, establishing the quadratic sequence-length bottleneck that motivates many later efficient-attention mechanisms.

Masks

  • Attention can restrict which query-key pairs are permitted by adding a mask \(M\) to the attention logits:

    \[\operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} + M\right)V\]

    • where disallowed connections receive a sufficiently negative value, conventionally conceptualized as:

      \[M_{ij} = -\infty\]
      • so that \(\operatorname{softmax}(\ldots,-\infty,\ldots)=0\) for the masked position.
Padding Mask
  • When variable-length sequences are padded into a common batch shape, padding tokens should not participate as meaningful keys or values. A padding mask sets their attention logits to negative infinity before softmax.
Causal Mask

  • Autoregressive language models must prevent position \(i\) from reading future tokens \(j>i\). The causal mask therefore has the form

    \[M_{ij} = \begin{cases} 0,&j\leq i\\ -\infty,&j>i \end{cases}\]
    • which produces a lower-triangular attention pattern.

  • Consequently, causal self-attention computes

    \[o_i = \sum_{j=1}^{i} \alpha_{ij}v_j\]
    • rather than summing over the entire sequence. The original Transformer applies this masking in the decoder so that predictions at a position depend only on previously generated outputs.

Multi-head and cross-attention variants

  • Multiple heads expand the types of relations an attention layer can represent. The surrounding encoder, decoder, and cross-attention arrangements differ primarily in the sources of their queries and key-value states.

Multi-Head Attention

  • A single attention operation produces one similarity structure between tokens. Multi-head attention instead learns multiple independent query, key, and value projections, allowing different heads to retrieve information using different representation subspaces.
Per-head projections and recombination

  • Each head computes attention in its own learned projection space. Their outputs are then concatenated and projected so that the layer can combine several token-interaction patterns.

  • For head \(r\):

    \[Q_r=QW_r^Q\] \[K_r=KW_r^K\] \[V_r=VW_r^V\]

    • and:

      \[\operatorname{head}_r = \operatorname{Attention}(Q_r,K_r,V_r)\]

  • The outputs of all heads are concatenated and projected:

    \[\operatorname{MHA}(Q,K,V) = \operatorname{Concat} ( \operatorname{head}_1, \ldots, \operatorname{head}_h ) W^O\]

    • where:

      \[W^O \in \mathbb{R}^{hd_v\times d_{\text{model}}}\]

  • Attention Is All You Need by Vaswani et al. (2017) motivates multi-head attention as a way to jointly attend to information from different representation subspaces and different positions, rather than forcing all relationships through one averaged attention operation.

  • In the original Transformer base configuration,

    \[h=8\]
    • and:
    \[d_k=d_v=\frac{d_{\text{model}}}{h}=64\]
    • with:
    \[d_{\text{model}}=512\]

  • Because each head operates in a lower-dimensional space, the aggregate computational cost of the heads remains similar to performing one full-dimensional attention operation.

  • The following figure (source) shows scaled dot-product attention on the left and multi-head attention on the right: queries and keys are multiplied and scaled before masking and softmax, while multi-head attention performs several projected attention operations in parallel before concatenation and output projection.

Why Multiple Heads Help

  • Consider two tokens whose relationship can be characterized along several dimensions. One head might learn a representation in which syntactic compatibility produces high query-key similarity, while another might emphasize entity identity or positional relationships. The heads need not be manually assigned such roles. Their projections are learned end-to-end.

  • Multi-head attention can therefore be interpreted as parallel learned retrieval systems:

    \[\operatorname{head}_1 \rightarrow \text{retrieval subspace 1}\] \[\operatorname{head}_2 \rightarrow \text{retrieval subspace 2}\] \[\vdots\] \[\operatorname{head}_h \rightarrow \text{retrieval subspace }h\]
    • followed by a learned combination through \(W^O\).

  • Visualizations in the Transformer paper illustrate distinct learned head patterns, including long-range dependencies and linguistic relations. This supports the representational motivation for multiple learned projections without implying that every head is always interpretable.

Self-Attention, Cross-Attention, and Masked Self-Attention

  • The original Transformer uses the same attention primitive in three configurations.
Encoder Self-Attention
  • Encoder self-attention obtains queries, keys, and values from the encoder’s current sequence representation:
\[Q=XW^Q\] \[K=XW^K\] \[V=XW^V\]
  • Every source position can attend to every other source position.
Decoder Masked Self-Attention
  • Decoder self-attention likewise derives \(Q\), \(K\), and \(V\) from the decoder sequence, but causal masking prevents access to future tokens:
\[\operatorname{Attention}_{\text{causal}} = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} + M_{\text{causal}} \right)V\]
  • This preserves autoregressive factorization:
\[p(y_1,\ldots,y_T) = \prod_{t=1}^{T} p(y_t\mid y_{<t})\]
Encoder-Decoder Cross-Attention
  • Cross-attention obtains queries from the decoder but keys and values from the encoder:
\[Q=H_{\text{dec}}W^Q\] \[K=H_{\text{enc}}W^K\] \[V=H_{\text{enc}}W^V\]
  • Cross-attention is therefore the Transformer’s direct descendant of classical encoder-decoder attention. The decoder asks what source-side information is useful for its current representation, but the recurrent Bahdanau alignment network is replaced by parallel scaled dot-product attention.

Position and complexity

  • A Transformer integrates attention with residual computation and positional information. Its dense all-pairs interactions motivate the long-context alternatives in the next chapter.

Transformer Attention Block

  • Attention is embedded inside a larger Transformer block rather than operating independently. The original encoder consists of repeated layers containing multi-head self-attention and a position-wise feed-forward network, together with residual connections and layer normalization.

  • Abstractly, the attention sublayer can be written as

    \[H' = \operatorname{LayerNorm} \left( H+\operatorname{MHA}(H) \right)\]
    • followed by:
    \[H'' = \operatorname{LayerNorm} \left( H'+\operatorname{FFN}(H') \right)\]
    • for the post-normalization architecture of the original Transformer.

  • The position-wise feed-forward network is applied independently to each token:

    \[\operatorname{FFN}(x) = \max(0,xW_1+b_1)W_2+b_2\]
    • so attention performs communication across token positions, while the feed-forward network performs nonlinear transformation within each token representation.

  • The following figure (source) shows the original Transformer encoder-decoder architecture, including stacked self-attention, masked decoder self-attention, encoder-decoder attention, feed-forward layers, residual connections, normalization, embeddings, and positional encodings.

Positional Information

  • Self-attention by itself does not encode sequence order. Permuting the input positions and applying the same permutation to the outputs leaves an unpositioned self-attention operation structurally unchanged. The Transformer therefore injects positional information into token representations.

  • The original model uses sinusoidal positional encodings:

    \[PE_{(pos,2i)} = \sin \left( \frac{pos} {10000^{2i/d_{\text{model}}}} \right)\] \[PE_{(pos,2i+1)} = \cos \left( \frac{pos} {10000^{2i/d_{\text{model}}}} \right)\]

    • and combines them with token embeddings:

      \[X_0=E_{\text{token}}+PE\]

  • Modern language models frequently replace these absolute encodings with relative or rotary positional mechanisms, but the fundamental requirement remains: attention needs an explicit mechanism for representing token position or relative displacement.

Computational Complexity

  • For sequence length \(n\) and representation dimension \(d\), computing the score matrix requires approximately \(O(n^2d)\) operations because every query interacts with every key.

  • The attention matrix itself requires \(O(n^2)\) memory if explicitly materialized. This becomes increasingly expensive as context grows:

    \[n=4{,}096 \Rightarrow n^2\approx16.8\text{ million}\]
    • whereas:
    \[n=128{,}000 \Rightarrow n^2\approx16.4\text{ billion}\]
    • query-key relationships per head.

  • This quadratic interaction structure is the central systems limitation of conventional full attention. It motivated several distinct lines of research that should not be conflated:

\[\text{Full Attention} \rightarrow \begin{cases} \text{Sparse/local attention}\\ \text{Low-rank attention}\\ \text{Linear attention}\\ \text{IO-efficient exact attention}\\ \text{KV-cache-efficient attention} \end{cases}\]
  • Later approaches address different bottlenecks: Longformer restricts connectivity, Linformer compresses sequence representations, linear attention changes the operator, and FlashAttention improves exact-attention execution.

Transition to Efficient Attention

  • Dense MHA established a general-purpose token-mixing baseline, but its quadratic prefill cost and persistent decoding cache motivate the long-context and cache-oriented mechanisms developed next.

  • This distinction explains why modern attention mechanisms modify different parts of the same underlying operation:

\[\operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V\]
  • Some methods reduce which elements of \(K\) are consulted. Others compress \(K\) and \(V\), approximate the attention matrix, reorganize the algebra, or preserve the exact operation while implementing it more efficiently on accelerators.

Efficient Long-Context Attention: Local, Sparse, Low-Rank, and Linear Approaches

  • Long sequences expose dense attention’s quadratic interaction cost, motivating several different approximations rather than one universal substitute. This chapter proceeds from fixed local patterns to global and dynamic sparsity, low-rank sequence projections, and linear/recurrent formulations. Their arithmetic scaling and access to distant history are compared separately.

The long-context scaling problem

  • Dense attention computes interactions for every query-key pair, which becomes expensive as sequence length grows. The following alternatives reduce computation by restricting connections, projecting the sequence, or changing the attention operator.

Why Full Attention Becomes Expensive

  • Full self-attention gives every query direct access to every key-value pair. For a sequence of length \(n\), the score matrix contains \(n^2\) query-key interactions:
\[S=\frac{QK^\top}{\sqrt{d_k}} \in\mathbb{R}^{n\times n}\]

  • This yields approximately \(O(n^2d)\) arithmetic complexity and, for a conventional materialized implementation, \(O(n^2)\) attention-state memory.

  • Efficient-attention research therefore asks whether every one of these \(n^2\) interactions is necessary. Several distinct answers emerged. Local and sparse attention compute only selected query-key pairs. Low-rank attention compresses the sequence dimension before computing attention. Linear attention replaces softmax attention with a kernelizable similarity function whose computation can be reassociated. These approaches change the attention operator itself, unlike FlashAttention, which preserves exact softmax attention and instead changes how it is executed.

Structured local and global patterns

  • Local windows impose a fixed neighborhood and offer predictable scaling. Longformer combines these windows with selected global positions to recover long-range information flow.

Sliding-Window and Local Attention

  • One of the simplest ways to reduce attention cost is to assume that most useful interactions are local. Instead of allowing position \(i\) to attend to every position, it attends only to positions within a window:
\[\mathcal{N}(i) = \{j:|i-j|\leq w\}\]

  • and computes:

    \[o_i = \sum_{j\in\mathcal{N}(i)} \alpha_{ij}v_j\]

  • If each token attends to only \(O(w)\) positions, the number of query-key interactions becomes \(O(nw)\) rather than \(O(n^2)\).

  • For a fixed \(w\), attention therefore scales linearly with sequence length.

Longformer

  • Longformer: The Long-Document Transformer by Beltagy et al. (2020) develops this idea into an attention mechanism for processing documents containing thousands of tokens. Longformer combines local sliding-window attention with a small number of globally attending positions, producing linear rather than quadratic sequence-length scaling.
Sliding-Window Attention

  • Longformer’s basic attention pattern assigns each token a fixed local window. If the total window contains approximately \(w\) positions, complexity is \(O(nw)\) rather than \(O(n^2)\).

  • The paper emphasizes that stacking multiple local-attention layers expands the effective receptive field. Even when a single layer communicates only locally, information can propagate increasingly far across the sequence as depth increases.

  • For a symmetric window, the neighborhood of position \(i\) can be written as:

    \[\mathcal{N}_i = \left[ i-\frac{w}{2}, i+\frac{w}{2} \right]\]

    • and attention is evaluated only for:

      \[j\in\mathcal{N}_i\]
      • rather than every \(j\in\{1,\ldots,n\}\).
Dilated Sliding-Window Attention
  • Longformer also supports dilation. Instead of attending to every token within the window, an attention head can skip positions using dilation \(d\):
\[\mathcal{N}_i^{(d)} = \{i-kd,\ldots,i,\ldots,i+kd\}\]

  • This expands the receptive field without proportionally increasing the number of attended positions.

  • Different heads can use different dilation configurations, allowing some heads to emphasize dense local information while others communicate across larger distances.

Global Attention

  • Pure local attention makes direct long-distance communication impossible within one layer. Longformer therefore introduces selected global tokens. A globally marked position attends to the complete sequence, and every ordinary token can attend to the global positions.

  • This produces an attention pattern conceptually consisting of:

\[\text{local edges} + \text{global edges}\]

  • If there are \(g\) global tokens, the approximate complexity becomes \(O(nw+ng)\) which remains linear in \(n\) when \(w\) and \(g\) are bounded.

  • Global positions depend on the task: classification may designate a classification token, while question answering may give question tokens global connectivity. Such positions connect otherwise distant local windows.

  • The following figure (source) shows the progression from full self-attention to Longformer’s sparse patterns, including sliding-window attention, dilated sliding-window attention, and global attention.

Low-rank sequence compression

  • Low-rank methods replace interactions with the full sequence by interactions with projected representations. The Linformer formulation makes the approximation explicit and introduces a rank-capacity tradeoff.

Low-Rank Attention

  • Sparse attention assumes that many token-pair interactions can be omitted. Low-rank approaches make a different assumption: the attention operation may contain substantial redundancy along its sequence dimension and can therefore be represented in a lower-dimensional space.

  • Standard attention computes

    \[\operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V\]

    • where:

      \[Q,K,V\in\mathbb{R}^{n\times d}\]

  • The expensive intermediate interaction remains \(n\times n\).

  • A low-rank approach instead compresses the sequence dimension from \(n\) positions to \(k\) projected positions, with \(k\ll n\) so that each query interacts with a much smaller representation of the sequence.

Linformer

  • Linformer: Self-Attention with Linear Complexity by Wang et al. (2020) is based on the empirical and theoretical observation that the attention matrix can be approximated by a low-rank matrix. It projects keys and values along the sequence dimension before attention, reducing both time and space complexity from quadratic to linear in sequence length when the projected dimension is fixed.

  • For one attention head, standard projected keys and values have shapes

\[K,V\in\mathbb{R}^{n\times d}\]

  • Linformer introduces projection matrices \(E,F\in\mathbb{R}^{k\times n}\) and constructs:

    \[\widetilde K=EK\] \[\widetilde V=FV\]

    • where:

      \[\widetilde K,\widetilde V \in \mathbb{R}^{k\times d}\]

  • Attention then becomes:

\[\operatorname{Linformer}(Q,K,V) = \operatorname{softmax} \left( \frac{Q(EK)^\top}{\sqrt{d_k}} \right) FV\]
  • Instead of an \(n\times n\) attention matrix, the score matrix is:
\[Q(EK)^\top \in \mathbb{R}^{n\times k}\]
Sequence scaling and projection sharing

  • Projecting the sequence length reduces the number of key-value positions each query must process. Linformer also considers how projection matrices can be shared across heads or layers to manage parameter cost.

  • Consequently, the dominant sequence-dependent complexity changes from \(O(n^2d)\) to \(O(nkd)\).

  • For fixed projected dimension \(k\), \(O(nkd)=O(n)\) with respect to sequence length.

  • The following figure (source) shows Linformer’s multi-head linear self-attention architecture and the projection of the sequence dimension into a smaller low-rank representation.

Projection Sharing

  • Linformer explores several ways of sharing the projection matrices \(E\) and \(F\) to further reduce parameter cost. Headwise sharing uses common projection matrices across attention heads within a layer. Key-value sharing additionally uses the same projection for keys and values, \(E=F\).

  • Layerwise sharing goes further and reuses a projection across layers, heads, keys, and values.

  • The paper also explores nonuniform projected dimensions because different heads and layers exhibit different spectral characteristics. Higher layers can sometimes be represented using smaller projected dimensions, suggesting that attention rank need not be constant throughout the network.

Limitation of the Low-Rank Assumption
  • Linformer’s efficiency comes from replacing the original sequence dimension with a learned compressed representation. The tradeoff is that the resulting operator is an approximation of full attention rather than an exact implementation of it. Information discarded by the projection cannot subsequently be recovered by the attention operation.

Kernelized and causal linear attention

  • Kernel feature maps allow attention products to be reassociated around an aggregated state. The causal case explains how this algebraic change supports efficient recurrent inference.

Linear Attention

Performer: Rethinking Attention with Performers

  • Performer: Rethinking Attention with Performers by Choromanski et al. (2021) introduces FAVOR+ (Fast Attention Via positive Orthogonal Random features), which approximates softmax attention using positive orthogonal random features. Unlike approaches that impose sparsity or low-rank structure, Performer approximates full-rank softmax attention with linear time and space complexity in sequence length for a fixed number of random features.

  • By combining positive random features with orthogonalization to reduce estimator variance, FAVOR+ enables stable kernel approximation while preserving the associative computation underlying linear attention.

  • The following figure (source) illustrates how Performer approximates conventional attention using random feature maps, changing the order of matrix multiplication to avoid explicitly constructing the quadratic attention matrix.

Kernel Formulation

  • Consider a generalized attention operation with a non-negative similarity function:
\[\operatorname{sim}(q_i,k_j)\]
  • The output is:
\[o_i = \frac {\sum_{j=1}^{n} \operatorname{sim}(q_i,k_j)v_j} {\sum_{j=1}^{n} \operatorname{sim}(q_i,k_j)}\]
  • Suppose the similarity can be expressed using a feature map \(\phi\):
\[\operatorname{sim}(q,k) = \phi(q)^\top\phi(k)\]
  • Substituting gives:
\[o_i = \frac { \sum_j \phi(q_i)^\top \phi(k_j)v_j^\top } { \sum_j \phi(q_i)^\top \phi(k_j) }\]
Reassociation and accumulated state

  • Associativity permits key-value products to be combined before processing queries. In causal attention, this leads to a recurrent summary that can be updated as new tokens arrive.

  • Because matrix multiplication is associative, terms independent of the query can be grouped:

\[o_i = \frac { \phi(q_i)^\top \left( \sum_j \phi(k_j)v_j^\top \right) } { \phi(q_i)^\top \left( \sum_j \phi(k_j) \right) }\]
  • Let’s define \(S = \sum_j \phi(k_j)v_j^\top\) and \(z = \sum_j \phi(k_j)\). Then each query requires only:
\[o_i = \frac {\phi(q_i)^\top S} {\phi(q_i)^\top z}\]
  • The critical difference is that the model no longer constructs the full \(n\times n\) matrix of pairwise interactions. The key-value summary \(S\) and key normalizer \(z\) are computed once and reused across queries, producing linear sequence-length complexity.

Feature Map

  • Exact softmax attention uses an exponential similarity:

    \[\operatorname{sim}(q,k) = \exp(q^\top k)\]

    • whose exact finite-dimensional feature map is unavailable. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention by Katharopoulos et al. (2020) therefore use the positive feature map:

      \[\phi(x) = \operatorname{ELU}(x)+1\]
      • which ensures non-negative similarity values while avoiding the zero-gradient behavior that a simple ReLU feature map could introduce for negative inputs.

  • This means the mechanism should not be interpreted as merely a faster implementation of ordinary softmax attention. It defines a different attention function chosen specifically so that the computation can be factorized.

Causal Linear Attention

  • Autoregressive attention introduces an additional constraint because position \(i\) may use only positions \(j\leq i\).

  • The linear-attention output becomes:

\[o_i = \frac { \phi(q_i)^\top \sum_{j=1}^{i} \phi(k_j)v_j^\top } { \phi(q_i)^\top \sum_{j=1}^{i} \phi(k_j) }\]
  • Let’s define running states as \(S_i = S_{i-1} + \phi(k_i)v_i^\top\) and \(z_i = z_{i-1} + \phi(k_i)\). Then:
\[o_i = \frac {\phi(q_i)^\top S_i} {\phi(q_i)^\top z_i}\]

  • The attention layer can therefore be evaluated recurrently. Instead of caching every historical key and value and rereading them at every decoding step, the model updates fixed-size summary states \(S_i\) and \(z_i\). This is the central reason the paper describes linear Transformers as RNNs.

  • The following figure (source) illustrates the correspondence between causal linear attention and recurrent computation, where accumulated key-value statistics act as the recurrent state.

Dynamic sparsity and Native Sparse Attention

  • Dynamic sparse methods select relevant interactions rather than relying entirely on fixed neighborhoods. NSA combines compression, selection, and local attention to connect algorithmic sparsity with hardware-friendly computation.

Sparse Attention as Dynamic Retrieval

  • Local attention uses a predetermined sparsity pattern. A natural extension is dynamic sparsity, where the model determines which portions of the context are useful for each query.

  • Conceptually, if \(\mathcal{I}(q_i) \subset \{1,\ldots,n\}\) is the set of context positions selected for query \(q_i\), sparse attention computes:

    \[o_i = \operatorname{Attention} \left( q_i, K_{\mathcal{I}(q_i)}, V_{\mathcal{I}(q_i)} \right)\]
    • rather than evaluating every query-key pair.

  • Sparse computation only accelerates hardware when selected interactions map efficiently to memory and matrix operations. Irregular token-level access can erase theoretical savings, motivating blockwise sparsity.

Native Sparse Attention

  • Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention by Yuan et al. (2025) addresses sparse attention as both an algorithmic and systems problem. NSA is trained with sparse attention from the outset and combines coarse-grained compression, fine-grained dynamic selection, and local sliding-window attention.

  • Rather than relying on one sparsity mechanism, NSA uses three parallel branches:

\[\text{NSA} = \text{Compression} + \text{Selection} + \text{Sliding Window}\]

  • Each branch captures a different type of contextual information.

  • The following figure (source) shows NSA’s hierarchical sparse-attention design, which combines compressed global context, selectively retrieved fine-grained blocks, and recent local context.

Compression Branch

  • The compression branch summarizes consecutive blocks of tokens into compressed key-value representations. Instead of retaining only individual token-level interactions, it provides a coarse global view over long context.

  • For a block of keys and values, a learned compression operator can be represented abstractly as

    \[\widetilde{k}^{\text{cmp}}_m = \operatorname{Compress}_K ( k_{s_m:s_m+l} )\] \[\widetilde{v}^{\text{cmp}}_m = \operatorname{Compress}_V ( v_{s_m:s_m+l} )\]
    • where \(l\) is the compression block size.

  • This branch allows a query to cheaply inspect information distributed throughout the historical context, even when individual token-level representations are not directly examined.

Selection Branch

  • Compression identifies broad regions that may be relevant, but coarse representations can discard fine-grained information. NSA therefore includes a selection branch that retrieves selected blocks of original token-level keys and values.

  • For each query, the mechanism obtains a set of selected blocks:

    \[\mathcal{B}_i = \operatorname{TopBlocks}(q_i)\]

    • and performs attention over their full-resolution token representations:

      \[o_i^{\text{slc}} = \operatorname{Attention} \left( q_i, K_{\mathcal{B}_i}, V_{\mathcal{B}_i} \right)\]

  • Block-level rather than arbitrary token-level selection is important for implementation efficiency because contiguous key-value regions can be loaded and processed more effectively on GPUs.

Sliding-Window Branch

  • The third branch guarantees high-resolution access to recent context. For window size \(w\), NSA retains:

    \[\widetilde K_t^{\text{win}} = K_{t-w:t}\] \[\widetilde V_t^{\text{win}} = V_{t-w:t}\]
    • and computes ordinary attention over these nearby tokens.

  • This branch captures the strong locality bias of language and ensures that dynamic selection is not responsible for recovering every short-range dependency.

Gated Combination

  • NSA computes the branches separately and combines their outputs through learned gating. Abstractly,

    \[o_t = g_t^{\text{cmp}}o_t^{\text{cmp}} + g_t^{\text{slc}}o_t^{\text{slc}} + g_t^{\text{win}}o_t^{\text{win}}\]
    • where the gates determine the relative contribution of compressed global context, selected fine-grained context, and recent local context.

  • The branches use independent key-value representations to reduce interference between their distinct roles. The design therefore separates coarse long-range summarization, precise retrieval, and local modeling rather than forcing one attention pattern to perform all three functions.

Hardware-Aligned Sparse Attention

  • NSA’s sparsity pattern is deliberately designed around blockwise GPU execution. During sparse selection attention, query heads belonging to the same GQA group share selected key-value blocks. The implementation therefore groups those queries together and loads the shared contiguous KV blocks into on-chip memory, avoiding redundant KV transfers.

  • This distinction is important:

    \[\text{fewer FLOPs} \not\Rightarrow \text{lower latency}\]
    • if the remaining operations require irregular memory accesses or underutilize the accelerator.

  • NSA’s implementation uses Triton kernels and arranges selection around contiguous blocks so that sparsity produces useful arithmetic intensity. The paper reports that at a sequence length of \(64\text{K}\), its evaluated configuration achieves substantial speedups over full attention across decoding, forward propagation, and backward propagation.

  • The following figure (source) compares NSA with full attention across model quality and computational efficiency, reporting speedups at a 64K sequence length for decoding, forward propagation, and backward propagation.

Comparisons and remaining limits

  • Local, low-rank, linear, and dynamic sparse mechanisms reduce different costs and impose different retrieval constraints. Their limits motivate attention designs that also consider KV-cache storage and accelerator execution.

Comparing the Main Approaches

  • The mechanisms in this family reduce quadratic attention in fundamentally different ways:
Mechanism Core idea Interaction pattern Approximation/change Sequence scaling
Full attention Every query attends to every key Dense Exact softmax \(O(n^2)\)
Sliding-window attention Attend only nearby tokens Static sparse Restricted receptive field \(O(nw)\)
Longformer Local windows plus global tokens Static sparse Restricted structured attention \(O(nw+ng)\)
Linformer Project sequence dimension to \(k\) Dense over compressed sequence Low-rank approximation \(O(nk)\)
Linear attention Kernelize and reassociate attention Implicit global Changes similarity function \(O(n)\) for fixed feature size
NSA Compression plus dynamic block selection plus local window Dynamic hierarchical sparse Learned sparse attention Subquadratic, configuration-dependent

A Useful Conceptual Distinction

  • These mechanisms can be organized according to what they remove from full attention.

  • Local and sparse attention remove interactions:

    \[n^2 \rightarrow |\mathcal E|\]
    • where \(\mathcal E\) is the selected set of query-key edges.

  • Linformer removes sequence dimensionality:

\[n \rightarrow k \qquad k\ll n\]

  • Linear attention removes explicit pairwise materialization by changing the algebra:

    \[(QK^\top)V \rightarrow Q(K^\top V)\]
    • after introducing a suitable feature-map formulation.

  • NSA combines multiple forms of selective context access while explicitly designing the resulting sparse computation for modern GPU hardware.

What These Methods Do Not Solve

  • Lower asymptotic complexity need not produce a faster GPU kernel. Memory traffic, launch overhead, arithmetic intensity, and matrix-unit utilization motivate IO-aware exact attention.

  • A second line of work therefore asks a different question: if exact dense softmax attention remains desirable, can the same mathematical result be computed without repeatedly materializing and moving the enormous attention matrix through GPU memory?

  • IO-aware exact attention provides a complementary approach: rather than reduce the computed token pairs, FlashAttention reorganizes memory movement and work partitioning. Its kernel principles reappear in the later implementation chapter.

IO-Aware Exact Attention: FlashAttention

  • FlashAttention improves the execution of exact dense softmax attention without changing which query-key interactions are permitted. This chapter follows the family from IO-aware tiling through work partitioning, asynchronous execution, and hardware-specific pipelining. It distinguishes mathematical equivalence from runtime and memory-efficiency claims.

FlashAttention: Exact Attention Without Quadratic Intermediate Storage

  • FlashAttention by Dao et al. (2022) recognizes that standard implementations materialize attention scores and probabilities in HBM even though the output is only sequence-by-head-dimension. For a head, the operator is unchanged:
\[O=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}}+M\right)V\]

IO-Aware Tiling

  • Small query blocks interact successively with key-value blocks loaded into on-chip SRAM. Fusing score multiplication, softmax updates, and value accumulation avoids writing the full score and probability tensors to HBM. Dense attention still requires \(O(N^2d_h)\) arithmetic, but its intermediate memory traffic decreases considerably.

  • The following figure (source) shows the GPU memory hierarchy and tiled IO-aware attention execution.

Online Softmax and Backward Recomputation

  • Online softmax maintains running row maxima and normalization statistics across tiles. For old statistics \(m,\ell\) and a new block \(\widetilde S\), the merged quantities are:
\[m'=\max(m,\max_j\widetilde S_j)\] \[\ell'=e^{m-m'}\ell+\sum_j e^{\widetilde S_j-m'}\]
  • The weighted output accumulator is rescaled correspondingly. Rather than storing quadratic probabilities for backward propagation, the kernel recomputes needed score blocks from saved inputs and compact normalization statistics. This trades additional arithmetic for lower memory movement.

FlashAttention-2: Parallelism and Work Partitioning

  • FlashAttention-2 by Dao (2023) reduces non-matrix-multiplication operations and improves occupancy by parallelizing across sequence blocks, not just batch and heads. Its forward pass splits queries across warps instead of splitting keys and values, reducing cross-warp shared-memory communication.

Query-Split Work Partitioning

  • Splitting the query block allows each warp to generate its own output region against shared key-value tiles. This makes long sequences and smaller batches better able to occupy GPU streaming multiprocessors.

  • The following figure (source) contrasts the warp work partitioning of the first and second versions.

FlashAttention-3: Hopper Asynchrony and Reduced Precision

  • FlashAttention-3 by Shah et al. (2024) uses producer-consumer warp specialization, asynchronous Tensor Core multiplication, and the Tensor Memory Accelerator to overlap loading, GEMM, and softmax. It also explores FP8 computation with block quantization and incoherent preprocessing to control quantization error.

Overlapping Compute and Data Movement

  • When one tile is consumed by matrix multiplication, another tile can be loaded asynchronously; softmax work for one block can overlap matrix multiplication for another. These implementation changes target Hopper’s execution model without changing the underlying softmax-attention objective.

FlashAttention-4: Blackwell-Specific Pipeline Co-Design

  • FlashAttention-4 by Zadouri et al. (2026) addresses asymmetric Blackwell scaling: Tensor Core throughput increased more rapidly than shared-memory bandwidth and exponentiation capacity. It combines fully asynchronous MMA pipelines, larger tiles, software-emulated exponential, and conditional softmax rescaling.

Forward and Backward Resource Balance

  • The forward pass shifts some exponentiation work to available arithmetic units, whereas the backward pass uses tensor memory and the two-CTA MMA mode to reduce shared-memory traffic and selected atomic reductions. Its CuTe-DSL implementation also improves iteration speed for hardware-specific kernels.

  • The following figure (source) illustrates the Blackwell-oriented asynchronous pipeline design.

Operator Versus Kernel

  • FlashAttention preserves exact dense softmax attention, whereas sparse attention changes the token-pair graph and recurrent linear attention changes the aggregation operator. Kernel improvements and KV sharing address distinct bottlenecks; the next chapter examines representations that shrink the state read during autoregressive decoding.

KV-Cache-Efficient Attention: MQA, GQA, and MLA

  • Autoregressive decoding repeatedly reads cached historical keys and values, introducing a bandwidth bottleneck distinct from prefill computation. This chapter compares Multi-Query and Grouped-Query Attention as KV-sharing strategies with Multi-Head Latent Attention as a representation-compression strategy. The derivations make cache-size reductions and their modeling tradeoffs explicit.

Decode-time memory bottlenecks

  • Autoregressive attention repeatedly accesses historical keys and values even when the model generates only one new query. The cache footprint and memory bandwidth motivate the head-sharing designs introduced next.

Why Autoregressive Decoding Has a Different Bottleneck

  • FlashAttention primarily improves how attention is executed, especially when many query tokens are processed together. Autoregressive decoding has a different performance profile. Once the prompt has been processed, an LLM typically generates one new token at a time. At decoding step \(t\), the model computes a new query but needs keys and values from all preceding positions \(q_t\):

    \[K_{1:t}=[k_1,\ldots,k_t]\] \[V_{1:t}=[v_1,\ldots,v_t]\]

  • The attention output for head \(i\) is:

\[o_{t,i}=\operatorname{softmax}\left(\frac{q_{t,i}K_{1:t,i}^{\top}}{\sqrt{d_h}}\right)V_{1:t,i}\]

  • Recomputing historical keys and values at every generation step would be wasteful, so inference systems store them in a KV cache. Each newly generated token appends its key and value representations to this cache.

  • For standard Multi-Head Attention (MHA) with \(H\) heads and head dimension \(d_h\), the number of cached elements per token per layer is \(2Hd_h\). For context length \(L\) and \(N_L\) Transformer layers, the cache therefore contains approximately \(2N_LLHd_h\) elements per sequence, before accounting for batch size and datatype size. DeepSeek-V2 by DeepSeek-AI (2024) uses the equivalent expression when motivating its KV-cache compression mechanism.

  • During token-by-token generation, attention often has relatively little arithmetic per byte loaded. The model repeatedly streams the historical KV cache from memory for a small number of new queries. Consequently, decoding can become memory-bandwidth-bound rather than compute-bound.

Head sharing: MQA and GQA

  • MQA reduces cached KV heads to one shared pair, whereas GQA provides an intermediate number of shared pairs. Their cache formulas and checkpoint-conversion procedures establish the practical tradeoff between head diversity and decoding bandwidth.

Multi-Query Attention

  • Fast Transformer Decoding: One Write-Head is All You Need by Shazeer (2019) introduced Multi-Query Attention (MQA) specifically to reduce this incremental-decoding bottleneck. The central observation is that query heads can remain independent while keys and values are shared across them.

  • In ordinary MHA, every head has separate projections:

    \[q_i=W_i^Qh\] \[k_i=W_i^Kh\] \[v_i=W_i^Vh\]

    • for \(i=1,\ldots,H\), and computes:

      \[o_i=\operatorname{Attention}(q_i,K_i,V_i)\]

  • MQA retains head-specific queries, \(q_i=W_i^Qh\), but constructs only one shared key and value:

    \[k=W^Kh\] \[v=W^Vh\]

    • so all query heads compute:

      \[o_i=\operatorname{Attention}(q_i,K,V)\]

What MQA Actually Shares

  • MQA does not collapse attention to a single head. There are still \(H\) different query heads. Each can therefore generate a different attention distribution against the shared keys:
\[A_i=\operatorname{softmax}\left(\frac{Q_iK^\top}{\sqrt{d_h}}\right)\]
  • What disappears is the duplication of key and value representations:
\[\begin{aligned}\text{MHA:}&\quad(K_1,V_1),\ldots,(K_H,V_H)\\\text{MQA:}&\quad(K,V)\end{aligned}\]
  • The query projections preserve head-specific retrieval behavior, while the shared keys and values dramatically reduce the state that must be cached and loaded during decoding.

KV-Cache Reduction with MQA

  • For MHA, the per-token cache contains \(2Hd_h\) elements. For MQA, it contains only \(2d_h\) elements. The idealized reduction factor is therefore \(\frac{2Hd_h}{2d_h}=H\) for the attention-layer KV state.

  • The same reduction applies to the amount of historical KV data that needs to be streamed for attention during decoding. This is why MQA’s primary advantage is often memory bandwidth, not merely fewer arithmetic operations.

The MQA Tradeoff

  • Sharing one key-value representation across all query heads is aggressive compression:
\[H\text{ KV heads}\rightarrow1\text{ KV head}\]
  • It substantially improves inference efficiency, but it also removes head-specific key and value projections. This naturally motivates an intermediate design \(1<G<H\), where \(G\) is the number of KV heads.

Grouped-Query Attention

MHA and MQA as Special Cases of GQA

  • When \(G=H\), every query head has its own KV head:
\[\operatorname{GQA}(G=H)=\operatorname{MHA}\]
  • When \(G=1\), all query heads share one KV head:
\[\operatorname{GQA}(G=1)=\operatorname{MQA}\]

  • Intermediate values \(1<G<H\) give ordinary GQA.

  • The following figure (source) shows Multi-Head Attention, Grouped-Query Attention, and Multi-Query Attention: MHA uses separate key and value heads for every query head, MQA shares one key and value head across all query heads, and GQA shares a key and value head within each query group.

KV-Cache Size with GQA

  • With \(G\) KV heads, the per-token KV cache contains \(2Gd_h\) elements rather than the MHA requirement \(2Hd_h\). Relative to MHA, the idealized attention KV-cache reduction factor is therefore \(\frac{H}{G}\).

Uptraining MHA into GQA or MQA

  • GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints by Ainslie et al. (2023) propose uptraining an MHA checkpoint into MQA or GQA. For MQA conversion, the original key and value projection matrices are mean-pooled:

    \[W_{\text{MQA}}^K=\frac{1}{H}\sum_{i=1}^{H}W_i^K\] \[W_{\text{MQA}}^V=\frac{1}{H}\sum_{i=1}^{H}W_i^V\]

  • For GQA, mean pooling occurs within each group. The converted model is then trained further so that it can adapt to the shared KV structure. The paper studies uptraining with approximately \(5\%\) of the original pretraining compute.

Sliding-Window KV Attention

  • KV-head sharing and sparse attention address orthogonal dimensions of the decoding problem. Sharing reduces how many KV representations exist per position, while sliding-window attention reduces how many historical positions remain directly visible.

  • For a causal window of size \(W\), token \(t\) attends only to:

    \[\mathcal{N}(t)=\{\max(1,t-W+1),\ldots,t\}\]
    • rather than all previous positions.

  • Combining local attention with shared KV heads attacks both dimensions: \(H\rightarrow G\) KV heads and \(L\rightarrow W\) attended positions. The resulting cache needed by a strictly windowed layer can scale approximately as \(O(GWd_h)\) rather than \(O(HLd_h)\) for full-context MHA.

  • The following figure illustrates attention architecture and implementation considerations in modern language models, providing additional context for combining locality constraints with KV-head sharing.

MHA vs. GQA vs. MQA

  • These three configurations retain multiple query heads but allocate different numbers of key and value heads. The comparison isolates changes in KV-cache storage and bandwidth from changes in the query-head count.
Property MHA GQA MQA
Query heads \(H\) \(H\) \(H\)
KV heads \(H\) \(G\) \(1\)
Queries per KV head \(1\) \(\frac{H}{G}\) \(H\)
KV elements/token \(2Hd_h\) \(2Gd_h\) \(2d_h\)
KV diversity Highest Intermediate Lowest
KV-cache bandwidth Highest Intermediate Lowest
Main goal Maximum head-specific representations Quality/efficiency compromise Maximum KV sharing

Beyond Sharing: Compressing the KV Representation

  • MQA and GQA reduce KV cache by sharing keys and values: \(H\rightarrow G\rightarrow1\) KV heads. A different possibility is to preserve rich multi-head behavior while caching a compressed latent representation from which the required key and value information can be derived.

  • This is the motivation behind Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model by DeepSeek-AI (2024). MLA uses low-rank joint compression of keys and values rather than simply sharing full-dimensional KV heads.

Multi-Head Latent Attention

  • Let the hidden representation for token \(t\) be \(h_t\in\mathbb{R}^{d}\). MLA first compresses the hidden representation into a much smaller joint KV latent:

    \[c_t^{KV}=W^{DKV}h_t\]
    • where \(c_t^{KV}\in\mathbb{R}^{d_c}\) and \(d_c\ll Hd_h\).

  • Full content keys and values can then be reconstructed through up-projections:

    \[k_t^C=W^{UK}c_t^{KV}\] \[v_t^C=W^{UV}c_t^{KV}\]

  • The model can cache the compact latent \(c_t^{KV}\) instead of the expanded multi-head content keys and values.

Joint KV Compression

  • MLA instead stores a compact latent representation from which multi-head KV behavior is derived:
\[h_t\xrightarrow{W^{DKV}}c_t^{KV}\xrightarrow{\{W^{UK},W^{UV}\}}\{k_t^C,v_t^C\}\]
  • The following figure (source) provides a simplified comparison of MHA, GQA, MQA, and MLA, showing that MLA jointly compresses keys and values into a latent vector that is cached during inference.

Absorbing the KV Up-Projection

  • Consider the content portion of the attention score, with \(k_j^C=W^{UK}c_j^{KV}\):
\[q_t^\top W^{UK}c_j^{KV}=\left((W^{UK})^\top q_t\right)^\top c_j^{KV}\]
  • Rather than explicitly reconstructing and caching the expanded key for every historical token, the query transformation can incorporate the key up-projection and interact directly with the compressed latent.

Query Compression

  • DeepSeek-V2 additionally applies low-rank compression to queries:

    \[c_t^Q=W^{DQ}h_t\] \[q_t^C=W^{UQ}c_t^Q\]
    • where \(d_c'\ll Hd_h\).

The RoPE Problem

  • Rotary Position Embedding applies a position-dependent transformation to queries and keys:
\[q_t'=R_tq_t\] \[k_j'=R_jk_j\]
  • If RoPE were applied directly to MLA’s compressed-content key after its up-projection, the position-dependent matrix would prevent the up-projection from being simply absorbed into the query projection.

Decoupled RoPE

  • MLA solves this using decoupled Rotary Position Embedding.

  • For queries:

    \[[q_{t,1}^R;\ldots;q_{t,H}^R]=\operatorname{RoPE}\left(W^{QR}c_t^Q\right)\]

  • For keys:

    \[k_t^R=\operatorname{RoPE}\left(W^{KR}h_t\right)\]

  • The final query and key are:

    \[q_{t,i}=[q_{t,i}^C;q_{t,i}^R]\] \[k_{t,i}=[k_{t,i}^C;k_t^R]\]

    • and attention uses:

      \[o_{t,i}=\sum_{j=1}^{t}\operatorname{Softmax}_j\left(\frac{q_{t,i}^{\top}k_{j,i}}{\sqrt{d_h+d_h^R}}\right)v_{j,i}^C\]

What MLA Caches

  • Conceptually, MLA retains the compressed content state \(c_t^{KV}\) plus the shared positional key \(k_t^R\) for each historical token. Instead of approximately \(2Hd_h\) elements per token as standard MHA does, the persistent representation scales approximately with \(d_c+d_h^R\).

  • Elements per token under MLA’s latent-cache formulation.

  • The following figure (source) shows the DeepSeek-V2 architecture and a detailed view of MLA, including the compressed KV latent, compressed query latent, content and RoPE components, and the representations cached during inference.

MLA’s Reported Inference Gains

  • In DeepSeek-V2, MLA is part of a larger architectural redesign that also includes DeepSeekMoE, so end-to-end model comparisons should not attribute every measured gain to attention alone. The paper reports that DeepSeek-V2 reduces KV-cache size by \(93.3\%\) relative to DeepSeek 67B and increases maximum generation throughput by \(5.76\times\) in its reported model-level comparison.

  • DeepSeek-V2 supports context lengths up to \(128\text{K}\) after its long-context extension procedure, making KV-cache reduction particularly consequential because cache memory grows linearly with context length.

Choosing and combining cache optimizations

  • Sharing heads, compressing per-token state, and limiting the number of resident positions address different components of memory use. Their interplay determines which attention state must be stored and loaded in a given inference regime.

Sharing vs. Compression

  • MHA, GQA, MQA, and MLA can be understood by the representation retained for historical tokens:
\[\boxed{\text{MHA}:\text{many full KV heads}}\] \[\boxed{\text{GQA}:\text{fewer shared full KV heads}}\] \[\boxed{\text{MQA}:\text{one shared full KV head}}\] \[\boxed{\text{MLA}:\text{compressed latent KV state}}\]
  • GQA and MQA reduce redundancy by sharing across heads. MLA instead attacks redundancy through low-rank joint compression.

KV-Cache Efficiency as a Separate Design Axis

  • Sliding-window attention reduces the number of positions, \(L\rightarrow W\). MQA and GQA reduce the number of KV heads, \(H\rightarrow G\). MLA reduces the representation stored for each historical position, \(2Hd_h\rightarrow d_c+d_h^R\). FlashAttention changes execution from HBM-heavy to tiled IO-aware attention.

  • These optimizations can therefore be combined. A modern attention stack can simultaneously use grouped or compressed KV representations, local or sparse context patterns, and hardware-efficient attention kernels.

From KV Compression to Hybrid Attention

  • Dense, sparse, shared-KV, and latent-KV designs target different resource constraints; they can be combined with IO-aware kernels.

  • Recurrent and hybrid mechanisms go further by summarizing historical interactions while retaining selected exact softmax layers. Lightning Attention provides the next example.

Lightning Attention and Hybrid Linear Attention

  • Linear attention can summarize previous tokens in recurrent state, while blockwise execution can recover accelerator-friendly parallelism. This chapter derives Lightning Attention’s intra-block and inter-block computation before contrasting it with exact softmax and hybrid stacks. The distinction between training-time throughput and token-by-token retrieval remains central.

Linear-attention foundations

  • Associativity permits certain attention operators to summarize the past into a recurrent state instead of revisiting every historical token. The causal computation and its left/right multiplication orders explain the motivation for blockwise execution.

From Quadratic Attention to Recurrent State

  • Standard causal attention explicitly forms interactions between each token and all preceding tokens:

    \[O=\left[(QK^\top)\odot M\right]V\]

    • where the causal mask is:

      \[M_{ts}=\begin{cases}1,&t\geq s\\0,&t<s\end{cases}\]

  • Linear attention exploits associativity to change the order of multiplication. Ignoring normalization and other architectural details, \((QK^\top)V\) can be reassociated as \(Q(K^\top V)\). The left-associated formulation first creates an \(n\times n\) interaction matrix, whereas the right-associated formulation first produces a fixed-size key-value summary.

  • In the causal setting, we can define the running state as :

    \[KV_t=\sum_{j=1}^{t}k_jv_j^\top\]
    • Then \(KV_t=KV_{t-1}+k_tv_t^\top\) and \(o_t^\top=q_t^\top KV_t\).

  • This gives a recurrent interpretation of causal linear attention: historical tokens are summarized into a fixed-dimensional matrix rather than retained as an ever-growing KV cache.

The Causal Linear-Attention Problem

Left Product vs. Right Product

  • The left product first evaluates query-key interactions:

    \[O_{\text{left}}=[(QK^\top)\odot M]V\]

  • The right product first aggregates keys and values:

    \[O_{\text{right}}=Q(K^\top V)\]

    • For causal attention, this becomes

      \[KV_t=KV_{t-1}+k_tv_t^\top\] \[o_t^\top=q_t^\top KV_t\]

  • Lightning Attention uses the left product for local computation and the right product for historical computation.

Blockwise Lightning Attention

  • Lightning Attention separates interactions within the current block from interactions represented by earlier blocks. This decomposition supports GPU-friendly dense operations locally while retaining a compact accumulated history across blocks.

Blockwise Decomposition

  • Divide the sequence into blocks of size \(B\):

    \[Q=[Q_1;Q_2;\ldots;Q_T]\] \[K=[K_1;K_2;\ldots;K_T]\] \[V=[V_1;V_2;\ldots;V_T]\]
    • where \(T=n/B\) and \(Q_t,K_t,V_t\in\mathbb{R}^{B\times d}\).

  • MiniMax-01: Scaling Foundation Models with Lightning Attention by MiniMax et al. (2025) describes the two pieces as inter-block and intra-block computation.

Intra-Block Attention

  • Interactions among tokens inside the current block use

    \[O_t^{\text{intra}}=\left[(Q_tK_t^\top)\odot M\right]V_t\]
    • where \(M\in\mathbb{R}^{B\times B}\) is a block-sized causal mask.

  • Across all blocks, the intra-block cost is approximately

    \[\frac{n}{B}O(B^2d)=O(nBd)\]
    • For fixed \(B\), this is linear in sequence length.

Inter-Block Attention

  • Information from completed blocks is summarized into \(KV_{t-1}\). For the current block \(O_t^{\text{inter}}=Q_tKV_{t-1}\), the state is updated as:

    \[KV_t=KV_{t-1}+K_t^\top V_t\]

  • The complete block output is:

    \[O_t=\left[(Q_tK_t^\top)\odot M\right]V_t+Q_tKV_{t-1}\]

Why the Decomposition Works

  • Keys belonging to earlier blocks are all causally valid and can be compressed into the accumulated state. Keys in the current block require an explicit causal mask because a token must not see later positions in its own block. Thus, past blocks use compressed state while the current block uses explicit causal attention.

Lightning Attention Forward Pass

  • The forward pass combines a local masked product within each block with a summary of all preceding blocks. Its state update allows the next block to reuse that accumulated history.
KV = 0

for each block t:
    load Q_t, K_t, V_t into on-chip memory
    O_intra = ((Q_t K_t^T) ⊙ M) V_t
    O_inter = Q_t KV
    O_t = O_intra + O_inter
    KV = KV + K_t^T V_t
    write O_t

  • The global state is only \(KV\in\mathbb{R}^{d\times d}\) rather than an attention matrix whose dimensions grow with sequence length.

  • The following figure (source) shows the IO-aware Lightning Attention algorithm, where each block computes causal intra-block attention and combines it with an accumulated inter-block \(KV\) state.

Efficiency, hardware mapping, and inference

  • The algebraic formulation is only part of the performance story: tiling, memory traffic, and scan/state-update behavior determine efficiency. These sections distinguish theoretical scaling from the execution regimes that matter in practice.

Computational Complexity

  • For every block, intra-block computation costs approximately \(O(B^2d)\), while inter-block multiplication and state updating involve approximately \(O(Bd^2)\) work. Across \(n/B\) blocks, the resulting complexity is:

    \[O(nBd+nd^2)\]

  • If \(B\) and \(d\) are fixed with respect to sequence length, this is \(O(n)\) in \(n\), compared with \(O(n^2d)\) for dense attention.

Memory Complexity

  • Lightning Attention operates on block-sized attention matrices plus a fixed-dimensional recurrent state, \(B\times B\) and \(d\times d\) respectively. Increasing sequence length therefore does not require constructing an increasingly large global attention matrix.

Why Lightning Attention Is Hardware-Aware

  • Token-by-token cumulative updates expose insufficient parallel work for GPUs. Blocking converts many small recurrent operations into matrix multiplications such as \(Q_tK_t^\top\), \(K_t^\top V_t\), and \(Q_tKV_{t-1}\) that map naturally onto GPU hardware. Lightning Attention-2 implements the tiled formulation in Triton for both forward and backward passes.

Lightning Attention vs. FlashAttention

  • FlashAttention computes exact softmax attention and reorganizes execution to minimize memory traffic, while its dense arithmetic complexity remains \(O(n^2d)\).

  • Lightning Attention changes the attention formulation so historical context can be represented through a recurrent matrix state and achieves sequence-linear asymptotic computation for fixed feature and block dimensions.

    \[\boxed{\text{FlashAttention}=\text{exact softmax attention with better execution}}\] \[\boxed{\text{Lightning Attention}=\text{linear attention with hardware-efficient execution}}\]

Lightning Attention vs. Earlier Linear Attention

  • Earlier kernelized linear attention already used associativity to transform pairwise attention into accumulated key-value statistics. Lightning Attention addresses the implementation bottleneck by performing recurrent accumulation at block granularity while recovering efficient parallel matrix multiplication inside each block.

Autoregressive Inference

  • Once the historical state \(KV_{t-1}\) has been computed, the new token contributes:

    \[KV_t=KV_{t-1}+k_tv_t^\top\]
    • and its output can be evaluated from the accumulated state without rereading an ever-growing KV cache.

  • Linear recurrent attention can retain a state whose dimensions depend on feature size rather than sequence length:

    \[KV_t\in\mathbb{R}^{d\times d}\]

What Is Lost by Compressing History?

  • Standard attention retains historical token representations separately, allowing a future query to assign an independently computed weight to each historical position. Linear recurrent attention compresses history into:

    \[KV_t=\sum_{j=1}^{t}k_jv_j^\top\]
    • so it does not retain the same explicit token-level retrieval structure as full softmax attention.

Hybrid architectures and tradeoffs

  • Compressed recurrent state does not provide the same arbitrary token-level retrieval as dense softmax attention. Hybrid models combine the two operators to distribute cost and representation capacity across layers.

Hybrid Linear-Softmax Attention

  • MiniMax-01: Scaling Foundation Models with Lightning Attention by MiniMax et al. (2025) provides a large-scale hybrid example. MiniMax-Text-01 uses Lightning Attention for most attention layers but periodically inserts conventional softmax-attention layers.

  • Its 80-layer architecture places one softmax-attention layer after every seven Lightning Attention layers:

    \[\underbrace{L,L,L,L,L,L,L}_{7\text{ Lightning layers}}\rightarrow S\rightarrow\underbrace{L,L,L,L,L,L,L}_{7\text{ Lightning layers}}\rightarrow S\rightarrow\cdots\]
    • where \(L\) denotes Lightning Attention and \(S\) denotes softmax attention.

Why Keep Softmax Layers?

  • Hybrid models use recurrent layers for economical long-range state propagation and occasional softmax layers for explicit token-level retrieval.

Scaling Lightning Attention in MiniMax-01

  • MiniMax-Text-01 scales the hybrid architecture to 456 billion total parameters, with 45.9 billion activated per token through its Mixture-of-Experts architecture. The report describes training contexts up to 1 million tokens and inference extrapolation to contexts up to 4 million tokens. These capabilities belong to the overall system rather than Lightning Attention alone.

A Broader View of Hybrid Attention

  • Hybrid designs distribute sequence-mixing work across operators with different retrieval and execution properties. Comparing those operators clarifies when recurrent history should be supplemented with explicit attention.
Mechanism Historical Representation Sequence Scaling Direct Token-Level Retrieval
Dense softmax attention Individual KV vectors \(O(n^2)\) Full
Sliding-window attention Recent individual KV vectors \(O(nW)\) Local
MQA/GQA Individual KV vectors with head sharing Still context-dependent Full over retained context
MLA Compressed latent per token Still context-dependent Full over latent-backed positions
Lightning Attention Accumulated matrix state + current block \(O(n)\) for fixed \(B,d\) Explicit within block
Hybrid Lightning + softmax Recurrent state plus periodic dense layers Architecture-dependent Periodically full

The Evolutionary Pattern

  • Earlier linear attention established the algebraic possibility \(Q(K^\top V)\). Lightning Attention makes the causal version practical by combining blockwise tiling, intra-block left products, inter-block right products, and IO-aware GPU kernels. Hybrid architectures then address the representational tradeoff by occasionally restoring full softmax attention.

  • Attention design has therefore evolved along several largely orthogonal axes: which tokens interact, how KV state is represented, whether history is compressed into recurrent state, and how the resulting computation is mapped onto hardware.

Attention in Modern LLMs: A Unified Design Taxonomy

  • Modern attention stacks combine decisions about token connectivity, KV storage, mathematical operators, and execution kernels. This chapter organizes those choices as largely independent axes instead of treating each named mechanism as a complete alternative. The resulting taxonomy separates resource savings during training, prefill, and decode.

A multidimensional attention taxonomy

  • Modern attention mechanisms modify different architectural and execution dimensions, so treating them as mutually exclusive alternatives obscures useful combinations. The baseline and design axes below separate connectivity, KV representation, operators, and kernels.

There Is No Single “Modern Attention”

  • Modern LLM attention combines independent choices about connectivity, KV organization, historical state, and accelerator execution. Dense MHA is one baseline, not the only valid combination.

  • It is therefore useful to decompose an attention architecture along four major axes:

    \[\boxed{ \text{Connectivity} \times \text{KV representation} \times \text{attention operator} \times \text{kernel implementation} }\]

  • These axes answer different questions:

Design axis Core question Representative approaches
Connectivity Which token pairs interact? Dense, sliding-window, sparse, NSA
KV representation What is stored for historical tokens? MHA, MQA, GQA, MLA
Attention operator How is history aggregated? Softmax, linear/recurrent, hybrid
Kernel implementation How is the operation executed efficiently? FlashAttention, Lightning Attention kernels
  • This distinction prevents several common category errors. FlashAttention, for example, is not an alternative to GQA: a model can use GQA and execute it with a FlashAttention-style kernel. Similarly, sliding-window attention can use MHA, MQA, or GQA.

Dense Multi-Head Attention as the Baseline

  • The original Transformer formulation in Attention Is All You Need by Vaswani et al. (2017) assigns separate query, key, and value projections to every head:

    \[Q_i=XW_i^Q\] \[K_i=XW_i^K\] \[V_i=XW_i^V\]

    • and computes:

      \[O_i = \operatorname{softmax} \left( \frac{Q_iK_i^\top}{\sqrt{d_h}} + M \right)V_i\]

      • followed by:

        \[O = \operatorname{Concat}(O_1,\ldots,O_H)W^O\]

  • Dense MHA offers the richest version of head-specific key-value representations among the standard MHA/MQA/GQA family. Its costs, however, appear along two different dimensions:

    • During prompt processing or training, the number of query-key interactions scales as \(O(n^2)\).

    • During autoregressive generation, the KV cache grows as \(O(nHd_h)\) per layer.

  • Much of the subsequent evolution of attention can be understood as targeting one of these two costs without unnecessarily sacrificing the other.

Connectivity and KV representation

  • Connectivity determines which positions can communicate, whereas KV organization determines what the model must keep in memory. These choices can be made independently and combined with different attention operators.

Axis 1: Which Tokens Interact?

  • Dense attention allows every query to interact with every causally valid preceding key:

    \[\mathcal{A}(t) = \{1,\ldots,t\}\]

  • This gives maximal direct connectivity but produces quadratic interaction growth.

  • Local attention restricts the accessible set:

    \[\mathcal{A}(t) = \{\max(1,t-W+1),\ldots,t\}\]
    • giving approximately \(O(nW)\) rather than \(O(n^2)\) attention interactions for fixed window size \(W\).

  • Longformer: The Long-Document Transformer by Beltagy et al. (2020) demonstrates a particularly influential combination of sliding-window local attention and selected global-attention positions. The local pattern provides inexpensive neighborhood processing, while global tokens provide direct long-range communication.

  • More recent sparse approaches make the connectivity pattern data-dependent. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention by Yuan et al. (2025) combines compressed coarse-grained context, selectively retrieved fine-grained blocks, and a local sliding window.

  • Connectivity therefore spans a spectrum:

    \[\text{dense} \rightarrow \text{structured sparse} \rightarrow \text{dynamic sparse}\]
    • with increasing selectivity over historical positions.

Axis 2: How Many KV Representations Are Stored?

Axis 3: Share the KV Cache or Compress It?

  • MQA and GQA reduce cache size by sharing representations across query heads. Multi-Head Latent Attention (MLA) introduces a different strategy: compress the KV representation itself.

  • DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model by DeepSeek-AI (2024) constructs a compressed latent:

    \[c_t^{KV} = W^{DKV}h_t\]

    • and derives content keys and values from it:

      \[k_t^C = W^{UK}c_t^{KV}\] \[v_t^C = W^{UV}c_t^{KV}\]

  • Rather than caching every expanded KV head, MLA can retain the compact latent together with the smaller positional component required by its decoupled RoPE formulation.

  • This creates another spectrum:

    \[\text{independent full KV heads}\] \[\downarrow\] \[\text{shared full KV heads}\] \[\downarrow\] \[\text{compressed latent KV state}\]

  • DeepSeek-V2 reports a \(93.3\%\) KV-cache reduction relative to DeepSeek 67B as part of its overall architecture, illustrating how aggressively this axis can affect inference memory.

Historical state and execution kernels

  • Attention may retain per-position state or aggregate history into a recurrent summary. Kernel algorithms then determine how the selected operator maps to memory and accelerator hardware without necessarily changing its mathematical output.

Axis 4: Is History Explicit or Compressed into State?

  • MHA, MQA, GQA, and MLA all preserve a representation associated with individual historical positions. Even when MLA compresses each token’s KV state, the model still retains token-indexed latent information.

  • Linear attention makes a more fundamental change by summarizing multiple historical tokens into a recurrent state:

    \[S_t = S_{t-1} + k_tv_t^\top\]

    • and querying that state with:

      \[o_t = q_t^\top S_t\]

  • Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention by Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention by Katharopoulos et al. (2020) (2020) formalizes this recurrent interpretation of causal linear attention.

  • Lightning Attention builds on this family with a blockwise, hardware-aware formulation:

    \[O_t^{\text{intra}} = [(Q_tK_t^\top)\odot M]V_t\] \[O_t^{\text{inter}} = Q_tS_{t-1}\] \[S_t = S_{t-1}+K_t^\top V_t\]

  • This changes the fundamental historical representation from \(\text{one state per token}\) to \(\text{one accumulated state}\) for the recurrent component.

Axis 5: Exact Operator vs. Efficient Execution

  • Another independent question is whether the attention operator itself changes.

  • Linformer changes the representation through low-rank projection:

    \[K,V \rightarrow EK,FV\]

  • Sparse attention changes the interaction graph:

    \[\text{all token pairs} \rightarrow \text{selected token pairs}\]

  • Linear attention changes the algebraic attention operator:

    \[\operatorname{softmax}(QK^\top)V \rightarrow \phi(Q) \left( \phi(K)^\top V \right)\]
    • with the appropriate normalization.

  • FlashAttention does none of these. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness by Dao et al. (2022) preserves exact softmax attention but changes the execution schedule:

    \[\boxed{ \text{same operator} + \text{better IO} }\]

  • FlashAttention-2, FlashAttention-3, and FlashAttention-4 continue this systems-oriented evolution through better parallelism, asynchronous execution, low precision, and architecture-specific pipelining.

Compositions and workload regimes

  • Architecture-level savings and kernel-level savings compound only when the workload exercises their respective bottlenecks. Separating training, prefill, and decode clarifies why a good combination depends on how and where attention executes.

Orthogonal Optimizations Can Be Composed

  • Because these mechanisms target different bottlenecks, a model need not choose exactly one.

  • Consider a model that uses GQA with sliding-window attention and a FlashAttention kernel. Its design can be decomposed as:

    \[\text{Connectivity} = \text{sliding window}\] \[\text{KV representation} = \text{GQA}\] \[\text{operator} = \text{softmax}\] \[\text{kernel} = \text{FlashAttention}\]

  • Another model could instead use:

    \[\text{Connectivity} = \text{dense}\] \[\text{KV representation} = \text{MLA}\] \[\text{operator} = \text{softmax}\] \[\text{kernel} = \text{optimized tiled attention}\]

  • A Lightning-based hybrid architecture occupies yet another point:

    \[\text{operator} = \text{mostly recurrent linear attention} + \text{periodic softmax attention}\]

  • Thinking in these dimensions is more accurate than describing each named technique as a mutually exclusive replacement for “attention.”

Training, Prefill, and Decode Are Different Regimes

  • Modern attention design is also shaped by that an LLM inference request contains computationally different phases.

  • During training, all sequence positions are usually processed together. Dense attention exposes a large amount of parallel matrix multiplication, but its quadratic interaction count and activation memory become expensive for long sequences.

  • During prefill, an entire prompt is processed in parallel:

    \[x_1,\ldots,x_n \rightarrow KV_{1:n}\]

  • This resembles training-time attention and benefits strongly from efficient tiled kernels such as FlashAttention.

  • During decode, only a small number of new tokens are typically processed per sequence:

    \[x_{n+1} \rightarrow x_{n+2} \rightarrow \cdots\]
    • while historical state must repeatedly be accessed.

  • The bottleneck can therefore shift:

    \[\text{training/prefill} \rightarrow \text{attention compute + memory traffic}\] \[\text{decode} \rightarrow \text{KV-cache capacity + memory bandwidth}\]

  • This explains why techniques such as FlashAttention and GQA solve complementary problems rather than competing ones.

Context Length Changes the Optimization Target

  • Let the context length be \(n\). For ordinary dense attention, the interaction matrix contains approximately \(n^2\) entries.

  • Doubling the context length therefore changes the interaction count by

    \[n^2 \rightarrow (2n)^2 = 4n^2\]

    • while a conventional KV cache grows only linearly:

      \[n \rightarrow 2n\]

  • Training and prefill thus face a quadratic-compute problem, whereas decode faces a linearly growing state-access problem.

  • Different attention innovations target different scaling laws:

Technique Primary quantity reduced
FlashAttention HBM IO
Sliding window Token interactions
Sparse attention Selected token interactions
MQA/GQA KV heads
MLA KV representation size
Linear/Lightning Sequence-dependent historical state

Scaling and architectural comparisons

  • Dense, local, sparse, and recurrent mechanisms make different choices about access to historical information. Comparing their asymptotic behavior alongside cache representation gives a more informative picture than a single efficiency ranking.

Dense vs. Local vs. Sparse vs. Recurrent History

  • Another useful taxonomy asks what information from the past remains directly addressable.

  • Dense attention preserves every historical position as an individually retrievable memory:

    \[\{(k_1,v_1),\ldots,(k_t,v_t)\}\]

  • Sliding-window attention preserves explicit retrieval only over recent history:

    \[\{(k_{t-W+1},v_{t-W+1}),\ldots,(k_t,v_t)\}\]

  • Dynamic sparse attention retains a potentially much larger context but selects only a subset for each query:

    \[\mathcal{S}_t \subset \{1,\ldots,t\}\]

  • Linear recurrent attention compresses earlier history into a state:

    \[S_t = \sum_{j\leq t}k_jv_j^\top\]

  • These mechanisms therefore make fundamentally different choices about memory:

    • Dense attention retains all historical positions.
    • Sliding-window attention retains only recent history.
    • Dynamic sparse attention retrieves selected history.
    • Linear recurrent attention compresses history into a fixed-size state.

A Unified Complexity View

  • The major mechanisms can be compared at a high level as follows. Exact constants and hardware behavior depend on implementation, so the table emphasizes the architectural scaling dimension each mechanism changes.
Mechanism Token Interaction Pattern Historical Representation Main Scaling Benefit
MHA Dense Per-token, per-head KV Baseline
MQA Dense Per-token shared KV Smaller KV cache
GQA Dense Per-token grouped KV Quality/cache tradeoff
Sliding Window Local Recent per-token KV \(O(nW)\) interactions
Longformer Local + global Structured sparse KV Linear interaction growth for fixed pattern
Linformer Low-rank projected Projected sequence representation Linear sequence scaling under fixed projection rank
NSA Dynamic hierarchical sparse Compressed + selected + local context Fewer attended positions
MLA Typically dense Compressed latent per token Much smaller KV cache
FlashAttention Dense Unchanged architecturally Lower IO; exact attention
Linear Attention Recurrent Fixed-dimensional state Linear sequence scaling
Lightning Attention Blockwise/recurrent Accumulated state + current block Hardware-efficient linear scaling
Hybrid Lightning + Softmax Mixed Recurrent + periodic explicit attention Scalability plus periodic full retrieval

Design history and selection

  • The historical overview illustrates why no single modern mechanism replaces every predecessor. Choosing an attention stack requires balancing retrieval fidelity, long-context scaling, memory footprint, and kernel practicality.

Attention Has Evolved Along Multiple Independent Directions

  • The historical progression is therefore not a simple sequence in which one mechanism replaces another.

  • One branch improves expressivity and alignment:

    \[\text{fixed context} \rightarrow \text{encoder-decoder attention} \rightarrow \text{self-attention} \rightarrow \text{multi-head attention}\]

  • Another reduces sequence complexity:

    \[\text{dense attention} \rightarrow \text{local/sparse attention} \rightarrow \text{linear/recurrent attention}\]

  • Another reduces inference state:

    \[\text{MHA} \rightarrow \text{MQA} \rightarrow \text{GQA}\]
    • with MLA introducing a separate low-rank compression direction.

  • Another improves hardware execution:

    \[\text{standard attention kernels} \rightarrow \text{FlashAttention} \rightarrow \text{FlashAttention-2} \rightarrow \text{FlashAttention-3} \rightarrow \\ \text{FlashAttention-4}\]

  • And another increasingly combines mechanisms:

    \[\text{single attention design} \rightarrow \text{hybrid attention stacks}\]

A Compact Evolutionary Timeline

  • The major milestones covered in this primer can be summarized chronologically:
Year Attention development
2015 Bahdanau additive attention and Luong global/local attention
2017 Transformer scaled dot-product MHA
2019 Multi-Query Attention
2020 Longformer, Linformer, linear attention
2022 FlashAttention
2023 GQA and FlashAttention-2
2024 MLA, Lightning Attention-2, FlashAttention-3
2025 Native Sparse Attention and large-scale Lightning hybrids
2026 FlashAttention-4
  • The trend is increasingly toward co-design rather than isolated algorithmic changes. Attention mechanisms are designed together with inference memory, accelerator characteristics, context length, sparsity patterns, and model architecture.

Choosing an Attention Design

  • There is no universally optimal attention mechanism because the appropriate design depends on the workload.

  • If preserving exact dense softmax attention is essential but memory traffic is the bottleneck, an IO-aware implementation such as FlashAttention addresses the systems problem without changing the operator.

  • If autoregressive KV-cache bandwidth dominates, MQA or GQA reduces duplication across heads, while MLA attacks the cache through latent compression.

  • If the context is extremely long and direct interaction with every historical token is unnecessary, local or sparse attention reduces the number of token pairs evaluated.

  • If maintaining a KV representation for every historical token is itself undesirable, recurrent linear mechanisms replace explicit history with fixed-dimensional state.

  • If neither full dense attention nor fully compressed recurrence provides the desired tradeoff, hybrid architectures combine mechanisms across layers or within an attention computation.

The Core Tradeoff

  • Nearly every efficient-attention method ultimately chooses where to compress information or computation.

    \[\boxed{ \begin{array}{ll} \text{Sparse attention} & \text{compresses connectivity} \\ \text{MQA/GQA} & \text{compresses KV-head diversity} \\ \text{MLA} & \text{compresses KV representation} \\ \text{Linear attention} & \text{compresses historical state} \\ \text{FlashAttention} & \text{compresses neither; optimizes execution} \end{array} }\]

  • This distinction is useful because reductions in computational cost are not free: they either exploit redundancy in the implementation or introduce an architectural assumption about what information does not need to remain explicitly represented.

From Attention Mechanisms to Attention Systems

  • The original attention question was primarily:

Which source positions should the model focus on?}

  • In modern large-scale models, the question has expanded:

Which positions should interact?}

What historical information should be stored?}

At what precision and representation?}

Which computations should be dense, sparse, or recurrent?}

How should those computations map onto accelerator hardware?}

  • Attention now spans representation, connectivity, memory, arithmetic complexity, and hardware execution. Workload-specific design must consider these dimensions together.

  • The enduring principle remains the same: give the model access to the information it needs while avoiding unnecessary representation and computation. What has changed is the scale at which that principle must now be engineered.

Ghost Attention and Instruction Persistence

  • Long conversational context introduces a behavioral challenge: a model may fail to maintain earlier instructions even when the tokens remain available. This chapter explains Ghost Attention as a training-data intervention, not a replacement attention kernel, and distinguishes instruction persistence from cache size and semantic long-term memory. Its limits become clearer when compared with structural attention methods.

Motivation: instruction persistence

  • A model can retain a long prompt yet fail to act on an instruction that appeared much earlier. Ghost Attention targets this behavioral failure through training examples, not through a replacement for softmax attention or the KV cache.

Why Long Conversations Create a Different Attention Problem

  • Efficient-attention mechanisms address computation, cache capacity, bandwidth, or utilization. Conversational LLMs introduce a different challenge: preserving earlier instructions despite intervening dialogue.

  • Consider a dialogue beginning with a system or user instruction \(I\) followed by multiple conversational turns:

    \[I,\ U_1,\ A_1,\ U_2,\ A_2,\ldots,U_T\]

  • Ideally, a persistent instruction should continue influencing every later response \(P(A_t\mid I,U_1,A_1,\ldots,U_t)\), even when \(t\gg1\).

Ghost Attention training procedure

  • The method reinforces persistent instruction following across dialogue turns by altering training conversations and supervision. Understanding the synthetic examples and loss masking makes its distinction from architectural attention precise.

Ghost Attention

  • Llama 2: Open Foundation and Fine-Tuned Chat Models by Touvron et al. (2023) introduces Ghost Attention (GAtt) as a data-generation and fine-tuning technique designed to improve instruction consistency across multiple dialogue turns.

  • Ghost Attention is a fine-tuning procedure, not a new attention operator: it does not change connectivity, KV storage, or asymptotic complexity. It changes the training examples and supervision for multi-turn instruction persistence.

The Core Training Idea

  • Suppose a dialogue begins with an instruction \(I\) and subsequently contains user messages \(U_1,U_2,\ldots,U_T\). During Ghost Attention data construction, the instruction is synthetically associated with later user turns. Conceptually:

    \[I,\ U_1,\ A_1,\ U_2,\ A_2,\ldots,U_T,\ A_T\]

    • becomes a structure resembling:

      \[(I,U_1),A_1,\ (I,U_2),A_2,\ldots,(I,U_T),A_T\]

  • The Llama 2 paper describes concatenating the instruction to all user messages in the synthetic training data, creating strong supervision for maintaining the instruction over multiple rounds.

Why It Is Called “Ghost” Attention

  • Ghost Attention uses repetition as a training scaffold. The model learns that later answers should continue behaving as though the original instruction remains salient even when the deployed conversation does not explicitly repeat it.

    \[\text{training: repeated instruction} \rightarrow \text{inference: normal instruction placement} \rightarrow \\ \text{desired persistent behavior}\]

Synthetic Instruction Generation

  • Llama 2 uses synthetic data to construct these multi-turn examples. The paper describes several classes of constraints, including instructions involving a hobby, language, or public figure.

  • The dialogue can move through unrelated topics while the expected responses continue satisfying the original persistent instruction, explicitly training the model to preserve constraints despite intervening conversational content.

Loss Masking

  • The Llama 2 GAtt procedure modifies the loss so the optimization does not become dominated by artificially duplicated instruction text. For multi-turn examples, the paper reports setting the loss to zero for tokens from earlier turns, including the repeated instruction, so training emphasizes the desired response.

  • Schematically:

    \[\mathcal{L}_{\text{GAtt}} = \sum_{j\in\mathcal{T}_{A_t}} -\log P(x_j\mid x_{<j})\]
    • where \(\mathcal{T}_{A_t}\) contains the target assistant tokens rather than every token in the synthetically expanded dialogue.

What Ghost Attention Teaches

  • Ordinary supervised fine-tuning often provides adjacent instruction-response examples:

    \[I\rightarrow A\]

  • Long conversations require the harder relationship:

    \[I \rightarrow \text{many intervening tokens} \rightarrow A_T\]

  • Ghost Attention strengthens this long-range dependency during training. The objective is not to increase the mathematical receptive field, but to improve whether the model uses that available context consistently.

    \[\boxed{ \text{context availability} \neq \text{context utilization} }\]

Effective context and competing instructions

  • Available context length describes the tokens that the architecture can process, not how faithfully the model prioritizes their instructions. Recency, competing turns, and prompt organization are relevant to that difference.

Attention Capacity vs. Instruction Persistence

  • Ghost Attention addresses a different axis from efficient attention. FlashAttention asks how dense attention can be computed with less IO; GQA reduces KV duplication across heads; MLA compresses per-token KV state; sparse attention reduces historical positions that require computation; Ghost Attention trains important instructions to remain behaviorally influential.

No Change to Inference-Time Attention Complexity

  • Because Ghost Attention is principally a fine-tuning technique, it does not by itself change the asymptotic cost of the underlying Transformer attention. If the model uses dense MHA, its attention computation remains approximately

    \[O(n^2d)\]

  • GAtt therefore belongs to a different layer of the stack:

    \[\text{attention architecture}\rightarrow\text{MHA/GQA/etc.}\] \[\text{attention implementation}\rightarrow\text{FlashAttention/etc.}\] \[\text{instruction-persistence training}\rightarrow\text{Ghost Attention}\]

Context Length Is Not the Same as Effective Context

  • Ghost Attention highlights a broader principle in long-context language models: a model’s nominal context window does not guarantee that every token inside that window has equal practical influence.

    \[\text{nominal context length} \neq \text{effective utilization of context}\]

  • Architectural improvements can increase the former, while training objectives, data construction, positional methods, retrieval behavior, and evaluation determine how effectively the model uses the latter.

Recency and Competing Context

  • In multi-turn dialogue, recent tokens are often highly predictive of the next response, creating competition between persistent instructions and immediately preceding conversational content.

    \[P(A_t) = P(A_t\mid I,H_{<t},U_t)\]
    • where \(H_{<t}\) denotes the intervening dialogue history.

  • Ghost Attention makes training examples explicitly exercise persistent instruction use rather than assuming standard next-token prediction will automatically produce robust multi-turn retention.

Relation to System Prompts and Persistent Instructions

  • The principle generalizes to conversational systems containing persistent instructions governing style, persona, formatting, allowed behavior, or task rules.

    \[\text{ordinary turn} \rightarrow \text{locally relevant information}\] \[\text{persistent instruction} \rightarrow \text{constraint across future turns}\]

  • Ghost Attention can therefore be viewed as an early explicit attempt to train a chat model to distinguish between information that merely appears earlier in the sequence and information whose scope extends across the dialogue.

Limitations and relation to other mechanisms

  • Ghost Attention changes the learned behavior of a model but does not create external memory or shrink attention computation. Its limitations and relationship to cache, sparse, and streaming techniques make its place in the taxonomy clearer.

Ghost Attention Is Not Explicit Memory

  • GAtt should be distinguished from persistent memory systems. Ghost Attention operates within the model’s current context, whereas a long-term memory system retrieves information that may no longer exist in the active context.

    \[\boxed{ \text{Memory} = \text{make information available again} }\] \[\boxed{ \text{Ghost Attention} = \text{train the model to keep using an available instruction} }\]

  • Once an instruction has fallen outside the model’s available context, Ghost Attention alone does not recover it.

Limitations

  • The Llama 2 paper reports that Ghost Attention improves consistency over multiple dialogue turns, but the effect does not persist indefinitely. The authors observe degradation as conversations become longer.

  • GAtt does not create a symbolic rule store or guaranteed persistent control channel; instruction persistence remains probabilistic. Its benefit also depends on the coverage of the synthetic training distribution.

Where Ghost Attention Fits in the Evolution of Attention

  • Ghost Attention is unusual because its name suggests a new neural attention operation, while its contribution actually belongs to post-training and conversational behavior.

  • It sits outside the main architectural progression:

    \[\text{Bahdanau} \rightarrow \text{MHA} \rightarrow \text{MQA/GQA} \rightarrow \text{MLA}\]

    • and outside the systems progression:

      \[\text{standard kernels} \rightarrow \text{FlashAttention} \rightarrow \text{later FlashAttention variants}\]

  • Instead, it represents a behavioral branch:

    \[\text{single-turn instruction following} \rightarrow \text{multi-turn instruction persistence}\]

A Complete Taxonomy Including Ghost Attention

  • Ghost Attention adds behavioral persistence as an independent training dimension. It does not change the structural costs or token connectivity of the underlying attention mechanism.
Technique What Changes? Primary Objective
Bahdanau/Luong Attention Encoder-decoder alignment Dynamic source selection
Self-Attention Interaction topology Token-to-token contextualization
MHA Representation views Multiple attention subspaces
Sliding-Window Attention Connectivity Reduce sequence interactions
Longformer Connectivity Local + selected global context
Linformer Sequence representation Low-rank sequence compression
Linear Attention Operator/state Avoid quadratic sequence interactions
MQA KV heads Reduce decode bandwidth
GQA KV groups Balance KV efficiency and capacity
MLA KV representation Compress per-token cache
Native Sparse Attention Dynamic connectivity Select useful long-context regions
FlashAttention Kernel execution Reduce memory I/O
Lightning Attention Operator + execution Hardware-efficient recurrent attention
Ghost Attention Fine-tuning data/objective Multi-turn instruction persistence

The Broader Lesson

  • Ghost Attention illustrates an important distinction:

    \[\boxed{ \text{being able to attend} \neq \text{learning what deserves sustained attention} }\]

  • Architecture determines which content is reachable, kernels determine execution cost, cache policies determine retained history, and instruction-oriented training affects how models use that context.

  • In this sense, the evolution of attention is not only about making longer contexts computationally feasible. It is also about making models use those contexts reliably, especially when some pieces of information should remain influential far longer than their position in the sequence would otherwise suggest.

Implementation Patterns and Tradeoffs

  • Implementing attention efficiently requires matching the numerical formulation to data layout, cache organization, and the execution regime. This chapter moves from stable softmax and mask fusion to MQA/GQA cache layouts, sparse and recurrent updates, precision, and validation. Its implementation guidance separates correctness checks from performance measurements.

Numerical foundations and reference attention

  • Before optimizing an attention kernel, it is useful to establish a correct reference implementation and stable normalization. These foundations identify which intermediate computations can be tiled or fused without changing the desired result.

Attention Design Is Now a Systems Choice

  • An attention formula does not determine its runtime characteristics. Tensor layout, masking, cache access, tiling, precision, and kernel scheduling shape actual latency, throughput, memory use, and numerical behavior.

  • A useful implementation view separates attention into five stages:

    \[X \rightarrow (Q,K,V) \rightarrow \text{score/state computation} \rightarrow \text{normalization or aggregation} \rightarrow O\]

  • Optimization can occur at every stage. MQA, GQA, and MLA alter the representation of \(K\) and \(V\); sparse attention changes which scores are computed; linear attention changes how history is aggregated; FlashAttention changes how the same dense softmax computation moves through the memory hierarchy.

Reference Scaled Dot-Product Attention

  • The simplest implementation of dense scaled dot-product attention closely follows the Transformer equations:
def attention(q, k, v, mask=None):
    scores = q @ k.transpose(-2, -1)
    scores = scores / math.sqrt(q.shape[-1])

    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))

    weights = torch.softmax(scores, dim=-1)
    return weights @ v

  • For tensors with shape \(Q,K,V \in \mathbb{R}^{B\times H\times N\times d_h}\) the score tensor has shape \(B\times H\times N\times N\).

  • That intermediate is the central systems problem with a naive implementation. Its size grows quadratically with sequence length:

    \[M_{\text{scores}} \propto BHN^2\]

  • At long contexts, explicitly writing the score and probability matrices to high-bandwidth memory can become more expensive than the arithmetic itself.

Numerically Stable Softmax

  • Attention implementations should not compute the exponential of raw scores directly. Given scores \(s_1,\ldots,s_n\), stable softmax subtracts the row maximum:

    \[m = \max_j s_j\] \[p_i = \frac{\exp(s_i-m)} {\sum_j\exp(s_j-m)}\]

  • Since softmax is invariant to adding or subtracting the same scalar from every logit,

    \[\operatorname{softmax}(s) = \operatorname{softmax}(s-m)\]
    • but the shifted representation avoids unnecessarily large exponentials.

  • This becomes especially important in low-precision training and in tiled algorithms such as FlashAttention, where a row’s maximum and normalization denominator must be maintained incrementally across blocks.

Online Softmax

  • Suppose an attention row is processed block by block. After processing some blocks, maintain a running maximum \(m\) and normalization statistic \(\ell\) with:

    \[\ell = \sum_j \exp(s_j-m)\]

  • When a new score block arrives with maximum \(m_{\text{new}}^{\text{block}}\), the updated maximum is:

    \[m' = \max \left( m, m_{\text{new}}^{\text{block}} \right)\]

  • The previous normalization statistic must then be rescaled:

    \[\ell' = e^{m-m'}\ell + \sum_{j\in\text{new block}} e^{s_j-m'}\]

  • The accumulated output is rescaled in the same way. This identity allows an exact softmax result to be produced without retaining the complete score matrix.

  • FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness by Dao et al. (2022) combines this online normalization with tiled loading of \(Q\), \(K\), and \(V\) so that intermediate attention matrices need not be materialized in high-bandwidth memory.

Masking and tensor organization

  • Mask semantics and tensor layout shape both correctness and memory efficiency. Organizing queries, heads, and sequence dimensions properly is also essential for efficient head sharing.

Masking Should Be Part of the Kernel

  • Causal masking is conceptually:

    \[M_{ij} = \begin{cases} 0,&j\leq i\\ -\infty,&j>i \end{cases}\]

    • followed by:

      \[A = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_h}}+M \right)\]

  • A naive implementation may explicitly allocate an \(N\times N\) mask. Optimized kernels instead infer whether a tile lies below, above, or across the causal diagonal and avoid materializing a full mask.

  • The same principle applies to sliding-window attention. Rather than first computing dense attention and masking almost all of it afterward, an efficient implementation should avoid computing blocks known to fall outside the window.

Multi-Head Attention Tensor Layout

  • A typical MHA projection begins with:

    \[X \in \mathbb{R}^{B\times N\times d_{\text{model}}}\]

    • and computes:

      \[Q=XW^Q\] \[K=XW^K\] \[V=XW^V\]

  • The projections are then reshaped into heads:

    \[Q,K,V \in \mathbb{R}^{B\times H\times N\times d_h}\]

    • with:

      \[d_{\text{model}} = Hd_h\]

  • Implementations often fuse the three projections into one matrix multiplication:

    \[[Q;K;V] = XW^{QKV}\]
    • rather than launching three independent operations.

  • The mathematical result is equivalent, but fusion reduces kernel-launch and memory-access overhead.

Head sharing and KV-cache implementation

  • MQA and GQA reduce KV duplication, but implementations can accidentally recover that overhead through unnecessary expansion. Cache layout, decode access patterns, and prefill parallelism determine whether architectural savings translate into throughput.

Implementing MQA

  • MQA keeps \(H_Q=H\) query heads but uses \(H_{KV}=1\) key-value head.

  • Conceptually:

      q = q_proj(x).view(batch, seq, num_q_heads, head_dim)
  k = k_proj(x).view(batch, seq, 1, head_dim)
  v = v_proj(x).view(batch, seq, 1, head_dim)

  • A naive implementation might physically repeat the shared keys and values:

      k = k.repeat_interleave(num_q_heads, dim=2)
  v = v.repeat_interleave(num_q_heads, dim=2)
    • but doing so destroys much of MQA’s memory advantage.

  • Efficient kernels instead interpret the single KV head as shared by multiple query heads without physically duplicating its stored representation.

Implementing GQA

  • For GQA, \(H_Q=H\) and \(H_{KV}=G\), where \(G<H\).

  • Each KV head serves \(R = \frac{H}{G}\) query heads.

  • The logical mapping is \(g(i)=\left\lfloor\frac{i}{R}\right\rfloor\), so query head \(i\) attends using \(K_{g(i)},V_{g(i)}\).

  • A simple conceptual implementation is:

q = q_proj(x).view(batch, seq, num_q_heads, head_dim)
k = k_proj(x).view(batch, seq, num_kv_heads, head_dim)
v = v_proj(x).view(batch, seq, num_kv_heads, head_dim)

group_size = num_q_heads // num_kv_heads
  • Production kernels should understand this grouped structure directly rather than expanding \(K\) and \(V\) to \(H_Q\) heads in memory.

KV-Cache Layout

  • During autoregressive decoding, each layer retains historical keys and values. A conceptual MHA cache has shape

    \[K_{\text{cache}}, V_{\text{cache}} \in \mathbb{R}^{B\times H_{KV}\times L\times d_h}\]

  • Its approximate storage in bytes is

    \[M_{\text{KV}} = 2 B N_L L H_{KV} d_h b\]
    • where \(b\) is bytes per stored element.

  • For BF16 or FP16,

    \[b=2\]

  • This equation makes several optimization levers explicit:

    \[B ,\quad L ,\quad H_{KV} ,\quad d_h ,\quad b\]

  • GQA reduces \(H_{KV}\). Sliding-window attention can reduce the retained effective \(L\). Quantized caches reduce \(b\). MLA changes the representation enough that the conventional \(H_{KV}d_h\) description is replaced by its compressed latent dimensions.

Decode Is Primarily a Data-Movement Problem

  • At decode step \(t\), the model usually generates only one or a small number of new query positions but must access historical state for potentially thousands of previous tokens.

  • The arithmetic intensity is therefore much lower than during prefill. Roughly speaking, each historical key and value is loaded to contribute to only a small amount of work for the current query.

  • This is why reducing

    \[H_{KV}\]

  • Through MQA or GQA can substantially improve generation throughput even though the query-head count remains unchanged.

  • The key engineering principle is:

    \[\boxed{ \text{Do not optimize decode as though it were prefill} }\]

  • The two phases expose different shapes, parallelism, and bottlenecks.

Prefill Is Highly Parallel

  • During prefill, the model processes \(N\) queries simultaneously against \(N\) keys.

  • This produces large matrix multiplications \(QK^\top\) and \(AV\), which can efficiently use accelerator matrix units.

  • FlashAttention is especially effective in this regime because it reorganizes these large dense operations to minimize HBM traffic while preserving exact softmax attention.

  • Decode kernels, in contrast, are commonly specialized around small query lengths and large cached-key lengths.

Avoid Materializing Repeated KV Heads

  • Consider GQA with \(H_Q=32\) and \(H_{KV}=8\).

  • Logically, each KV head serves four query heads. Expanding the cache to 32 KV heads would multiply its apparent storage and bandwidth by \(\frac{32}{8} = 4\) and largely eliminate the benefit of GQA.

  • Broadcasting should therefore occur logically inside the computation:

    \[Q_i \leftrightarrow K_{g(i)},V_{g(i)}\]
    • rather than physically through duplicated tensors.

Sparse, local, and recurrent implementations

  • Locality and sparsity help only when the kernel avoids work on excluded positions. Recurrent attention requires a different set of state updates and blockwise computations that should be evaluated on their own terms.

Sliding-Window Cache Eviction

  • For a strictly local attention layer with window size \(W\), token \(t\) can only access positions \([t-W+1,t]\).

  • Once a cached position falls permanently outside this range, it no longer needs to remain in that layer’s active KV cache.

  • A ring buffer can therefore store at most \(W\) positions:

slot = position % window_size
k_cache[:, :, slot] = k_new
v_cache[:, :, slot] = v_new

  • This changes cache growth from \(O(L)\) to \(O(W)\) for a strictly windowed layer.

  • Care is required when positional encodings, global tokens, hybrid layers, or retrieval mechanisms make older positions accessible through another route.

Sparse Attention Should Avoid Dense Work

  • Sparse attention only provides computational savings if the implementation avoids evaluating discarded interactions.

  • An inefficient pattern is

    \[\text{dense scores} \rightarrow \text{sparse mask} \rightarrow \text{discard most scores}\]
    • because the expensive dense computation has already occurred.

  • The desired pattern is

    \[\text{sparse selection} \rightarrow \text{compute selected blocks only}\]

  • Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention by Yuan et al. (2025) explicitly emphasizes hardware-aligned blockwise sparsity for this reason. Sparse structure should map onto efficient accelerator operations rather than merely reducing the number of mathematically nonzero entries.

Linear Attention State Updates

  • Causal linear attention can be implemented through a running state:

    \[S_t = S_{t-1} + \phi(k_t)v_t^\top\]

    • and, depending on the formulation, a normalization state:

      \[z_t = z_{t-1} + \phi(k_t)\]

  • The output becomes:

    \[o_t = \frac{ \phi(q_t)^\top S_t }{ \phi(q_t)^\top z_t }\]

  • During decoding, neither state grows with sequence length. This gives linear attention its characteristic constant-state recurrent interpretation.

  • During training, however, computing all prefix-dependent states efficiently requires specialized scan or blockwise algorithms.

Lightning Attention’s Blockwise Implementation

  • Lightning Attention handles causal linear computation by dividing the sequence into blocks. For block \(t\):

    \[O_t^{\text{intra}} = [(Q_tK_t^\top)\odot M]V_t\] \[O_t^{\text{inter}} = Q_tS_{t-1}\]

    • and:

      \[S_t = S_{t-1} + K_t^\top V_t\]

  • The implementation therefore converts token-wise recurrent work into accelerator-friendly matrix multiplications.

  • Block size \(B\) becomes an engineering hyperparameter. Very small blocks increase overhead and recurrent synchronization; very large blocks increase the cost of \(B\times B\) intra-block attention.

Memory, precision, and benchmarking

  • GPU performance depends on the tradeoff between recomputation and memory traffic, as well as the precision and capabilities of the target hardware. Correct benchmarking must separate prefill from decode and include numerical validation.

Recompute vs. Store

  • Backpropagation introduces another important systems tradeoff. Saving every intermediate from the forward pass reduces recomputation but consumes memory.

  • Attention kernels often instead save compact statistics and reconstruct intermediate quantities during the backward pass.

  • FlashAttention is a prominent example: it avoids storing the full attention matrix and recomputes blocks of scores as needed during backward propagation.

  • This exchanges additional arithmetic for substantially lower memory traffic:

    \[\text{more FLOPs} \quad\leftrightarrow\quad \text{less memory IO}\]

  • On modern accelerators, this can be favorable because matrix arithmetic has become much faster relative to off-chip memory movement.

Precision Is Part of Attention Design

  • Attention combines dot products, exponentiation, normalization, and weighted accumulation, which have different numerical sensitivities.

  • Low-precision implementations may use reduced precision for matrix multiplications while maintaining selected statistics or accumulations in higher precision.

  • FlashAttention-3 explores FP8 attention while using techniques including block quantization and incoherent processing to control numerical error.

  • The general engineering pattern is mixed precision:

    \[\text{high-throughput low-precision matmul} + \text{higher-precision critical reductions}\]
    • rather than requiring every attention operation to use the same datatype.

Benchmark the Correct Regime

  • Reporting a single “attention speed” can hide important differences. At minimum, benchmarks should separately evaluate training, prefill, and decode.

  • Useful independent variables include batch size, sequence length, query length, head dimension, number of query heads, number of KV heads, datatype, causal versus bidirectional attention, and sparsity/window configuration.

  • Prefill and single-token decoding have different bottlenecks, so a kernel or cache optimization must be benchmarked in its intended regime.

Measure Memory and Bandwidth, Not Just FLOPs

  • The theoretical operation count is only one component of performance.

  • For an attention implementation, useful metrics include latency, tokens per second, HBM bytes transferred, peak memory usage, achieved FLOPs per second, and hardware utilization.

  • FlashAttention’s success is an important demonstration that two algorithms computing mathematically identical attention can have substantially different runtime because one performs far less off-chip memory traffic.

Validate Numerical Equivalence

  • When replacing a reference implementation with an optimized exact kernel, correctness should be tested independently from speed.

  • For exact-attention implementations, compare outputs and gradients \(\Delta_O = \|O_{\text{optimized}}-O_{\text{reference}}\|\) and \(\Delta_{\nabla} = \|\nabla_{\text{optimized}}-\nabla_{\text{reference}}\|\) under appropriate precision-dependent tolerances.

  • Tests should include causal boundaries, padding, variable sequence lengths, extreme logits, different head dimensions, and grouped-query configurations.

  • Approximate mechanisms such as Linformer, sparse attention, or linear attention require a different evaluation because mathematical equality with dense softmax attention is not expected. Their efficiency gains must instead be considered jointly with downstream quality.

Putting the implementation choices together

  • The decision tree consolidates the architectural and systems considerations into workload-specific choices. The checklist that follows turns those choices into validation steps rather than assuming theoretical savings imply production gains.

A Practical Decision Tree

  • When attention is too expensive, first identify which resource is actually limiting the workload.

  • If the bottleneck is the quadratic score computation during long training or prefill:

    \[\text{consider} \rightarrow \text{sparse/local/linear attention}\]
    • or retain dense attention and optimize its execution with FlashAttention when exactness is required.

  • If the bottleneck is decode-time KV bandwidth:

    \[\text{consider} \rightarrow \text{MQA/GQA/MLA}\]

  • If the bottleneck is KV-cache capacity at very long context:

    \[\text{consider} \rightarrow \text{GQA/MLA/windowing/cache quantization}\]

  • If the goal is sequence-length-independent recurrent state:

    \[\text{consider} \rightarrow \text{linear or Lightning-style attention}\]

  • If the issue is failure to preserve conversational instructions rather than computational cost:

    \[\text{consider training/post-training behavior}\]
    • rather than changing the attention kernel itself.

Implementation Checklist

  • A robust attention implementation should answer all of the following questions before optimization:

    • What is the attention operator: softmax, sparse softmax, linear, or hybrid?

    • What is the connectivity pattern: dense, causal, local, global-local, or dynamically sparse?

    • How many query and KV heads are present?

    • Is historical state stored per token or compressed recurrently?

    • What is the prefill workload?

    • What is the decode workload?

    • Is the implementation compute-bound, bandwidth-bound, or memory-capacity-bound?

    • Which intermediates are stored versus recomputed?

    • Which tensors use reduced precision?

    • Does the kernel exploit sparsity and KV sharing without physically expanding tensors?

    • Are numerical outputs and gradients validated against a reference implementation?

The Engineering Principle

  • Attention optimization is ultimately a problem of matching mathematical structure to hardware structure.

  • Dense softmax attention exposes highly parallel matrix multiplication but expensive all-pairs interactions. FlashAttention preserves those interactions while improving their data movement.

  • MQA and GQA recognize that storing independent KV heads can be unnecessarily expensive during decoding. MLA goes further by changing what representation is cached.

  • Sparse attention recognizes that not every token pair necessarily deserves computation. Linear and Lightning Attention recognize that explicit per-token historical state itself can sometimes be replaced by a recurrent summary.

  • The practical objective is therefore not simply:

    \[\min \text{FLOPs}\]

    • but something closer to:

      \[\boxed{ \text{maximize model quality and useful throughput} \quad \text{subject to compute, memory, bandwidth, and latency constraints} }\]

  • This systems perspective explains much of the evolution of attention: successful mechanisms increasingly co-design the mathematical operator, model architecture, memory representation, and accelerator execution strategy.

Attention Sinks and StreamingLLM

  • Bounded-memory streaming inference must decide which historical KV entries remain resident after the context grows beyond the cache budget. This chapter explains attention sinks, StreamingLLM’s sink-plus-recent-token policy, and the resulting limits on accessible old information. It compares eviction-based streaming with trained local attention, KV compression, and recurrent history.

Streaming and the sink phenomenon

  • A bounded rolling cache appears sufficient for streaming inference, but ordinary pretrained attention can assign disproportionate mass to initial tokens. The sink phenomenon explains why naive eviction can cause output degradation.

The Streaming-Inference Problem

  • Autoregressive Transformers normally retain the keys and values of every previous token so that a new query can attend to the complete history:

    \[K_{1:t},V_{1:t}\]

  • Consequently, the KV cache grows linearly with generated sequence length:

    \[M_{\text{KV}} \propto t\]

  • For an indefinitely running dialogue, agent, or streaming application, retaining every historical KV vector eventually becomes impractical. A natural solution is to retain only the most recent \(W\) tokens:

    \[K_{t-W+1:t},V_{t-W+1:t}\]

    • giving a fixed-size cache:

      \[M_{\text{KV}} = O(W)\]

  • However, Efficient Streaming Language Models with Attention Sinks by Xiao et al. (2023; ICLR 2024) showed that simply discarding old KV states causes the performance of pretrained LLMs to collapse once the sequence exceeds the cache window.

The Attention-Sink Phenomenon

  • Efficient Streaming Language Models with Attention Sinks by Xiao et al. (2023; ICLR 2024) observed that attention often concentrates on the first few positions even when their token content is not important. These attention sinks explain why discarding all but recent tokens can destabilize streaming inference.

  • Suppose attention for query \(q_t\) is

    \[a_{tj} = \frac{ \exp(q_t^\top k_j/\sqrt{d}) }{ \sum_{\ell\leq t} \exp(q_t^\top k_\ell/\sqrt{d}) }\]

    • with:

      \[\sum_{j\leq t}a_{tj}=1\]

  • Some early positions can repeatedly receive substantial attention mass \(a_{t1}\gg a_{tj}\) for many later positions \(t\), even when token \(1\) carries little direct semantic information for the current prediction.

  • The phenomenon is therefore different from an important early token being retrieved because of its meaning. The initial positions can serve as stable destinations for attention mass—a sink.

Why the First Tokens Become Sinks

  • Causal attention gives the earliest tokens a structural advantage. The first token is visible to essentially every later position:

    \[1 \in \mathcal{A}(t) \qquad \forall t\geq1\]
    • whereas a later token \(j\) is only visible when \(t\geq j\).

  • Early tokens therefore participate in far more attention computations during autoregressive training. Efficient Streaming Language Models with Attention Sinks by Xiao et al. (2023; ICLR 2024) argue that this makes them natural locations for the model to learn sink behavior.

  • Softmax normalization is also important because attention probabilities must sum to one. Subsequent empirical work, When Attention Sink Emerges in Language Models: An Empirical View by Gu et al. (2024), found evidence that sinks can behave like key biases that absorb otherwise uninformative attention scores and that the phenomenon is closely tied to normalized softmax attention.

Why Naive Window Attention Fails

  • Consider ordinary sliding-window inference with cache size \(W\). At time \(t\), only \(\{t-W+1,\ldots,t\}\) remain available.

  • Eventually, the initial sink tokens are evicted:

    \[\{1,2,\ldots,S\} \notin \text{KV cache}\]

  • The model was not necessarily trained under this altered attention distribution. Removing the positions that learned to absorb substantial attention mass can therefore perturb attention normalization throughout the network.

  • The StreamingLLM experiments show that retaining only recent tokens performs poorly once the sequence exceeds the cache size, whereas retaining a few initial tokens largely restores stable language modeling.

Sink-preserving streaming inference

  • StreamingLLM retains a small set of initial positions alongside the most recent tokens. Its cache structure and position-handling requirements explain how a bounded cache can process long streams without promising full historical recall.

StreamingLLM

  • StreamingLLM turns this observation into a simple inference strategy. Instead of retaining only the most recent window, preserve:

    \[\boxed{ \text{attention sinks} + \text{recent-token window} }\]

  • If \(S\) initial sink tokens and \(W\) recent tokens are retained, the active cache is conceptually

    \[\mathcal{C}_t = \{1,\ldots,S\} \cup \{t-W+1,\ldots,t\}\]

  • The total retained context is therefore approximately \(S+W\) tokens regardless of how long the stream becomes.

  • The paper reports that retaining only four initial tokens was sufficient in its experiments to stabilize several tested pretrained LLM families.

Constant-Memory Streaming

  • Conventional full-context decoding stores:

    \[\text{KV cache size} = O(t)\]

  • A pure sliding window stores \(O(W)\), but can destabilize models when sink tokens disappear.

  • StreamingLLM stores \(O(S+W)\) where \(S\) and \(W\) are fixed.

  • Thus, for total stream length, \(\boxed{M_{\text{StreamingLLM}} = O(1)}\) for a fixed sink count and local window.

  • This does not mean the model has direct access to an infinite amount of historical content. It means inference can continue indefinitely without KV memory increasing indefinitely.

What “Infinite Sequence Length” Means Here

  • This distinction is important. StreamingLLM enables an indefinitely long stream, but it does not turn a finite-window Transformer into a model that can directly retrieve every token ever observed.

  • At step \(t\), the model sees only

    \[\text{initial sinks} + \text{recent window}\]
    • rather than \(x_1,x_2,\ldots,x_t\) in full.

  • Therefore:

    \[\boxed{ \text{unbounded streaming length} \neq \text{unbounded retrievable context} }\]

  • Old semantic information that has left the recent window is not magically preserved by the sink tokens. StreamingLLM is primarily a stable streaming-inference mechanism, not a general long-term memory system.

Cache Structure

  • A conceptual cache implementation separates persistent sink slots from a rotating recent-token region:
sink_k = k_cache[:, :, :num_sink_tokens]
sink_v = v_cache[:, :, :num_sink_tokens]

recent_k = recent_window_k
recent_v = recent_window_v

active_k = concat(sink_k, recent_k, dim=sequence_dim)
active_v = concat(sink_v, recent_v, dim=sequence_dim)

  • As new tokens arrive, the sink portion remains unchanged:

    \[K_{\text{sink}}^{(t+1)} = K_{\text{sink}}^{(t)}\]

    • while the recent window behaves like a bounded queue \([K_{t-W+1},\ldots,K_t]\) becomes:

      \[[K_{t-W+2},\ldots,K_t,K_{t+1}]\]

  • Memory therefore remains bounded while the stream continues.

Position Handling Matters

  • KV eviction creates another subtle problem: positional encodings.

  • Suppose the active cache contains the initial sink tokens and recent tokens whose original absolute positions are very large:

    \[1,2,3,4,99997,99998,99999,100000\]

  • Feeding those large absolute positions directly into a model trained only over much shorter position ranges can introduce an out-of-distribution positional regime.

  • StreamingLLM therefore assigns positional information based on positions inside the current cache rather than allowing position indices to grow indefinitely. The project description specifically identifies relative cache positions as part of its streaming recipe.

  • Conceptually:

    \[\text{stream positions} : [1,2,3,4,99997,\ldots,100000]\]

    • are mapped into a bounded positional configuration such as:

      \[\text{cache positions} : [0,1,2,3,4,\ldots,S+W-1]\]
      • according to the positional scheme used by the model.

Interpreting sinks and comparing windows

  • Attention sinks stabilize normalization and behavior; they should not be mistaken for an explicit semantic summary of distant content. Comparing sink-preserving caches with ordinary sliding windows makes the distinction operational.

Attention Sinks Are Not Summary Tokens

  • It is tempting to interpret sink tokens as learned summaries of the entire preceding context, but that is not what the original StreamingLLM result establishes.

  • Attention sinks help stabilize learned attention distributions; they are not summaries of the discarded sequence.

  • This distinguishes attention sinks from recurrent linear-attention states:

    \[S_t = S_{t-1}+k_tv_t^\top\]
    • which explicitly aggregate information from historical tokens.

  • StreamingLLM’s sink states are retained \(K_{\text{sink}},V_{\text{sink}}\) but they are not continuously updated to summarize every evicted token.

Dedicated Sink Tokens

  • The StreamingLLM study also explored introducing a dedicated placeholder token during pretraining whose purpose is to act as an attention sink.

  • If the model is explicitly provided a token \(x_{\text{sink}}\) that is always visible, it can learn to allocate otherwise unnecessary attention there instead of implicitly assigning that role to ordinary initial tokens.

  • The paper reports that introducing such a dedicated sink token during pretraining further improves streaming behavior.

  • This suggests that attention sinks are not merely accidental artifacts that must always be eliminated. Once understood, the behavior can be deliberately incorporated into model design.

Experimental Scale

  • Efficient Streaming Language Models with Attention Sinks by Xiao et al. (2023; ICLR 2024) evaluated StreamingLLM across model families including Llama 2, MPT, Falcon, and Pythia. They report stable language modeling over streams of up to 4 million tokens and beyond, without fine-tuning the tested pretrained models for the streaming mechanism.

  • Again, this result refers to continuous processing using the bounded sink-plus-window cache, not direct full-attention access over a four-million-token KV cache.

  • In the streaming setting, the paper reports up to a \(22.2\times\) speedup relative to a sliding-window recomputation baseline.

Sliding Window vs. StreamingLLM

  • Both approaches bound cache residency, but StreamingLLM preserves initial sink tokens alongside recent positions. That distinction changes inference stability without restoring access to evicted content.
Property Sliding Window StreamingLLM
Initial tokens retained No Yes
Recent tokens retained Yes Yes
Cache grows with stream No No
Sink states preserved No Yes
Direct access to old evicted content No No
Requires model fine-tuning No No
Streaming memory \(O(W)\) \(O(S+W)\)
  • The computational difference is small: \(W \rightarrow W+S\) with \(S\ll W\), but preserving those few positions can substantially change model stability.

Relationship to other memory strategies

  • A streaming cache policy acts at inference time and can coexist with shared KV heads or compressed representations. Comparisons with linear attention, external memory, and training methods clarify which limitations the cache policy does and does not address.

StreamingLLM vs. Sliding-Window Attention During Training

  • StreamingLLM should also be distinguished from models architecturally trained with local sliding-window attention.

  • In architectural sliding-window attention, locality is part of the model’s learned connectivity:

    \[\mathcal{A}(t) = [t-W+1,t]\]

  • During training and inference.

  • StreamingLLM instead asks how a model originally trained with a finite context can be deployed on an indefinitely continuing stream while maintaining bounded KV storage.

  • Therefore:

    \[\boxed{ \text{Sliding-window architecture} = \text{attention connectivity choice} }\]

    • whereas:

      \[\boxed{ \text{StreamingLLM} = \text{KV-cache/inference strategy exploiting attention sinks} }\]

StreamingLLM vs. GQA and MLA

  • GQA reduces the head dimension of KV storage:

    \[H \rightarrow G\]

  • MLA reduces the representation dimension stored per token:

    \[K_t,V_t \rightarrow c_t^{KV}\]

  • StreamingLLM reduces the number of historical positions retained:

    \[t \rightarrow S+W\]

  • These approaches therefore optimize different dimensions of the cache.

  • Conceptually, KV memory has several multiplicative factors:

    \[M_{\text{KV}} \propto \underbrace{L}_{\text{positions}} \times \underbrace{H_{KV}}_{\text{KV heads}} \times \underbrace{d}_{\text{representation}} \times \underbrace{b}_{\text{precision}}\]

  • StreamingLLM primarily attacks \(L\), while GQA attacks \(H_{KV}\) and cache quantization attacks \(b\) with MLA changing the representation term more fundamentally.

StreamingLLM vs. Linear Attention

  • Linear attention takes a fundamentally different approach. Instead of retaining selected historical KV positions, it recursively aggregates history:

    \[S_t = S_{t-1} + \phi(k_t)v_t^\top\]

  • Thus:

    \[\text{StreamingLLM} \rightarrow \text{retain selected explicit tokens}\]

    • whereas:

      \[\text{linear attention} \rightarrow \text{compress history into recurrent state}\]

  • Both can have sequence-length-independent memory during streaming, but they achieve it through very different representations.

StreamingLLM vs. Long-Term Memory

  • Consider information appearing at token \(j\) where

    \[S<j<t-W\]

  • That token is neither an initial sink nor part of the current recent window, so its KV state has been discarded.

  • Consequently, StreamingLLM alone cannot directly retrieve it.

  • An external memory or retrieval system could instead store and later restore relevant information:

    \[x_j \rightarrow M \rightarrow \operatorname{retrieve}(x_j) \rightarrow C_t\]

  • The distinction is:

    \[\boxed{ \text{StreamingLLM} = \text{bounded working context} }\] \[\boxed{ \text{external memory} = \text{recoverable long-term information} }\]

Attention Sinks vs. Ghost Attention

  • Attention sinks and Ghost Attention sound related but describe almost completely different phenomena.

  • Attention sinks are learned attention patterns in which early positions attract unusually high attention mass:

    \[\text{early token} \leftarrow \text{large attention weight}\]

  • Ghost Attention is a post-training technique designed to improve persistence of conversational instructions.

  • Their purposes are therefore:

    \[\text{Attention sinks} \rightarrow \text{streaming stability and attention behavior}\] \[\text{Ghost Attention} \rightarrow \text{multi-turn instruction consistency}\]

  • Neither should be confused with a new basic attention operator such as MHA or linear attention.

Systems implications

  • The central insight is that not all cached positions play the same role in a pretrained model. Cache policy must respect those learned attention behaviors while remaining explicit about its limits for long-term information retrieval.

A Deeper Interpretation

  • Subsequent research has continued investigating why sink behavior emerges. When Attention Sink Emerges in Language Models: An Empirical View by Gu et al. (2024) report that attention sinks appear during effective language-model pretraining and are strongly connected to the normalization behavior of softmax attention.

  • More recent theoretical and empirical work has proposed additional interpretations, including that sink behavior may help Transformers control excessive information mixing across layers. These interpretations remain an active research topic rather than a single settled explanation.

  • The safest architectural conclusion is therefore empirical:

    \[\boxed{ \text{many pretrained causal LMs rely strongly on a few early positions} }\]
    • and removing those positions from the KV cache can substantially disturb model behavior.

The Systems Lesson

  • StreamingLLM illustrates a recurring theme in efficient attention: seemingly redundant states can play important roles because models adapt to the exact computational structure present during training.

  • A naive systems optimization might reason:

    \[\text{old token} \Rightarrow \text{probably irrelevant} \Rightarrow \text{evict}\]

  • Attention sinks show why this inference can fail. Semantic importance and computational importance are not identical:

    \[\boxed{ \text{low semantic content} \not\Rightarrow \text{low architectural importance} }\]

  • By preserving only a few structurally important initial KV states alongside a bounded recent window, StreamingLLM turns this observation into a practical constant-memory streaming strategy.

  • Together with GQA, MLA, sparse attention, FlashAttention, and linear attention, it demonstrates another independent axis of attention optimization:

    \[\boxed{ \text{not only how attention is computed or represented, but which KV states must remain resident} }\]

Alternative Attention Formulations: Talking-Heads and Synthesizer

  • Multi-head attention normally computes query-key scores separately for each head, but that structure is not mathematically compulsory. This chapter examines Talking-Heads as cross-head mixing and Synthesizer as an alternative to explicit pairwise query-key alignment. Their results motivate a broader view of attention as learned token routing rather than one fixed score formula.

Why modify attention internals?

  • Ordinary multi-head attention assumes independent head distributions and explicit query-key interactions. Talking-Heads and Synthesizer probe these assumptions separately while keeping token mixing as their common computational purpose.

Why Revisit Standard Multi-Head Attention?

  • The Transformer established scaled dot-product multi-head attention as the dominant attention formulation. For head \(h\):

    \[S_h = \frac{Q_hK_h^\top}{\sqrt{d_k}}\] \[A_h = \operatorname{softmax}(S_h)\] \[O_h = A_hV_h\]

    • and the outputs of the heads are concatenated:

      \[O = \operatorname{Concat}(O_1,\ldots,O_H)W^O\]

  • This design makes two strong structural choices. First, each attention head independently computes its own attention distribution. Second, the attention matrix is dynamically produced from pairwise query-key interactions.

  • Talking-Heads tests communication between heads while retaining pairwise QK scoring; Synthesizer tests whether pairwise QK scoring is needed at all.

Independent Heads in Standard MHA

  • In ordinary MHA, head \(h\) produces its logits independently:

    \[S_h = Q_hK_h^\top\]

  • Softmax is also independently applied within each head:

    \[A_h = \operatorname{softmax}(S_h)\]

  • Information from different heads is mixed only after each head has already completed its attention operation, through concatenation and the output projection:

    \[O = \operatorname{Concat}(A_1V_1,\ldots,A_HV_H)W^O\]

  • Thus, head \(h_1\) cannot directly influence the attention distribution computed by head \(h_2\).

  • Talking-Heads Attention by Shazeer et al. (2020) modifies precisely this property by adding learned linear projections across the head dimension before and after softmax.

Talking-Heads: communication across heads

  • Cross-head projections before and after softmax allow head interactions during attention-weight formation. This changes representational flexibility without automatically solving quadratic sequence scaling.

Letting Attention Heads Communicate

  • Instead of treating each head’s attention matrix as completely independent, Talking-Heads Attention introduces information exchange across heads at two points:

    \[\boxed{ \text{QK logits} \rightarrow \text{mix heads} \rightarrow \text{softmax} \rightarrow \text{mix heads} \rightarrow \text{apply values} }\]

  • This allows one set of heads to specialize in generating attention logits while another representation of heads can be used when applying those probabilities to values.

Pre-Softmax Head Mixing

  • Let the raw attention logits be represented as a tensor

    \[S \in \mathbb{R}^{N\times N\times H_k}\]
    • where \(H_k\) denotes the number of key/query attention heads.

  • Standard MHA independently normalizes each slice:

    \[A_{ijh} = \operatorname{softmax}_j(S_{ijh})\]

  • Talking-Heads Attention instead learns a projection across the head dimension:

    \[\widetilde{S}_{ijr} = \sum_h S_{ijh} W^{L}_{hr}\]
    • where \(W^L\) is a learned projection over attention heads.

  • Softmax is then applied to the mixed logits:

    \[\widetilde{A}_{ijr} = \operatorname{softmax}_j ( \widetilde{S}_{ijr} )\]

  • A head’s normalized attention pattern can therefore depend on logits produced by multiple query-key heads.

Post-Softmax Head Mixing

  • Talking-Heads Attention performs another projection after normalization:

    \[A'_{ijv} = \sum_r \widetilde{A}_{ijr} W^W_{rv}\]
    • where \(W^W\) is another learned head-mixing matrix.

  • These resulting weights are then applied to value heads:

    \[O_{iv} = \sum_j A'_{ijv}V_{jv}\]

  • Conceptually, the architecture therefore contains three notions of heads:

    \[\text{key/query heads} \rightarrow \text{softmax heads} \rightarrow \text{value heads}\]
    • rather than requiring one fixed head identity throughout the entire attention operation.

Why Head Mixing Can Help

  • Standard MHA implicitly assumes that the same decomposition into heads is appropriate for:

    • computing query-key similarities,

    • defining normalized attention distributions, and

    • retrieving value representations.

  • Talking-Heads Attention relaxes this assumption.

  • One head might identify a useful relation:

    \[\text{subject} \leftrightarrow \text{verb}\]

    • while another captures:

      \[\text{entity} \leftrightarrow \text{coreference}\]

  • Mixing their logits before softmax allows the normalized attention distribution to combine evidence from both relations.

  • The original paper reports improvements in masked-language-model perplexity and downstream language-understanding and question-answering transfer experiments while adding comparatively few parameters.

Talking-Heads vs. GQA

  • Talking-Heads Attention and Grouped-Query Attention both manipulate the concept of an attention “head,” but for very different reasons.

  • GQA reduces the number of KV heads:

    \[H_{KV}<H_Q\]
    • primarily to reduce KV-cache storage and memory bandwidth.

  • Talking-Heads Attention instead adds learned communication across the head dimension \(H\xrightarrow{\text{projection}}H'\) to increase representational flexibility.

  • Thus:

    \[\boxed{ \text{GQA} \rightarrow \text{share KV representations} }\]

  • whereas

    \[\boxed{ \text{Talking-Heads} \rightarrow \text{mix attention-head information} }\]

Synthesizer: alternative alignment construction

  • Synthesizer tests whether attention patterns must arise from pairwise query-key comparisons. Its dense, random, and factorized variants illustrate different ways to construct mixing structures.

Is Query-Key Interaction Necessary?

  • Standard self-attention constructs its alignment matrix from content-dependent token interactions:

    \[S = QK^\top\]

  • Every entry

    \[S_{ij}\]
    • therefore explicitly measures an interaction between token \(i\) and token \(j\).

  • Synthesizer: Rethinking Self-Attention for Transformer Models by Tay et al. (2021) asks whether these pairwise dot products are actually necessary. The paper constructs attention-like alignment matrices without explicit token-token query-key interactions and finds that several such variants remain surprisingly competitive.

  • The central idea is:

    \[\boxed{ A = \operatorname{softmax} ( \text{synthesized alignment matrix} ) }\]

    • instead of necessarily requiring

      \[A = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)\]

Dense Synthesizer

  • Dense Synthesizer generates attention weights directly from each token representation rather than comparing that token against every other token using query-key dot products.

  • For token representation

    \[x_i\]

  • A learned function produces a vector of logits across sequence positions:

    \[s_i = W_2 \sigma ( W_1x_i )\]

  • The resulting row is normalized:

    \[a_i = \operatorname{softmax}(s_i)\]

    • and applied to values:

      \[o_i = \sum_j a_{ij}v_j\]

  • The attention distribution remains dependent on the current token representation \(x_i\), but the explicit pairwise comparison \(q_i^\top k_j\) is absent.

  • This changes the conceptual model from \(\text{query asks which keys match}\) to \(\text{token directly predicts where to attend}\).

Random Synthesizer

  • Synthesizer pushes the idea further with an alignment matrix that does not depend on token content.

  • Let \(R \in \mathbb{R}^{N\times N}\) be a learned matrix.

  • Attention can then be computed as

    \[A = \operatorname{softmax}(R)\] \[O = AV\]

  • The same learned positional attention structure can therefore be used regardless of the actual token identities in the sequence.

  • Tay et al. found that learned or even random alignment structures could perform surprisingly competitively in several experimental settings, suggesting that some of the Transformer’s effectiveness comes from components beyond explicit content-based query-key matching.

Fixed Random Attention

  • An even stronger ablation fixes the random alignment matrix rather than learning it.

  • If \(R\sim\mathcal{D}\) is sampled once, the model can use \(A = \operatorname{softmax}(R)\) throughout training.

  • In this case, the attention structure itself contributes no learned token-token matching.

  • The surrounding network—value projections, feed-forward layers, residual connections, normalization, and other layers—must adapt around this fixed communication pattern.

  • The competitive behavior of such variants was evidence against an overly narrow interpretation in which Transformer performance is attributed solely to sophisticated pairwise attention maps.

Factorized Synthesizer

  • A full synthetic matrix requires storage proportional to \(N^2\) for maximum sequence length \(N\).

  • Synthesizer therefore also studies factorized forms that represent the matrix through lower-dimensional components.

  • Conceptually:

    \[R \approx AB^\top\]
    • where \(A,B \in \mathbb{R}^{N\times r}\) and \(r\ll N\).

  • The number of parameters can then decrease from approximately \(O(N^2)\) to \(O(Nr)\).

  • The paper reports that factorized Synthesizer variants can be competitive with efficient low-rank attention alternatives on encoder tasks.

Mixtures of Synthetic and Dot-Product Attention

  • The Synthesizer results do not establish that content-dependent attention is useless. In fact, the paper reports that combining synthetic attention with ordinary dot-product attention consistently improved over the corresponding Transformer baselines in its experiments.

  • A hybrid formulation can be viewed conceptually as

    \[S = \alpha S_{\text{dot}} + \beta S_{\text{synth}}\]
    • where \(S_{\text{dot}} = QK^\top\) and \(S_{\text{synth}}\) is generated by a synthesizing function.

  • After combination:

    \[A = \operatorname{softmax}(S)\]

  • This suggests that content-dependent retrieval and learned structural attention patterns can provide complementary signals.

Comparing routing assumptions

  • The two approaches address different design assumptions: head independence versus the source of alignment weights. Viewing them alongside content dependence clarifies their relationship to efficient attention methods and modern LLM kernels.

Content-Based vs. Content-Independent Attention

  • Standard self-attention is strongly content based:

    \[A_{ij} = f(x_i,x_j)\]
    • because both queries and keys depend on token representations.

  • Dense Synthesizer is only partially pairwise-content dependent:

    \[A_{i,:} = f(x_i)\]

  • Random Synthesizer can instead use:

    \[A = f(\theta)\]
    • where \(\theta\) represents learned positional parameters but not the current input tokens.

  • These variants reveal a useful spectrum:

    \[\text{pairwise content} \rightarrow \text{query-conditioned structure} \rightarrow \text{learned static structure} \rightarrow \text{fixed structure}\]

  • Attention therefore need not always be understood exclusively as dynamic semantic retrieval.

What These Experiments Tell Us About Attention

  • The original Transformer encouraged the interpretation:

    \[\text{attention} \approx \text{content-based token retrieval}\]

  • Synthesizer demonstrates that useful token mixing can also arise when the alignment matrix is generated without explicit query-key comparisons.

  • This suggests a broader interpretation:

    \[\boxed{ \text{attention-like layer} = \text{learned mechanism for routing information across positions} }\]

  • Content similarity is one powerful method for determining that routing, but it is not the only possible one.

Attention as Token Mixing

  • Consider a generic layer:

    \[O = AV\]

  • Much of the layer’s behavior is determined by the mixing matrix:

    \[A \in \mathbb{R}^{N\times N}\]

  • Different attention mechanisms primarily disagree about how \(A\) should be constructed.

Mechanism Construction of Mixing Structure
Dot-product attention (A=\operatorname{softmax}(QK^\top))
Additive attention Learned nonlinear query-key scoring
Sparse attention Content/structure over selected entries
Linformer Low-rank projected sequence interactions
Dense Synthesizer Token-conditioned generated weights
Random Synthesizer Learned content-independent matrix
Talking-Heads Dot-product attention + cross-head mixing
  • From this perspective, self-attention is one member of a broader family of token-mixing operators.

Talking-Heads and Synthesizer Solve Different Problems

  • Talking-Heads Attention accepts the query-key paradigm but questions independent heads:

    \[QK^\top \rightarrow \boxed{\text{cross-head mixing}} \rightarrow \operatorname{softmax}\]

  • Synthesizer questions the query-key paradigm itself:

    \[\cancel{QK^\top} \rightarrow \boxed{\text{synthesized alignment}} \rightarrow \operatorname{softmax}\]

  • The distinction can be summarized as:

Mechanism Pairwise QK? Dynamic Content Dependence? Cross-Head Mixing?
MHA Yes Yes No
Talking-Heads Yes Yes Yes
Dense Synthesizer No explicit pairwise QK Yes, from source token Optional
Random Synthesizer No No Optional

Relationship to Efficient Attention

  • These approaches are historically important because they broadened the attention design space, but they should not automatically be grouped with modern long-context efficiency techniques.

  • Talking-Heads primarily changes head interaction, not sequence complexity:

    \[O(N^2) \rightarrow O(N^2)\]

  • A basic Synthesizer may avoid query-key dot-product computation but can still produce an \(N\times N\) alignment matrix.

  • Therefore, eliminating \(QK^\top\) does not by itself eliminate quadratic sequence storage or token mixing.

  • This is an important distinction:

    \[\boxed{ \text{removing dot-product attention} \neq \text{removing quadratic attention matrices} }\]

  • Longformer, Linformer, Native Sparse Attention, linear attention, and related methods specifically target long-sequence scaling. Talking-Heads and Synthesizer primarily investigate what an attention layer needs to represent.

Relationship to Modern LLM Attention

  • Modern autoregressive LLM attention development has largely emphasized other constraints:

    \[\text{KV-cache bandwidth}\] \[\text{long-context computation}\] \[\text{memory capacity}\] \[\text{accelerator utilization}\]

  • This led to mechanisms such as:

    \[\text{MQA/GQA} \rightarrow \text{fewer KV heads}\] \[\text{MLA} \rightarrow \text{compressed KV representation}\] \[\text{FlashAttention} \rightarrow \text{IO-aware exact computation}\] \[\text{sparse/local attention} \rightarrow \text{fewer token interactions}\] \[\text{linear attention} \rightarrow \text{recurrent historical state}\]

  • Talking-Heads and Synthesizer instead belong to an earlier but conceptually important branch asking what the internal structure of attention itself should look like.

Generalized routing and conclusions

  • A shared routing formulation captures dot-product, synthetic, sparse, and implicit linear operators. The concluding discussion uses that view to separate architectural expressivity from computational savings.

A More General Attention Equation

  • The various mechanisms discussed throughout this primer can be described through a generalized routing formulation:

    \[O = \mathcal{R}(X)V\]
    • where \(\mathcal{R}(X)\) is the routing or mixing operator.

  • For ordinary self-attention:

    \[\mathcal{R}(X) = \operatorname{softmax} \left( \frac{XW^Q(XW^K)^\top} {\sqrt{d_k}} \right)\]

  • For a Dense Synthesizer:

    \[\mathcal{R}(X) = \operatorname{softmax} ( f_\theta(X) )\]

  • For a learned Random Synthesizer:

    \[\mathcal{R}(X) = \operatorname{softmax}(R_\theta)\]

  • For sparse attention:

    \[\mathcal{R}(X) = \operatorname{SparseSoftmax} ( S(X) )\]

  • For linear attention, the matrix \(\mathcal{R}(X)\) may never be explicitly constructed at all; associativity allows the equivalent interaction to be evaluated through aggregated state.

  • This provides a useful conceptual progression:

    \[\boxed{ \text{attention} \rightarrow \text{routing} \rightarrow \text{token mixing} }\]

The Broader Lesson

  • Talking-Heads Attention shows that attention heads need not be isolated computational channels. Learned communication across heads can enrich the attention distribution before values are retrieved.

  • Synthesizer goes further and demonstrates experimentally that useful Transformer token mixing can emerge even without explicit pairwise query-key similarity, while also finding that content-based dot products remain valuable when combined with synthetic mechanisms.

  • Together, these approaches reveal that the familiar equation \(\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\) is not the definition of attention itself. It is one particularly effective design point in a much larger space of learned information-routing mechanisms.

  • Treating routing as a design choice lets researchers independently change connectivity, head mixing, cached history, alignment generation, and hardware execution.

The Evolution of Attention

  • The evolution of attention is a branching set of solutions to different bottlenecks, not a single sequence of replacements. This concluding chapter connects classical alignment, long-context approximation, state compression, and hardware-oriented kernels in one framework. It closes by separating arithmetic, memory capacity, bandwidth, and accelerator utilization.

Historical foundations

  • Attention progressed from decoder-side retrieval to sequence-wide token interaction and then to specialized architecture and systems optimizations. The historical stages below track those questions rather than implying that every newer mechanism supersedes its predecessors.

From a Bottleneck Fix to a Systems Primitive

  • Attention began as a solution to a specific weakness of recurrent encoder-decoder models: forcing an entire input sequence through a single fixed-dimensional representation. It has since evolved into one of the principal computational primitives of modern foundation models.

  • The progression can be understood as a sequence of increasingly broad questions:

\[\boxed{ \begin{gathered} \text{What information should the decoder retrieve?}\\ \downarrow\\ \text{Can every token dynamically retrieve from every other token?}\\ \downarrow\\ \text{How should multiple retrieval heads cooperate?}\\ \downarrow\\ \text{How can attention scale to long sequences?}\\ \downarrow\\ \text{How can KV state be reduced during generation?}\\ \downarrow\\ \text{How should attention map efficiently onto modern hardware?} \end{gathered} }\]
  • Modern attention design therefore spans model architecture, sequence connectivity, state representation, numerical algorithms, memory systems, and accelerator kernels rather than being defined by a single equation.

Stage 1: Encoder-Decoder Attention

\[c_i = \sum_j \alpha_{ij}h_j\]

  • where:

    \[\alpha_{ij} = \frac{\exp(e_{ij})} {\sum_k\exp(e_{ik})}\]

  • Each decoding step could therefore retrieve different information from the encoder.

  • Effective Approaches to Attention-based Neural Machine Translation by Luong et al. (2015) further explored global and local attention along with alternative alignment functions.

  • The central conceptual transition was:

    \[\boxed{ \text{compress everything once} \rightarrow \text{retrieve relevant information dynamically} }\]

Stage 2: Self-Attention Replaces Recurrence

  • Attention Is All You Need by Vaswani et al. (2017) generalized the retrieval idea by allowing tokens within the same sequence to attend directly to one another.

  • Scaled dot-product attention computes:

    \[\operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V\]

  • This removed recurrence as the primary sequence-mixing mechanism and enabled substantially more parallel computation during training.

  • Instead of propagating information sequentially \(x_1 \rightarrow h_1 \rightarrow h_2 \rightarrow \cdots \rightarrow h_n\) self-attention creates direct interaction paths \(x_i \leftrightarrow x_j\) between token positions.

  • The Transformer consequently transformed attention from an auxiliary encoder-decoder component into the central sequence-processing operation.

Stage 3: Multi-Head Attention

  • Multi-head attention introduced multiple independently projected query, key, and value spaces:

    \[\operatorname{head}_h = \operatorname{Attention} ( QW_h^Q, KW_h^K, VW_h^V )\]

    • followed by:

      \[\operatorname{MHA}(Q,K,V) = \operatorname{Concat} ( \operatorname{head}_1,\ldots,\operatorname{head}_H ) W^O\]

  • Different heads can learn different interaction patterns, effectively creating multiple parallel routing mechanisms over the same sequence.

  • This established the architecture that became the baseline for most subsequent attention research:

    \[\boxed{ \text{many Q heads} + \text{many K heads} + \text{many V heads} }\]

Stage 4: Rethinking the Attention Operator

  • Once MHA became established, researchers began testing which aspects of the formulation were actually necessary.

  • Talking-Heads Attention by Shazeer et al. (2020) allowed information to mix across heads before and after softmax:

    \[QK^\top \rightarrow \text{head mixing} \rightarrow \operatorname{softmax} \rightarrow \text{head mixing}\]

  • Synthesizer: Rethinking Self-Attention for Transformer Models by Tay et al. (2021) went further by showing that useful token mixing could be obtained without explicitly constructing attention through pairwise query-key dot products.

  • These experiments broadened the interpretation of attention from simply \(\text{query-key similarity}\) toward \(\boxed{\text{learned information routing between representations}}\).

Long-context algorithms and exact kernels

  • Quadratic token interactions drove local, sparse, low-rank, and linear alternatives, while IO-aware kernels improved exact attention execution. These are distinct answers to different long-context bottlenecks.

Stage 5: The Quadratic-Scaling Problem

  • Dense self-attention requires every query to interact with every key.

  • For sequence length \(N\), \(QK^\top \in \mathbb{R}^{N\times N}\) giving approximately \(O(N^2d)\) arithmetic and \(O(N^2)\) attention-score storage in a straightforward implementation.

  • As context lengths increased, this became one of the Transformer’s central scaling constraints. Efficient-Transformer research consequently branched into several approaches rather than converging on a single replacement.

Stage 6: Local and Sparse Attention

  • One response was to question whether every token actually needs to interact with every other token.

  • Local attention restricts connectivity to a neighborhood:

    \[j \in [i-W,i+W]\]

  • Reducing the number of interactions from approximately

    \[N^2\]

  • To

    \[NW\]
    • for fixed window size \(W\).

  • Longformer: The Long-Document Transformer by Beltagy et al. (2020) combines sliding-window attention with selected global positions, providing local processing while preserving direct long-range communication through designated tokens.

  • More recent mechanisms make sparsity content dependent. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention by Yuan et al. (2025) combines compressed representations, token selection, and local windows while designing the sparse computation around hardware-efficient block operations.

  • The fundamental change is:

    \[\boxed{ \text{all token pairs} \rightarrow \text{selected token pairs} }\]

Stage 7: Low-Rank Attention

  • A different line of work asks whether the dense attention structure contains substantial redundancy.

  • Linformer: Self-Attention with Linear Complexity by Wang et al. (2020) projects keys and values along the sequence dimension into a lower-dimensional representation.

  • Conceptually:

    \[K \rightarrow EK\] \[V \rightarrow FV\]
    • where the projected sequence dimension \(k\) satisfies \(k\ll N\).

  • Attention then operates against a compressed representation of the sequence:

    \[\operatorname{softmax} \left( \frac{Q(EK)^\top}{\sqrt{d_k}} \right) FV\]

  • The underlying idea is:

    \[\boxed{ \text{compress the sequence representation before all-pairs attention} }\]

Stage 8: Linear Attention

  • Linear attention takes a more fundamental approach by rearranging attention computation so that an explicit \(N\times N\) matrix is unnecessary.

  • Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention by Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention by Katharopoulos et al. (2020) (2020) expresses suitable attention formulations through feature maps:

    \[\operatorname{sim}(q,k) = \phi(q)^\top\phi(k)\]

  • Associativity then permits:

    \[\phi(Q) \left( \phi(K)^\top V \right)\]

    • instead of:

      \[\left( \phi(Q)\phi(K)^\top \right)V\]

  • In the causal setting, history can be accumulated into recurrent state:

    \[S_t = S_{t-1} + \phi(k_t)v_t^\top\]

  • This reveals an important connection:

    \[\boxed{ \text{Transformer-style parallel training} \leftrightarrow \text{RNN-style recurrent inference} }\]

Stage 9: IO-Aware Exact Attention

  • Algorithmic complexity alone does not determine wall-clock performance.

  • FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness by Dao et al. (2022) showed that dense softmax attention could be made dramatically more efficient without approximating the attention operator.

  • Rather than materializing the complete score and probability matrices in high-bandwidth memory, FlashAttention tiles the computation and maintains online softmax statistics in faster on-chip memory.

  • Importantly:

    \[\boxed{ \text{FlashAttention output} = \text{exact dense softmax-attention output} }\]
    • up to ordinary numerical precision.

  • The innovation is therefore not:

    \[\text{compute fewer attention relationships}\]

  • But:

    \[\boxed{ \text{move dramatically less intermediate data} }\]

Stage 10: Hardware-Co-Designed Attention

  • FlashAttention-2, FlashAttention-3, and FlashAttention-4 progressively adapt attention execution to changing GPU architectures.

  • The progression reflects a broader shift from mathematical optimization alone toward hardware-algorithm co-design:

    \[\text{algorithm} + \text{memory hierarchy} + \text{matrix units} + \text{asynchronous execution} + \text{precision}\]

  • As accelerator matrix multiplication becomes increasingly fast, supporting operations such as memory movement, synchronization, normalization, and exponentiation can become relatively more important.

  • This leads to a crucial systems principle:

    \[\boxed{ \text{minimum FLOPs} \neq \text{minimum runtime} }\]

  • FlashAttention itself demonstrated this clearly: exact dense attention can outperform nominally lower-complexity alternatives when the latter map poorly onto hardware.

Decode-oriented KV and state optimizations

  • Autoregressive decoding exposes memory-bandwidth and cache-capacity problems that training-oriented optimizations do not fully address. Head sharing, latent compression, recurrent state, and sink-preserving caches target different forms of that burden.

Stage 11: The KV-Cache Bottleneck

  • Training and prefill emphasize attention computation, but autoregressive decoding exposes a different problem.

  • Every layer traditionally stores historical keys and values:

    \[K_{\text{cache}}, V_{\text{cache}}\]

  • For MHA, approximate cache storage scales as:

    \[M_{\text{KV}} \propto 2N_LL H d_h\]
    • where \(N_L\) is the number of layers and \(L\) is cached sequence length.

  • During single-token decoding, repeatedly reading this state can make inference strongly memory-bandwidth constrained.

  • Attention research therefore began optimizing not only attention computation, but also what historical state must be stored.

Stage 12: Multi-Query Attention

  • Fast Transformer Decoding: One Write-Head is All You Need by Shazeer (2019) introduced Multi-Query Attention.

  • MHA uses:

    \[H_Q=H_K=H_V=H\]

    • while MQA uses:

      \[H_Q=H\] \[H_K=H_V=1\]

  • All query heads therefore share one key head and one value head.

  • The idealized KV-head reduction is:

    \[H \rightarrow 1\]
    • which substantially reduces cache storage and bandwidth.

  • MQA established another independent optimization dimension:

    \[\boxed{ \text{attention heads need not imply independent KV heads} }\]

Stage 13: Grouped-Query Attention

Stage 14: Multi-Head Latent Attention

  • DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model by DeepSeek-AI (2024) introduced Multi-Head Latent Attention.

  • Instead of primarily reducing the number of KV heads, MLA compresses the historical representation itself:

    \[c_t^{KV} = W^{DKV}h_t\]

  • Keys and values can be reconstructed from the latent representation when needed:

    \[k_t^C = W^{UK}c_t^{KV}\] \[v_t^C = W^{UV}c_t^{KV}\]

  • Thus:

    \[\boxed{ \text{MHA} \rightarrow \text{store many full KV heads} }\] \[\boxed{ \text{GQA/MQA} \rightarrow \text{store fewer full KV heads} }\] \[\boxed{ \text{MLA} \rightarrow \text{store compressed latent KV state} }\]

  • This reframes KV-cache optimization as a representation-learning problem rather than only a head-sharing problem.

Stage 15: Lightning Attention

  • Lightning Attention revisits the linear-attention branch with stronger emphasis on hardware-efficient causal computation.

  • A sequence is partitioned into blocks, combining explicit intra-block interactions:

    \[O_t^{\text{intra}} = [(Q_tK_t^\top)\odot M]V_t\]

    • with recurrent inter-block state:

      \[O_t^{\text{inter}} = Q_tS_{t-1}\] \[S_t = S_{t-1} + K_t^\top V_t\]

  • This yields:

    \[O_t = O_t^{\text{intra}} + O_t^{\text{inter}}\]

  • The design combines two desirable computational properties:

    \[\boxed{ \text{matrix-multiplication-friendly blocks} + \text{sequence-linear recurrent history} }\]

  • Hybrid architectures can additionally interleave Lightning-style layers with conventional softmax attention, retaining occasional explicit token-level retrieval while reducing the cost of most sequence mixing.

Stage 16: Attention Sinks and Streaming Inference

  • StreamingLLM exposed yet another independent question:

    \[\text{Which historical KV states actually need to remain resident?}\]

  • Pure sliding-window eviction can destabilize pretrained models because a few early positions frequently behave as attention sinks.

  • A streaming cache can instead retain:

    \[\boxed{ \text{initial sink tokens} + \text{recent window} }\]

  • For fixed sink count \(S\) and window \(W\):

    \[M_{\text{stream}} = O(S+W) = O(1)\]
    • with respect to total stream length.

  • This does not create unlimited semantic memory; it creates bounded-memory continuous inference.

System-level synthesis

  • The historical trajectory branches into independent architectural choices and workload-specific implementation decisions. The final comparisons consolidate those choices into a common framework for understanding attention systems.

The Independent Axes of Modern Attention

  • The historical progression can be misleading if interpreted as a sequence in which each new mechanism replaces the previous one.

  • Modern attention mechanisms instead modify largely independent dimensions.

Design Axis Representative Mechanisms Primary Question
Alignment Bahdanau, Luong Which source information matters?
Token interaction Self-attention Which tokens communicate?
Head structure MHA, Talking-Heads How are interaction subspaces organized?
Connectivity Local, Longformer, NSA Which token pairs are computed?
Sequence compression Linformer Can interactions use compressed sequence representations?
Operator Softmax, linear attention How is routing mathematically computed?
KV sharing MQA, GQA How many KV heads are stored?
KV representation MLA What historical representation is cached?
Kernel execution FlashAttention How is exact attention mapped to hardware?
Historical state Lightning Attention Can history become recurrent state?
Cache residency StreamingLLM Which historical states remain resident?
Behavioral persistence Ghost Attention How does the model learn to preserve instructions?

A Unified Decomposition

  • A modern attention layer can be viewed as the composition of several choices:

    \[\boxed{ \mathcal{A} = \mathcal{K} \circ \mathcal{R} \circ \mathcal{H} \circ \mathcal{C} \circ \mathcal{S} }\]

    • where conceptually:

      \[\mathcal{S} = \text{state representation}\] \[\mathcal{C} = \text{connectivity pattern}\] \[\mathcal{H} = \text{head organization}\] \[\mathcal{R} = \text{routing/attention operator}\] \[\mathcal{K} = \text{hardware kernel}\]

  • These dimensions can often be combined rather than chosen exclusively.

  • For example, a model could conceptually use:

    \[\boxed{ \text{GQA} + \text{sliding-window connectivity} + \text{FlashAttention kernel} }\]

  • Another model could combine:

    \[\boxed{ \text{MLA} + \text{dense causal attention} + \text{IO-aware fused execution} }\]

  • A hybrid long-context architecture might instead combine:

    \[\boxed{ \text{linear recurrent layers} + \text{occasional softmax-attention layers} }\]

  • The term “attention mechanism” therefore increasingly describes a stack of design decisions rather than one isolated mathematical formula.

Four Different Meanings of Attention Efficiency

  • Much confusion around efficient attention comes from treating efficiency as a single property.

  • Arithmetic efficiency asks:

    \[\text{How many operations are performed?}\]
    • Sparse and linear attention primarily target this dimension.

  • Memory-capacity efficiency asks:

    \[\text{How much state must remain stored?}\]
    • GQA, MLA, windowing, cache quantization, and recurrent-state methods target this dimension.

  • Memory-bandwidth efficiency asks:

    \[\text{How much data must move during computation?}\]
    • MQA, GQA, MLA, and FlashAttention can substantially affect this dimension.

  • Hardware-utilization efficiency asks:

    \[\text{How effectively can accelerator resources execute the algorithm?}\]
    • FlashAttention, hardware-aligned sparse attention, and blockwise Lightning-style computation directly emphasize this dimension.

  • Thus:

    \[\boxed{ \text{efficient attention} \neq \text{one complexity class} }\]

Training, Prefill, and Decode Are Different Problems

  • Attention should also be evaluated separately across three execution regimes.

Training

  • Training processes many query positions simultaneously and requires backward propagation.

  • Major concerns include:

    \[\text{quadratic interactions}\] \[\text{activation memory}\] \[\text{HBM traffic}\] \[\text{backward-pass efficiency}\]

Prefill

  • Prefill evaluates an existing prompt and similarly exposes large parallel matrix operations.

  • Major concerns include:

    \[\text{long-sequence latency}\] \[\text{attention-kernel throughput}\]

Decode

  • Decode typically processes one or a few new queries against a large historical state.

  • Major concerns become:

    \[\text{KV-cache size}\] \[\text{KV-cache bandwidth}\] \[\text{small-query kernel efficiency}\]

  • Consequently:

    \[\boxed{ \text{the best training optimization} \neq \text{the best decoding optimization} }\]

The Central Tradeoff: Explicit Access vs. Compression

  • Many attention mechanisms can be organized around one fundamental tension.

  • Dense attention retains explicit access to every historical token:

    \[\{k_1,v_1\}, \ldots, \{k_t,v_t\}\]

  • Windowed attention discards distant positions:

    \[\{k_{t-W+1},v_{t-W+1}\}, \ldots, \{k_t,v_t\}\]

  • Sparse attention retains selected interactions:

    \[\mathcal{S}_t \subset \{1,\ldots,t\}\]

  • MLA compresses each token’s KV representation.

  • Linear attention goes further and compresses many historical tokens into recurrent state:

    \[S_t = f(S_{t-1},k_t,v_t)\]

  • Increasing compression generally improves efficiency but changes how explicitly historical information remains accessible.

  • The design question becomes:

    \[\boxed{ \text{How much historical structure can be compressed without losing the retrieval behavior the task requires?} }\]

Attention Has Evolved Along Multiple Directions

  • The development of attention is therefore better represented as a branching tree than a single chronological line:

    \[\text{Attention} \rightarrow \begin{cases} \text{better alignment}\\ \text{self-attention}\\ \text{head specialization}\\ \text{sparse connectivity}\\ \text{low-rank compression}\\ \text{linear/recurrent state}\\ \text{KV sharing}\\ \text{KV latent compression}\\ \text{IO-aware kernels}\\ \text{streaming cache management}\\ \text{behavioral persistence} \end{cases}\]

  • Each branch attacks a different limitation.

The Evolution in One View

  • The timeline below consolidates the attention mechanisms discussed in the preceding chapters. Each milestone names its primary architectural or execution contribution rather than implying it replaces earlier methods.
Era Mechanism Core Innovation
2015 Bahdanau Attention Dynamic source alignment
2015 Luong Attention Global/local alignment variants
2017 Transformer MHA Self-attention replaces recurrence
2019 MQA Shared KV heads for faster decoding
2020 Talking-Heads Cross-head communication
2020 Longformer Local + global sparse connectivity
2020 Linformer Low-rank sequence projection
2020 Linear Attention Associative/recurrent attention
2021 Synthesizer Attention without explicit pairwise QK
2022 FlashAttention IO-aware exact attention
2023 GQA Intermediate KV sharing
2023 StreamingLLM Attention sinks for bounded streaming cache
2023 FlashAttention-2 Improved parallelism and work partitioning
2024 MLA Compressed latent KV state
2024 FlashAttention-3 Hopper-aware asynchronous and low-precision execution
2024+ Lightning Attention Hardware-efficient blockwise linear attention
2025 Native Sparse Attention Trainable hardware-aligned dynamic sparsity
2026 FlashAttention-4 Blackwell-oriented algorithm/kernel co-design

From Attention Mechanisms to Attention Systems

  • The earliest attention papers primarily asked:

    \[\boxed{ \text{Where should the model look?} }\]

  • The Transformer expanded the question to:

    \[\boxed{ \text{How should every representation communicate with every other representation?} }\]

  • Long-context research then asked:

    \[\boxed{ \text{Which of those interactions are actually necessary?} }\]

  • MQA, GQA, and MLA asked:

    \[\boxed{ \text{What historical state actually needs to be stored?} }\]

  • FlashAttention asked:

    \[\boxed{ \text{How should the same mathematical operation move through hardware?} }\]

  • Linear and Lightning Attention asked:

    \[\boxed{ \text{Does history need to remain as explicit per-token KV state at all?} }\]

  • StreamingLLM asked:

    \[\boxed{ \text{Which cached positions must remain resident for stable continuous inference?} }\]

  • The trajectory therefore moves from an isolated neural-network mechanism toward a full systems problem:

    \[\boxed{ \text{Attention} = \text{routing} + \text{representation} + \text{memory} + \text{computation} + \text{hardware} }\]

Final Perspective

  • No single attention mechanism dominates every workload because the underlying bottlenecks differ.

  • No single attention mechanism optimizes every workload. Dense attention supports explicit retrieval; local, sparse, and recurrent alternatives reduce interactions or historical state; MQA, GQA, and MLA lower cache costs; and FlashAttention improves exact-attention execution. Streaming and instruction-persistence methods address separate constraints.

  • The central lesson of attention’s evolution is therefore not that one mechanism has replaced another. It is that attention has become increasingly factorized into independently optimizable components.

The future of attention is not merely a better attention equation; it is better co-design of routing, state, memory, and computation.

References

Classical alignment and encoder-decoder attention

Multi-head attention variants and alternative formulations

Local, sparse, and content-routed attention

Low-rank and approximate attention

Linear and recurrent attention

KV-cache-efficient attention

IO-aware exact attention and FlashAttention

Attention sinks, streaming, and cache residency

Instruction persistence through attention-oriented training

Broader efficient-attention literature

Attention implementation and open-source resources

Citation

If you found our work useful, please cite it as:

@article{Chadha2021DistilledAttention,
  title   = {Attention},
  author  = {Jain, Vinija and Chadha, Aman},
  journal = {Aman's AI Journal},
  year    = {2021},
  note    = {\url{https://aman.ai}}
}