Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction

TL;DR

This paper introduces Token Statistics Self-Attention (TSSA), a linear-time attention mechanism for transformers that replaces pairwise similarity computations, achieving comparable performance with improved efficiency across vision and language tasks.

Contribution

The paper presents a novel linear-time attention module derived from a variational form of the MCR$^2$ objective, enabling more efficient transformer architectures.

Findings

TSSA achieves linear computational complexity with comparable accuracy.

Replacing standard attention with TSSA reduces computational cost significantly.

TSSA challenges the necessity of pairwise similarity in transformer success.

Abstract

The attention operator is arguably the key distinguishing factor of transformer architectures, which have demonstrated state-of-the-art performance on a variety of tasks. However, transformer attention operators often impose a significant computational burden, with the computational complexity scaling quadratically with the number of tokens. In this work, we propose a novel transformer attention operator whose computational complexity scales linearly with the number of tokens. We derive our network architecture by extending prior work which has shown that a transformer style architecture naturally arises by "white-box" architecture design, where each layer of the network is designed to implement an incremental optimization step of a maximal coding rate reduction objective (MCR$^2$). Specifically, we derive a novel variational form of the MCR$^2$ objective and show that the architecture that results from unrolled gradient descent of this variational objective leads to a new attention module called Token Statistics Self-Attention (TSSA). TSSA has linear computational and memory complexity and radically departs from the typical attention architecture that computes pairwise similarities between tokens. Experiments on vision, language, and long sequence tasks show that simply swapping TSSA for standard self-attention, which we refer to as the Token Statistics Transformer (ToST), achieves competitive performance with conventional transformers while being significantly more computationally efficient and interpretable. Our results also somewhat call into question the conventional wisdom that pairwise similarity style attention mechanisms are critical to the success of transformer architectures. Code will be available at https://github.com/RobinWu218/ToST.