Linear Attention: Removing the Softmax Buys Associativity

Attention computes, for each query, a weighted average of value vectors, with weights given by the softmax of query-key similarities. Written as matrices, the output is the softmax of Q times K transposed, all multiplied by V.

Q, K and V each have n rows and d columns, where n is the sequence length and d the head dimension. In language modelling n is thousands or hundreds of thousands and d is typically 64 or 128, so n is enormously larger than d.

Now look at where the cost sits. Q times K transposed produces an n-by-n matrix. That is the quadratic term, in both compute and memory.

The question worth asking is whether that matrix has to exist. Matrix multiplication is associative, so the product of three matrices can be bracketed either way. Multiplying K transposed by V first gives a d-by-d matrix, which is tiny and independent of sequence length, and the total work becomes linear in n.

So the quadratic cost is not inherent in attention's shape. It is imposed by one thing sitting between the two multiplications.

The obstacle is that the expression is not a product of three matrices. It is the product of a nonlinear function of two of them with the third.

Softmax is applied to the n-by-n score matrix, row by row: exponentiate every entry, then divide each row by its own sum. Both operations require the matrix to exist. You cannot exponentiate entries you have not computed, and you cannot normalise by a row sum without having summed the row.

Associativity applies to matrix products. It does not apply across a nonlinearity, so there is no rebracketing available while the softmax is there.

That isolates the target precisely. If the similarity between a query and a key could be written as a dot product of transformed vectors, with no nonlinearity applied to the matrix as a whole, the three-matrix product would be genuine and could be reassociated.

That is exactly what linear attention does. Replace the exponential of a dot product with a dot product of feature maps applied to each vector separately.

Choose a feature map, a function applied to each query and each key vector independently. Define the similarity between a query and a key as the dot product of their feature maps.

The crucial property is that the nonlinearity now acts on individual vectors before any interaction, rather than on the matrix of interactions. Applying it to Q and to K separately leaves a genuine product of three matrices: feature-mapped Q, times feature-mapped K transposed, times V.

Now reassociate. Compute feature-mapped K transposed times V first. Because feature-mapped K has n rows and some feature dimension, and V has n rows and d columns, the product has feature dimension by d entries, with n summed away.

That quantity is a fixed-size matrix. It does not grow with the sequence.

Every query then reads from it with a single small matrix-vector product, and the normaliser is handled by maintaining a running sum of feature-mapped keys alongside it.

Total cost is linear in n and the state is bounded. The trilemma's third corner, reached by deleting one operation.

The fixed-size matrix has an interpretation that reframes the whole subject.

It is a sum over positions of outer products between each key's feature vector and its value vector. A sum can be built incrementally: start from zero, and at each position add the outer product of that token's feature-mapped key with its value.

Written that way, the state at time t equals the state at time t minus 1 plus an outer product. That is a recurrence. The state is a matrix rather than a vector, and the update is additive rather than passing through a nonlinearity, but it is unambiguously a recurrent network.

So linear attention and a linear recurrent network are the same object. One arrived by removing the softmax from attention; the other by generalising an RNN's state from a vector to a matrix.

The reading step has an interpretation too. Multiplying a query's feature vector by the state matrix retrieves a weighted combination of every value ever written, weighted by how well the query matches each key. It is an associative memory: write with outer products, read with a matrix-vector product.

The most useful consequence is that the same model has two computation modes producing identical results, and you pick per situation.

The parallel mode brackets the product the original way, materialising the n-by-n matrix. Cost is quadratic in sequence length, and every position is computed simultaneously as dense matrix multiplications that saturate the hardware. This is training mode.

The recurrent mode brackets the other way, maintaining the fixed-size state and stepping through positions. Cost is linear in sequence length, memory is constant, and each step is small. This is generation mode.

They compute the same function, so a model can be trained in one and served in the other with no conversion.

That is precisely what classical RNNs could not do, and it is the resolution of the trilemma from the previous lesson. Parallel training because the parallel mode exists; constant-cost inference because the recurrent mode exists; compressed rather than perfect recall because the state is finite.

Production implementations use a chunked hybrid: parallel within blocks to exploit tensor cores, recurrent between blocks to keep memory bounded.

The substitution was not free, and being precise about the loss explains every empirical result in the last lesson.

Softmax is sharp. The exponential means a score slightly higher than its neighbours receives a disproportionately larger weight, so attention can concentrate almost all its mass on one position. That is what makes exact retrieval possible: attending to token 7 and essentially ignoring everything else.

A dot product of feature maps is comparatively flat. Without the exponential, the weight distribution over positions is smoother, so the mechanism averages where softmax would select.

And there is a capacity limit that no feature map removes. The state is a fixed-size matrix holding a sum of outer products. Each write adds to the same accumulator, so writes interfere. Past a number of stored associations related to the state dimension, retrieving any one cleanly becomes impossible, because the others are superimposed on it.

That is the compression from the previous lesson, made concrete. It is not a defect of the implementation. A finite state holds finitely many things.

The plain formulation has a second problem, and fixing it turns out to be where most of the recent progress lives.

The state is a running sum with no decay. Every token ever seen contributes equally and permanently. A token from a hundred thousand positions ago has exactly the weight of the one just processed.

For language that is clearly wrong. Relevance decays, topics change, and a bounded state filled with everything is a state that has resolved nothing.

The fix is to multiply the state by a decay factor before each additive update, so older contributions shrink geometrically. That is a forget gate, and it is the same idea long short-term memory introduced decades ago, reappearing because the same problem reappeared.

The design question is what the decay depends on. A fixed constant is simple and treats all content alike. Making it a learned function of the current token lets the model decide, per position, how much of the past to retain: a delimiter can flush, a salient entity can persist.

That input-dependence is the pivotal idea, and the next lesson shows the state space line arriving at exactly the same conclusion from the opposite direction.

The state's dimensions determine both cost and capacity, and the arithmetic is worth carrying.

The state is a matrix with one dimension set by the feature map's output size and the other by the value dimension. Enlarging either increases how much can be stored before writes interfere destructively, and increases memory and per-token compute proportionally.

This is the recall-throughput tradeoff, and it is a dial rather than a discrete choice. A small state gives fast, memory-light inference and poor exact retrieval. A large state approaches attention's recall and approaches its cost. Attention itself is the limiting case where the state grows with the sequence, which is why it never saturates.

Two further observations follow.

The tradeoff is per layer, so different layers can carry different state sizes, and there is no requirement for uniformity.

And it composes with the previous point about decay. A small state with good input-dependent forgetting can outperform a larger state that retains indiscriminately, because capacity spent on irrelevant history is capacity wasted.

Linear attention was proposed well before it worked, and the intervening history is instructive rather than merely chronological.

Early formulations chose a feature map to approximate the exponential kernel, reasoning that if the approximation were good the model would behave like softmax attention at linear cost. The approximations were mediocre and the models were noticeably worse, which is why the idea was dismissed as an efficiency hack with a quality penalty.

The reframing that changed things was to stop approximating. Once the state is understood as an associative memory with a write rule and a decay rule, the design question becomes how to write, retain and read well, rather than how to imitate a softmax.

From there the useful ideas follow directly: input-dependent decay so the model controls retention, more careful write rules that remove a stale association before writing a new one rather than merely adding, normalisation that keeps the state well-conditioned over long sequences, and larger states where recall matters.

That is the shift from a cheaper approximation of attention to a different mechanism with its own strengths, and it is why the last two years of results look different from the first few.

This lesson reduces to four statements, all of which the next one will meet again from the other direction.

Softmax is what forces the quadratic cost, because a nonlinearity applied to the score matrix blocks reassociating the product. Remove it and the matrices can be bracketed so a fixed-size quantity is maintained instead.

That fixed-size quantity is a recurrent state, updated by outer products and read by a matrix-vector product. Linear attention and a linear recurrent network are the same object.

The same parameters support two computation modes, quadratic-parallel for training and linear-recurrent for generation, and they compute the same function. This is what classical RNNs lacked and it is the resolution of the trilemma.

And the losses are specific rather than general: softmax's sharpness, which enabled exact selection, and a finite capacity beyond which stored associations interfere. Decay and larger states mitigate both without removing either.

The next lesson starts from a linear dynamical system, adds discretisation and a principled initialisation, and arrives at the same recurrence with an input-dependent forget gate.