Subquadratic Attention with a Hierarchical Tree-Structured Indexer

📌 Highlights

ZHA is a hierarchical sparse-attention mechanism that preserves a full KV cache but searches it through recursively compressed key summaries. Its central claim is asymptotic: with suitable tree-depth and top-p scaling, routing plus final attention can grow subquadratically rather than merely improving dense attention by a constant factor.

The strongest empirical evidence is on Qwen3.6-35B-A3B. Deep C4 L6 and wide C256 L1 routers with R=16 remained within noise of dense attention on most reported reasoning, agentic coding, and retrieval evaluations; the stated exception was RULER-13, where the deep router scored 93.3% versus 96.0% for dense attention while attending to 1%-6% of the original KV cache.

The key caveat is that the largest speed claims are modeled rather than measured. The paper predicts a crossover against dense attention near 48K tokens and against MSA near 0.48M tokens, but does not empirically validate quality or latency at the ultra-long contexts required for its billion-token extrapolation.

• Query-dependent top-p routing avoids the fixed-budget tradeoff of top-K selection.
• Only reference queries are learned; the base model remains frozen during conversion.
• Six-level trees did not intrinsically cause the retrieval collapse expected from a shallow-only design.
• The main open question is whether sparsity and routing laws remain valid beyond the evaluated context range.

🎯 Introduction

Dense attention compares every query with every prior key. During prefill this creates quadratic computation, while during decode each new token rereads a KV cache whose storage and bandwidth cost grow linearly with context length. Fixed-state sequence mixers avoid these costs but can saturate, and existing fine-grained sparse methods often relocate much of the work into an indexer that still scores too many KV entries.

The paper argues that an effective replacement needs both query-dependent sparsity and an active state that grows with predictive information but remains sublinear in sequence length. Six fiction books provide the motivating observation: for Qwen3.6-35B-A3B, the smallest top-p KV subsets grow sublinearly as power laws for p<1. However, the union across future queries can cover a substantial fraction of a 256K cache, so permanently discarding low-attention keys would be unsafe.

ZHA addresses these two facts together. It keeps the cache, learns a searchable hierarchy of summaries, and routes each query independently from coarse nodes to raw KV leaves. The objective is to preserve dense-attention expressiveness where it matters while reducing the number of query-key comparisons.

🔬 Methodology

Tree construction is incremental during decode. With chunk size C, a first-level summary is added every C decode steps, a second-level summary every C^2 steps, and so on. The paper's roofline model estimates amortized construction overhead at around 1% or less. Each reference query produces a compressed key and bias that approximate the chunk's log unnormalized attention mass; summaries are recursively merged at higher levels.

Routing starts at the top of the tree. For each query head, candidate nodes are scored, sorted, and truncated to the smallest prefix whose estimated mass reaches p. Selected nodes are expanded recursively until leaf chunks are reached, with parentless and incomplete chunks handled explicitly. The final leaf union is passed to ordinary scaled dot-product attention. Query heads sharing a KV head jointly load the union of selected leaves but retain independent routing and attention masks.

The only learned components are shared reference queries. The frozen dense model supplies target attention mass pooled into chunks at every tree level, and a short KL warmup trains the references. The reported Qwen and Nemotron conversions use 4B tokens at 32K context, and the method supports grouped KV heads and positional encodings including RoPE-related configurations.

📊 Experiments

The main evaluation converts Qwen3.6-35B-A3B using deep C4 L6 and wide C256 L1 trees, each with R=16 reference queries and 4B-token KL warmup. Routing thresholds p are evaluated at 0.8, 0.9, and 0.95. The suite covers AIME 2026, MMLU-Pro, GPQA Diamond, LongBench v2, RULER, and agentic coding evaluations including SWE-bench Verified; the paper also converts Nemotron-3.5-Lightning-30B-A3B to test a hybrid-NoPE architecture.

ZHA stays close to dense attention on AIME 2026, MMLU-Pro, and GPQA Diamond, and both geometries are reported to remain competitive on agentic coding and long-context retrieval. Against MSA, which selects 16 blocks of 128 tokens and receives 40B tokens of KL warmup, ZHA's deep router attends to a similar 1.1% of decode keys on LongBench v2 but scores 14 points higher. The deep ZHA exception on RULER-13 is 93.3% versus 96.0% for dense attention.

The geometry sweep contains 28 configurations evaluated on RULER at 64K, 128K, and 256K with p values 0.7, 0.8, 0.9, 0.95, and 0.99. Across 13 tasks at 256K, retrieval is robust when 10%-30% of keys are attended; deeper trees at fixed C^L=4096 are more effective at a given sparsity. The four-parameter accuracy model transfers across held-out lengths with RMSE 1.44 percentage points overall and 1.65 percentage points on the 256K subset.

For systems scaling, the paper fits per-level selection counts on AIME 2026 responses of 16K-64K tokens using C=4, R=16, and depths L=2-6. Its roofline extrapolation for Qwen 35B in BF16 with R=1, C=8, p=0.75 predicts that ZHA becomes faster than dense attention near 48K tokens and faster than MSA near 0.48M tokens; at 10^9 tokens, the modeled speedups are about 373-417 over dense attention and 52-58 over MSA, with selected depth L=5. These are modeled attention-subsystem decode times at batch size 1, not measured end-to-end results.

fig:dense-top-p-growth

Caption: Dense top-$p$ set size versus causal context length in Qwen3.6-35B-A3B on six fiction books (averaged across books, layers, and heads). We observe that at all $p < 1$ (shades of blue), the top-$p$ set size increases sublinearly as a power law, indicating that longer sequences become increasingly compressible.

Why It Matters: This is the empirical motivation for replacing exhaustive attention with a hierarchy: the query-relevant subset grows more slowly than the full context.

fig:top-p-query-union

Caption: Fraction of KV entries of a 256K sequence covered by the union of top-$p$ attention subsets as the number of future queries increases, for several values of $p$. We observe that for the higher $p$ values we use in practice, a substantial fraction of the total KV cache is included in the union of the queries.

Why It Matters: It explains why ZHA retains the full KV cache and performs query-dependent routing instead of permanently evicting a fixed subset.

fig:architecture-overview

Caption: Overview of the ZHA architecture: (a) tree construction and (b) sparse routing. During tree construction, for each chunk of contiguous key nodes, we construct a compressed key node at the higher level by an attention-weighted sum of the keys, where the attention is computed with respect to the learned reference queries. In practice, we use multiple reference queries and store the key nodes for each in parallel. This process continues recursively with higher level key nodes being constructed from the summaries of lower levels. During routing, we start at the top of the tree and compute dot-product similarity scores between our query and the key nodes of that level. We select the set of nodes...

Why It Matters: The figure captures the complete data path from online tree construction through recursive routing to ordinary attention over selected KV entries.

fig:evaluation-qwen

Caption: (a) Reasoning and (b) retrieval evaluation scores with dense attention, MSA, and deep (C4~L6) or wide (C256~L1) ZHA ($p=0.95$, $R=16$, converted from Qwen3.6-35B-A3B). Appendix~app:training-evaluation gives evaluation details and results across routing thresholds. We observe that ZHA with both wide and deep geometries maintains performance even at high sparsity levels, while MSA suffers performance degradation especially on retrieval tasks such as RULER and LongBench~v2.

Why It Matters: It is the main visual comparison showing that both one-level and six-level ZHA can preserve benchmark quality under sparse routing.

fig:geometry-ruler-model

Caption: Geometry and sparsity have predictable effects on RULER accuracy (mean across tasks). Greater sparsity without degrading performance is achieved by: (a) decreasing compression $C^L$, (b) increasing depth at fixed compression ($C^L=4096$), or (c) making key summaries more expressive (with increasing reference-query number $R$) at fixed chunk size and depth. (The horizontal dotted lines in panels (a)--(c) show dense attention.) (d) For the simple geometry-to-performance model, predictions versus measurements across 64K, 128K, and 256K, holding out each length in turn (RMSE: 1.44 percentage points).

Why It Matters: It provides the clearest ablation evidence for how geometry controls the quality-sparsity tradeoff and how well the proposed accuracy model transfers across lengths.

fig:speedup-latency

Caption: For a fixed top-p, and hence sparsity level, ZHA (here, $R=1$, $C=8$, $p=0.75$) scales subquadratically according to our decode roofline model, eventually dramatically outperforming both dense attention and the MSA sparse attention variant. Importantly, we have not empirically validated the scaling at these ultra-long context lengths due to the lack of existing LLMs which can operate at these contexts. Because of this, while we can obtain subquadratic latency, it is unclear to what extent reasoning and retrieval performance degrades at the level of depths required for these ultra-long contexts. Latencies are modeled per-token decode times for the attention subsystem of one layer at batch si...

Why It Matters: It summarizes the paper's central systems claim while making the empirical-validation boundary explicit.

🔮 Conclusion

The paper presents ZHA as a practical hierarchical indexer for sparse attention: it retains query-dependent access to the full KV cache, can be attached to frozen pretrained models through a short reference-query distillation phase, and remains robust across wide and deep tree geometries in evaluations up to 256K context.

The empirical results support the central mechanism but not the full asymptotic deployment claim. ZHA preserves benchmark quality at substantial sparsity and outperforms MSA on the reported LongBench v2 comparison, while the roofline and power-law models indicate how subquadratic behavior could emerge at much longer contexts.

The authors also identify a fundamental tuning problem: repeated top-p pruning can make cumulative retained mass shrink with depth, while compensating by increasing per-level p can make upper tree levels effectively dense. Long-context systems therefore need both accurate summaries and a routing schedule whose sparsity law remains stable as depth grows.

🛠️ Future Research Improvements

The paper directly proposes learning chunk boundaries rather than fixing them, introducing recency bias so the tree is more detailed near the current token, and finetuning the base model on long-context data with an aggressive ZHA geometry. These changes could reduce the quality loss observed at high sparsity, but they require matched evaluations to separate better routing from model adaptation.

A second direction is to attach compressed value summaries to pruned subtrees. That would allow pruned regions to contribute coarse attention outputs while selected leaves receive token-level attention, potentially mitigating the loss caused by hard pruning.

The most important validation step is evaluation at 1M tokens and beyond on models capable of such contexts, with measured kernels and end-to-end latency. NoPE and other positional-encoding schemes also require dedicated extrapolation studies rather than assuming the Qwen p-RoPE behavior transfers.

🏭 Potential Industry Use Scenarios

If the modeled scaling is confirmed in optimized kernels, ZHA could support long-context document assistants that retain large archives or multi-document workspaces without reducing every future query to a fixed global memory state. The query-dependent union property is especially relevant when different questions need different portions of the archive.

Agentic coding systems are another credible target because the evaluation includes SWE-bench Verified and the paper reports that relevant information can be distributed across code, tool outputs, and earlier reasoning. Adaptive top-p routing could allocate more keys to diffuse software-engineering contexts and fewer to concentrated lookups.

Long-running conversational or research agents could benefit from a full retained history with sparse per-token retrieval, but deployment decisions should wait for empirical evidence beyond 256K context and for kernel measurements that include launch overhead, memory bandwidth, and worker-wave utilization.

💬 Critical Analysis

The strongest aspect of ZHA is the alignment between its mechanism and the observed structure of attention: top-p subsets are sublinear, query unions remain broad, and recursive summaries provide a plausible search structure. The geometry sweep is unusually useful because it shows that depth itself is not necessarily destructive at fixed compression, while reference-query capacity and sparsity provide interpretable controls. The reported 1.44-percentage-point held-out RMSE further suggests that benchmark degradation varies smoothly over the tested design space.

The main concern is extrapolation. The latency model depends on power-law fits collected from AIME 2026 responses of 16K-64K tokens and uses limited depth data to estimate behavior at 10^9 tokens. Even if the routing cost is subquadratic, the paper does not establish that the selected leaves retain enough information for reasoning or retrieval at those lengths. The authors also report that the deepest C4 L6, R=16 configuration at p=0.7 reached 44.0% on AIME versus 93.1% for dense attention, illustrating how an apparently favorable sparsity law can coincide with severe quality loss.

Reproducibility will depend on implementation details that the paper's roofline model treats explicitly but the experiments do not measure with released optimized kernels. A serious reproduction should therefore report end-to-end latency, per-level selected-node counts, KV-cache bandwidth, launch overhead, and quality by task rather than relying on the modeled crossover points alone. It should also test whether the learned reference-query structure remains stable across model scales, positional encodings, and training distributions.