ZHA: How a Tree-Structured Indexer Achieves Sub-Quadratic Attention

Large language models' long-context capabilities are being held back by the quadratic complexity of attention. Existing approaches either sacrifice expressiveness (e.g., sliding windows or fixed-state linear attention) or deliver only constant-factor speedups (e.g., sparse attention whose indexer is itself quadratic). Zyphra's proposed Hierarchical Attention (ZHA) attempts to break through on asymptotic complexity: using a recursively compressed tree-structured indexer, the total cost grows sublinearly with sequence length, yielding overall sub-quadratic (but still super-linear) attention.

Why existing sparse attention isn't "sparse" enough

The original intent of sparse attention is to compute only a subset of token pairs, but in most methods the indexer (which determines which keys are selected) itself carries quadratic overhead. In other words, the attention computation you save gets eaten up by routing costs. Another class of linear attention variants compresses history into a fixed-size state, which inevitably loses information on long sequences, so their expressiveness cannot match full attention. ZHA's entry point is straightforward: make the construction and query costs of the indexing structure itself grow sublinearly with length, rather than merely optimizing by a constant factor.

Quadratic attention complexity refers to standard self-attention, where each query computes similarity against all keys, producing an L×L attention matrix, with compute and memory costs growing with the square of sequence length L. When context scales from 8K to 128K, this cost grows by roughly 256×, becoming the primary bottleneck for long-context inference. Linear attention instead uses kernel functions or recurrent states to compress historical information into a fixed-size memory, reducing complexity to linear, but the fixed state has limited capacity, so early information in long sequences gets overwritten or diluted, degrading performance on retrieval tasks. Sparse attention attempts a compromise: computing only a subset of query-key pairs, but in most methods the "indexer" still needs to score the full sequence, thus retaining quadratic overhead—which is exactly what the article means by "the attention you save gets eaten up by routing costs."

How the tree-structured indexer works

ZHA builds a recursively compressed key summary tree over the KV cache. An incoming query traverses from the top of the tree downward, filtering layer by layer, and ultimately retrieves only a small subset of sparse keys. The key design is that the tree's depth can grow iteratively: as the sequence lengthens, the tree deepens, and the combined cost of routing and final attention grows sublinearly relative to sequence length. The paper's conclusion is that overall attention thus achieves a sub-quadratic but super-linear scaling curve.

This hierarchical filtering approach does not aim for one-shot localization; instead, it uses multiple layers of coarse-to-fine summaries to progressively narrow the candidate set. Compared with a flat sparse index, the tree structure spreads the complexity of "finding relevant keys" across logarithmic depth.

KV cache is the cached historical key/value tensors during autoregressive inference, avoiding recomputation at each step; its size grows linearly with the number of generated tokens and is also the main source of memory usage in long-context scenarios. ZHA does not perform flat retrieval directly on the original keys; instead, it maintains an additional summary tree over this cache: leaf nodes correspond to original key blocks, and upper-level nodes are compressed representations of their children (e.g., means or learned pooling vectors). During querying, the query is first compared against top-level summaries to select several relevant subtrees, then refined layer by layer downward, ultimately performing exact attention on only a small number of leaf nodes. This way, routing cost depends on the number of tree nodes rather than the number of tokens; as the tree deepens, the candidate set at each layer shrinks exponentially, compressing the total cost of "finding relevant keys" to sublinear.

Low-cost migration to existing models

ZHA's practical value lies in not requiring pretraining from scratch. The paper shows that a brief KL distillation training phase on the order of a few hundred billion tokens is enough to cheaply apply ZHA to existing pretrained models, while maintaining performance on long-context retrieval as well as reasoning, math, and agentic coding benchmarks. This is crucial for real-world adoption—most teams don't have the budget to retrain base models, and short distillation plus structural replacement is a more realistic path.

KL distillation here means using the original full-attention model as a teacher, so that the student model with the ZHA structure replaced imitates the teacher's output distribution on the same inputs, with the loss being the KL divergence between the teacher's and student's next-token distributions. Compared with pretraining from scratch, distillation requires only a small number of tokens to align the new structure's attention behavior with the original model, avoiding catastrophic forgetting. The "order of a few hundred billion tokens" mentioned in the paper may seem large, but relative to the typically trillions of tokens used to pretrain base models it is still more than an order of magnitude smaller, and it does not require preparing large-scale corpora anew—hence "cheaply applied." The implicit premise of this path is that ZHA's tree-routing capability can be learned via distillation from the teacher's full attention, rather than having to emerge on its own through large-scale pretraining.

Robustness to tree geometry

An easily overlooked engineering question: are the tree's structural parameters (branching factor, depth) sensitive? The paper reports that the method is highly robust to tree geometry, and can even stably support trees as deep as 6 layers. This means that in actual deployment, there is no need for fine-grained tuning for different sequence lengths, and the extra routing layers introduced by deep trees do not break performance.

Performance model and long-context extrapolation

The authors build a detailed performance model with three conclusions: first, the method exhibits sub-quadratic scaling in measurements; second, it significantly outperforms competing methods beyond 1 million context; third, it has the potential to extend the effective context window to hundreds of millions or even billions of tokens. It should be noted that the latter two points come from model projections and experimental extrapolation, and are directional judgments rather than already-delivered product metrics.

Looking at the research lineage, ZHA is in line with recent ideas of "using hierarchical structures to reduce long-context costs," but it explicitly takes asymptotic complexity as a design goal rather than merely optimizing constant factors. For readers interested in long-context inference costs, sparse attention redesign, and model distillation migration, this paper offers an architectural direction worth following.

One-line Summary

Zyphra proposes tree-structured hierarchical attention ZHA, achieving sub-quadratic long-context attention with sublinear indexing.

Paragraph Summary

Zyphra's Hierarchical Attention (ZHA) addresses the problem of long-context attention being dragged down by quadratic complexity by building a recursively compressed key summary tree over the KV cache. The query filters from top to bottom layer by layer, retrieving only a small number of sparse keys, so that the combined cost of routing and attention grows sublinearly with sequence length, yielding overall sub-quadratic but super-linear scaling. Compared with sparse attention whose indexer is itself quadratic and fixed-state linear attention with limited expressiveness, ZHA avoids both defects. The method is highly robust to tree geometry and can support trees as deep as 6 layers; a short KL distillation on the order of a few hundred billion tokens is enough to migrate it to existing pretrained models, while maintaining performance on retrieval, reasoning, math, and agentic coding benchmarks. The performance model shows it significantly outperforms competing methods beyond 1 million context and has the potential to extrapolate to hundreds of millions to billions of tokens.