Expanding Attention: A Taylor-Series Route to Linear-Time Transformers with a Fixed Number of Recurrent States

summary

The gist

The gist The expansion of an exponential kernel into a truncated Taylor series allows for linear-time transformers with exactly n recurrent states, which dramatically reduces memory footprint and

In short

The paper expands the exponential attention kernel into a truncated Taylor series of $n$ terms to create linear-time transformers with fixed recurrent states. This method reduces memory and increases decoding speed by replacing quadratic attention costs with a fixed number of recurrent states, achieving near-exact results for small truncation orders.

Key concepts

Exponential Kernel Expansion: The paper approximates the complex exponential function used in attention by writing it as a truncated Taylor series. This allows the expensive $O(N^2)$ attention calculation to be reformulated into a sum of $n$ simpler, manageable terms, effectively linearizing the computational complexity with respect to sequence length N.
Recurrent States: Instead of storing all past keys and values (KV cache growing as O(N)), the method collapses the causal history into exactly $n$ recurrent states. Each token update involves a single rank-one outer product, meaning the state size remains constant regardless of how long the sequence is.
Tensor-Power Inner Product Factorization: A key technique used is factoring each Taylor term into a tensor-power inner product. This factorization allows the sum over all previous positions to be moved inside the calculation, enabling incremental accumulation of states and decoupling position dependence from the kernel itself.

Terminology used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Sample Complexity — explained in 20 episodes
Retrieval-Augmented Generation (RAG) — explained in 18 episodes
Federated Learning — explained in 17 episodes
Quantum Fisher Information (QFI) — explained in 16 episodes
Flow Matching — explained in 15 episodes

This episode discusses

The paper

Expanding Attention: A Taylor-Series Route to Linear-Time Transformers with a Fixed Number of Recurrent States · Read on arXiv

Zicong Huang

We expand the exponential kernel of attention into a truncated Taylor series of n terms and factorize each term into a tensor-power inner product, which collapses the causal history into exactly n recurrent states, each updated by one rank-one outer product per token and independent of context length N. The recursion is lossless with respect to the truncated kernel (2.2 times 10-16 in float64), and the states are additive across chunks (5.7 times 10-14), so prefill admits a parallel scan. Measured exponents are t proportional to N 2.28 for softmax against N 0.91 for states, with a crossover at N about 728 and 20.8 times at N = 8192; per-token decoding latency is constant at 7.3,μ s against N 1.04 for softmax, giving 76.0 times at N = 65536. Accuracy improves exponentially in n while cost grows combinatorially, since the m-th state is m wide---a bound confirmed by SVD---so the optimum is second order, n = 3, at 6.5 times 10-4 error and 9, 945 elements per head. Query and key normalization is a prerequisite: without it the argument exceeds the convergence radius and adding terms stops helping.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Expanding Attention: A Taylor-Series Route to Linear-Time Transformers with a Fixed Number of Recurrent States".

Jane: The gist The expansion of an exponential kernel into a truncated Taylor series allows for linear-time transformers with exactly n recurrent states,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, to wrap up what we just covered about expanding that exponential kernel into Taylor terms, they are essentially showing us how to transform attention from being quadratic time complexity into something linear time with a fixed number of states.

Jane: That’s the core idea—they take that big messy exponential formula and break it down into a series, and by using tensor-power factorization, they get exactly n recurrent states.

Lu: What this means for me is that instead of storing every single past key and value pair, which grows with context length N, you only need to maintain these fixed number of state vectors per token.

Meng: And the cost per token for those states is independent of how long the sequence t gets, which sounds like a massive win for scaling inference.

Tom: Exactly. The paper confirms that this setup allows for autoregressive decoding with a constant time latency per token, which they measured at about seven point three microseconds against the standard softmax attention.

Jane: That constant time is really what matters when you're dealing with very long sequences, like those sixty-five thousand tokens we’ve been testing on <ref:2610.20637#pg1>.

Lu: And they show that for prefill, the scaling is much better than before; they get a growth exponent of zero point nine one against N, compared to two point two eight for the softmax baseline.

Meng: So when I think about deploying this to production, that means we can handle much longer context windows on our existing hardware without needing a massive memory jump just to hold the attention history.

Tom: Right, and they’re pretty confident in their accuracy here too; they showed that the recursion is exact with respect to the truncated kernel up to one-hundred-millionth of a percent in float64 precision.

Jane: That level of numerical stability on top of achieving linear time is what makes this result really solid for me, Tom.

Lu: But there are still trade-offs; accuracy improves exponentially with n, but the cost grows combinatorially, so they found that the second order, meaning n equals three is the sweet spot for performance and error rate <ref:2610.20637#pg1>.

Meng: Three states sounds manageable to me; it’s not throwing away all our memory just to chase a theoretical maximum accuracy gain.

Tom: The paper also lays out some clear prerequisites for using this method; you need your queries and keys to be normalized so that their ratio is less than or equal to one, and you have to stop the expansion at an even highest order.

Jane: So the takeaway is that this technique offers a way to get near-linear time performance by fixing the number of recurrent states, but we still have some strict rules about how you prepare your input data.

Lu: And they admit their measurements were done on a CPU setup with random synthetic data, so those exact millisecond figures won't transfer perfectly to a GPU environment yet.

Meng: That’s fair; I wouldn't deploy based only on synthetic CPU numbers, but the underlying mechanism for linear scaling is definitely something worth exploring.

Tom: So we’ve seen how they get there—linear time with fixed states, constant latency per token, and a clear optimal order of three—but now we need to see how this actually plays with real-world data distributions.

The paper's summary: Tom: So, we’re moving on to what these authors are suggesting as improvements on this Taylor series approach to attention. They aren't just stopping at linear time; they are pushing the limits of how much we can squeeze out of this fixed-state model.

Jane: That means they’re looking at how to make that constant-time token decoding even faster, or maybe how to get more accuracy without increasing those n states too much.

Lu: One key improvement is related to the state update itself; they show that because the states are additive, you can use a parallel scan for prefill.

Meng: A parallel scan sounds good for throughput, but I’m worried about fidelity—how sure are we that adding up chunks of work will still give us the exact same result as building it all at once?

Tom: The authors tested this aggressively; they verified this additivity with a high fidelity agreement of five point six eight four times ten to the negative fourteen when comparing it against a much larger context length <ref:2610.20637#pg1>.

Jane: That’s strong evidence that we can actually scale prefill time by just chunking the work, which is huge for handling massive inputs.

Lu: And they confirmed that this parallel scanning technique reproduces the full state built on a larger context length with very high accuracy.

Meng: So if we use this to handle longer documents, it should be much more efficient than just running a giant sequential loop over the whole thing.

Tom: Exactly; when you look at prefill scaling, they hit an exponent of zero point nine one against N, which is a significant improvement over the two point two eight we saw with the standard softmax attention <ref:2610.20637#pg2>.

Jane: And for decoding, they show that this state model keeps its growth exponent near one point zero four against N, which is much closer to linear than the softmax scaling.

Lu: This means that as context length grows, you’re seeing better efficiency gains in both phases of the process compared to traditional methods.

Meng: So for me at startup, that means we could potentially cut down on our inference server costs significantly if we can actually implement this parallel scan reliably without introducing errors.

Tom: Right, and they also confirmed that the states are exactly additive, which supports the idea that these computations can be done in parallel across different segments of input.

Jane: So what this changes for us is that we might not need to increase our hardware just to handle slightly longer context windows anymore; we just need to manage those n states better.

Lu: The main limitation they point out is still the precision issue; they had to project down the dimension d to sixteen just so the state size stayed manageable for their tests.

Meng: That’s a real caveat for me; if we have a very high-dimensional embedding space, that projection cost could become significant, even if the recursion itself is fast.

Tom: And they also mentioned that these measurements were single-threaded on a CPU using random synthetic data, so we need to be careful when applying these exact numbers to our GPU benchmarks.

Jane: So the next logical step is seeing how this performs when we feed it real, messy data instead of just random keys and queries.

The paper's improvements: Tom: So, to wrap up this discussion on "Expanding Attention: A Taylor-Series Route to Linear-Time Transformers with a Fixed Number of Recurrent States," we’ve seen how they transform quadratic attention into a fixed number of recurrent states for linear time inference.

Jane: That means we can run much longer sequences without the memory exploding, which is huge for real-world applications where context length is everything.

Lu: The core idea is that by using n Taylor terms, you replace a dynamic sum over all past positions with a fixed set of recurrent states that are updated once per token.

Meng: It’s good that it’s constant time per token, but what I still need to know is how much memory this model actually consumes across different types of tasks when we move away from random synthetic data.

Tom: The authors are pretty confident in the accuracy; they confirmed the recursion is exact with respect to the truncated kernel up to one-hundred-millionth of a percent in float64 precision.

Jane: And they found that the second order expansion, n equals three, gives them an error rate around six point five times ten to the negative four, which is pretty good for practical use.

Lu: They also showed that the states are additive, which means they can be parallelized across chunks of input during prefill.

Meng: That parallelization is what makes it appealing for throughput; if we can prove that additivity holds on real traffic, this could really cut down on our prefill costs.

Tom: And the scaling exponents they reported—zero point nine one against N for prefill and one point zero four against N for decoding—show a much better scaling profile than standard softmax attention.

Jane: So, it’s not just about being fast; it’s about having a predictable, superior scaling curve when dealing with long sequences.

Lu: The main caveat they flagged is that the measurements were done on a CPU using random synthetic data, so we still need to validate these numbers on actual production workloads.

Meng: I agree; synthetic data doesn't catch the kinds of memory leaks or precision issues you see when you run something on real, messy traffic.

Tom: So, while this paper gives us a solid mathematical framework for achieving linear time with fixed states using the Taylor series expansion, we still have to work on validating those practical performance claims.

Jane: Exactly; it’s a very promising structural change for attention mechanisms that points toward much more efficient models down the road.

Lu: This work opens up a whole new avenue for how we represent attention history, moving away from the dense matrix approach entirely.

Meng: I’m curious to see if this fixed state idea can be integrated into larger Mixture-of-Experts setups, like the DynFS-MoE framework we were just looking at.

Tom: That’s a great thought; connecting these two ideas could lead to something even more efficient than what this paper alone suggests.

Conclusion: Tom: So we’ve looked at "Expanding Attention: A Taylor-Series Route to Linear-Time Transformers with a Fixed Number of Recurrent States," and essentially, they’re showing us how to turn that quadratic attention problem into something linear time with a fixed number of states.

Jane: That means we can run much longer sequences without the memory exploding, which is huge for real-world applications where context length is everything.

Lu: The core idea is that by using n Taylor terms, you replace a dynamic sum over all past positions with a fixed set of recurrent states that are updated once per token.

Tom: The authors are pretty confident in the accuracy; they confirmed the recursion is exact with respect to the truncated kernel up to one-hundred-millionth of a percent in float64 precision.

Jane: And they found that the second order expansion, n equals three gives them an error rate around six point five times ten to the negative four, which is pretty good for practical use.