Fast Gauss Sums via Flash Attention

2026-09-04

AI Summary

本文通过使用flash attention方法解决了高斯核求和的计算效率问题，提高了速度、降低了内存开销并保持了准确性。

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.

Problem

• Gaussian kernel sums
• flash attention
• computational efficiency
• memory overhead
• accuracy

Innovation

• Flash Attention
• Gaussian Kernel Sums
• Softmax Reduction
• Memory Scaling
• Hardware-aware Optimization