Zyphra's ZHA Tree Router Scans Just 2.8% of Keys at 256K Tokens

Dense self-attention compares each query with every available key. For a sequence of n tokens, that produces O(n²) score calculations during training or prefill. During autoregressive decoding, each new token scans an O(n) key-value (KV) cache. That scaling makes million-token contexts impractical. Zyphra’s paper, Subquadratic Attention with a Hierarchical Tree-Structured Indexer, proposes a sparse router that organizes the KV cache as a tree and searches only selected branches.

Zyphra Hierarchical Attention, or ZHA, retains the original KV cache and adds recursively compressed key summaries above it. Future queries traverse those summaries from the root to selected leaves, where attention runs over a sparse set of original key-value pairs. Skipping one node prunes every key below it. The reported efficiency gains focus on decode, where repeated cache scans become increasingly expensive.

Why sparse attention can still scan everything

Sliding-window attention limits each token to nearby context. Linear-attention variants and state-space models compress history into fixed-size state, reducing compute but potentially weakening retrieval of specific distant tokens. Sparse attention can preserve the full cache, yet many sparse selectors still score every stored key before discarding most of them. The resulting attention matrix is sparse, but the selector retains an O(n) scan for each decoded token.

ZHA replaces that scan with hierarchical routing. Each query begins at the root and scores only the children of selected nodes. Total work depends on the tree depth and the number of branches expanded. Depth alone does not guarantee subquadratic scaling; the paper’s argument also relies on the selected set growing sublinearly as context expands.