Stochastic Sparse Attention for Memory-Bound Inference

• URL: http://arxiv.org/abs/2605.01910v1
• Date: Sun, 03 May 2026 14:44:14 GMT
• Title: Stochastic Sparse Attention for Memory-Bound Inference
• Abstract summary: We present Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling $S ll n_k$ indices from the post-softmax distribution.<n>We also propose Bernoulli $qKmathsfT$ sampling as a complementary technique to sparsify the score stage.
• Score: 19.301894658575502
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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 sampling to design variance-reduced, GPU-friendly variants, demonstrating $1.5\times$ decode-step attention kernel speedup over FlashInfer and FlashDecoding on an NVIDIA RTX 6000 Ada while matching baseline accuracy at 32k-token contexts. 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 orthogonal to upstream techniques such as ternary quantization, low-rank projections, and KV-cache compression. Together, they point toward sparse, multiplier-free, and energy-efficient inference. We open-source our kernels at: https://github.com/OPUSLab/SANTA.git

Related papers

• The Sample Complexity of Quantum Entanglement Allocation [0.13537117504260623]
  We show that the answer depends on the allocation choices created by the queries.<n>A larger memory can require no more data.<n>We derive an exact tradeoff with extra fresh detector calls and transfer the learning law to structured transaction co-location.
  arXiv  Detail & Related papers  (2026-09-09T13:19:27Z)
• SSTQ:Privacy-Preserving Vector Quantization via Subsampled Stochastic TurboQuant [79.24089819400126]
  Subsampled TurboQuant (SSTQ) is a framework that combines overcomplete equal-width tight frames, coordinate subsampling, and privacy-aware quantization.<n>SSTQ achieves mean squared error scaling while using only $lceil log N il + b$ bits per client, where $N = (d)$ is the frame size.<n>We also derive a surrogate privacy-aware codebook objective that reduces the codebook-dependent MSE scaling from $O(4b)$ to $O(2b)$.
  arXiv  Detail & Related papers  (2026-08-05T17:51:25Z)
• Accelerating Discrete Diffusion Models with Parallel-In-Time Sampling [55.388363730120325]
  We parallelize the mainstream $$-leaping algorithm for absorbing discrete diffusion in a Continuous-Time Markov Chain (CTMC) framework.<n>By leveraging the continuous-time integral form of the $$-leaping algorithm and the Picard method, we achieve parallel-in-time sampling acceleration.<n>Our research expands the potential of discrete diffusion models for efficient parallel inference, with broader implications for applications such as molecular structure and language generation.
  arXiv  Detail & Related papers  (2026-07-01T10:59:33Z)
• Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering [3.845089876398699]
  Flash-GMM achieves a textbf20$times$ speedup over existing implementations.<n>We show that soft GMM clustering is now a viable drop-in replacement for $k$-means.
  arXiv  Detail & Related papers  (2026-06-09T14:07:18Z)
• Nectar: Neural Estimation of Cached-Token Attention via Regression [36.709992864731866]
  We show how a compact neural network can be used to evaluate softmax attention over a fixed long context.<n>Nectar fits two networks per layer and KV-head: a target network that predicts the attention output and a score network that predicts the log-normalizer.<n>We report experiments on models from 1.7B to 8B parameters across five long-context datasets.
  arXiv  Detail & Related papers  (2026-05-10T21:51:36Z)
• Scaling Attention via Feature Sparsity [50.64995497733461]
  Scaling Transformers to ultra-long contexts is bottlenecked by the $O(n2 d)$ cost of self-attention.<n>We propose Sparse Feature Attention, where queries and keys are represented as $k$-sparse codes that preserve high-dimensional expressivity.<n>Across GPT-2 and Qwen3 pretraining, SFA matches dense baselines while improving speed by up to $2.5times$ and reducing FLOPs and KV-cache by nearly 50%.
  arXiv  Detail & Related papers  (2026-03-17T08:41:50Z)
• FlashSampling: Fast and Memory-Efficient Exact Sampling [62.5203057469482]
  FlashSampling is an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in.<n>Across H100, H200, B200, and B300 GPUs, FlashSampling speeds up kernel-level decode workloads.<n>In end-to-end vLLM experiments, it reduces time per output token by up to $19%$ on the models we test.
  arXiv  Detail & Related papers  (2026-03-16T19:37:08Z)
• Parallel Sampling via Autospeculation [13.643401888306398]
  We present algorithms to accelerate sampling via counting in two settings: any-order autoregressive models and denoising diffusion models.<n>We show that, by issuing oracle calls in parallel, the expected sampling time can be reduced to $widetildeO(n1/2)$.<n>We introduce a novel technique to obtain our results: speculative rejection sampling.
  arXiv  Detail & Related papers  (2025-11-11T06:09:44Z)
• Fast Rates for Bandit PAC Multiclass Classification [73.17969992976501]
  We study multiclass PAC learning with bandit feedback, where inputs are classified into one of $K$ possible labels and feedback is limited to whether or not the predicted labels are correct. Our main contribution is in designing a novel learning algorithm for the agnostic $(varepsilon,delta)$PAC version of the problem.
  arXiv  Detail & Related papers  (2024-06-18T08:54:04Z)
• Debiased Distribution Compression [30.600795754425775]
  We introduce a new suite of compression methods suitable for compression with biased input sequences. We provide succinct and accurate posterior summaries while overcoming biases due to burn-in, approximate Markov chain Monte Carlo, and tempering.
  arXiv  Detail & Related papers  (2024-04-18T16:11:16Z)
• Stochastic Approximation Approaches to Group Distributionally Robust Optimization and Beyond [89.72693227960274]
  This paper investigates group distributionally robust optimization (GDRO) with the goal of learning a model that performs well over $m$ different distributions. To reduce the number of samples in each round from $m$ to 1, we cast GDRO as a two-player game, where one player conducts and the other executes an online algorithm for non-oblivious multi-armed bandits. In the second scenario, we propose to optimize the average top-$k$ risk instead of the maximum risk, thereby mitigating the impact of distributions.
  arXiv  Detail & Related papers  (2023-02-18T09:24:15Z)
• Linear Time Sinkhorn Divergences using Positive Features [51.50788603386766]
  Solving optimal transport with an entropic regularization requires computing a $ntimes n$ kernel matrix that is repeatedly applied to a vector. We propose to use instead ground costs of the form $c(x,y)=-logdotpvarphi(x)varphi(y)$ where $varphi$ is a map from the ground space onto the positive orthant $RRr_+$, with $rll n$.
  arXiv  Detail & Related papers  (2020-06-12T10:21:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.