PoM: A Linear-Time Replacement for Attention with the Polynomial Mixer

TL;DR

PoM is a new token mixing mechanism that replaces self-attention with linear complexity, maintaining performance across various domains while significantly reducing computational costs.

Contribution

Introducing Polynomial Mixer (PoM), a linear-time token mixing method that preserves transformer universality and matches attention-based models' performance.

Findings

PoM achieves comparable results to attention models in multiple domains.

PoM drastically reduces computational cost for long sequences.

PoM satisfies the universal sequence-to-sequence approximation property.

Abstract

This paper introduces the Polynomial Mixer (PoM), a novel token mixing mechanism with linear complexity that serves as a drop-in replacement for self-attention. PoM aggregates input tokens into a compact representation through a learned polynomial function, from which each token retrieves contextual information. We prove that PoM satisfies the contextual mapping property, ensuring that transformers equipped with PoM remain universal sequence-to-sequence approximators. We replace standard self-attention with PoM across five diverse domains: text generation, handwritten text recognition, image generation, 3D modeling, and Earth observation. PoM matches the performance of attention-based models while drastically reducing computational cost when working with long sequences. The code is available at https://github.com/davidpicard/pom.