Hierarchical attention enables faster processing of very long texts
Subquadratic Attention with a Hierarchical Tree-Structured Indexer
Computation and LanguageArtificial Intelligence
Summary
Large language models often struggle to handle very long texts because the usual way they pay attention to words gets slower as the text gets longer. The authors introduce a new method called Zyphra Hierarchical Attention that organizes information in a tree-like structure to find important parts quickly, making the process faster than before. This approach can work with existing models after some short additional training and still performs well on tasks that require understanding long contexts, math problems, and coding. Their method also works well even when the tree is very deep and could help models process texts with millions or even billions of words.
What this means in practice
- •For large language model engineers: Extend context windows of LLMs efficiently to handle extremely long inputs beyond millions of tokens with manageable computation cost.
- •For cloud ai service providers: Deploy large models supporting longer user inputs and queries for better reasoning or coding assistance in cloud platforms.$Commercial implications: Enables cloud services to offer advanced AI models processing ultra-long documents for enterprise customers, differentiating products.
Authors
Leon Lufkin, Cutter Dawes, Alexander Speicher, Kamesh Krishnamurthy, Tomás Figliolia, Beren Millidge
Abstract
The quadratic scaling of attention is a fundamental bottleneck limiting the long-context capabilities of large language models (LLMs). A variety of approaches have attempted to circumvent this, but they either fail to match attention's expressiveness (e.g., sliding-window or linear attention variants with fixed-size state) or only achieve constant-factor speedups (e.g., sparse attention methods whose indexers themselves scale quadratically). In this paper, we derive an asymptotically subquadratic sparse attention method based on a hierarchical, tree-structured indexer, which we call Zyphra Hierarchical Attention (ZHA). ZHA constructs a tree of recursively compressed key summaries over the KV cache, which future queries traverse from the top down to retrieve a sparse subset of keys. Crucially, by iteratively growing the tree's depth, the total cost including routing and final attention grows sublinearly in sequence length, resulting in overall subquadratic (but superlinear) scaling for attention. We demonstrate that ZHA can be cheaply applied to existing pretrained models with only a short KL distillation training stage over a few billion tokens, and maintains performance on long-context retrieval and challenging reasoning, mathematical, and agentic coding evaluation tasks. Moreover, we show that our method is highly robust to tree geometry including very deep trees of up to 6 levels. Finally, we develop a detailed performance model and show that our approach scales subquadratically, dramatically outperforms competing methods beyond 1 million context, and ultimately has the potential to extend effective context windows to the hundreds of millions or billions of tokens.