# Subquadratic Attention with a Hierarchical Tree-Structured Indexer

> ZHA is a hierarchical sparse attention method that achieves subquadratic scaling with query-dependent sparsity, matching dense performance on long-context tasks while attending to only 1–6% of the KV cache.

- **Source:** [arXiv](https://arxiv.org/abs/2610.42112)
- **Published:** 2026-11-14
- **Permalink:** https://picx.dev/p/i11Eut
- **Whiteboard:** https://picx.dev/p/i11Eut/image

## Summary

## Summary (Overview)

- **Zyphra Hierarchical Attention (ZHA)** is a novel sparse attention method that uses a hierarchical, tree-structured indexer to achieve **asymptotically subquadratic** attention scaling, overcoming the quadratic bottleneck of standard attention while maintaining query-dependent sparsity.
- ZHA constructs a tree of recursively compressed key summaries over the KV cache; queries traverse this tree top-down, selecting a sparse subset of keys via **adaptive top-p routing** at each level.
- The method can be applied to **existing pretrained models** (up to 35B parameters) with only a **1-billion-token KL distillation phase**, freezing the base model and learning only a small set of reference queries.
- ZHA maintains performance on **long-context retrieval (RULER, LongBench v2), mathematical reasoning (AIME), and agentic coding (SWE-bench)** tasks, while attending to only 1–6% of the KV cache.
- A detailed **roofline performance model** predicts ZHA becomes faster than dense attention at ~48K tokens and faster than MiniMax Sparse Attention (MSA) at ~0.48M tokens, with potential speedups of 373–417× over dense attention at $10^9$ tokens.

## Introduction and Theoretical Foundation

### The Quadratic Bottleneck of Attention

Standard attention computes an all-to-all mapping between sequence elements, where each query compares against every key in the sequence. This leads to **quadratic computational cost** in prefill and **linear memory bandwidth costs** during decode (re-reading the KV cache per token). These costs prohibit extremely long context windows despite their demonstrated benefits for reasoning and retrieval tasks.

### Limitations of Existing Approaches

| Method Family | Approach | Limitation |
|---|---|---|
| State-space models / sliding-window | Fixed-size state | State saturates at long contexts, degrading performance |
| Hybrid methods | Interleave fixed-state mixers with full attention | Few full-attention layers still dominate at long contexts |
| Sequence compression (CSA, NSA, HNet) | Online compression of multiple tokens | Constant-factor speedups, brittle in practice |
| KV cache eviction | Offline eviction of low-value tokens | Impossible to know token utility ahead of time |
| Sparse attention (DSA, MSA) | Query-dependent top-K selection | Indexer itself scales quadratically (60–90% of time at ~1M tokens) |
| Chunk-sparse (PISA, HiP) | Multi-level chunking | Log-linear complexity but retrieval degradation at long contexts |

### Key Insight: Natural Sparsity of Attention

The authors demonstrate that attention's effective state size grows **sublinearly** with sequence length. They operationalize this via the **top-p set**: the smallest subset of keys whose attention probabilities sum to at least $p$. Attention can be decomposed as:

$$
\begin{array}{c} \text {Attention} (q, K, V) = p \cdot \text {Attention} (q, K _ {p}, V _ {p}) \\ + (1 - p) \cdot \text {Attention} (q, K _ {\bar {p}}, V _ {\bar {p}}), \end{array}\tag{1}
$$

where $(K_p, V_p)$ is the top-p subset and $(K_{\bar{p}}, V_{\bar{p}})$ contains the remaining KV pairs.

**Empirical finding:** The top-p set size grows as a **sublinear power law** in context length (exponent depending on $p$), suggesting hierarchical, scale-invariant structure in attention representations.

### Why Hierarchical Search?

The authors show that while individual queries are sparse, the **union** of top-p sets across future queries covers the entire cache:

$$
U _ {p} (W) = \frac {1}{T - W} \left| \bigcup_ {q = T - W} ^ {T - 1} S _ {p} (q) \right|
$$

where $S_p(q)$ is the top-p subset for query at position $q$. As $W$ grows, $U_p(W) \to 1$, meaning the full KV cache must be retained. This motivates a **query-dependent retrieval mechanism** that identifies relevant keys without scanning all of them—hence the need for a hierarchical tree indexer.

## Methodology

### Tree Construction

ZHA recursively creates key summaries that compress contiguous chunks of keys. For a chunk $K_C = [k_0, \dots, k_{C-1}]$ with chunk size $C$, the **log unnormalized attention mass** (Log-Sum-Exp, LSE) for query $q$ is:

$$
F _ {C} (q) := \mathrm{LSE} _ {j \in [ C ]} \left(q ^ {\top} k _ {j} / \sqrt {d}\right) = \log \sum_ {j \in [ C ]} \exp (q ^ {\top} k _ {j} / \sqrt {d})
$$

The first-order Taylor approximation at reference query $\hat{q}_r$ gives a tangent $t_r(q) = q^{\top}\hat{k}_r/\sqrt{d} + \hat{b}_r$, where:

$$
\begin{array}{r l} & {\hat {k} _ {r} = \mathrm{Attn} (\hat {q} _ {r}, K _ {C}, K _ {C}) = \sum_ {j \in [ C ]} \frac {\exp (\hat {q} _ {r} ^ {\top} k _ {j} / \sqrt {d}) k _ {j}}{\sum_ {j ^ {\prime} \in [ C ]} \exp (\hat {q} _ {r} ^ {\top} k _ {j ^ {\prime}} / \sqrt {d})},} \\ & {\hat {b} _ {r} = F _ {C} (\hat {q} _ {r}) - \hat {q} _ {r} ^ {\top} \hat {k} _ {r} / \sqrt {d}.} \end{array}
$$

The chunk's LSE for a future query is approximated by $\max_{r \in [R]} t_r(q)$ over $R$ reference queries. Higher-level nodes (spanning $C^{\ell}$ original keys) are constructed recursively from their children's summaries.

### Training

The **only learnable parameters** are the reference queries. Training uses a KL-divergence loss matching predicted chunk probabilities to true attention distributions at every tree level:

$$
\mathcal {L} _ {\mathrm{KL}} = \sum_ {\ell = 1} ^ {L} \mathbb {E} _ {q} [ D _ {\mathrm{KL}} (p _ {\ell} (q) \| \hat {p} _ {\ell} (q)) ]
$$

where $p_{\ell,\mathcal{C}}(q)$ is the target chunk probability and $\hat{p}_{\ell,\mathcal{C}}(q)$ is the predicted probability from the router's approximate log attention mass. Training requires only **1 billion tokens** with the backbone frozen.

### Routing

At each level, ZHA selects the **top-p nodes**—the smallest prefix of nodes (ordered by score) whose cumulative estimated attention mass reaches fraction $p$:

$$
m _ {\ell} (q) = \min \left\{m \in [ n _ {\ell} ]: \frac {\sum_ {i = 1} ^ {m} \exp s _ {\ell , \mathcal {C} _ {i}} (q)}{\sum_ {i = 1} ^ {n _ {\ell}} \exp s _ {\ell , \mathcal {C} _ {i}} (q)} \geq p \right\}
$$

Selected nodes' children are passed to the next level, and the process repeats until leaf chunks. The final token indices are:

$$
\mathcal {I} (q) = \{0 \} \cup \left(\bigcup_ {\mathcal {C} \in \mathcal {S} _ {1} (q)} \mathcal {C}\right) \cup \mathcal {C} _ {\text { cur }} (q)
$$

with standard attention computed over this sparse subset: $o(q) = \mathrm{Attn}(q, K_{\mathcal{I}(q)}, V_{\mathcal{I}(q)})$.

## Empirical Validation / Results

### Reasoning and Retrieval Performance

ZHA was evaluated on Qwen3.6-35B-A3B (converted with both deep C4-L6 and wide C256-L1 geometries, R16 reference queries, $p = 0.95$):

| Benchmark | Dense | ZHA Deep (C4-L6) | ZHA Wide (C256-L1) | MSA |
|---|---|---|---|---|
| AIME 2026 | ~baseline | ~baseline | ~baseline | ~baseline |
| MMLU-Pro | ~baseline | ~baseline | ~baseline | ~baseline |
| GPQA Diamond | ~baseline | ~baseline | ~baseline | ~baseline |
| RULER-13 | 96.0% | 93.3% | ~96% | Significantly lower |
| LongBench v2 | ~baseline | **12.7 pts higher than MSA** | ~baseline | Lower |

Key findings:
- Both ZHA geometries perform **within noise of dense** on reasoning and coding tasks.
- The deep router attends to only **1–6% of the KV cache** with just a 3% RULER drop.
- ZHA **outperforms MSA** on long-context retrieval despite MSA being trained on 40× more tokens.

### Robustness to Geometry

- **Depth at fixed compression** ($C^L = 4096$): Deeper geometries sustain better retrieval at higher sparsity levels, demonstrating depth itself is not the limiting factor.
- **Reference count R**: Increasing $R$ improves sparsity-performance tradeoff.
- **Predictive model**: A simple 4-parameter model predicts RULER accuracy from $(R, C, L, p)$:

$$
\begin{array}{c} {d _ {\ell} (R, C) = \log \biggl (1 + \frac {C ^ {\ell}}{R ^ {r}} \biggr),} \\ {\widehat {A} = A _ {\mathrm{dense}} (N) \prod_ {\ell = 1} ^ {L} \exp \left[ - a (1 - p) ^ {b} d _ {\ell} ^ {h} \right].} \end{array}\tag{2}
$$

This model achieves **RMSE of 1.44 percentage points** in leave-one-length-out cross-validation.

### Variable Sparsity Across Tasks

ZHA's top-p routing adapts sparsity to the task:
- **RULER** (narrow lookup): 0.44% of chunks retained
- **SWE-bench** (agentic coding): 5.0% of chunks retained
- **Fiction books** (diffuse): 7.4% of chunks retained

### Latency Scaling

The roofline model predicts per-token decode latency:

$$
\begin{array}{l} \boldsymbol {t} _ {\mathrm{ZHA}} (T) = t _ {\mathrm{tree}} (T) + \lambda_ {A} + \boldsymbol {\Phi} _ {\mathrm{attn}} (T) \\ \qquad + \sum_ {\ell = 1} ^ {L (T)} \left[ \lambda_ {\ell} + \boldsymbol {\Phi} _ {\mathrm{routing}, \ell} (T) + \boldsymbol {\Phi} _ {\mathrm{scan}, \ell} (T) \right] \end{array}\tag{3}
$$

with each kernel following a three-term roofline model (compute, memory bandwidth, worker waves):

$$
\boldsymbol {\Phi} _ {x} (T) = \max \left\{ \begin{array}{l} \mathbf {F L O P s} _ {x} (T) / P _ {x}, \\ \mathbf {b y t e s} _ {x} (T) / \mathrm{BW} _ {x}, \\ \boldsymbol {t} _ {x, \text { worker waves }} (T) \end{array} \right\}.\tag{4}
$$

**Key predictions** (R1, C8, $p = 0.75$):
- Faster than dense attention: **~48K tokens** (L1)
- Faster than MSA: **~0.48M tokens** (L2)
- At $10^9$ tokens: **373–417× speedup over dense**, 52–58× over MSA (L5)

Per-level selection counts follow power laws with exponents $\widehat{S}_\ell(T) = c_\ell T^{\alpha_\ell}$, which are modeled cumulatively from the top down:

$$
\begin{array}{r l} & {\varepsilon_ {\ell} = 2 \varepsilon (1 - 2 ^ {- (L - \ell)}),} \\ & {\beta_ {\ell} = \beta + \varepsilon_ {\ell}, \qquad \varepsilon_ {L} = 0,} \\ & {\widehat {\alpha} _ {\ell} = \prod_ {k = \ell} ^ {L} \beta_ {k}, \qquad \widehat {\alpha} _ {L} = \beta .} \end{array}\tag{7}
$$

## Theoretical and Practical Implications

### Theoretical Significance

1. **Subquadratic scaling with query-dependent sparsity**: ZHA satisfies both requirements for an effective attention replacement—query-dependent retrieval AND sublinear active state growth—while retaining the full KV cache.

2. **Power-law structure of attention**: The finding that top-p set size grows as a sublinear power law suggests natural language has scale-invariant, self-similar structure that hierarchical indexers can exploit.

3. **Depth robustness**: ZHA is the first hierarchical attention variant to demonstrate convincing scaling to depths > 1 while maintaining parity with dense attention on challenging tasks.

4. **Positional encoding dependence**: The sublinear power-law structure depends on positional encoding (partial RoPE works well; NoPE shows larger retrieval losses), suggesting architectural interactions worth further study.

### Practical Implications

1. **Cheap adaptation**: ZHA can be applied to existing pretrained models with only 1B tokens of KL distillation and negligible parameter overhead (reference queries only), without finetuning the base model.

2. **Extreme context scaling potential**: The performance model suggests ZHA could extend effective context windows to **hundreds of millions or billions of tokens**, where it dramatically outperforms both dense attention and sparse attention baselines.

3. **Adaptive sparsity**: Top-p routing automatically adjusts retrieval budget based on task demands—narrow for targeted lookups, broad for diffuse aggregation—impossible with fixed top-K methods.

## Conclusion

ZHA introduces a principled approach to subquadratic attention via hierarchical tree-structured indexing. Key takeaways:

- **Sublinear state growth** matches the natural compressibility of attention representations, validated by power-law top-p scaling.
- **Cheap conversion** of existing models (1B tokens, frozen backbone) with maintained performance on reasoning, coding, and retrieval.
- **Robust to geometry**, including deep trees up to 6 levels, with predictable accuracy-sparsity tradeoffs.

### Future Directions

1. **Empirical validation at ultra-long contexts**: The primary limitation is lack of validation beyond 256K contexts; the authors plan to scale to frontier open-source models at 1M+ tokens.
2. **Dynamic chunk boundaries**: Learning chunk boundaries instead of fixed contiguous chunks.
3. **Recency bias in tree structure**: Allocating more detail near the current token.
4. **Compressed value summaries**: Allowing pruned subtrees to contribute coarse attention outputs.
5. **Dynamic top-p scaling with depth**: The critical challenge is scaling $p$ per-level to track the decrease in predictive information: fixing $p < 1$ causes $p_{\text{eff}} = p^L \to 0$ (state asymptotes to constant), while rescaling to $p^{1/L}$ causes $p \to 1$ (returns to quadratic). The optimal scaling depends on the intrinsic statistical structure of ultra-long context data, which remains an open question.

---

_Markdown view of https://picx.dev/p/i11Eut, served by PicX — AI-generated visual whiteboard summaries of research papers._