SANTA++: Sampling Attention through Representative Keys

Kyle Lee; Christian Z. Pratt; Ruoyu Fang; Heekyung Lee; Avinash Lohitsa; Ryan Modafe; Kerem Y. Camsari

cs.LGSep 28, 2026

SANTA++: Sampling Attention through Representative Keys

Authors: Kyle Lee, Christian Z. Pratt, Ruoyu Fang, Heekyung Lee, Avinash Lohitsa, Ryan Modafe, Kerem Y. Camsari

Organizations: Department of Electrical and Computer Engineering University of California, Santa Barbara, CA, USA · Flucta, San Francisco, CA, USA

Abstract

Attention often concentrates on a small subset of tokens in the context, but which subset matters changes from one query to the next. To exploit this changing structure, we introduce SANTA++, a training-free stochastic attention method that uses representative keys for memory-efficient selection without scanning the entire key-value (KV) cache. Cached keys are organized into teams, and the query scores one representative from each team to decide which teams to sample. We compute exact attention scores within the sampled teams and reweight each team's contribution by the inverse of its inclusion probability. This importance sampling correction estimates attention over the full cache, with a sampling budget that lets us trade memory reads for accuracy. Remarkably, with 32 or 64 sampled teams, SANTA++ uses 16% to 22% of dense attention's KV reads and retains 94% to 99% of the dense-attention baseline's scores on LongBench v2 and HELMET's retrieval-augmented generation subset, and 85% to 91% on RULER, with Qwen2.5-7B-Instruct at 32K context. With 31 sampled teams, our GPU implementation delivers a $1.69\times$ attention speedup over the dense FlashAttention baseline at 32K context. By reducing the number of cache entries read, SANTA++ in principle complements architectures with compressed KV representations, such as multi-head latent attention. Our kernels are available at: https://github.com/OPUSLab/santapp-kernel-demo.git.

Figures & tables

Figure 1: Representative-guided attention sampling. After prefill, parent groups cover the prompt, and actual member keys represent the teams built within each parent. For each decoding query, we form a routing weight for each team from its representative score and team size, and use these weights to sample teams

\mathcal{S}

. Selected teams’ keys and values are read in full, and their members’ exact attention contributions are computed. Each team’s unnormalized attention mass and weighted value sum are divided by its conditional inclusion probability

c_{g}

before the corrected sums are combined and normalized (Section 2 ).

Figure 2: A mean-key mass surrogate for 16 constructed logits. One high score dominates the exact mass, which is approximately

760\times

the surrogate.

Task	Teams (k-means) $P=16,R=4$	Centroid-guided sampling $B=4$	Centroid-guided sampling $B=16$
CWE	1.621	2.806	3.429
S-NIAH-1	1.248	2.076	2.827
QA-1	1.613	2.429	3.495

Table 1: Routing approximation to dense attention. Mean

\mathrm{KL}(p\|q)

in nats, where

p

is dense attention over prompt tokens and

q

spreads each group’s routing probability equally across its members. All variants use the same dense SDPA queries and cache states. No sampling is performed in this diagnostic. Appendix B.2 gives the definitions and averaging procedure.

Method	$S$	KV (%)	Score (norm) (%)	Score (abs) (%)
SDPA	–	100.00	100.00	57.40 $\pm$ 2.00
SANTA	128	50.43	99.13	56.90 $\pm$ 2.05
	256	50.75	99.04	56.85 $\pm$ 2.05
	512	51.30	100.17	57.50 $\pm$ 2.02
	1024	52.17	99.48	57.10 $\pm$ 2.07
	2048	53.53	99.65	57.20 $\pm$ 2.08

Table 2: HELMET RAG: full 32K budget sweep with Qwen2.5-7B-Instruct. Score (abs) is the equal-weight mean substring exact match on Natural Questions, TriviaQA, PopQA, and HotpotQA. Score (norm) is the absolute score divided by the corresponding SDPA score at the same context length, expressed as a percentage. KV is the equal-weight mean of the four RAG task-level total-access percentages, including representative-key routing reads.

S

is the nominal budget. The

\pm

entries are 95% bootstrap confidence interval half-widths.

Method	$S$	KV (%)	Score (norm) (%)	Score (abs) (%)
SDPA	—	100.00	100.00	31.91 $\pm$ 2.04
SANTA	128	50.36	95.95	30.62 $\pm$ 2.06
	256	50.59	102.80	32.80 $\pm$ 1.96
	512	50.97	106.85	34.10 $\pm$ 2.04
	1024	51.56	99.84	31.86 $\pm$ 1.91
SANTA ++ (k-means)	128	16.68	99.07	31.61 $\pm$ 2.01

Table 3: LongBench v2: full budget sweep with a 32K context cap. Qwen2.5-7B-Instruct is evaluated on 503 prompts using the brief-reasoning prompt and middle-truncation procedure in Appendix C.3 . Score (abs) is accuracy. Score (norm) is the absolute score divided by the corresponding SDPA score, expressed as a percentage. KV is reported GQA-aware logical access, and

S

is the nominal budget. The

\pm

entries are the reported 95% bootstrap confidence interval half-widths.

Method	$S$	KV (%)	Score (norm) (%)	Score (abs) (%)
SDPA	—	100.00	100.00	83.51 $\pm$ 0.66
SANTA	128	50.32	99.15	82.80 $\pm$ 0.64
	256	50.54	99.75	83.31 $\pm$ 0.67
	512	50.91	100.11	83.61 $\pm$ 0.65
	1024	51.52	99.92	83.45 $\pm$ 0.67
	2048	52.49	99.99	83.50 $\pm$ 0.66

Table 4: RULER: full 32K budget sweep with Qwen2.5-7B-Instruct. Score (abs) is the equal-weight mean of all 13 task scores over 500 completed prompts per task. Score (norm) is the absolute score divided by the corresponding SDPA score at the same context length, expressed as a percentage. KV is pooled GQA-aware logical access, including representative-key routing reads, and

S

is the nominal budget. Absolute-score

\pm

entries are 95% task-stratified bootstrap confidence interval half-widths.

Method	GPU launches	Median latency ( $\mu$ s)	Speedup vs. Flash
SANTA ++	5	63.52	1.69 $\times$
Forced Flash SDPA	2	107.33	1.00 $\times$

Table 5: Complete attention-operator latency on the captured 32K, batch-one FP16 workload. Measurements include selection and all launches. Preparation is excluded.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Additional Jensen examples. Each row compares 16 raw logits

s_{i}

, their exponentials, and the exact-to-surrogate mass ratio

\sum_{i}e^{s_{i}}/\bigl[n\exp(n^{-1}\sum_{i}s_{i})\bigr]

. The displayed ratios are approximately (a)

1.03

for near-uniform logits, (b)

30

for bimodal logits, (c)

135

for the heavy-tailed example, and (d)

760

for the sparse example. The exponentiated-score axes are logarithmic.

		S-NIAH-1		S-NIAH-2		S-NIAH-3		MK-NIAH-1		MK-NIAH-2
Method	$S$	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)
SDPA	—	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	98.00 $\pm$ 2.50
SANTA ++ (Teams)	128	22.54	100.00 $\pm$ 0.00	25.37	100.00 $\pm$ 0.00	25.50	100.00 $\pm$ 0.00	24.95	99.00 $\pm$ 1.50	25.57	94.00 $\pm$ 4.50
	256	30.33	100.00 $\pm$ 0.00	34.75	99.00 $\pm$ 1.50	34.77	100.00 $\pm$ 0.00	34.05	100.00 $\pm$ 0.00	34.60	98.00 $\pm$ 2.50
	512	42.60	100.00 $\pm$ 0.00	48.49	100.00 $\pm$ 0.00	48.19	100.00 $\pm$ 0.00	47.63	100.00 $\pm$ 0.00	47.78	95.00 $\pm$ 4.50
Centroid-guided ( $B=4$ )	128	20.90	100.00 $\pm$ 0.00	22.06	100.00 $\pm$ 0.00	22.53	98.00 $\pm$ 2.50	21.98	99.00 $\pm$ 1.50	22.87	98.00 $\pm$ 2.50

Appendix

Table 6: Task-variant evaluation on RULER at 8K context. Qwen2.5-7B-Instruct with 100 prompts per task. SANTA ++ (teams) uses whole-team sampling with

P=16,R=4

; Centroid-guided (

B=4

) and (

B=16

) denote the two centroid-guided baselines.

S

is the nominal budget. KV is the reported theoretical GQA-aware access relative to dense SDPA. Accuracy is reported as mean

\pm

95% bootstrap confidence interval half-width. All Tasks accuracy is the equal-weight mean over all 13 tasks, with a stratified bootstrap that resamples prompts within each task. This is a separate evaluation from the 20-prompt dense-trajectory KL diagnostic.

Figure 4: Key-to-centroid distance across RULER tasks. Mean distance for key clusters, contiguous chunks, and random chunks on 100 prompts per task, using Qwen2.5-7B-Instruct at 8,192-token context and nominal group size

P=16

. Contiguous chunks lie between key clusters and random chunks on every plotted task.

Task	Context length	WER threshold	Heads (%)
CWE	8192	$\geq 0.50$	82.1 [79.8, 84.9]
CWE	32768	$\geq 0.50$	80.2 [76.9, 84.8]
CWE	8192	$\geq 0.75$	54.0 [48.4, 60.7]
CWE	32768	$\geq 0.75$	49.5 [41.7, 60.7]
CWE	8192	$\geq 0.90$	22.1 [18.4, 27.0]
CWE	32768	$\geq 0.90$	23.6 [16.2, 36.4]

Appendix

Table 7: Oracle prefill-window effectiveness. Percentage of heads across all layers meeting each WER threshold from Equation 12 . For each prompt, layer, and head, WER compares the attention mass in the top

K_{\mathrm{w}}=32

disjoint windows of width

w=16

with the mass in the

wK_{\mathrm{w}}=512

highest-attention tokens. Results use 10 prompts per task at 8K and 32K context. Brackets give 95% bootstrap confidence intervals from resampling prompts.

Figure 5: Key versus probe-fingerprint clustering. Accuracy on four RULER tasks with 50 prompts per task and nominal budgets $S=128,256,512$ . Bars compare key clustering, probe-fingerprint clustering, SANTA, and dense SDPA; error bars denote the displayed 95% confidence intervals.

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	–	100.00	100.00	62.65 $\pm$ 2.02
SANTA	128	51.62	100.24	62.80 $\pm$ 2.02
	256	52.70	99.60	62.40 $\pm$ 2.03
	512	54.39	99.60	62.40 $\pm$ 2.03
	1024	56.77	100.16	62.75 $\pm$ 2.05
	2048	60.14	100.00	62.65 $\pm$ 2.07

Appendix

Table 8: HELMET RAG: full 8K budget sweep with Qwen2.5-7B-Instruct. Acc (abs) is the equal-weight mean substring exact match on Natural Questions, TriviaQA, PopQA, and HotpotQA. Acc (norm) is the absolute score divided by the corresponding SDPA score at the same context length, expressed as a percentage. KV is the equal-weight mean of the four RAG task-level total-access percentages, including representative-key routing reads.

S

is the nominal budget. The

\pm

entries are the reported score-interval half-widths.

Method	$S$	Natural Questions		TriviaQA		PopQA		HotpotQA		MS MARCO
Method	$S$	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)
SDPA	–	100.00	53.20 $\pm$ 4.30	100.00	82.20 $\pm$ 3.30	100.00	64.40 $\pm$ 4.10	100.00	50.80 $\pm$ 4.40	100.00	58.51 $\pm$ 3.89
SANTA	128	51.62	54.40 $\pm$ 4.40	51.55	82.40 $\pm$ 3.40	51.79	63.00 $\pm$ 4.20	51.51	51.40 $\pm$ 4.30	51.72	59.16 $\pm$ 3.81
	256	52.73	53.80 $\pm$ 4.40	52.58	81.80 $\pm$ 3.40	52.97	63.80 $\pm$ 4.20	52.52	50.20 $\pm$ 4.50	52.75	58.03 $\pm$ 3.75
	512	54.46	53.00 $\pm$ 4.30	54.21	82.20 $\pm$ 3.30	54.79	64.20 $\pm$ 4.20	54.09	50.20 $\pm$ 4.50	54.27	59.06 $\pm$ 3.84
	1024	56.90	54.40 $\pm$ 4.30	56.49	82.20 $\pm$ 3.30	57.30	64.00 $\pm$ 4.10	56.39	50.40 $\pm$ 4.40

Appendix

Table 9: HELMET task results at 8K context length. Each accuracy is reported as the mean

\pm

the half-width of its 95% confidence interval.

Method	$S$	Natural Questions		TriviaQA		PopQA		HotpotQA		MS MARCO
Method	$S$	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)
SDPA	–	100.00	46.20 $\pm$ 4.30	100.00	82.00 $\pm$ 3.50	100.00	62.80 $\pm$ 4.30	100.00	38.60 $\pm$ 4.20	100.00	34.39 $\pm$ 4.36
SANTA	128	50.43	46.40 $\pm$ 4.30	50.41	81.00 $\pm$ 3.50	50.46	62.20 $\pm$ 4.20	50.41	38.00 $\pm$ 4.30	50.46	37.63 $\pm$ 4.60
	256	50.77	45.60 $\pm$ 4.40	50.71	81.00 $\pm$ 3.40	50.81	62.80 $\pm$ 4.20	50.71	38.00 $\pm$ 4.40	50.77	35.79 $\pm$ 4.58
	512	51.33	46.00 $\pm$ 4.40	51.23	81.40 $\pm$ 3.30	51.40	62.80 $\pm$ 4.20	51.23	39.80 $\pm$ 4.20	51.29	34.49 $\pm$ 4.35
	1024	52.23	46.40 $\pm$ 4.40	52.05	80.40 $\pm$ 3.40	52.32	62.80 $\pm$ 4.20	52.06	38.80 $\pm$ 4.10

Appendix

Table 10: HELMET task results at 32K context length. Each accuracy is reported as the mean

\pm

the half-width of its 95% confidence interval.

		S-NIAH-1		S-NIAH-2		S-NIAH-3		MK-NIAH-1		MK-NIAH-2
Method	$S$	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)
SDPA	—	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	98.80 $\pm$ 0.90
SANTA	128	50.97	100.00 $\pm$ 0.00	51.11	100.00 $\pm$ 0.00	51.06	99.80 $\pm$ 0.30	51.13	100.00 $\pm$ 0.00	51.25	98.80 $\pm$ 0.90
	256	51.56	100.00 $\pm$ 0.00	51.81	100.00 $\pm$ 0.00	51.72	99.80 $\pm$ 0.30	51.83	100.00 $\pm$ 0.00	52.08	99.00 $\pm$ 0.90
	512	52.52	100.00 $\pm$ 0.00	52.93	100.00 $\pm$ 0.00	52.76	100.00 $\pm$ 0.00	52.94	100.00 $\pm$ 0.00	53.41	98.20 $\pm$ 1.10
	1024	53.98	100.00 $\pm$ 0.00	54.62	100.00 $\pm$ 0.00	54.35	100.00 $\pm$ 0.00	54.62	100.00 $\pm$ 0.00	55.43	98.20 $\pm$ 1.10

Appendix

Table 11: Per-task accuracy and KV memory access on RULER (8K context) on Qwen2.5-7B-Instruct. KV is the theoretical GQA-aware KV memory access (%), and

S

denotes the nominal budget. Accuracy values are reported as mean

\pm

95% bootstrap confidence interval half-width.

		S-NIAH-1		S-NIAH-2		S-NIAH-3		MK-NIAH-1		MK-NIAH-2
Method	$S$	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)
SDPA	—	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	98.60 $\pm$ 1.00	100.00	99.00 $\pm$ 0.90	100.00	93.40 $\pm$ 2.20
SANTA	128	50.28	100.00 $\pm$ 0.00	50.31	100.00 $\pm$ 0.00	50.30	99.20 $\pm$ 0.70	50.31	99.20 $\pm$ 0.70	50.37	93.20 $\pm$ 2.20
	256	50.47	100.00 $\pm$ 0.00	50.53	100.00 $\pm$ 0.00	50.51	98.80 $\pm$ 0.90	50.53	99.00 $\pm$ 0.90	50.64	92.80 $\pm$ 2.20
	512	50.81	100.00 $\pm$ 0.00	50.91	100.00 $\pm$ 0.00	50.87	98.80 $\pm$ 0.90	50.90	99.00 $\pm$ 0.90	51.12	93.60 $\pm$ 2.10
	1024	51.37	100.00 $\pm$ 0.00	51.53	100.00 $\pm$ 0.00	51.45	98.60 $\pm$ 1.00	51.51	99.00 $\pm$ 0.90	51.92	93.40 $\pm$ 2.10

Appendix

Table 12: Per-task accuracy and KV memory access on RULER (32K context) on Qwen2.5-7B-Instruct. KV is the theoretical GQA-aware KV memory access (%), and

S

denotes the nominal budget. Accuracy values are reported as mean

\pm

95% bootstrap confidence interval half-width.

Directly timed region	Median latency ( $\mu$ s)
Grouped representative scorer	20.32
Parallel per-query top- $K$	15.51
Fused exact merge + GQA union	16.16
Grouped selected attention	19.97
Split-softmax reducer	9.95
Complete five-launch SANTA ++	63.52

Appendix

Table 13: Profiler-off SANTA ++ timings for independently measured regions of the captured 32K workload. Entries are medians of round medians; the complete operator is measured directly.

Kernel	CTAs	NCU dur. ( $\mu$ s)	Occ. (%)	Elig.	DRAM (%)	Reg./thr.
SANTA ++ representative scorer	248	14.14	24.49	0.08	63.85	56
SANTA ++ top- $K$	868	10.78	84.51	2.40	10.61	30
SANTA ++ merge + union	4	10.94	16.62	0.38	1.59	44
SANTA ++ selected attention	128	14.30	12.88	0.20	19.46	144
SANTA ++ reducer	28	4.80	8.33	0.05	11.07	40
Flash split-KV forward	128	89.92	8.33	0.10	83.95	200

Appendix

Table 14: Nsight Compute measurements for the captured 32K, batch-one workload. Entries are per-metric medians of three independent captures of a fixed sampling realization; durations use kernel replay.

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	86.22 $\pm$ 4.56
SANTA ++ (Latent k-means)	128	56.28	80.80	69.67 $\pm$ 5.92
	256	70.09	91.56	78.94 $\pm$ 5.14
	512	86.18	98.07	84.56 $\pm$ 4.72
	1024	101.89	97.81	84.33 $\pm$ 5.00

Appendix

Table 15: Token sparsity in a shared-latent cache at 8K. DeepSeek-V2-Lite-Chat, with 50 matched prompts per task on niah_multivalue , fwe , and qa_1 . Absolute accuracy is the equal-weight task mean; normalized accuracy divides it by the SDPA mean, times 100. KV denotes logical shared-row access, including routing, under Equation 14 . Throughout these tables,

\pm

gives half the span of a 95% percentile-bootstrap interval using 10,000 resamples. Prompts are resampled within each task, retaining task strata for the aggregate. Interval endpoints need not be symmetric about the point estimate.

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	96.67 $\pm$ 2.67
SANTA ++ (Latent k-means)	128	47.57	84.82	82.00 $\pm$ 7.00
	256	62.25	94.48	91.33 $\pm$ 5.33
	512	80.34	100.00	96.67 $\pm$ 3.00
	1024	98.45	99.31	96.00 $\pm$ 3.00

Appendix

Table 16: RULER fwe at 8K. Accuracy on 50 matched prompts per method. Normalized accuracy divides the reported two-decimal task score by SDPA’s score on this task, times 100. Access and uncertainty follow Table 15 .

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	94.00 $\pm$ 4.00
SANTA ++ (Latent k-means)	128	57.88	77.66	73.00 $\pm$ 8.75
	256	71.89	93.09	87.50 $\pm$ 5.50
	512	87.67	98.94	93.00 $\pm$ 4.75
	1024	102.64	92.55	87.00 $\pm$ 7.75

Appendix

Table 17: RULER niah_multivalue at 8K. Accuracy on 50 matched prompts per method. Normalized accuracy divides the reported two-decimal task score by SDPA’s score on this task, times 100. Access and uncertainty follow Table 15 .

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	68.00 $\pm$ 13.00
SANTA ++ (Latent k-means)	128	58.41	79.41	54.00 $\pm$ 14.00
	256	72.81	85.29	58.00 $\pm$ 14.00
	512	88.50	94.12	64.00 $\pm$ 13.00
	1024	103.87	102.94	70.00 $\pm$ 13.00

Appendix

Table 18: RULER qa_1 at 8K. Accuracy on 50 matched prompts per method. Normalized accuracy divides the reported two-decimal task score by SDPA’s score on this task, times 100. Access and uncertainty follow Table 15 .

Explore similar work

May 3, 2026cs.LG

Stochastic Sparse Attention for Memory-Bound Inference

Autoregressive decoding becomes bandwidth-limited at long contexts, as generating each token requires reading all

n_k

key and value vectors from KV cache. We present Stochastic Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling

S \ll n_k

indices from the post-softmax distribution and aggregates only those value rows. This yields an unbiased estimator of the post-softmax value aggregation while replacing value-stage multiply-accumulates with gather-and-add. We introduce stratified and systematic sampling to design variance-reduced, GPU-friendly variants. Evaluated on Llama-3.1-8B-Instruct at 32k-token contexts, S

^2

ANTA matches baseline accuracy while achieving up to

1.5\times

decode-step attention-kernel speedup over FlashInfer and FlashDecoding on an NVIDIA RTX 6000 Ada. In batched long-context generation, these kernel gains translate to up to

1.25\times

end-to-end decode-latency speedup. Finally, we propose Bernoulli

qK^\mathsf{T}

sampling as a complementary technique to sparsify the score stage, reducing key-feature access through stochastic ternary queries. Both methods are complementary to upstream quantization, low-rank projection, KV-cache compression, and KV-cache selection methods. Together, they point toward sparse, multiplier-free, and energy-efficient inference. We open-source our kernels at: https://github.com/OPUSLab/SANTA.git

Kyle Lee, Corentin Delacour, Kevin Callahan-Coray +5

Department of Electrical and Computer Engineering, University of California, Santa Barbara, CA, USA · Department of Electrical and Computer Engineering, University of Michigan, Ann Arbor, MI, USA · Department of Electrical and Computer Engineering, Carnegie Mellon University, Pittsburgh, PA, USA

Dynamic Sparse Attention Autoregressive Decoding Key-Value Cache Key-Value Cache Compression Gumbel-Softmax Relaxation Stochastic

Ask PDF

Jun 20, 2026cs.CL

Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers

We propose Keyless Attention, an attention mechanism that eliminates the key projection entirely, operating over queries and values only. This yields a Value-Only Cache that reduces KV cache memory and access overhead by exactly 50% over standard attention, while matching or exceeding standard attention's decode throughput. Beyond efficiency, we introduce Depth-

m

Attention Factorization: standard attention computes a depth-2 factorization of the attention bilinear form, while Keyless Attention realizes a depth-

m

instance of this family. At m=3, Keyless Attention matches the projection matrix count of standard attention via a value-space routing matrix that replaces the key projection and introduces a coupling between routing and retrieval. Experiments across five models and four architectures (GPT-2 280M, GPT-2 557M, Pythia 410M, Qwen2 1.5B, and Llama 3.2 1B) show that Keyless Attention matches or outperforms standard QKV attention on perplexity in 4 out of 5 models. On downstream zero-shot evaluation (GPT-2 557M), Keyless Attention outperforms on 4 out of 5 commonsense reasoning benchmarks, while achieving 50% KV cache reduction throughout.

Xin Gao

Department of Mathematics and Statistics York University Toronto, M3J 1P3

Transformer Architectures Reasoning Benchmark

Ask PDF

Apr 22, 2026cs.LG

Simplified Sparse Attention via Gist Tokens

Sparse attention can reduce the cost of long-context inference, but most variants introduce new architectural components. We introduce Simplified Sparse Attention (SSA), a simpler approach to sparse attention that requires no architectural changes. Concretely, we first perform continued pretraining on sequences interleaved with gist tokens. We optimize the standard next-token loss as usual, but the gist tokens use an attention mask to restrict what parts of the context the language model can attend to; this teaches the model to pack each chunk's important information into the gist tokens. At inference time, SSA scores chunks via attention between the current query and the small set of gist tokens, selectively unfolding the top-k chunks by reintroducing their corresponding raw tokens. Since the query is scored only against the gist tokens, we avoid the memory-bandwidth cost associated with naive scoring against the full KV cache, without requiring the auxiliary KV cache approach used by sparse attention methods. On LongBench, SSA consistently outperforms compression and inference-time sparse-attention baselines under the same compression ratio. More strikingly, in retrieval-augmented generation, SSA can even outperform full attention after continued pretraining by over 5.7 points. We attribute this to the ability of SSA's selective unfolding, which concentrates attention on the query-relevant chunks and effectively filters out noise. SSA further extends to a hierarchical gist-of-gist variant (H-SSA) that achieves log-linear decoding complexity while maintaining or improving accuracy at high compression ratios up to 32x. The code is available at https://github.com/yuzhenmao/simplified-sparse-attention/.

Yuzhen Mao, Michael Y. Li, Emily B. Fox

Stanford University

Dynamic Sparse Attention Efficient Long-Context Inference

Ask PDF