Subquadratic Attention with a Hierarchical Tree-Structured Indexer
Leon Lufkin, Cutter Dawes, Alexander Speicher, Kamesh Krishnamurthy, Tomás Figliolia, Beren Millidge
arXiv:2610.42112v1Today’s paper tackles one of the biggest bottlenecks in modern language models: attention gets expensive very fast as context grows. In a standard transformer, every new token can potentially interact with every earlier token. That all-to-all flexibility is a huge part of why transformers work so well, but it also means the cost of processing long sequences grows quadratically in prefill, and during decoding the model has to keep rereading a growing key-value cache. So if you want models that can truly reason over very long documents, codebases, conversations, or agent traces, the usual attention mechanism becomes a practical wall. A lot of prior work has tried to get around that wall. Some methods replace full attention with fixed-size states, like sliding-window or state-space style approaches. Those can be much cheaper, but the tradeoff is that the state saturates: once the sequence gets long enough, the model no longer has room to represent everything it might need. Hybrid systems mix cheap layers with a few full-attention layers, but at very long contexts those expensive layers still dominate. Other methods compress tokens, compress heads, or evict parts of the cache, but they often give only modest speedups and can be brittle, because it is hard to know in advance which tokens will matter later. This paper argues that the better path is not to throw away the cache, but to search it more intelligently. The key idea is query-dependent sparsity: different queries should be allowed to retrieve different subsets of the context. That preserves the flexibility of attention while avoiding the cost of comparing every query to every key. The challenge is that the usual sparse-attention methods move the cost somewhere else. They still store the full cache, and the indexer that decides what to attend to can itself become expensive. The authors’ goal is to make that indexing step asymptotically cheaper than quadratic too. Their answer is a hierarchical sparse attention method called Zyphra Hierarchical Attention, or ZHA. The core intuition is simple: instead of flatly scoring every key for every query, build a tree of compressed summaries over the key-value cache. A query starts at the top of the tree and recursively descends through the most promising branches. At each level, the model evaluates coarse summaries first, then only drills down into the branches that look relevant. In effect, the query searches the cache from coarse to fine, rather than brute-forcing every token. Why should this work at all? The motivation comes from two empirical observations. First, the authors study full attention in pretrained language models on long fiction and find that the “top-p” subset of keys, meaning the smallest set whose attention probabilities sum to most of the mass, grows sublinearly with context length. In other words, the amount of truly relevant information seems to increase, but not as fast as the raw sequence length. That suggests attention has an effective state that is smaller than the full cache. Second, if you look across many future queries, the union of their top-p subsets eventually covers almost the entire cache. So even if a token is unimportant for one query, it may matter later. That argues against permanently evicting tokens, and in favor of keeping the full cache but searching it selectively. The authors interpret this as evidence for a hierarchical structure in attention itself. If the most relevant keys can be organized in a way that has power-law, sublinear growth, then a tree-like indexer may be able to route queries efficiently. That is the conceptual leap behind ZHA. Mechanically, ZHA builds recursively compressed summaries of the cache. A query does not score every token directly. Instead, it traverses the tree top-down, using the summaries to decide which branches to keep. Only at the leaves, or near the leaves, does it perform the more expensive attention over the selected subset of keys. The important point is that both the routing work and the final attention work grow sublinearly with sequence length when the tree is grown appropriately. So the overall attention cost becomes superlinear, but still asymptotically subquadratic. That is the paper’s main computational claim. A practical strength of the method is that it can be added to existing pretrained models without changing the base weights. The authors use a short KL distillation stage over only a few billion tokens to teach the model to use the hierarchical indexer. They report converting models up to 35 billion parameters this way, with virtually no degradation on difficult evaluations, including mathematical reasoning, agentic coding, and long-context retrieval. The experimental setup has two main parts. First, they study the structure of attention and sparsity patterns in pretrained models on long-context fiction books, to justify the tree assumption. They measure top-p subset sizes across layers and heads, and also examine how the union of relevant keys grows across many future queries. Second, they train and evaluate ZHA itself on large pretrained models, then test it on a mix of long-context retrieval tasks and more demanding reasoning and coding benchmarks. They also vary the tree geometry quite a bit, including very deep trees with up to six levels, to see how sensitive the method is to the exact structure of the indexer. One of the striking findings is robustness. ZHA is reported to work well across a range of tree geometries, including deep trees, which matters because a hierarchical method can easily become fragile if it depends on one very specific branching pattern. The paper emphasizes that performance does not collapse when the tree gets deeper, and that the method remains effective even as the geometry changes. On the efficiency side, the authors develop a detailed performance model and compare ZHA against dense attention and against existing sparse methods such as MSA. Their roofline model suggests that ZHA’s scaling remains subquadratic when extrapolated to very long sequences, including tens of millions to billions of tokens. In their model, ZHA can dramatically outperform competing methods beyond about one million tokens, and for a performant geometry with p around 0.75, they predict a crossover against MSA around half a million tokens. The exact numbers here are model-based rather than universal guarantees, but they support the central claim that a hierarchical indexer can move sparse attention beyond constant-factor gains. What do these results actually show? First, that full attention patterns in language models are not arbitrary noise: they have enough structure to be compressed in a query-dependent way. Second, that a tree-structured search procedure can exploit that structure without destroying model quality, at least for the tasks studied. Third, that the bottleneck in prior sparse methods really is the indexer, and that making the indexer hierarchical can change the scaling story in a meaningful way. There are also important caveats. The paper is careful to note that the sublinear structure they observe is empirical, not a theorem about all possible sequences. Adversarial or highly random sequences would likely not be compressible this way. The results also depend on the data distribution and, in some cases, on positional encoding choices. The authors mention that one model with partial RoPE shows stronger power-law behavior than a NoPE model, although the NoPE model still performs comparably with ZHA. So the exact geometry of attention and the representation of position may matter. Another limitation is that the method is still an approximation. It aims to recover the important parts of full attention, not to compute exact attention everywhere. That means the quality of the tree summaries and the routing policy is crucial. If the summaries fail to preserve relevance, the query could miss key information. The paper’s empirical robustness is encouraging, but the approach will likely need careful tuning as context lengths become even larger and tasks become more adversarial. The broader takeaway is that the paper reframes long-context attention as a search problem over a growing memory. Instead of asking how to make every query look at less, it asks how to organize memory so that each query can find what it needs with much less work. If the empirical scaling trends hold up, that could be a big step toward practical context windows in the hundreds of millions or even billions of tokens. In that sense, ZHA is not just another sparse attention variant. It is an attempt to change the asymptotic shape of long-context inference while keeping the expressive, query-specific behavior that makes attention so powerful in the first place.
Also spotted that day
Previous daily papers