FLuRKA: Fast and accurate unified Low-Rank & Kernel Attention

TL;DR

FLuRKA introduces a unified low-rank and kernel attention mechanism that significantly accelerates transformer training and inference while maintaining or improving accuracy across diverse tasks.

Contribution

The paper proposes FLuRKA, a novel fusion of low-rank and kernel attention methods, offering a training-efficient transformer with theoretical and empirical speed and quality guarantees.

Findings

FLuRKA achieves up to 3.3x speedup over low-rank methods.

FLuRKA attains up to 20x speedup over flash-attention models.

FLuRKA maintains or surpasses the accuracy of constituent methods across various tasks.

Abstract

Many efficient $\textit{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 strengths. We observe these strengths synergistically complement each other and exploit them to fuse low-rank and kernel methods, producing a new class of transformers: FLuRKA ($\textbf{F}$ast $\textbf{L}$ow-$\textbf{R}$ank & $\textbf{K}$ernel$ \textbf{A}$ttention). FLuRKA are highly $\textit{training-efficient}$ with faster model speeds $\textit{and}$ similar model qualities compared to constituent low-rank and kernel methods. We theoretically and empirically evaluate the speed and quality of FLuRKA. Our model speed analysis posits a variety of parameter configurations where FLuRKA exhibit speedups over low-rank and kernel approximations and our model quality analysis bounds the error of FLuRKA with respect to full-attention. Empirically, we instantiate three FLuRKA variants which experience speedups of up to 3.3x and 1.7x over low-rank and kernel methods respectively. This translates to speedups of up to 20x over models with flash-attention. Across a diverse set of tasks spanning language modeling, language understanding, long sequence modeling, machine translation, and image classification, FLuRKA achieve comparable accuracy with underlying low-rank and kernel approximations, occasionally surpassing both.