Scaling Attention via Feature Sparsity

Yan Xie; Tiansheng Wen; Tangda Huang; Bo Chen; Chenyu You; Stefanie Jegelka; Yifei Wang

cs.LGMar 17, 2026

Scaling Attention via Feature Sparsity

Authors: Yan Xie, Tiansheng Wen, Tangda Huang, Bo Chen, Chenyu You, Stefanie Jegelka, Yifei Wang

Organizations: School of Electronic Engineering, Xidian University · Stony Brook University · TUM · MIT · Amazon AGI SF Lab

Abstract

Scaling Transformers to ultra-long contexts is bottlenecked by the $O(n^2 d)$ cost of self-attention. Existing methods reduce this cost along the sequence axis through local windows, kernel approximations, or token-level sparsity, but these approaches consistently degrade accuracy. In this paper, we instead explore an orthogonal axis: feature sparsity. We propose Sparse Feature Attention (SFA), where queries and keys are represented as $k$ -sparse codes that preserve high-dimensional expressivity while reducing the cost of attention from $Θ(n^2 d)$ to $Θ(n^2 k^2/d)$ . To make this efficient at scale, we introduce FlashSFA, an IO-aware kernel that extends FlashAttention to operate directly on sparse overlaps without materializing dense score matrices. Across GPT-2 and Qwen3 pretraining, SFA matches dense baselines while improving speed by up to $2.5\times$ and reducing FLOPs and KV-cache by nearly 50%. On synthetic and downstream benchmarks, SFA preserves retrieval accuracy and robustness at long contexts, outperforming short-embedding baselines that collapse feature diversity. These results establish feature-level sparsity as a complementary and underexplored axis for efficient attention, enabling Transformers to scale to orders-of-magnitude longer contexts with minimal quality loss. Code is available at https://github.com/YannX1e/Sparse-Feature-Attention.

Figures & tables

Figure 1: Overview of our proposed method. (a) Trade-off between performance and speed. Compared to directly reducing dimensionality with short embeddings , our method achieves a more favorable balance, delivering a 259% speedup over the original dimensionality while improving performance by 21.4% relative to the short-embedding baseline. (b) Computational and memory efficiency comparison. Our method reduces KV-cache memory usage by 41% and FLOPs by 49%.

Figure 2: Three paradigms of attention. Left: Standard attention computes all

N\times N

query–key interactions in the full feature dimension

d

. Middle: Sparse attention reduces cost by selecting, for each query

i

, a small subset of keys

\Omega_{i}

and masking the remaining logits before softmax, but each retained interaction still spans all

d

features. Right: Sparse Feature Attention (ours) keeps all tokens but sparsifies along the feature axis by selecting the top-

k

channels in

Q

and

K

(

\tilde{Q}=\mathrm{Topk}_{k}(Q),\tilde{K}=\mathrm{Topk}_{k}(K)

). Attention is then computed only over overlapping selected features with sparse matrix multiplication. This shifts sparsity from the token axis (

N\times N

) to the feature axis, achieving efficiency while preserving token coverage.

Table 2: Long context pretraining results. Comparison of NIAH accuracy rates for different lengths under various training lengths. (a) Models are trained on 8k synthetic NIAH data, and the accuracy rate on test lengths from 1k to 8k. (b) Models are trained on 32k synthetic NIAH data, and the accuracy rate on test lengths from 1k to 32k.

Panel of Figure 3 — Figure 3: Latency vs. feature sparsity. Latency Comparison of dense attention and SFA (ours) at different modular levels in Transformers under 16k context length. Higher sparsity brings substantial decrease in latency.

Panel of Figure 4 — Figure 4: Latency vs. feature sparsity with various config. Latency Comparison of dense attention and SFA (ours) at different head dimensions and context lengths. Notably, the latency of SFA can be much lower than dense attention under high dimension per head and long context, e.g., Figure 4(d) .

Panel of Figure 5 — Figure 5: Scaling dense attention and SFA with context length. SFA can consistently reduce both the computatin cost and KV cache size by a constant factor of at least $2$ .

Model	Variant	General Tasks			NIAH Acc
Model	Variant	GSM-8K	Arxiv	PubMed	4096	8192	16384	32768
Qwen3-0.6B	Base	59.59	13.65	10.48	90	87	77	52
	Fine-tune	63.42	41.17	40.54	94	92	79	55
	SFA ( $k=16$ )	61.46	39.14	39.03	95	93	77	53
Qwen3-4B	Base	75.44	31.52	29.19	97	95	90	81
	Fine-tune	76.18	49.31	49.05	99	96	92	84

Table 3: Evaluation on general reasoning tasks and synthetic retrieval (NIAH). Accuracy is in %.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Size	#Parameters	hidden_size	num_layers	num_heads	short_hidden	position_embedding
Small	124M	768	12	12	384	APE
Medium	350M	1024	24	16	512	APE
Large	596M	1024	28	16/8	512	RoPE

Appendix

Table 4: Base model configurations. “Short” refers to halving the hidden size for

Q/K

Model	Dataset	Epoch	Batch_Size	Lr	Warmup_Ratio	Gradient_Ckecpointing
Qwen3-0.6B	GSM8K	3	256	6e-4	0.1	False
	Arxiv	2	256	1e-5	0.05	False
	PubMed	2	256	2e-5	0.05	False
	NIAH	3	256	2e-5	0.05	False
Qwen3-4B	GSM8K	3	256	6e-6	0.1	True
	Arxiv	2	256	2e-6	0.1	True

Appendix

Table 5: Configurations for fine-tuning MoE models

Panel of Figure 6 — Figure 6: Comparison of latencies in different attention scenarios: (a) full attention (Time-to-first-token) and (b) kv-cache attention (Time-to-next-token) .

	8192		16384		32768		65536
Config	TFLOPs	INOPs	TFLOPs	INOPs	TFLOPs	INOPs	TFLOPs	INOPs
Dense_128	$2.23$	/	$8.92$	/	$35.67$	/	$142.67$	/
Sparse_32/128	$1.20$	$28.31$	$4.79$	$58.72$	$19.17$	$121.63$	$76.70$	$251.67$
Sparse_16/128	$1.15$	$18.87$	$4.59$	$39.85$	$18.35$	$83.89$	$73.40$	$176.17$
Sparse_8/128	$1.13$	$13.63$	$4.54$	$29.36$	$18.14$	$62.91$	$72.57$	$134.22$
Dense_64	$1.12$	/	$4.48$	/	$17.94$	/	$71.75$	/

Appendix

Table 6: Operation counts for standard flash attention and flash attention sparse. The number of floating-point operations (FLOPs) and integer operations (INOPs) were counted separately for feature dimensions of 64 and 128 under different context lengths.

Kernel	Dense	Dense w/o compute	FlashSFA	FlashSFA w/o compute
HBM Bandwidth (GB/s)	14.22	1194.34	17.14	919.38

Appendix

Table 7: HBM Bandwidth Comparison. “w/o compute” denotes measuring memory throughput with computation logic disabled.

Context Length ( $N$ )	1024	4096	8192	16384	32768	65536
torch.topk	0.730	2.701	5.336	10.684	21.205	42.374
RTop- $k$ (Ours)	0.221	0.589	1.089	2.080	4.057	8.080
Ratio of RTop- $k$ (%)	10.51	2.10	1.96	1.90	1.03	0.51

Appendix

Table 8: Latency comparison (ms) between standard torch.topk and RTop-

k

kernel across different context lengths. The Ratio indicates the percentage of time RTop-

k

consumes relative to the total attention forward pass.

	Context Length
Variant	1024	4096	8192	16384	32768	65536
Dense_256	$10.981$	$176.202$	$712.983$	$2894.459$	$11\,772.470$	$49\,197.700$
Sparse_32/256	$3.209$	$45.893$	$180.375$	$715.016$	$2886.208$	$11\,529.740$
Sparse_24/256	$2.771$	$33.959$	$154.789$	$612.812$	$2488.042$	$8309.513$
Sparse_16/256	$2.307$	$29.149$	$128.523$	$510.183$	$2079.824$	$7388.784$
Sparse_12/256	$1.969$	$21.284$	$109.036$	$431.497$	$1769.727$	$5063.700$

Appendix

Table 9: Latency (ms) versus context length.

	Token-Level Operation
Model	Variant	Latency@128k $\downarrow$		PPL $\downarrow$	Acc $\uparrow$
Model	Variant	Decode	Forward	OWT/Pile	PiQA	LAMBDA	ARC-e	ARC-c	HellaS	Avg
GPT2 124M	Dense (full)	17.08	16.86	17.29	56.34	22.78	28.35	14.32	19.61	28.28
	Longformer	6.75	7.93	18.73	54.25	21.27	28.02	13.01	18.92	28.10
	+SFA ( $k=8$ )	5.23	6.18	19.30	52.81	20.54	26.39	12.59	17.24	25.91
	Feature-Level Operation

Appendix

Table 10: Latency, Perplexity and Accuracy results comparison with various compression and acceleration techniques, categorized into Token-Level and Feature-Level Operations. For token-level operations, "Longforemer" ( Beltagy et al., 2020 ) denotes fixed token sparsity pattern, "NSA" ( Yuan et al., 2025 ) denotes dynamic token sparsity pattern. "Dense (full)" baselines use full hidden size and uncompressed KV cache; "Short (

d=X

)" denotes baselines with half feature dimensions; "Quant" denotes 8-bit quantization aware training (QAT ( Liu et al., 2024b ) ) on weights and activations; "Low-Rank" denotes PCA-based projection matrix fine-tuning; "MLA" denotes multi head latent attention ( Liu et al., 2024a ) , and "MLA + SFA" combines SFA with latent key/value. "Latency@128k" is measured by "Decoding with KV cache (TTNT) (ms)" and "Prefilling with full attention (TTFT) (s)". PPL is evaluated on OpenWebText for GPT-2 and Pile for Qwen3.

Panel of Figure 7 — Figure 7: Entropy of Top-K feature selection across layers and heads. We plot the normalized entropy of the TopK index distribution for each attention head and layer of Qwen3-0.6B when applying SFA. (a) Entropy of the TopK positions of query vectors $Q$ (16 heads, due to GQA). (b) Entropy of the TopK positions of key vectors $K$ (8 heads). Each cell corresponds to one (layer, head) pair, and brighter colors indicate higher entropy (more balanced use of feature dimensions) .

	Token Sparse (Training)
Model		Latency@128k		PPL	Acc
Model		Decode	Forward	OWT/Pile	PiQA	LAMBDA	ARC-e	ARC-c	HellaS	Avg
GPT2 124M	Dense (full)	17.08	16.86	17.29	42.74	22.78	28.35	8.12	19.61	24.32
	SFA	14.12	9.41	18.17	41.62	21.03	28.41	7.39	19.26	23.54
	Routing	7.92	8.37	18.64	41.39	21.08	28.31	7.11	18.89	23.35
	Longformer	6.75	7.93	18.73	41.28	21.27	28.02	7.01	18.92	23.30

Appendix

Table 11: Comparison results with token-sparse, KV-pruning, low-rank, and kernel baselines on GPT-2. Token-sparse training methods include Routing ( Roy et al., 2021 ) and Longformer ( Beltagy et al., 2020 ) ; KV-pruning (training-free) methods include H 2 O ( Zhang et al., 2023 ) , Quest ( Tang et al., 2024 ) , and SnapKV ( Li et al., 2024 ) ; Loki ( Singhania et al., 2024 ) is a low-rank key compression method (training-free); Performer ( Choromanski et al., 2021 ) is a kernel-based approximation. Rows marked “+SFA (

k=8

)” apply our feature-sparse SFA to Longformer and SnapKV, showing that SFA is orthogonal to these approaches and can be combined with them.

Panel of Figure 8 — Figure 8: Ablation of sparsity $k$ on GPT-2 124M with fixed head dimension $d=64$ . Perplexity on OpenWebText (left) and latency at 32K context (right) as a function of the Top- $k$ sparsity level used by SFA. The dashed gray line denotes the dense (full) attention baseline; the red curve shows SFA with different $k$ .

Panel of Figure 9 — Figure 9: Ablation of head dimension $d_{\text{head}}$ on GPT-2 124M with fixed sparsity $k=8$ . Perplexity on OpenWebText (left) and latency at 32K context (right) as a function of the head dimension $d_{\text{head}}$ used by SFA. The dashed gray line denotes the dense (full) attention baseline; the red curve shows SFA with different $d_{\text{head}}$ .

Figure 10: Validation loss curves of SFA on GPT-2 (124M) pre-training. We compare varying sparsity levels

k\in\{2,4,8,16\}

. The curves decrease smoothly and monotonically without divergent spikes, demonstrating that SFA maintains training stability even under aggressive sparsity (

k=2

Context Length	1k	2k	3k	4k	Speedup@4k
Dense(full)	93	87	79	62	1.0x
SFA( $k=8$ )	95	90	80	66	1.5x
SFA( $k=16$ )	96	90	83	71	1.2x

Appendix

Table 12: NIAH accuracy (%) within 4k Context Length. Qwen3-0.6B trained on Pile dataset with 4k window, and the accuracy rate on NIAH test lengths from 1k to 4k.

Panel of Figure 11 — Figure 11: Eigenvalue spectrum analysis for Qwen3-0.6B model. Layer-wise effective dimension of (a) query and (b) key activation with normalized cumulative eigenvalue of 0.9, evaluated on the same sampled subset of the Pile validation set in Appendix F .

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
Jun 22, 2026Huzama Ahmad, Se-Young YunLong-Context Language Model Inference LLM Inference Acceleration
KAIST
Block Sparse Flash Attention
Dec 7, 2025Daniel Ohayon, Itay Lamprecht, Itay Hubara +3GPU Acceleration Self-Attention
Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel
AdaSplash-2: Faster Differentiable Sparse Attention
Apr 16, 2026Nuno Gonçalves, Hugo Pitorro, Vlad Niculae +4GPU Kernel Optimization Self-Attention
Instituto Superior Técnico, Universidade de Lisboa · Instituto de Telecomunicações · Carnegie Mellon University +4

Jun 22, 2026J/K move · Enter open · S save

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Huzama Ahmad, Se-Young Yun

KAIST

Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention Block-Sparse Attention

Open paper AskPDF arXiv

Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.

Abstract

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Block Sparse Flash Attention

AdaSplash-2: Faster Differentiable Sparse Attention