KDEformer: 通过核密度估计加速Transformer
机器学习
2023-06-30 v2 计算机视觉与模式识别
数据结构与算法
摘要
点积注意力机制在现代深度架构（如Transformer）中对于序列建模起着关键作用，然而，该模型的朴素精确计算在序列长度上具有二次时间和内存复杂度，阻碍了长序列模型的训练。关键瓶颈在于softmax函数分母中配分函数的计算以及softmax矩阵与值矩阵的乘法。我们的关键观察是，前者可以简化为核密度估计（KDE）问题的一个变体，并且高效的KDE求解器可以进一步通过基于子采样的快速矩阵乘积来加速后者。我们提出的KDEformer可以在次二次时间内近似注意力，并具有可证明的谱范数界，而所有先前的结果仅提供逐元素误差界。实验上，我们验证了KDEformer在各种预训练模型上，在准确性、内存和运行时间方面优于其他注意力近似方法。在BigGAN图像生成中，我们实现了比精确计算更好的生成分数，加速超过4倍。对于使用T2T-ViT的ImageNet分类，KDEformer实现了超过18倍的加速，而准确率下降小于0.5%。
引用
@article{arxiv.2302.02451,
title = {KDEformer: Accelerating Transformers via Kernel Density Estimation},
author = {Amir Zandieh and Insu Han and Majid Daliri and Amin Karbasi},
journal= {arXiv preprint arXiv:2302.02451},
year = {2023}
}
备注
26 pages, 7 figures