Prism: Spectral-Aware Block-Sparse Attention

Paper: 2602.08426 Authors: Xinghao Wang, Pengyu Wang, Xiaoran Liu, Fangxu Liu, Jason Chu, Kai Song, Xipeng Qiu Categories: cs.CL, cs.LG

Abstract

This paper addresses fundamental inefficiencies in block-sparse attention mechanisms for long-context LLMs. The authors identify that the common approach of mean pooling combined with RoPE (Rotary Position Embeddings) creates “blind spots” for local positional information by filtering out high-frequency signals. Prism solves this by separating block selection into high-frequency and low-frequency components using energy-based temperature calibration. The method is training-free and achieves accuracy parity with full attention while delivering up to 5.1x speedup.

Key Contributions

• Identification of the blind spot problem: Mean pooling over RoPE-encoded representations filters high-frequency positional signals
• Spectral decomposition for block selection: Separating attention block scoring into frequency components
• Energy-based temperature calibration: Adaptive calibration that accounts for frequency-dependent energy distributions
• Training-free deployment: No fine-tuning required, directly applicable to existing models

The Problem: Blind Spots in Block-Sparse Attention

Block-Sparse Attention Background

For long sequences, full attention with $O(n^2)$ complexity becomes prohibitive. Block-sparse attention divides the key-value sequence into blocks and selects only the most relevant blocks to attend to, reducing computation significantly.

The standard approach:

1. Divide keys into blocks of size $B$: $\{K_1, K_2, \ldots, K_{n/B}\}$
2. Compute a representative for each block (typically mean pooling): $\bar{K}_i = \frac{1}{B}\sum_{j \in \text{block}_i} K_j$
3. Score each block: $s_i = Q \cdot \bar{K}_i^T$
4. Select top-$k$ blocks and compute full attention only within those blocks

The RoPE + Mean Pooling Problem

RoPE encodes position by rotating embedding dimensions at different frequencies:

$\text{RoPE}(x, m) = x \odot \cos(m\theta) + \text{rotate}(x) \odot \sin(m\theta)$

where $m$ is the position and $\theta_d = 10000^{-2d/D}$ defines frequency for dimension $d$.

When mean pooling is applied over a block of RoPE-encoded keys, the averaging acts as a low-pass filter. High-frequency components — which encode fine-grained local positional information — cancel out due to destructive interference:

$\bar{K}_i = \frac{1}{B}\sum_{m \in \text{block}_i} \text{RoPE}(K, m) \approx \text{low-freq components only}$

This creates “blind spots” where the block selection mechanism cannot distinguish between blocks that differ primarily in local positional patterns, leading to incorrect block selection and degraded accuracy.

The Prism Solution

Spectral Decomposition

Prism separates the block selection into two parallel pathways:

1. Low-frequency pathway: Captures semantic similarity (content-based relevance)
2. High-frequency pathway: Captures local positional patterns (position-based relevance)

For each query-block pair, two scores are computed:

$s_i^{\text{low}} = Q^{\text{low}} \cdot \bar{K}_i^{\text{low},T}$ $s_i^{\text{high}} = Q^{\text{high}} \cdot \bar{K}_i^{\text{high},T}$

where the superscripts denote frequency-filtered versions of the representations.

Frequency Separation

The dimensions of the RoPE-encoded representations are partitioned based on their rotation frequency $\theta_d$:

• Low-frequency dimensions: $d$ where $\theta_d < \theta_{\text{threshold}}$ (slow rotation, semantic content)
• High-frequency dimensions: $d$ where $\theta_d \geq \theta_{\text{threshold}}$ (fast rotation, positional detail)

This partition is natural because RoPE assigns monotonically decreasing frequencies across dimensions.

Energy-Based Temperature Calibration

The two frequency components have different energy scales, so their scores are not directly comparable. Prism introduces adaptive temperature calibration:

$s_i = \frac{s_i^{\text{low}}}{\tau^{\text{low}}} + \frac{s_i^{\text{high}}}{\tau^{\text{high}}}$

where the temperatures $\tau$ are computed based on the energy distribution of each frequency band:

$\tau^{\text{low}} = \sqrt{\mathbb{E}[\|Q^{\text{low}}\|^2] \cdot \mathbb{E}[\|\bar{K}^{\text{low}}\|^2]}$

$\tau^{\text{high}} = \sqrt{\mathbb{E}[\|Q^{\text{high}}\|^2] \cdot \mathbb{E}[\|\bar{K}^{\text{high}}\|^2]}$

This ensures both components contribute proportionally to the final block selection score.

Experimental Results

Accuracy Parity with Full Attention

Across multiple long-context benchmarks:

• Prism achieves accuracy within 1% of full attention on most tasks
• Significantly outperforms standard block-sparse attention (mean pooling baseline)
• Particularly strong on tasks requiring precise positional retrieval (needle-in-a-haystack)

Speedup

• Up to 5.1x speedup over full attention on long sequences
• Overhead of spectral decomposition is negligible compared to attention computation savings
• Speedup scales with sequence length as the sparsity ratio increases

Comparison with Other Sparse Attention Methods

Prism outperforms existing block-sparse methods that use:

• Simple mean pooling (the blind spot problem)
• Learned block selection (requires training, less generalizable)
• Fixed sparsity patterns (cannot adapt to content)

Ablation Studies

• Frequency threshold: The optimal split point balances semantic and positional information
• Temperature calibration: Removing calibration significantly degrades performance, confirming the importance of energy normalization
• Block size: Prism’s advantage is most pronounced with larger blocks where mean pooling loses more high-frequency information

Analysis

Why High Frequencies Matter

High-frequency RoPE components encode the precise position within a local window. For tasks like:

• Retrieval: Finding a specific piece of information requires knowing exactly where it is
• Coreference: Resolving references depends on relative position
• Recency: Distinguishing recent from distant context requires positional precision

Without high-frequency signals, block selection becomes “position-blind” within local regions, selecting blocks based only on semantic similarity.

Visualization of Block Selection

The authors visualize block selection patterns showing:

• Standard mean pooling misses blocks containing positionally relevant information
• Prism correctly identifies these blocks through the high-frequency pathway
• The two pathways are complementary, each catching blocks the other misses

Significance

Prism elegantly solves a fundamental problem in efficient long-context attention. By identifying the spectral blind spot created by mean pooling over RoPE representations and providing a principled decomposition solution, it achieves the best of both worlds: the efficiency of block-sparse attention with the accuracy of full attention. The training-free nature makes it immediately deployable, and the theoretical grounding in signal processing provides a clear understanding of why it works.