FLuRKA: Fast fused Low-Rank & Kernel Attention
Ahan Gupta Yueming Yuan Yanqi Zhou Charith Mendis
Many efficient approximate self-attention techniques have become prevalent
since the inception of the transformer architecture. Two popular classes of
these techniques are low-rank and kernel methods. Each of these methods has its
own strengths. We observe these strengths synergistically complement each other
and exploit these synergies to fuse low-rank and kernel methods, producing a
new class of transformers: FLuRKA (Fast Low-Rank and Kernel Attention). FLuRKA
provide sizable performance gains over these approximate techniques and are of
high quality. We theoretically and empirically evaluate both the runtime
performance and quality of FLuRKA. Our runtime analysis posits a variety of
parameter configurations where FLuRKA exhibit speedups and our accuracy
analysis bounds the error of FLuRKA with respect to full-attention. We
instantiate three FLuRKA variants which experience empirical speedups of up to
3.3x and 1.7x over low-rank and kernel methods respectively. This translates to
speedups of up to 30x over models with full-attention. With respect to model
quality, FLuRKA can match the accuracy of low-rank and kernel methods on GLUE
after pre-training on wiki-text 103. When pre-training on a fixed time budget,
FLuRKA yield better perplexity scores than models with full-attention.