---
title: Second-Order Linearized Cross-Attention
url: https://www.emergentmind.com/topics/second-order-linearized-cross-attention
type: topic
---

# Second-Order Linearized Cross-Attention

Second-order linearized cross-attention is an approximation technique for the cross-attention mechanism in neural networks, designed to achieve sub-quadratic memory and computational complexity while more closely mirroring the curvature of the standard softmax attention. By leveraging a Taylor expansion of the exponential function in the softmax, it introduces additional terms beyond the first-order (linear) kernel feature map used by standard linear transformers. This approach organizes higher-order interactions between queries and keys into summary structures that are tractable for large-scale inference when the key/query dimension is small relative to the sequence length [2010.14816].

## 1. Mathematical Foundation

The standard softmax cross-attention computes, for queries $Q \in \mathbb{R}^{n \times d_k}$, keys $K \in \mathbb{R}^{m \times d_k}$, and values $V \in \mathbb{R}^{m \times d_v}$:

$$
A = \mathrm{softmax}(Q K^T) V
$$

Linearized attention replaces the exponential kernel $\exp(q \cdot k)$ with an inner product in a feature space $\varphi$, yielding:

$$
A \approx \varphi(Q) (\varphi(K)^T V)
$$

Second-order linearized cross-attention, as introduced by Mercat, applies a second-order Taylor expansion to the exponential:

$$
\exp(x) \approx 1 + x + \frac{x^2}{2}
$$

where $x = (\widetilde Q \widetilde K^T) / (\alpha\sqrt{d_k})$, $\widetilde Q$, $\widetilde K$ are layer-normalized queries and keys, and $\alpha > 1$ is a scalar controlling the scale of $x$ to maintain approximation accuracy. Both numerator and denominator in attention are approximated in this manner.

The second-order term is efficiently formulated:

$$
\sum_{j=1}^m (q_i \cdot k_j)^2 v_j = \sum_{m=1}^{d_k}\sum_{\ell=1}^{d_k}q_i^m q_i^\ell \left[\sum_{j=1}^m k_j^m k_j^\ell v_j\right]
$$

defining the second-order key summary $S^{(2)}_{m,\ell} = \sum_{j=1}^m k_j^m k_j^\ell v_j$.

## 2. Algorithmic Implementation

The computation proceeds without forming the explicit $n \times m$ attention matrix, relying on intermediate key–value summary matrices and tensors. The steps are:

1. **Layer Normalization and Scaling:** $\widetilde Q \leftarrow \text{LayerNorm}(Q)$, $\widetilde K \leftarrow \text{LayerNorm}(K)$, scale $\leftarrow 1/(\alpha\sqrt{d_k})$.
2. **First-Order Summary:** $A^{(1)} = \widetilde K^T V$.
3. **First-Order Numerator:** $N^{(1)} = (\widetilde Q A^{(1)}) \cdot \text{scale}$.
4. **Second-Order Summary:** For $m,\ell=1\ldots d_k$, $S^{(2)}_{m\ell} = \sum_{j=1}^m \widetilde K_{j,m} \widetilde K_{j,\ell} V_j$.
5. **Second-Order Numerator:** For each $i = 1\ldots n$, $N^{(2)}_i = \frac{1}{2}\cdot \text{scale}^2 \cdot \sum_{m,\ell} \widetilde Q_{i,m} \widetilde Q_{i,\ell} S^{(2)}_{m\ell}$.
6. **Zeroth-Order Numerator:** $N^{(0)} = \mathbf{1}_n (\sum_{j=1}^m V_j)^T$.
7. **Combine Numerators:** $\text{Num} = N^{(0)} + N^{(1)} + N^{(2)}$.
8. **Compute Denominator:** Repeat steps 2–7 with $V_j$ set to $1$ for all $j$.
9. **Final Attention Output:** For $i = 1\ldots n$, $\text{Attention}_i = \text{Num}_i / \text{Denom}_i$.

Both $A^{(1)}$ and $S^{(2)}$ can be constructed in a single streaming pass over the keys; chunking is possible for memory constraints.

## 3. Computational Complexity and Efficiency

| Approximation | Time Complexity                      | Memory Complexity           |
|---------------|-------------------------------------|----------------------------|
| First-order   | $O(n d_k d_v + m d_k d_v)$          | $O(n d_k + m d_k)$         |
| Second-order  | $O(m d_k^2 d_v + n d_k^2 d_v)$      | $O(d_k^2 d_v)$             |
| Softmax       | $O(n m d_v)$                        | $O(n m)$                   |

Second-order linearized attention induces an extra factor of $d_k$ in both time and memory over first-order, due to the $d_k \times d_k \times d_v$ summary tensor, but remains more efficient than softmax when $d_k \ll m$.

## 4. Kernel Feature Maps

Mercat’s formulation for the first-order mapping is $\varphi(x) = [1; x /( \alpha\sqrt{d_k} )]$, covering both constant and linear contributions in one structure. The second-order term builds on all pairwise products of query vector components multiplied by the corresponding pairwise key–value summaries.

This feature map approach allows for tractable computation of higher-order approximations. A plausible implication is that additional orders could be implemented, but with quickly rising cost.

## 5. Stability, Normalization, and Limitations

Stability requires careful normalization: both queries and keys are layer-normalized (affine-free) and scaled by $\alpha$ (empirically, $\alpha=3$ is used) to ensure that dot-products $x$ remain in the $O(1)$ regime so that the Taylor expansion remains accurate. If $|x|$ grows beyond approximately $0.5$, the Taylor approximation degrades, and reverting to softmax or increasing $\alpha$ is advised. Higher-order expansions are theoretically possible but become computationally intractable for $o > 2$ unless $d_k$ is very small compared to $m$.

## 6. Empirical Evaluation and Recommendations

No empirical evaluation on real-world data sets (such as MT or image–text tasks) is reported; existing tests are on random data, and no performance benchmarks are provided versus either softmax or first-order linear attention. A plausible implication is that further research is required to assess impact on accuracy, perplexity, and runtime/memory usage in practical settings [2010.14816].

Recommended usage scenarios are those in which $d_k$ is small to moderate ($d_k \leq 64$) and $m$ (number of keys) is very large, so that $O(d_k^2 d_v)$ remains much less than $O(m d_k d_v)$. For strict memory requirements or where only coarse approximation to the softmax is tolerable, first-order schemes are preferred. Monitoring the scale of $q_i \cdot k_j$ is necessary to avoid Taylor approximation breakdown.

## 7. Context and Significance

Second-order linearized cross-attention extends techniques from "linear transformers" motivated by prior work (Katharopoulos et al., Shen et al.), which sought to deploy efficient attention via kernel tricks and low-rank approximations [2010.14816]. The addition of second-order terms recovers more of the nonlinearity of softmax, potentially yielding more faithful approximations in long-sequence contexts. This approach holds particular promise when system constraints make quadratic complexity in sequence length infeasible, though the ultimate trade-off between accuracy and resource consumption for real-world tasks remains to be fully characterized.

Source: https://www.emergentmind.com/topics/second-order-linearized-cross-attention