Linear Attention Reassociation

Theory Overview

Standard self-attention operations compute similarities between queries and keys utilizing a softmax function. This incurs a quadratic time complexity with respect to the sequence length. Linear Attention provides a way to reduce this to linear time by replacing the softmax with a non-negative feature map φ(·) applied directly to queries and keys.

Instead of evaluating V_(out) = softmax(QK^T)V, we evaluate:

V_(out) = (φ(Q) · (φ(K)^T V))/(φ(Q) · (φ(K)^T 1))

Crucially, because matrix multiplication is associative, we can compute the sum over keys and values first, creating compact running summaries:

S = \sum_i φ(k_i) v_i^T
Z = \sum_i φ(k_i)

This "reassociation" enables recurrent computation of attention over sequences, making it computationally efficient for long contexts.

Experimenting with the Module

The interactive Linear Attention Reassociation module allows you to explore this mechanism with 2D queries and keys, testing how feature maps change the running sums and the final normalized attention.

To investigate these behaviors:

Adjust the Preset to explore regular or adversarial conditions (e.g., all negative values).
Switch the Feature map to observe how ELU+1 vs. ReLU impacts denominator stability, particularly where the denominator becomes exactly zero.
Change the Prefix to dynamically update the running sums and verify the observable attention output against the provided reference.

Worked Example

Let us manually calculate the "General" preset using the "ELU plus one" feature map and a Prefix of 1.

The query vector is q = [0.5, 1]. Using ELU+1 on non-negative values, we simply add 1: φ(q) = [1.5, 2].
With Prefix 1, we consider a single key k_0 = [1, 0] and value v_0 = 1. Applying the feature map gives φ(k_0) = [2, 1].

The running sums are:

Z = φ(k_0) = [2, 1] S = φ(k_0) * v_0 = [2, 1] * 1 = [2, 1]

The kernel denominator is φ(q) · Z = (1.5 × 2) + (2 × 1) = 3 + 2 = 5.0.
The numerator is φ(q) · S = (1.5 × 2) + (2 × 1) = 5.0.
The kernel attention output is 5.0/5.0 = 1.0.

This result matches both the interactive module output and the theoretical expectation for a single item (where attention must be exactly 1).

Limitations

Linear attention approximations do not perfectly replicate the standard softmax attention distribution. If a feature map like ReLU causes all key representations to become zero, the denominator goes to zero, rendering the output undefined ("Unavailable" in the module). Additionally, the feature maps chosen must be strictly non-negative to mimic probabilities, which sometimes requires compromises in expressiveness.

References

The mathematical basis for reassociating linear attention is derived in the foundational work on linear transformers. Specifically, replacing the softmax operator with a non-negative feature map over dot-products enables the computationally efficient running-sum formulation presented here:

Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020).Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. Proceedings of the 37th International Conference on Machine Learning (ICML).

Linear Attention Reassociation