PADRe: A Unifying Polynomial Attention Drop-in Replacement for Efficient Vision Transformer

TL;DR

PADRe introduces a polynomial-based attention mechanism that unifies various efficient attention methods, offering faster computation with comparable accuracy across vision tasks.

Contribution

It proposes a unifying polynomial attention framework that replaces traditional self-attention, improving efficiency and encompassing recent alternative attention mechanisms.

Findings

PADRe is 11x to 43x faster than standard self-attention.

Maintains similar accuracy to traditional self-attention.

Effective across diverse vision tasks like classification and detection.

Abstract

We present Polynomial Attention Drop-in Replacement (PADRe), a novel and unifying framework designed to replace the conventional self-attention mechanism in transformer models. Notably, several recent alternative attention mechanisms, including Hyena, Mamba, SimA, Conv2Former, and Castling-ViT, can be viewed as specific instances of our PADRe framework. PADRe leverages polynomial functions and draws upon established results from approximation theory, enhancing computational efficiency without compromising accuracy. PADRe's key components include multiplicative nonlinearities, which we implement using straightforward, hardware-friendly operations such as Hadamard products, incurring only linear computational and memory costs. PADRe further avoids the need for using complex functions such as Softmax, yet it maintains comparable or superior accuracy compared to traditional self-attention. We assess the effectiveness of PADRe as a drop-in replacement for self-attention across diverse computer vision tasks. These tasks include image classification, image-based 2D object detection, and 3D point cloud object detection. Empirical results demonstrate that PADRe runs significantly faster than the conventional self-attention (11x ~ 43x faster on server GPU and mobile NPU) while maintaining similar accuracy when substituting self-attention in the transformer models.