Item

Abstract

Dot-product attention mechanism plays a crucial role in modern deep
architectures (e.g., Transformer) for sequence modeling, however, na\"ive 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\times$ speedup. For ImageNet classification with T2T-ViT, KDEformer
shows over $18\times$ speedup while the accuracy drop is less than $0.5\%$.