Linearized Cross-Attention Mechanism

• Linearized cross-attention is a kernel-based variant of softmax attention that approximates the exponential kernel to enable constant or linear compute scaling.
• The mechanism employs explicit feature maps and Taylor series expansions to balance efficiency with model expressivity and accuracy.
• It is pivotal for architectures managing long-context or external-memory tasks, achieving near-optimal performance with reduced computational overhead.

A linearized cross-attention mechanism replaces the quadratic complexity and storage costs of softmax-based cross-attention with linear-algebraic or kernel-based approximations that enable constant or linear scaling in sequence length or number of lookups. This is typically achieved by removing or approximating the softmax normalization and rewriting the cross-attention as a kernel operation or a matrix-vector product over fixed-size statistics. Linearized cross-attention is central to efficient architectures for long-context or external-memory models and has seen numerous variants targeting practical, theoretical, or expressivity-driven improvements.

1. Foundations: From Softmax to Linearized Cross-Attention

Standard cross-attention, as used in Transformers, is defined by

$$A_{\mathrm{soft}}(Q, K, V) = \mathrm{Softmax}(Q K^{\top})\, V,$$

where queries $Q \in \mathbb{R}^{n_q \times d}$, keys $K \in \mathbb{R}^{n_k \times d}$, and values $V \in \mathbb{R}^{n_k \times d_v}$. The softmax application requires building the full $n_q \times n_k$ attention matrix, demanding $O(n_q n_k d)$ compute and $O(n_q n_k)$ memory per head.

Linearized cross-attention replaces the softmax with either an explicit linear kernel, a feature map, or an approximation to the exponential kernel: $$A_{\mathrm{lin}}(Q, K, V) = \phi(Q)\bigl[\phi(K)^{\top} V\bigr],$$ for some feature map $\phi: \mathbb{R}^d \rightarrow \mathbb{R}^r$ (Barnfield et al., 4 Feb 2026). When $\phi$ is the identity, this yields a bilinear attention mechanism; other choices (such as $\phi(x) = \exp(x)$ or random Fourier features) more closely match the exponential kernel at the heart of softmax attention (Brébisson et al., 2016, Heinsen, 2024).

Elimination or linearization of the softmax nonlinearity radically alters both the computational scaling and the expressivity of the resulting attention layers.

2. Core Mechanisms and Feature Map Interpretations

The simplest linearized cross-attention directly removes the softmax, as in: $R(D, Q) = (H^\top H) q,$ with $H$ the matrix of key (and value) vectors and $q$ a single query. Here, $C = H^\top H$ acts as a covariance-like summary matrix, allowing any number of queries to be answered as $Cq$ at $O(k^2)$ cost (where $k$ is the hidden size), decoupling lookup cost from sequence length $n$ (Brébisson et al., 2016).

A general kernel perspective frames softmax attention as operating with $\exp(q^\top k)$ similarities. Linearized mechanisms substitute a feature map $\phi$ such that $$\exp(q^\top k) \approx \langle \phi(q), \phi(k) \rangle$$. Choices include:

• Identity ($\phi(x) = x$), yielding strictly linear attention (Brébisson et al., 2016).
• Elementwise exponential ($\phi(x) = \exp(x)$), yielding an exponential kernel (Heinsen, 2024).
• Higher-order Taylor expansions ($\phi$ contains constant, linear, quadratic, etc., terms) to better approximate the softmax kernel (Mercat, 2020).

The omission of explicit normalization means output magnitudes are not inherently bounded or probabilistic, so additional gating, normalization, or downstream processing is commonly required.

3. Higher-Order and Exponential-Kernel Linearization

Moving beyond first-order (strictly linear) mechanisms, second-order and kernel-based schemes more accurately approximate softmax:

• Second-order linearized cross-attention (Mercat, 2020): Expands $\exp(q^\top k)$ via a Taylor series to include zeroth- ($1$), first- ($q^\top k$), and second-order ($(q^\top k)^2/2$) terms. These are contracted with keys and values using precomputable matrices or tensors (e.g., $\sum_j k_j v_j^\top$, $\sum_j k_j^m k_j^\ell v_j$). This maintains linear dependence on sequence length but introduces quadratic dependence on the hidden dimension $d$ due to the quadratic term.
• Exponential kernel feature map (Heinsen, 2024): For each query,
  1. Precompute a summary $m_Z$ for the denominator and $m_V$ for the numerator via log-sum-exp reductions over all keys and values.
  2. For each incoming query $q$, output is computed with log-sum-exp mixing with constant (in key count) compute and memory.
  3. Updates to the memory with new keys/values require in-place log-sum-exp operations, still at constant memory.

These mechanisms offer tunable trade-offs between accuracy and cost. First-order (linear) models sacrifice expressivity but achieve maximal efficiency. Second-order and exponential variants restore some curvature and selectivity at moderate additional cost.

4. Empirical Behavior, Accuracy, and Trade-Offs

Empirical analyses demonstrate that:

• Purely linear attention (identity map) underperforms softmax attention in tasks requiring selective focus or probabilistic normalization, but still outperforms no-attention baselines in QA (Brébisson et al., 2016).

• Gated linear attention, where each hidden state is modulated by a learned gate (e.g., $f_t = \mathrm{sigmoid}(W h_t + b) \odot h_t$), bridges much of the performance gap with modest computational overhead (Brébisson et al., 2016).

• Second-order mechanisms offer improved approximation to softmax, especially in regimes where moderate correlations are informative, but incur quadratic scaling in the hidden dimension and more complex tensor contractions (Mercat, 2020).

• Exponential-kernel and log-sum-exp approaches allow $O(1)$ per-query cost in the fixed-key/value, non-autoregressive setting, with output precision comparable to standard softmax on language and standard benchmarks (Heinsen, 2024).

Accuracy generally follows: softmax > gated linear > higher-order linear > basic linear. The speedup of linearized mechanisms is most pronounced in real-time or high-query-load regimes, long-context applications, or memory-constrained deployments.

5. Theoretical Analysis and Provable Optimality

Recent work on in-context learning has established several facts regarding the expressivity and optimality of linearized cross-attention:

• Expressivity limitation: Single-layer linear self-attention fails to recover Bayes-optimal predictors in multi-modal, latent-factor models when the task distribution varies covariances from prompt to prompt (Barnfield et al., 4 Feb 2026).

• Provable optimality with depth: A deep stack of linearized cross-attention layers, with residual updates and trained via gradient flow, can "whiten" the context covariance and provably recover Bayes-optimal predictions as the number of layers and context length grows (Barnfield et al., 4 Feb 2026). This optimality emerges from repeated application of the update $$F_t = F_{t-1} + \alpha X + (\beta/L) X X^\top F_{t-1}$$, so that the effective kernel converges to the prompt-covariance inverse needed for Bayes-optimal prediction.

• Role of feature maps: The specific choice of feature map $\phi$ and stacking depth dictates whether the architecture can invert prompt-dependent covariances or only learn a global average, thus affecting its theoretical limits in multi-modal or variable-task settings.

6. State-Based and Specialized Linearized Cross-Attention

The CrossWKV architecture in RWKV-7 generalizes linearized cross-attention to a state-space, recurrent setting (Xiao et al., 19 Apr 2025):

• Weighted Key-Value (WKV) recurrence: Maintains a state matrix summarizing all past key-value contributions. At each step, the state is updated with a vector-valued decay and a rank-one correction, allowing for compressed yet expressive context representations.

• CrossWKV: Extends WKV to cross-modal fusion, where a (possibly long) sequence of image features (keys/values) is fused with a text sequence (queries) in a single unidirectional sweep. Vector-valued gating, low-rank adaptation (LoRA), and non-diagonal input-dependent transition matrices expand the mechanism's expressive power to simulate all regular languages and explicit permutation-tracking tasks—capabilities not attainable in purely diagonal or standard state-space approaches.

• Practical outcome: Achieves state-of-the-art text-to-image synthesis performance at lower compute and memory (FID=2.88, CLIP=0.33 on ImageNet 256x256), with linear inference scaling and robust behavior on long prompts or resource-constrained devices (Xiao et al., 19 Apr 2025).

7. Limitations, Practical Guidelines, and Extensions

Linearized cross-attention presents several trade-offs:

• Limitations:

  • Fixed-size summaries risk information loss in highly structured or extremely long sequences (Brébisson et al., 2016).
  • Absence of explicit normalization (unless specifically constructed) can lead to unstable outputs and degraded probabilistic interpretation (Brébisson et al., 2016, Mercat, 2020).
  • Higher-order mechanisms can greatly increase time and memory pressure when the hidden dimension grows (Mercat, 2020).
  • Practical implementations require careful attention to numerical stability (use of float32, log-sum-exp tricks) (Mercat, 2020, Heinsen, 2024).
• Guidelines:
  • Choose hidden sizes to balance representational power against $O(k^2)$ or $O(d^2)$ costs (Brébisson et al., 2016).
  • Employ simple gating or nonlinear feature maps to restore selectivity (Brébisson et al., 2016, Mercat, 2020).
  • Reserve linearization for use-cases with long context, large query volume, or memory constraints where classic softmax attention becomes prohibitive (Brébisson et al., 2016, Heinsen, 2024).
  • For mutable key sets, use incremental updates to summary statistics as in log-sum-exp-based schemes (Heinsen, 2024).
• Extensions: Linearized cross-attention is a flexible paradigm; variants include kernel-parameterized mechanisms, state-based recurrences suitable for multimodal streams, high-order (quadratic or beyond) Taylor expansions for enhanced expressivity, and hybrid attention stacks to mediate between efficiency and accuracy. It is applicable to external-memory retrieval, fast encoder-decoder architectures, multimodal data fusion, and on-device or real-time inference.

---

References:

• "A Cheap Linear Attention Mechanism with Fast Lookups and Fixed-Size Representations" (Brébisson et al., 2016)
• "Higher Order Linear Transformer" (Mercat, 2020)
• "Multi-layer Cross-Attention is Provably Optimal for Multi-modal In-context Learning" (Barnfield et al., 4 Feb 2026)
• "Cross-attention for State-based model RWKV-7" (Xiao et al., 19 Apr 2025)
• "Softmax Attention with Constant Cost per Token" (Heinsen, 2024)