Why exact attention does not need a quadratic matrix

Abstract

Does exact attention require keeping every token-to-token weight in memory? The familiar matrix formula suggests it does, but the output only needs a weighted sum. This review asks how that distinction permits an exact computation with much smaller intermediate storage. The answer rests on combining partial sums under a shared normalization scale. FlashAttention makes this algebra useful on GPUs by arranging computation around memory movement. The question is deliberately about evaluating a fixed attention function: it does not concern which tokens a trained model chooses to attend to.

Related Work

Dao et al.’s widely cited FlashAttention is the primary source. It combines tiled attention, recomputation during differentiation, and an analysis of transfers between large GPU memory and small, fast on-chip memory. Its dense algorithm preserves the attention operation while avoiding a stored matrix of all pairwise weights. [1]

Two earlier contributions clarify the mechanism. Milakov and Gimelshein show how to update the softmax normalizer alongside a running maximum, reducing memory passes without replacing softmax. [2] Rabe and Staats show that attention's quadratic intermediate storage is avoidable, including when differentiating it, while retaining quadratic time complexity. Their work separates a mathematical storage requirement from the behavior of a conventional implementation. [3]

Method/Mechanism

For one query, let s_j be its scaled score for key j, and let v_j be that key's value vector. The desired output is the sum of exp(s_j)v_j, divided by the sum of exp(s_j). Both sums can be accumulated. The complication is numerical: directly exponentiating a large score can overflow.

Subtracting the largest score solves that problem without changing the ratio. For each processed block, retain three quantities: its maximum m, its denominator l = Σ exp(s_j − m), and its vector numerator u = Σ exp(s_j − m)v_j. The scalar normalizer follows the online-softmax construction; accumulating the numerator makes it an attention computation. [2] [3]

Here is the merge identity, written for two nonempty blocks A and B with finite scores. Set m = max(m_A, m_B). Multiply each block's denominator and numerator by exp(m_block − m), then add corresponding quantities. The final answer is u/l. Multiplication simply changes both blocks to the same exponential scale; cancellation in the ratio establishes correctness.

Example 1: why averaging block outputs fails. This is an illustrative calculation. Block A has scores (0, 0) and scalar values (0, 2). Its summary is (m, l, u) = (0, 2, 2), so its local output is 1. Block B has one score ln(2) and value 4. Its summary is (ln(2), 1, 4), and its local output is 4.

The combined maximum is ln(2). Rescale A by one half, giving denominator 1 and numerator 1. B needs no rescaling. The merged output is therefore (1 + 4)/(1 + 1) = 2.5. Direct evaluation gives weights (1, 1, 2)/4 and the same result.

In this particular example, averaging the two local outputs also gives 2.5, but only because their unnormalized masses happen to match. Replace B's score with ln(4). The local outputs remain 1 and 4, yet the correct merged answer becomes (2 + 16)/(2 + 4) = 3. The normalizers carry essential information that the local outputs alone lose.

Key Findings

FlashAttention computes score tiles in fast on-chip memory and consumes them before discarding them. During the backward pass, it reconstructs needed tiles using saved normalization statistics instead of retaining the full attention matrix. The extra arithmetic can pay for itself through reduced memory transfers. [1]

Example 2: the size of an unnecessary intermediate. For an illustrative single head with 16,384 queries and equally many keys, one dense matrix contains 268,435,456 entries. At two bytes per entry, that is 512 MiB. An output with 64 value coordinates per query contains 1,048,576 entries, or 2 MiB at the same precision. These are calculated tensor sizes, not measured peak memory: inputs, gradients, scratch space, and accumulation precision add further costs.

Normalization can be deferred. Partial numerators and denominators preserve information that independently normalized block outputs discard.
Storage and arithmetic scale differently. Avoiding the matrix does not remove the dense pairwise score calculations.
Recomputation can improve speed. Whether storing an intermediate is worthwhile depends on the cost of moving it as well as calculating it.
Exactness needs a numerical qualifier. An algebraically identical result need not be bitwise identical after floating-point operations are reordered.

Limitations

The memory calculation assumes the consumer needs only the attention output. If an analysis explicitly requests every attention weight, writing that output itself requires quadratic space. A compact internal computation cannot make a requested quadratic artifact disappear. Similarly, reducing temporary attention storage does not eliminate the key-value cache used during autoregressive generation.

Dense FlashAttention retains O(N²d) arithmetic for sequence length N and head dimension d. Its reported speed gains depend on the evaluated hardware and comparison implementation; they are not a universal multiplier. [1] Single-token decoding also has a different workload from processing all query positions together, so conclusions should specify the phase being measured.

The merge calculation above excludes fully masked blocks. An implementation must handle those explicitly rather than subtracting negative infinity from itself. Precision, causal masking, and dropout also belong in a correctness comparison. “Same formula” is a necessary starting point, not a substitute for checking a working kernel.

Future Directions

A useful experiment would hold queries, keys, values, masks, and precision fixed while varying tile size and sequence length. Compare outputs and gradients against a high-precision reference, then record peak allocated memory, memory traffic, and elapsed time separately. Disable dropout initially; introduce matched dropout masks in a second comparison. This design would distinguish numerical disagreement from stochastic variation.

For performance, include both prompt processing and single-token decoding, and report the baseline implementation precisely. A memory reduction can be real even when elapsed time barely changes. That outcome would help identify when arithmetic, launch overhead, or another part of the model becomes the limiting cost.

Open question: Can tile selection minimize memory traffic while enforcing a useful worst-case bound on output and gradient error across long sequences in low precision?

Summary

Exact attention needs the correctly normalized weighted sum, not a persistent record of every weight. Online normalization provides the algebra; tiling makes it useful on hardware. The central distinction is between eliminating an intermediate and eliminating interactions. The former saves memory while preserving the dense computation.