[KongchangAI]
· 2 min read· 1,366 words

Latent Sparse Attention Breaks Through the Million-Token Context Bottleneck

Latent Sparse Attention Breaks Through the Million-Token Context Bottleneck

LSA uses hierarchical clustering in latent space to escape the sparse-attention cost-accuracy trap at extreme context lengths.

Sparse attention is key to million-token contexts, but existing indexers are either still quadratic or lose accuracy at extreme lengths. Latent Sparse Attention (LSA) hierarchically clusters keys in a learned latent space, having queries score cluster nodes along a tree instead of individual keys—making indexing and attention sub-quadratic and speeding up prefill, decode, and training. Distilled-LSA injects LSA into existing models; Native-LSA trains natively with next-token loss alone. N-LSA matches full-attention accuracy at 128k tokens with 1.6× faster single-layer training; at 10M tokens on GLM-5.3-FLASH, D-LSA achieves 22.7×/4.6× prefill/decode speedups over sparse attention and 72.4×/42.5× over full attention, with comparable accuracy.

As agentic workflows evolve toward long-horizon tasks, models increasingly need ultra-long contexts. But once context length exceeds a million tokens, the computational cost of full attention becomes untenable. Sparse attention is seen as a promising solution—using a dedicated indexer to select a subset of keys for each query, thereby reducing overhead. Yet existing indexing methods either remain quadratic in context size or lose accuracy at extreme lengths. Latent Sparse Attention (LSA) is proposed precisely to resolve this dilemma with an efficient attention method.

Screenshot of the Latent Sparse Attention paper

The Accuracy-Efficiency Dilemma of Sparse Attention

At the heart of sparse attention is the indexer: it selects, from a massive pool of keys, the subset most relevant to the current query, so that attention computation happens only on that subset. Ideally, the indexer's own cost should be far lower than full attention, while its selection must be accurate enough not to lose critical information.

In reality, current mainstream indexing methods struggle to achieve both at ultra-long contexts. One class of methods reduces attention computation, but the indexing process itself still grows quadratically with context length—at the million-token level, this overhead becomes the bottleneck again. Another class keeps indexing cost down but shows clear accuracy degradation as context lengthens, weakening the model's ability to leverage distant information. The LSA authors take this tension as their design starting point, aiming to push both indexing and attention toward sub-quadratic complexity while preserving accuracy at ultra-long contexts.

Full attention and quadratic complexity: In standard self-attention, every query computes similarity against all keys, at a cost of O(n²d) (n = sequence length, d = head dimension). As n grows from 128k to 1M, the number of attention matrix elements increases by roughly 64×, and memory and compute costs explode accordingly. Sparse attention reduces complexity to O(nkd) by having each query interact with only k keys, but introduces an additional indexing cost—"how to choose these k keys"—which is exactly the core problem LSA aims to solve.

The difference between prefill and decode: The prefill stage processes the entire prompt at once; it is a large-scale matrix multiplication that is compute-intensive and highly parallel. The decode stage generates token by token, handling only one query per step but repeatedly reading the KV Cache; it is memory-bound and latency-sensitive. Because the two stages have different bottlenecks, the same indexing method often yields significantly different speedups across them—which is why LSA reports two separate sets of numbers.

How LSA Sidesteps Per-Key Scoring with Hierarchical Clustering

LSA's core idea is to hierarchically cluster keys in a learned latent space. Instead of scoring every key individually, the query scores the hierarchy formed by the clusters. This way, the query only needs to evaluate a small number of cluster nodes top-down along the hierarchy tree to locate the subset of keys that should participate in attention, avoiding a linear scan over all keys.

The direct benefit shows up in three areas: prefill, decode, and training. At long contexts, all three accelerate thanks to lower indexing cost. Because the cluster structure is learned in latent space, the indexer can adapt to the data distribution at low cost rather than relying on fixed heuristics. The paper emphasizes that this lets LSA maintain accuracy at ultra-long contexts instead of trading accuracy for speed.

Latent space and hierarchical clustering: Latent space refers to a low-dimensional representation space obtained via projection inside the model. LSA first maps keys into this space, then groups them layer by layer using a tree structure (similar to hierarchical k-means or a learnable routing tree): the root node represents all keys, and divisions become finer toward the leaves. The query starts at the root, scores only a small number of candidate clusters at each level and selects a few branches to descend, ultimately reaching the key subset corresponding to leaf nodes. This coarse-to-fine retrieval reduces the number of scoring operations from O(n) to the order of O(log n), and the tree structure can be precomputed and updated as the data distribution is learned, avoiding a linear scan over all keys.

Two Deployment Paths: Distillation and Native Training

To suit different model usage scenarios, LSA offers two variants. Distilled-LSA (D-LSA) targets already-pretrained models, injecting LSA capability into existing models via distillation without training from scratch. This is a pragmatic path for teams that want to upgrade long-context capability at low cost.

Native-LSA (N-LSA) is trained natively using only the next-token loss, without relying on distillation signals. This approach is closer to designing a long-context model from scratch. Experiments show that N-LSA reaches accuracy comparable to full attention at 128k-token contexts while achieving a 1.6× speedup in single-layer training. This result demonstrates that hierarchical indexing is not only an inference-time optimization but can also serve as a native training-time structure.

Knowledge distillation: Distillation lets the output distribution of a well-trained "teacher" model (here, full attention or a stronger model) guide a "student" model (the model with LSA structure injected). D-LSA uses this mechanism to calibrate the new indexer with the teacher's attention behavior without altering the bulk of the original pretrained weights, giving existing models long-context capability at low cost without rerunning expensive pretraining.

Next-token loss: This is the standard autoregressive language modeling objective, requiring only that the model predict the next token in a sequence. N-LSA is trained solely with this loss, meaning the indexer and attention structure are learned end-to-end entirely from the language modeling signal, without additional distillation supervision. It is therefore closer to designing a long-context model from scratch and better tests the learning capability of the structure itself.

Measured Speedups: Significant Gains in Prefill and Decode

The paper reports comparison data at 10M-token contexts on GLM-5.3-FLASH (QWEN-3.8-27B). Compared with its original sparse attention, D-LSA achieves a 22.7× speedup in single-layer prefill and a 4.6× speedup in decode; compared with full attention, prefill speeds up by 72.4× and decode by 42.5×. Alongside these speedups, the model maintains comparable accuracy across a variety of tasks.

What makes these numbers significant is the context scale—10M tokens far exceeds typical evaluation ranges and is exactly the regime where full attention costs are most unbearable. One detail: the speedup factors differ markedly between the two baselines—several-fold to tens-of-fold relative to sparse attention, and tens-of-fold relative to full attention. This itself reflects that the original sparse attention still carries considerable indexing overhead at ultra-long contexts, and LSA compresses that cost further.

Implications for Ultra-Long Context Applications

LSA's value is not limited to speeding up a single attention layer. Tasks in agentic workflows—open-ended scientific research, long-document understanding, multi-turn tool calling—all depend on the model maintaining information coherence across extremely long contexts. If attention layers can run efficiently at million-token scale, model designers gain more room to keep extending context windows without repeatedly compromising between cost and accuracy.

Methodologically, LSA reframes indexing as hierarchical retrieval in latent space, offering an extensible framework for future research: how cluster structures adapt to tasks, how distillation and native training complement each other, and how to trade off hierarchy depth against accuracy are all directions worth exploring. The paper's two variants also show that upgrading long-context capability can proceed along two tracks—serving existing models and supporting native design of new ones.

One-line Summary

LSA uses hierarchical clustering in latent space as an index, resolving the cost-accuracy dilemma of sparse attention at ultra-long contexts.

Paragraph Summary

Sparse attention is seen as a key solution for million-token contexts, but existing indexers either still grow quadratically with context or degrade in accuracy at extreme lengths. Latent Sparse Attention (LSA) hierarchically clusters keys in a learned latent space; queries score cluster nodes along a hierarchy tree rather than every key, pushing both indexing and attention toward sub-quadratic complexity and benefiting prefill, decode, and training alike. The paper offers two variants: Distilled-LSA injects LSA into existing pretrained models via distillation, while Native-LSA trains natively using only next-token loss. At 128k context, N-LSA matches full-attention accuracy while speeding up single-layer training by 1.6×; on GLM-5.3-FLASH at 10M tokens, D-LSA speeds up prefill by 22.7× and decode by 4.6× versus the original sparse attention, and by 72.4× and 42.5× versus full attention, with accuracy essentially on par.

#Latent Sparse Attention#LSA#sparse attention#long context#hierarchical clustering#knowledge distillation#million-token context
Share:

Related articles