SchoenbAt: Rethinking Attention with Polynomial basis

TL;DR

SchoenbAt tackles the inefficiency of kernelized attention by replacing the Fourier-based expansion with a polynomial basis derived from Schoenberg’s theorem, enabling a Random Maclaurin Feature (RMFA) representation of dot-product kernels. A two-stage Pre-Post Scaling Batch Normalization (ppSBN) keeps inputs within a bounded space to satisfy the polynomial-basis assumptions while preserving output scale, yielding an unbiased estimator of kernelized attention. The authors prove unbiasedness and provide a finite-sample concentration bound, and validate the approach with experiments on the Long Range Arena showing substantial speedups and competitive accuracy across NLP and vision tasks. This work offers a drop-in replacement for attention that scales more favorably on long sequences without sacrificing performance, with practical implications for efficient Transformer variants.

Abstract

Kernelized attention extends the attention mechanism by modeling sequence correlations through kernel functions, making significant progresses in optimizing attention. Under the guarantee of harmonic analysis theory, kernel functions can be expanded with basis functions, inspiring random feature-based approaches to enhance the efficiency of kernelized attention while maintaining predictive performance. However, current random feature-based works are limited to the Fourier basis expansions under Bochner's theorem. We propose Schoenberg's theorem-based attention (SchoenbAt), which approximates dot-product kernelized attention with the polynomial basis under Schoenberg's theorem via random Maclaurin features and applies a two-stage regularization to constrain the input space and restore the output scale, acting as a drop-in replacement of dot-product kernelized attention. Our theoretical proof of the unbiasedness and concentration error bound of SchoenbAt supports its efficiency and accuracy as a kernelized attention approximation, which is also empirically validated under various random feature dimensions. Evaluations on real-world datasets demonstrate that SchoenbAt significantly enhances computational speed while preserving competitive performance in terms of precision, outperforming several efficient attention methods.

Figures (6)

• Figure 1: The module design structure of SchoenbAt.
• Figure 1: The dot-product kernels we study in this work along with their non-negative Maclaurin coefficients, where $\mathcal{K}(\cdot)=f(\langle x,y \rangle )$.
• Figure 2: Computation graphs for kernelized attention and RMFA. In each figure, the data on the left represents the input to the attention layer. Here, operators $(\cdot)$ and $\otimes$ respectively denote matrix multiplication and outer product. Each step's main computation cost is marked on the left side of the operators, and the dimensions of data and intermediate results are indicated with superscripts.
• Figure 3: The loss, perplexity (lower is better), and Bleu scores (higher is better) of the traditional Transformer with and without ppSBN across training epochs on machine translation task of Multi30k dataset. In each plot, solid lines represent the Transformer with ppSBN, while dashed lines represent the Transformer without ppSBN.
• Figure 4: The average absolute differences between SchoenbAt and kernelized attention on five dot-product kernels. Plots of different colors represent different dimensions of attention input.

Theorems & Definitions (8)

• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• proof
• proof
• proof
• proof