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

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$):

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.