Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Attention via Pre-Scoring: Prioritizing Informative Keys in Transformers

Zhexiang Li, Haoyu Wang, Yutong Bao, David Woodruff

TL;DR

This paper tackles the quadratic cost of transformer self-attention by introducing a clustering- or leverage-based pre-scoring step that prioritizes informative keys before applying HyperAttention. By selecting a small, representative set of keys (often $k=d{+}1$ centroids) and optionally using a leverage-score sketch, the method preserves heavy-attention signals while reducing computation, achieving up to 20× speedups over FlashAttention in long-context scenarios and substantial perplexity improvements on ChatGLM2/GLM3 and competitive accuracy in Vision Transformers. The authors provide theoretical guarantees under a planted-subspace model, showing that clustering-based pre-scoring can recover all heavy keys with high probability, and extend the approach to $\,\ell_p$-norms via Minkowski-$k$-means. Empirical results across language and vision tasks demonstrate robust improvements in efficiency-accuracy trade-offs, supported by reproducible experiments and open-source code.

Abstract

Recent advances in transformer architectures deeply enhanced long-context language modeling. Among them, HyperAttention achieves competitive efficiency by combining a single-level LSH-based clustering with uniform residual sampling. However, HyperAttention fails to find all significant keys, which in turn raises the overall perplexity. We propose a pre-scoring mechanism that prioritizes significant keys before applying HyperAttention. We introduce three scoring methods: $k$-means and kernel $k$-means clustering, $k$-median clustering, and leverage score-based ranking (inspired by LevAttention) to filter keys effectively. We further replace HyperAttention's original uniform residual sampling, relying exclusively on our pre-scoring mechanism. Experiments on ChatGLM2 (131k token context) reduce perplexity from 12 to 8.3, which outperforms standard HyperAttention. Moreover, when running on the Vision-Transformer (ViT), our method shows that it can guarantee similar accuracy compared with LevAttention, and will surpass LevAttention given specific parameters. Although this method introduces some computational overhead, its combination with HyperAttention achieves up to 20 times faster than FlashAttention, providing a balanced trade-off between speed and modeling accuracy. Our results highlight the effectiveness of integrating pre-scoring into hierarchical attention mechanisms, significantly improving transformer efficiency.

Efficient Attention via Pre-Scoring: Prioritizing Informative Keys in Transformers

TL;DR

This paper tackles the quadratic cost of transformer self-attention by introducing a clustering- or leverage-based pre-scoring step that prioritizes informative keys before applying HyperAttention. By selecting a small, representative set of keys (often centroids) and optionally using a leverage-score sketch, the method preserves heavy-attention signals while reducing computation, achieving up to 20× speedups over FlashAttention in long-context scenarios and substantial perplexity improvements on ChatGLM2/GLM3 and competitive accuracy in Vision Transformers. The authors provide theoretical guarantees under a planted-subspace model, showing that clustering-based pre-scoring can recover all heavy keys with high probability, and extend the approach to -norms via Minkowski--means. Empirical results across language and vision tasks demonstrate robust improvements in efficiency-accuracy trade-offs, supported by reproducible experiments and open-source code.

Abstract

Recent advances in transformer architectures deeply enhanced long-context language modeling. Among them, HyperAttention achieves competitive efficiency by combining a single-level LSH-based clustering with uniform residual sampling. However, HyperAttention fails to find all significant keys, which in turn raises the overall perplexity. We propose a pre-scoring mechanism that prioritizes significant keys before applying HyperAttention. We introduce three scoring methods: -means and kernel -means clustering, -median clustering, and leverage score-based ranking (inspired by LevAttention) to filter keys effectively. We further replace HyperAttention's original uniform residual sampling, relying exclusively on our pre-scoring mechanism. Experiments on ChatGLM2 (131k token context) reduce perplexity from 12 to 8.3, which outperforms standard HyperAttention. Moreover, when running on the Vision-Transformer (ViT), our method shows that it can guarantee similar accuracy compared with LevAttention, and will surpass LevAttention given specific parameters. Although this method introduces some computational overhead, its combination with HyperAttention achieves up to 20 times faster than FlashAttention, providing a balanced trade-off between speed and modeling accuracy. Our results highlight the effectiveness of integrating pre-scoring into hierarchical attention mechanisms, significantly improving transformer efficiency.
Paper Structure (35 sections, 5 theorems, 18 equations, 6 figures, 7 tables, 2 algorithms)

This paper contains 35 sections, 5 theorems, 18 equations, 6 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Algorithm
  4. Structural Guarantees
  5. Matrix Structure and Assumptions
  6. Remark
  7. Connection to real key matrices.
  8. Connection to the planted model and the constant gap $D_{\min}$.
  9. Pre-Scored HyperAttention
  10. Experiments and Results
  11. Speed Comparison on Flash Attention
  12. Accuracy Comparison
  13. Results on ChatGLM3-6B-32k.
  14. Monkey Patching Vision-Transformer
  15. Conclusion
  16. ...and 20 more sections

Key Result

Lemma 1

In this model, for each row $i\in S_0$, consider that $||A_i||^2=1$ we have $h_i \le \frac{\|A_i\|^2}{\sigma_{\min}^2} = \frac{1}{\Theta(1/\varepsilon)} = O(\varepsilon).$

Figures (6)

  • Figure 1: Single layer speed test by only forward pass
  • Figure 2: Single layer speed test by forward and backward pass
  • Figure 3: GLM2: The perplexity performance of various key selection strategies—K-means, K-median, and Leverage Score (Lev)—under different values of $k$ (the number of selected keys), where each $k$ is sampled as a power of $2$ (e.g., $128$, $2048$, $8192$).
  • Figure 4: GLM3: perplexity vs. top-$k$ for K-means/K-median/Leverage, with/without residual sampling.
  • Figure 5: K-Means: Median percentage vs. sampled keys.
  • ...and 1 more figures

Theorems & Definitions (11)

  • Lemma 1: Upper Bound on Noise Leverage
  • Lemma 2: Lower Bound on Signal Leverage
  • Theorem 1: Leverage‐Score Separation
  • proof : Proof sketch
  • Theorem 2: K-means Clustering
  • Corollary 1: Singleton case of Theorem 2
  • proof : Proof Sketch
  • Claim 1: $\ell_p$-Generalization
  • proof
  • proof
  • ...and 1 more