[论文解读] KDEformer: Accelerating Transformers via Kernel Density Estimation

Amir Zandieh, In‐Su Han+2 more|arXiv (Cornell University)|Feb 5, 2023

Advanced Neural Network Applications被引用 5

一句话总结

KDEformer 通过将点积注意力重新表述为核密度估计（KDE）问题，将 Transformer 的计算加速，实现了次二次计算复杂度，并具备可证明的谱范数误差界。在 ImageNet 上，其推理速度最高可提升 18.3 倍，且精度下降不足 0.5%；在 BigGAN 上，FID 分数优于精确注意力，同时 FLOPs 减少 4.14 倍。

ABSTRACT

Dot-product attention mechanism plays a crucial role in modern deep architectures (e.g., Transformer) for sequence modeling, however, naïve exact computation of this model incurs quadratic time and memory complexities in sequence length, hindering the training of long-sequence models. Critical bottlenecks are due to the computation of partition functions in the denominator of softmax function as well as the multiplication of the softmax matrix with the matrix of values. Our key observation is that the former can be reduced to a variant of the kernel density estimation (KDE) problem, and an efficient KDE solver can be further utilized to accelerate the latter via subsampling-based fast matrix products. Our proposed KDEformer can approximate the attention in sub-quadratic time with provable spectral norm bounds, while all prior results merely provide entry-wise error bounds. Empirically, we verify that KDEformer outperforms other attention approximations in terms of accuracy, memory, and runtime on various pre-trained models. On BigGAN image generation, we achieve better generative scores than the exact computation with over $4 imes$ speedup. For ImageNet classification with T2T-ViT, KDEformer shows over $18 imes$ speedup while the accuracy drop is less than $0.5\%$.

研究动机与目标

• 为解决 Transformer 自注意力机制中的二次计算与内存瓶颈问题。
• 开发一种具备可证明谱范数误差保证的注意力近似方法，而非仅依赖逐元素误差界。
• 通过降低计算复杂度，实现长序列模型的高效推理与训练。
• 相较于先前的近似注意力方法，提升视觉与语言模型在生成性能与效率方面的表现。

提出的方法

• 将注意力计算简化为对对角缩放矩阵 D 和注意力矩阵 A 的广义核密度估计（KDE）问题。
• 使用基于 KDE 的概率进行重要性采样，选取少量列用于近似注意力计算。
• 通过变换将指数核密度问题转化为高斯 KDE 问题，从而利用最先进的 KDE 求解器。
• 通过 KDE 估计对角缩放矩阵 D，避免显式计算完整注意力矩阵 A。
• 利用估计的采样矩阵 Π 和对角矩阵 eD，通过子采样矩阵乘积高效计算输出。
• 提供谱范数误差界（公式 1），确保注意力输出的稳定性和保真度。

实验结果

研究问题

• RQ1 注意力机制能否被重新表述为核密度估计问题，以实现次二次计算？
• RQ2 基于 KDE 的采样与缩放能否在保持模型性能的同时降低注意力的计算复杂度？
• RQ3 注意力近似能否实现谱范数误差界，从而提供强于逐元素误差界的理论保证？
• RQ4 所提方法在长序列任务中是否在精度、速度与内存效率方面优于现有近似注意力机制？
• RQ5 KDEformer 是否能在减少 FLOPs 的同时，提升 BigGAN 等模型的生成质量？

主要发现

• 在 BigGAN 图像生成任务中，KDEformer 的 FID 分数低于精确注意力（生成质量更优），同时 FLOPs 减少 4.14 倍。
• 在 T2T-ViT 的 ImageNet 分类任务中，KDEformer 达到 82.08% 的 top-1 准确率，仅比精确计算低 0.47%，同时 FLOPs 减少 18.3 倍。
• 在 Long Range Arena 的端到端训练中，KDEformer 在 Text 数据集上比精确注意力快 8 倍，且精度下降可忽略。
• KDEformer 在 LRA 数据集上的峰值内存平均为 2.28 GB，显著低于精确方法（9.21 GB）及其他竞争方法。
• 在 LRA 上，该方法在平均测试准确率（57.15%）中位列第二，同时是所有近似方法中最快且内存效率最高者。
• 实证结果表明，KDEformer 在多个基准测试中，均优于 Reformer、Performer 和 ScatterBrain，在精度与效率方面表现更优。