aiqbench analysis
This research demonstrates how Flash Attention mechanisms can be repurposed to accelerate Gaussian kernel sums, which are fundamental to various kernel methods like MMDs and SVGD. By leveraging existing hardware-aware optimizations for softmax attention, the authors achieve significant speedups in memory usage and accuracy without requiring custom GPU code. This work bridges the gap between transformer-specific optimizations and broader machine learning computational needs.
Fast Gauss Sums via Flash Attention
arXiv:2609.04910v1 Announce Type: new Abstract: Gaussian kernel sums are the computational core of maximum mean discrepancies (MMDs), kernel gradient flows, Stein variational gradient descent (SVGD), and many other kernel methods. At the same time, softmax attention has received an extraordinary amount of hardware-aware code engineering, culminating in flash attention. We show that Gauss kernel sums with arbitrary, signed weights can be evaluated via flash attention: two small input augmentations turn the normalized softmax reduction into the unnormalized Gauss sum, without writing a single line of custom GPU code. For feature dimension D>8 in fp16, this approach beats compiled PyTorch code as well as PyKeOps kernels (often significantly) in speed, memory-overhead and accuracy. Indeed, its memory scaling remains linear.
Abstract:Gaussian kernel sums are the computational core of maximum mean discrepancies (MMDs), kernel gradient flows, Stein variational gradient descent (SVGD), and many other kernel methods. At the same time, softmax attention has received an extraordinary amount of hardware-aware code engineering, culminating in flash attention. We show that Gauss kernel sums with arbitrary, signed weights can be evaluated via flash attention: two small input augmentations turn the normalized softmax reduction into the unnormalized Gauss sum, without writing a single line of custom GPU code. For feature dimension D>8 in fp16, this approach beats compiled PyTorch code as well as PyKeOps kernels (often significantly) in speed, memory-overhead and accuracy. Indeed, its memory scaling remains linear.
Machine Learning (cs.LG); Numerical Analysis (math.NA)
G.1.0; G.4; I.2.6
arXiv:2609.04910 [cs.LG]
(or arXiv:2609.04910v1 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2609.04910 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 213 words
Related coverage
Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers
Hugging Face · Aug 18, 2026
Enabling DeepSeek-V4-Flash Training on AMD Instinct MI355X GPUs with Primus
AMD · Sep 3, 2026
AMD Instinct MI455X vs MI355X: A Technical Look at the Advancing AI 2026 Inference Numbers
AMD · Oct 8, 2026
Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference
NVIDIA · Jul 31, 2026