Vollständige Aufmerksamkeit schlägt zurück: Überführung von vollständiger Aufmerksamkeit in sparse Aufmerksamkeit innerhalb von hundert Trainingsschritten

Zusammenfassung

Die Inferenz bei langen Kontexten in großen Sprachmodellen wird durch die quadratischen Kosten der vollständigen Aufmerksamkeit (Full Attention) zum Engpass. Bestehende effiziente Alternativen basieren häufig entweder auf nativem spärlichem Training oder auf heuristischem Token-Eviction, was einen unerwünschten Kompromiss zwischen Effizienz, Trainingskosten und Genauigkeit schafft. In dieser Arbeit zeigen wir, dass Full-Attention-LLMs bereits intrinsisch spärlich sind und mit nur minimaler Anpassung in hochspärliche Modelle transformiert werden können. Unser Ansatz beruht auf drei Beobachtungen: (1) nur eine kleine Teilmenge der Attention-Heads benötigt tatsächlich eine vollständige Verarbeitung langer Kontexte; (2) der Abruf über große Reichweiten wird hauptsächlich durch einen niedrigdimensionalen Unterraum bestimmt, sodass relevante Token effizient mit einem 16-dimensionalen Indexer abgerufen werden können; (3) das nutzbare Token-Budget ist stark abfrageabhängig, was eine dynamische Top-p-Auswahl besser geeignet macht als eine feste Top-k-Verschlankung. Basierend auf diesen Erkenntnissen schlagen wir RTPurbo vor, das den vollständigen KV-Cache nur für Retrieval-Heads beibehält und einen leichten Token-Indexer für spärliche Aufmerksamkeit einführt. Durch die Ausnutzung der intrinsischen Spärlichkeit des Modells erreicht RTPurbo eine Verschlankung mit nur wenigen hundert Trainingsschritten. Experimente mit Long-Context-Benchmarks und Reasoning-Aufgaben zeigen, dass RTPurbo eine nahezu verlustfreie Genauigkeit beibehält und gleichzeitig erhebliche Effizienzsteigerungen erzielt, darunter eine bis zu 9,36-fache Beschleunigung der Prefill-Phase bei einem Kontext von 1 Mio. Token und eine etwa 2,01-fache Beschleunigung der Decode-Phase. Diese Ergebnisse legen nahe, dass eine starke spärliche Inferenz durch standardmäßiges Full-Attention-Training ohne teures natives spärliches Vortraining erreicht werden kann.

English

Long-context inference in large language models is bottlenecked by the quadratic cost of full attention. Existing efficient alternatives often rely either on native sparse training or on heuristic token eviction, creating an undesirable trade-off among efficiency, training cost, and accuracy. In this work, we show that full-attention LLMs are already intrinsically sparse and can be transformed into highly sparse models with only minimal adaptation. Our approach is built on three observations: (1) only a small subset of attention heads truly requires full long-context processing; (2) long-range retrieval is governed primarily by a low-dimensional subspace, allowing relevant tokens to be retrieved efficiently with a 16-dimensional indexer; and (3) the useful token budget is strongly query-dependent, making dynamic top-p selection more suitable than fixed top-k sparsification. Based on these insights, we propose RTPurbo, which retains the full KV cache only for retrieval heads and introduces a lightweight token indexer for sparse attention. By exploiting the model's intrinsic sparsity, RTPurbo achieves sparsification with only a few hundred training steps. Experiments on long-context benchmarks and reasoning tasks show that RTPurbo preserves near-lossless accuracy while delivering substantial efficiency gains, including up to a 9.36times prefill speedup at 1M context and about a 2.01times decode speedup. These results suggest that strong sparse inference can be obtained from standard full-attention training without expensive native sparse pretraining.