用闪速注意力加速高斯核求和，性能超越现有代码。
Fast Gauss Sums via Flash Attention
- 通过输入微调，用闪速注意力实现高斯核求和
- 在fp16下特征维数大于8时速度超PyTorch与PyKeOps
- 内存线性增长，适合大规模核方法应用
高斯核求和是最大均值差异（MMD）、核梯度流、斯坦因变分梯度下降（SVGD）等核方法的核心计算。与此同时，软最大注意力已获得大量硬件优化的代码工程支持，形成闪速注意力。本文展示：通过两个小的输入变换，可将归一化的软最大注意力约简为非归一化的高斯求和，无需编写任何自定义GPU代码。在fp16精度下，当特征维度D>8时，该方法在速度、内存开销和精度上均优于编译后的PyTorch代码及PyKeOps内核（常显著领先），且内存增长保持线性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。