Spectraformer: A Unified Random Feature Framework for Transformer

TL;DR

Spectraformer introduces a unified random feature framework for Transformer attention, achieving state-of-the-art performance among random feature methods and offering various trade-offs in efficiency and accuracy.

Contribution

We present Spectraformer, a systematic framework that unifies and advances random feature-based attention approximation in Transformers, setting new performance benchmarks.

Findings

Achieves performance comparable to top sparse and low-rank methods on Long Range Arena.

Establishes a new state-of-the-art for random feature-based Transformers.

Offers multiple variants with different accuracy, training time, and memory trade-offs.

Abstract

Linearization of attention using various kernel approximation and kernel learning techniques has shown promise. Past methods used a subset of combinations of component functions and weight matrices within the random feature paradigm. We identify the need for a systematic comparison of different combinations of weight matrices and component functions for attention learning in Transformer. Hence, we introduce Spectraformer, a unified framework for approximating and learning the kernel function in the attention mechanism of the Transformer. Our empirical results demonstrate, for the first time, that a random feature-based approach can achieve performance comparable to top-performing sparse and low-rank methods on the challenging Long Range Arena benchmark. Thus, we establish a new state-of-the-art for random feature-based efficient Transformers. The framework also produces many variants that offer different advantages in accuracy, training time, and memory consumption. Our code is available at: https://github.com/cruiseresearchgroup/spectraformer .