Source-linked AI summary
Fast Gauss Sums via Flash Attention
Nicolaj Rux, Sebastian Neumayer
TL;DR
Gaussian kernel sums are central but expensive reductions in MMD, gradient flows, and SVGD. The paper reduces them to flash attention through two input augmentations, then finds that for D > 8 the flash variants are fastest with competitive accuracy and linear memory scaling, subject to fp16 and differentiation constraints.
Problem
Exact Gaussian kernel sums are computational bottlenecks in MMD gradient flows and SVGD, with naive evaluation requiring O(MN) memory.
Method
Two reductions transform normalized softmax attention into unnormalized Gaussian kernel sums, including a differentiable reweight variant using public PyTorch attention APIs.
Results
For D > 8, flash variants are fastest, while for D ≥ 16 they run 2–21 times faster than PyKeOps forward and 3–10 times faster with gradients, with similar linear memory scaling.
Takeaways & Limitations
Flash attention provides a fast, memory-efficient drop-in implementation for unnormalized Gaussian kernel sums with signed weights in fp16.
Abstract
from arXiv · showhide
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.
1 Introduction
Gaussian kernel sums underpin major kernel methods but are expensive to evaluate exactly at scale. The paper shows that flash attention can evaluate them after two small input augmentations, without custom GPU code.
- Gaussian kernel sums appear in essentially every Gaussian kernel method, including MMD gradient flows and Stein variational gradient descent.
- Naive evaluation requires O(MN(C+D)) operations and O(MN) memory.
- Classical fast summation improves asymptotics only in restricted regimes by trading exactness, whereas PyKeOps performs an exact fused GPU reduction.
- Two small input augmentations reduce Gaussian kernel sums to flash attention without custom CUDA or Triton code.
- The paper contributes differentiable and logits-based attention reductions, fp16 safeguards, and benchmarks against compiled PyTorch and PyKeOps.
2 Gaussian kernel sums from flash attention
The paper converts normalized softmax attention into an unnormalized Gaussian sum using prescaling or reweighting. The reweighting route is differentiable through public PyTorch APIs, but fp16 underflow imposes a bound on κ and N.
- Flash attention computes the normalized reduction with O(MN(D+C)) operations and O((M+N)(D+C)) memory, while its logits are available as an fp32 by-product.
- Prescale: Prescale removes softmax normalization by multiplying attention outputs by exp(lseτ(q, k)m), but the implementation cannot handle gradients.
- Reweight: Reweight avoids nondifferentiable logits by adding two entries to queries and keys, enabling automatic differentiation through public SDPA.
- Numerical stability: Because fp16 has dynamic range [2^-14, 2^16], avoiding β underflow restricts κ ≤ 2^14 or N ≤ 268 435 455.
- Reweight: The reweight construction augments queries, keys, and values before recovering the Gaussian sum from attention outputs via καm/βm.
3 Applications
MMD and particle-flow updates can be expressed as Gaussian kernel sums, including query gradients and the kernel component of SVGD. This makes the flash-based reductions applicable to these methods.
- Maximum mean discrepancy: For discrete measures, MMD uses a Gaussian kernel sum with q = k and C = 1, followed by the inner product vT s.
- MMD flows and Stein variational gradient descent: Particle flows require gradients of the Gaussian sum with respect to query points.
- MMD flows and Stein variational gradient descent: The query gradient is another kernel sum with values (vnkn, vn) ∈ R^(D+1), covering both terms of the SVGD update.
- MMD flows and Stein variational gradient descent: Plain autograd through the reweight algorithm can compute these particle-flow gradients.
4 Numerical results
The benchmarks compare flash-based Gaussian summation with compiled PyTorch and PyKeOps across input sizes and dimensions. Flash variants are especially competitive for D = 32 and D > 8, while PyKeOps remains strongest in low dimensions.
- Benchmark setup: The benchmark measures elapsed time, memory overhead, and relative L2 error against an fp64 reference for fp16 Gaussian sums.It compares PyTorch, PyKeOps, reweight, and prescale across forward and gradient computations.
- Sweep over N: For D = 3, PyKeOps is marginally faster than the flash variants at large N and uses consistently 3 times less memory, but has worse error.The PyTorch implementation quickly runs out of memory because it allocates an entire N × N matrix.
- Sweep over N: For D = 32, the flash variants outperform PyTorch and PyKeOps in both speed and memory overhead while maintaining healthy relative accuracy.The forward comparison is shown in Figure 1, and the gradient comparison in Figure 2 is similar; prescale is unavailable for gradients.
- Sweep over D: 2.7×10^-3 is PyKeOps's approximate error across dimensions, compared with around 3.5 × 10^-4 for the other methods.The flash variants have error around 4.0 × 10^-4, slightly above PyTorch at 3.3 × 10^-4.
- Sweep over D: For D ≤8, PyKeOps is best in speed and memory, whereas the flash variants are fastest for D > 8 and have roughly linear memory overhead.PyKeOps's speed degenerates as dimension increases, while padding makes reweight worse than prescale for some powers-of-two dimensions.
5 Conclusion
Two small input augmentations turn flash attention into a fast, memory-efficient method for unnormalized Gaussian kernel sums with signed weights. The approach excels in fp16, while its quadratic complexity leaves combining it with subquadratic approximations as future work.
- Conclusion: Two small input augmentations make flash attention a drop-in method for evaluating unnormalized Gaussian kernel sums with signed weights.Both proposed variants retain quadratic complexity despite their implementation-level speedups.
- Conclusion: For D ≥16, flash-based sums run 2−21 times faster than PyKeOps forward and 3−10 times faster with gradients, with up to 7 times lower error.The flash variants also retain similar linear memory scaling.
- Conclusion: The method excels at fp16, which often suits sampling, flows, and testing but not ill-conditioned solvers.The conclusion identifies subquadratic approximations as a promising direction for combining with these constant-factor gains.
6 Appendix
In fp32, the memory_efficient backend is compared across feature dimensions and problem sizes, with PyKeOps remaining the state-of-the-art implementation. The proof establishes bounds for the auxiliary quantities α and β.
- Proof of Proposition 1: The proof bounds β between κ(N + 1)^−1 and κ, while the bound on α follows from convex-combination and norm-convexity properties.
- Single precision: For fp32, PyKeOps remains the state-of-the-art implementation because flash is unavailable and memory_efficient is comparatively slow at low dimensions.For D ≤32, fp32 memory_efficient is around 10 times slower than flash-fp16; PyKeOps retains the lowest relative L2 error.
- Single precision: Figure 4 sweeps D for the forward Gauss sum with B = 64, C = 1, and N = 16384 in fp32 using memory_efficient.
- Single precision: Figure 5 sweeps N for the forward Gauss sum with B = 4, D = 32, and C = 32 in fp32 using memory_efficient.