Zyphra adds tree-based routing to Hierarchical Attention for sparse retrieval

• ZHA makes token selection sparse too, routing queries through the learned summary tree in Zyphra’s announcement.
• Existing models can be adapted with a frozen backbone: Zyphra’s distillation results report just 81,920 added trainable parameters for Qwen3.6-35B-A3B.
• A six-level router matched dense attention on SWE-bench and exceeded it on LongBench v2, according to Zyphra’s evaluations.
• The headline speedups are projections: Zyphra’s roofline model estimates roughly 400× faster attention than dense at one billion tokens, rather than measured serving gains.

The paper builds summaries with learned reference queries and constructs its index as tokens arrive. Across 28 tested tree geometries, deeper trees often retrieved better at the same compression.

Many sparse-attention methods still scan every candidate with a lightweight indexer, as ZyphraAI explains. Pruning entire subtrees before their keys are scored is ZHA’s best trick.

Its retrieval path has three stages:

1. Group contiguous keys into chunks, then recursively summarize those chunks into higher tree levels.
2. Compare the query against the highest-level summaries and expand promising branches, potentially in widely separated regions.
3. Descend to the original KV entries and perform ordinary attention over the selected subset.

At each level, top-p selects the smallest set of nodes reaching a target estimated attention mass. Precise lookups can explore few branches; diffuse queries can expand many, without a fixed top-k budget.

The subquadratic claim depends on increasing tree depth with context and keeping routing plus final attention work sublinear per query. In Zyphra’s measurements, dense attention’s top-p sets grew sublinearly on natural-language text.

The paper’s analysis explicitly allows quadratic behavior for adversarial random-token retrieval. Zyphra also hypothesizes that suppressing irrelevant attention could preserve signal at ultra-long contexts.

ZHA’s adaptation recipe, described in the official write-up, separates the language model from its new indexer:

• Backbone: all original model weights remain frozen.
• Trainable components: reference queries, totaling 81,920 parameters for Qwen3.6-35B-A3B, approximately 0.00024% of the model.
• Objective: KL distillation teaches the indexer to reproduce dense attention’s probability mass across regions at each tree level.
• Data: a few billion tokens, trained at 32K context length.

Zyphra’s Qwen evaluations compare the dense model with a wide, single-level router and a deep, six-level router:

• RULER-13 at 256K: dense 96.0% → deep 93.3%, −2.7 points, while attending to 2.8% of decode keys. The wide router scored 96.2%.
• LongBench v2: dense 58.0% → deep 63.4%, +5.4 points, while attending to approximately 1.1% of decode keys.
• SWE-bench Verified: dense 67.4% → deep 67.4%, unchanged. The wide router scored 68.4%.

The same write-up reports more noticeable RULER degradation on Nemotron-3.5-Lightning-30B-A3B, whose attention layers use no positional encoding.

Zyphra’s GPU roofline model compares attention latency against dense attention and MiniMax Sparse Attention (MSA), accounting for compute, memory bandwidth, kernel launches, and hardware parallelism.

For a lightweight ZHA configuration, it predicts:

• Crossover points: approximately 48K tokens against dense attention and 480K against MSA.
• 10 million tokens: 32–35× faster than dense and 4.7–5.1× faster than MSA.
• One billion tokens: 373–417× faster than dense and 52–58× faster than MSA.

Empirical quality evaluation at million-token contexts and beyond remains future work in the paper.

ZHA retains the complete KV cache plus the summary tree, leaving storage proportional to context length. In the paper’s cache-coverage experiment, the union of top-p key sets across later queries in 256K fiction sequences expanded toward the entire cache, even though individual queries were sparse.