ZHA: Zyphra Hierarchical Attention

Why Is Efficient Attention So Difficult?

One of the most remarkable properties of transformers is their ability to retrieve information from anywhere in a sequence.

When a language model generates its next token, attention allows it to look back at everything that came before and decide which pieces of information are relevant. For example, a coding model might need to retrieve a function definition from thousands of lines earlier, while a reasoning model might need to revisit an intermediate conclusion from a long chain of thought.

Crucially, the model does not have to decide in advance which information will be important. Every new query can independently retrieve whichever parts of the context it needs.

This flexibility is enormously powerful, but it is also expensive.

In standard attention, each token is compared against every preceding token. This results in quadratic complexity. Doubling the context length therefore increases the total attention computation by approximately four times. During generation, each additional token also requires reading an increasingly large key-value (KV) cache containing representations of the model's preceding context.

As we push toward contexts of tens of millions of tokens, computational costs quickly become prohibitive.

A substantial amount of research has therefore focused on making attention more efficient. However, existing approaches tend to fall into two broad categories, each with an important limitation.

The first approach is to compress the past into a fixed-size state.

State-space models (SSMs), linear attention variants, and sliding-window attention mechanisms can achieve highly efficient scaling by maintaining a bounded representation of the preceding sequence. The problem is that the amount of potentially useful information in a sequence does not stop growing.

Imagine reading an increasingly large collection of books while being allowed to maintain only a fixed number of notes. Initially, your notes might capture most of what matters. But eventually, new information must overwrite or interfere with older information, regardless of how useful that information might become later.

The second approach is sparse attention.

Sparse attention methods preserve the ability to retrieve arbitrary information by selecting a small subset of the KV cache for every query. Instead of attending to every preceding token, the model might attend to only the most relevant 1%.

However, there is a catch: how does the model determine which 1% of the tokens are relevant?

Many sparse attention approaches use a lightweight indexer that scores every preceding token to identify the subset worth retrieving. Although the final attention operation is sparse, the indexer itself must still search the entire context. This effectively moves the quadratic computation from attention into the indexer. Making the indexer cheaper can provide large constant-factor improvements, but the underlying scaling problem remains.

Some methods address this by indexing compressed chunks rather than individual tokens. However, searching every chunk at a single resolution still requires work that grows linearly with the context, again delivering only a constant-factor improvement in asymptotic scaling.

These limitations suggest that an effective attention replacement needs two properties simultaneously.

First, retrieval must remain query-dependent. Different questions require different information, so we cannot simply discard a fixed set of supposedly unimportant tokens.

Second, the amount of actively retrieved information must be allowed to grow with context length, but more slowly than the context itself. A fixed-size memory will eventually become insufficient, while linearly growing active memory brings us back to quadratic attention.

Our work starts from the observation that natural language appears to have precisely this property. When we examine attention distributions in pretrained LLMs, we find that the number of tokens accounting for a fixed fraction of attention probability tends to grow sublinearly with sequence length. In our experiments on long-form fiction, this relationship is often well described by a power law.

In other words, as a document becomes longer, the amount of information the model needs to retrieve grows, but it does not necessarily grow as quickly as the document itself. This suggests an important possibility: perhaps we can preserve attention's flexibility without paying the cost of searching every token.

Hierarchical Attention: Searching the Context with a Tree

The key insight behind ZHA is that finding relevant information should not require examining every piece of information individually.

Consider searching for a particular passage in a very large book. You would not normally read every sentence to determine where to look. Instead, you might start with the table of contents, identify a promising chapter, look through its sections, and finally locate the relevant paragraph.

This is the basic intuition behind hierarchical attention. Rather than maintaining a flat index over all preceding tokens, ZHA builds a tree of increasingly compressed representations of the KV cache.

Building the tree

At the lowest level, we begin with the original keys from the model's KV cache.

We divide these keys into contiguous groups, or chunks, and construct a compressed representation of each group. We then group these summaries together and compress them again, repeating this process recursively to build higher levels of the tree.

For example, with a chunk size of four, one level summarizes four original keys, two levels summarize 16 keys, and six levels can summarize regions of 4,096 tokens.

The result is a hierarchy in which the highest levels describe large regions of the context, while lower levels provide progressively finer detail.

An important challenge is making these summaries useful for retrieval.

Simply averaging the keys in a chunk would lose much of the information necessary to determine whether that chunk is relevant to a particular query. Different queries might care about completely different aspects of the same chunk.

Instead, ZHA constructs summaries using a small set of learned reference queries. Each reference query provides a different way of examining the keys in a chunk. These reference queries allow ZHA to approximate how much total attention a future query would assign to the chunk.

Intuitively, rather than asking a summary to represent everything inside a region, we train it to answer a much more specific question:

How relevant is this region likely to be to the query I am currently trying to answer?

These summaries are then recursively combined to construct the higher levels of the tree while the original KV cache is still retained. ZHA compresses the keys for the purpose of searching, rather than permanently replacing the information that might need to be retrieved later.

Routing through the tree

Once the tree has been constructed, retrieving information becomes a coarse-to-fine search.

When a new query arrives, ZHA first compares it against the summaries at the highest level of the tree. Rather than examining all the tokens represented by these summaries, it determines which regions appear most relevant and expands only those regions. At the next level, it repeats the process, comparing the query against the finer-grained summaries within the selected regions. This continues until the search reaches the original KV entries, at which point the model performs ordinary attention over the selected subset.

Importantly, ZHA does not simply follow one branch of the tree. It can expand multiple branches located anywhere in the context, allowing it to retrieve information from widely separated regions. The decision about which branches to expand uses a mechanism called top-p routing. Rather than selecting a fixed number of nodes, top-p selects the smallest group of nodes accounting for a specified fraction of the estimated attention probability.

This makes the amount of retrieval adaptive to the query. If a query is looking for a single specific fact, the attention distribution might be highly concentrated and ZHA will explore only a few branches. If a query requires synthesizing information from many different parts of the context, the attention distribution may be more diffuse, and ZHA will expand substantially more branches.

Adaptive retrieval is critical. With a fixed top-k budget, we face a strict trade-off between having a large enough k to answer diffuse queries, and not wasting compute when only a small highly-concentrated subset is necessary. Top-p instead allows the amount of computation to adapt naturally to the task.

How does this achieve subquadratic scaling?

The most important advantage of a hierarchical index is that we no longer need to search the entire index at the same resolution. At each level, large regions of the context can be eliminated from consideration without examining their individual contents. Provided that the number of regions explored grows sublinearly with context length, both the routing computation and the final attention computation can remain sublinear per query. This is fundamentally different from sparse attention methods that still need to scan a flat index over the entire context. The tree makes the process of finding relevant tokens itself sparse.

It is also fundamentally different from a fixed-size state. ZHA does not impose a fixed upper bound on how much information can be stored or retrieved. The active state can continue growing with the context, while remaining substantially smaller than the complete sequence. Moreover, the full KV-cache is maintained at the leaf level of the tree so that the original keys and values remain available for retrieval.

Distilling ZHA into Pretrained Models

An especially attractive property of ZHA is that it does not require training a new language model from scratch. Instead, we can take an existing pretrained transformer and teach a lightweight hierarchical indexer to approximate its original dense attention behavior.

We do this through a short KL-distillation phase. The original model remains frozen throughout training. We observe how its dense attention distributes probability across different regions of the context, and train ZHA's reference queries to reproduce these probabilities at each level of the tree.

In effect, the pretrained model teaches the indexer how to search its own memory. Because the only learnable components are the reference queries, the number of additional parameters is extremely small. For our Qwen3.6-35B-A3B experiments, ZHA adds only 81,920 trainable parameters, approximately 0.00024% of the base model's parameter count.

We find that training on just a few billion tokens at a context length of 32K is sufficient to obtain strong performance, including evaluations at much longer contexts. This is particularly important because it means hierarchical attention can potentially be deployed on top of existing capable models, without repeating their expensive pretraining process or changing the weights responsible for their reasoning and language capabilities.

Is model quality preserved with ZHA?

Efficient attention methods have often been evaluated primarily on relatively simple retrieval tasks, such as finding a single hidden fact in a long document. However, these benchmarks do not fully capture the demands placed on modern reasoning models.

A long-context reasoning or coding task may require repeatedly revisiting different parts of the context, combining information from many locations, and maintaining coherence across tens of thousands of generated tokens.

We therefore evaluate ZHA on a much broader and more challenging collection of tasks, including AIME 2026 mathematical reasoning, MMLU-Pro, GPQA Diamond, SWE-bench Verified, Terminal-Bench 2.1, RULER, and LongBench v2.

We compare the original dense model with two ZHA configurations: a wide, single-level tree and a deep, six-level tree.

The results are encouraging.

On RULER-13 at 256K tokens, our wide-tree ZHA model achieves 96.2% accuracy, compared with 96.0% for the original dense model. A much deeper six-level tree achieves 93.3%, despite attending to only a small fraction of the original KV cache.

On LongBench v2, the deep-tree model scores 63.4%, compared with 58.0% for dense attention, while attending to approximately 1.1% of the KV cache during decoding.

And on SWE-bench Verified, a demanding agentic coding benchmark involving real-world software engineering tasks, ZHA achieves performance comparable to the dense baseline: 67.4% for the deep tree and 68.4% for the wide tree, compared with 67.4% for dense attention.

We also compare against MiniMax Sparse Attention (MSA), which uses a fixed number of blocks selected by a learned indexer. The MSA indexer underwent only the KL warm-up phase, using 40B tokens, following the training recipe described in the MSA paper.

Although MSA remains competitive on several reasoning benchmarks, it suffers substantially larger degradation on long-context retrieval. On RULER-13, for example, MSA achieves 77.1%, compared with 93.3% and 96.2% for our deep and wide ZHA configurations.

On LongBench v2, our deep router outperforms MSA by approximately 12.7 percentage points at a similar fraction of attended keys, despite requiring only one-tenth as many training tokens for the indexer.

These results are especially significant because they demonstrate that hierarchical routing can preserve performance on tasks requiring much more than retrieving one isolated fact.

We also test ZHA on Nemotron-3.5-Lightning-30B-A3B, which uses no positional encoding (NoPE) in its attention layers. The method continues to work well on reasoning and agentic coding tasks, although retrieval performance on RULER degrades more noticeably, in line with findings on sparse adaptation of NoPE attention from other work.

Overall, our experiments suggest that hierarchical attention can preserve much of the expressive power of dense attention while avoiding the need to examine most of the context.

Robustness to Tree Geometry and the Potential for Dramatic Speedups

One natural concern with hierarchical attention is that the tree itself might introduce a new bottleneck. If information must pass through multiple levels of compression, wouldn't increasingly deep trees become less accurate?

This is a crucial question because the subquadratic scaling of ZHA ultimately depends on our ability to increase tree depth as context lengths grow. To investigate this, we train and evaluate 28 different tree configurations, varying the number of reference queries, the number of tokens represented by each chunk, and the depth of the tree.

Remarkably, we find that ZHA is highly robust across a wide range of these choices. At moderate sparsity levels, retrieval performance is similar across many different tree geometries. More surprisingly, when we hold the overall compression factor fixed, deeper trees often perform better than shallower trees at the same level of sparsity.

This suggests that progressively refining the search through multiple levels of abstraction can be more effective than making one aggressive compression decision.

How fast can ZHA become?

Reducing the number of attended tokens is only part of the performance story.

Actual GPU performance also depends on the cost of constructing the tree, routing through its different levels, selecting nodes, and finally performing attention on the retrieved tokens. To study these costs, we develop a detailed GPU roofline performance model accounting for computation, memory bandwidth, kernel launch overhead, and hardware parallelism.

We use this model to compare the predicted performance of ZHA against dense attention and a representative sparse attention method with a quadratic indexer.

The results show why changing the asymptotic scaling matters. For a lightweight ZHA configuration, our model predicts that hierarchical attention becomes faster than dense attention at approximately 48K tokens and faster than MSA at approximately 480K tokens.

Beyond these crossover points, the advantage continues to increase as the context grows. At a context length of ten million tokens the model projects an approximate speedup of 32–35× over dense and 4.7–5.1× over MSA. At one billion tokens, these projected speedups reach approximately 373–417× and 52–58×, respectively.

At sufficiently long contexts, even very large constant-factor improvements become less important than avoiding quadratic scaling altogether.

The Narrow Path to Truly Ultra-Long Contexts

Beyond the immediate architecture and experimental results, our work raises a more fundamental question:

What does it actually take for a language model to use an arbitrarily long context effectively?

It is tempting to think of long-context scaling as purely a computational problem. If we could make dense attention sufficiently fast and inexpensive, perhaps we could simply keep expanding the context window indefinitely.

However, our analysis suggests the situation may be more subtle. To understand why, it is useful to consider two opposing failure modes.

Attention becomes too diffuse

As the context grows, dense attention must compare a query against an increasing number of tokens. But not all of these tokens contain useful predictive information. As more irrelevant or weakly relevant tokens are added, their collective contribution to the attention distribution can become increasingly significant.

Prior work has suggested that softmax attention can become dispersive at large sequence lengths: probability spreads across an increasingly large number of distractors, potentially obscuring the information that is actually relevant.

In other words, even if we could afford to compute dense attention over an enormous context, that would not necessarily guarantee that the model could retrieve the correct information from it.

This suggests that sparsity might be useful for more than computational efficiency. By filtering out irrelevant parts of the past, sparse attention may act as a form of denoising, allowing the model to focus more precisely on the information that matters.

Attention becomes too sparse

However, making attention increasingly sparse introduces the opposite problem.

If the model retrieves too little information, eventually it will discard useful context. This is particularly relevant for hierarchical attention because sparsity compounds as we descend the tree.

Suppose we retain a fixed fraction of attention probability at every level. As we add more levels, the fraction of the original attention mass retained after the full traversal can progressively shrink toward zero.

In the limit, this can cause the active retrieval state to effectively collapse toward a fixed size.

We would then recreate the same fundamental limitation as fixed-state architectures: the model would no longer be able to retain an increasing amount of relevant information as the sequence grows.

One might try to solve this by becoming increasingly conservative at each level, retaining almost all the attention probability as the tree gets deeper. But this leads to the opposite problem. If higher levels of the tree retain nearly every branch, the hierarchy stops providing meaningful pruning, and the computational cost approaches that of dense attention.

We are therefore faced with a narrow path between two undesirable extremes:

Too little sparsity, and attention becomes computationally prohibitive and potentially overwhelmed by irrelevant information. Too much sparsity, and the information the model needs to reason is lost.

Matching compression to the predictive information in the data

We believe the way through this tradeoff is to recognize that the appropriate amount of compression should be determined by the statistical structure of the sequence itself.

A long natural language sequence contains substantial redundancy. The millionth token can still provide useful new information, but the amount of additional predictive information it contributes may be much smaller than that of earlier tokens.

The total amount of relevant information can therefore continue growing without needing to grow linearly with the total sequence length. This is the regime we ultimately want:

• The active state continues growing, preventing the information bottleneck of fixed-size memories.
• It grows sublinearly, avoiding the prohibitive cost of dense attention.
• Retrieval remains query-dependent, so different questions can access different parts of the complete history.

Achieving this regime asymptotically will require more sophisticated control of sparsity than simply applying a fixed top-p threshold at every tree level. In particular, the routing threshold must evolve with tree depth and context length so that the cumulative compression tracks the information actually needed for prediction. Crucially, this scaling is not guaranteed by the attention architecture itself. It depends on the distribution of the underlying data.

The broader implication is that truly ultra-long context may require something more fundamental than simply making dense attention faster.

At sufficiently large scales, the challenge is not just storing or processing more tokens. It is determining how much of the past remains useful, and retrieving precisely that information without being overwhelmed by everything else.

Our results suggest that hierarchical attention provides a promising architectural foundation for doing this. By combining query-dependent retrieval with a dynamically growing, sublinear active state, ZHA demonstrates that attention need not be restricted to either fixed-size memory or quadratic computation.

We believe the most exciting possibility is that the right form of sparsity could simultaneously address both the computational and informational challenges of extremely long contexts.

If we can learn to match the growth of active attention to the growth of predictive information in real-world sequences, contexts of hundreds of millions or even billions of tokens may eventually become practical. The future of long-context language modeling may therefore depend on giving models the ability to find exactly what matters, at every scale.