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

## Abstract

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 ( 11 × \!\times × ∼ similar-to \,\sim\, ∼ 43 × \times × faster on server GPU and mobile NPU) while maintaining similar accuracy when substituting self-attention in the transformer models.

## 1 Introduction

Transformers have been pivotal in recent advancements across natural language processing, computer vision, and speech processing. In computer vision, Vision Transformers (ViTs) [] are particularly significant, driving innovations in areas such as open-world recognition [], image and video generation [, ], and large multi-modal models [, ].

However, the computational and memory costs associated with self-attention operations in transformers can increase quadratically relative to input size. This computationally-demanding challenge is particularly pronounced in real-world computer vision applications, which often involve inherently large input sizes, such as high-resolution images and videos in virtual reality and camera applications, or large 3D point clouds in autonomous driving. The need to run transformers on hardware platforms with limited computational and memory resources further intensifies these challenges.

To address these computational hurdles, we propose PADRe, a unifying Polynomial Attention Drop-in Replacement scheme. PADRe is a novel approach that can effectively substitute standard self-attention at low (linear computational and memory) costs without sacrificing accuracy. PADRe’s unifying framework provides a guide for designing alternative transformer architectures. In particular, we observe and demonstrate that many recently-proposed attention replacement mechanisms (e.g., Hyena [], Mamba [], Conv2Former [], SimA [], Castling-ViT []), as well as the standard attention itself [], may in fact be interpreted as specific instances of the PADRe framework (See Table 1).

Our framework relies on polynomial approximants in the input to replace the standard self-attention mechanism. Polynomials are well-studied objects known to be powerful and efficient multi-dimensional function approximants over compact sets []. We use these existing results, together with the aforementioned observations, as guiding principles. We also leverage simple, mobile-friendly operations, such as Hadamard products, to induce nonlinearities, providing the necessary modeling capacity of the polynomial functions. Moreover, our design allows for easily and efficiently scaling up the modeling capacity of existing network architectures by increasing the degree, while maintaining linear computational complexity and memory cost.

In addition to discussing the theoretical underpinnings and unifying properties of PADRe, we further experimentally demonstrate the significant advantages of PADRe over the standard self-attention through a variety of experiments involving computer vision tasks in practical contexts. More specifically, by replacing the self-attention operation with our design, we maintain the model performance while only incurring linear computational and memory costs. Further, by profiling on actual hardware platforms, we demonstrate the significantly better on-device efficiency of PADRe as compared to the standard self-attention operation.

We summarize the main contributions of the paper as follows:

- •
We propose PADRe, a unifying framework for replacing the attention mechanism based on polynomial approximants, which provides an efficient, linear-complexity alternative to the standard self-attention without sacrificing accuracy.
- •
We show that PADRe provides a unifying formulation for many of the recently-proposed attention replacement schemes, including Hyena [], Mamba [], Conv2Former [], SimA [], and Castling-ViT [].
- •
We implement a specific instance of PADRe following the proposed design principles. We demonstrate the performance and advantage of PADRe through multiple experiments related to diverse, practical computer vision applications. Our PADRe-based models achieve similar or better accuracy than the original ones, while incurring significantly lower costs. Our on-device profiling further validates the computational efficiency of PADRe.

## 2 Related Work

Efficient Vision Transformer Backbones:
A majority of existing research proposes holistically designed vision transformer backbones. These backbones are capable of extracting features for various vision tasks, such as image classification, object detection, and segmentation.
Most of these methods employ a hierarchical architecture, where the spatial dimension of the image feature map is progressively downsampled []. This strategy allows the deeper transformer layers to operate on smaller inputs, thereby reducing computational and memory costs.
Specifically, many recent designs adopt convolution layers in the early stages of the network and apply transformers only to significantly downsampled feature maps later in the network [, , , ]. While this approach considerably improves efficiency, it still suffers from the inherent quadratic complexity when input size becomes large.
Building upon hierarchical and/or convolution-transformer hybrid architectures, some studies propose more efficient, alternative attention schemes, e.g., ReLU-based attention [], transposed attention [], convolutional modulation [], additive attention [], shift-add attention [], linear-angular attention []. However, all these models are designed to process a single 2D image and cannot be directly applied to more complex visual inputs, such as 3D point clouds. Many of them also rely on jointly designing the entire backbone to achieve good performance on 2D tasks and do not work well as drop-in replacements. In fact, using some of these efficient attention methods as drop-ins can degrade the model’s performance, as reported by [].

Efficient Attention Drop-in Replacements: To enhance the computational efficiency of vision transformers, various alternative attention mechanisms have been introduced. For instance, Swin [] proposes windowed attention, where self-attention computation is limited to local neighborhoods of an image. Although this theoretically reduces computation, the intricate reshaping and indexing operations pose practical challenges on resource-constrained hardware platforms. Several works propose linear-cost attentions [, , , ]. Some of them (e.g., [, ]) need to compute normalization factors at inference, which requires summing over all the tokens. This is inefficient on memory-constrained mobile platforms, especially when the number of tokens is large. Others (e.g., [, ]) are special cases of our proposed PADRe framework (see Section 3.5). Furthermore, it is worth noting that many of these alternative attentions still require multiple heads to maintain accuracy, a requirement that PADRe eliminates.

## 3 PADRe Framework Approach

We describe our framework’s technical details; at a high level, PADRe may be described as follows:

*Each element of the output tensor is a polynomial function of the elements of the input tensor, the coefficients of which are polynomial functions of the parameters (weights). Further, this dependency is such that the coefficients and the output itself may be computed efficiently (e.g., in linear time).*

For $N\in\mathbb{N}$ tokens and embedding / feature channel dimension $D\in\mathbb{N}$, consider an input $X\in\mathbb{R}^{N\times D}$. We denote $d$ as the degree of our framework; this parameter determines the strengths of the nonlinearities as discussed below. PADRe operations are performed on each batch element identically.

We omit batch dimension for brevity and consider a single head; if more than one head is used, each is treated independently using copies of the framework. PADRe consists of three main components: (1) Linear transformations; (2) Nonlinearities; (3) Optional operations (e.g., output resizing, normalization). Further details regarding each step are discussed below in the following sections.

### 3.1 Linear Transformations

Linear transformations consist of the following operations: for $i=1,...,d$, compute

$$ $Y_{i}=A_{i}\,X\,B_{i},$ (1) $$

where $A_{i}$ and $B_{i}$ are $N\!\times\!N$ and $D\!\times\!D$ weight matrices, and $Y_{i}\in\mathbb{R}^{N\times D}$. We impose *structure* on the matrices $A_{i}$ and $B_{i}$ for efficient multiplication; common structures include convolutions, sparse patterns, low-rank and hierarchical matrices, which can be optimized and parallelized on existing platforms. Note that left and right multiplications create a form of *mixing* among tokens (in the spatial dimension in the case of visual input) and in the embedding/channel dimension, respectively.
Coupled with nonlinearities (Section 3.2) this leads to the presence of *cross-terms* (i.e., multivariate monomials) in the expansion that are key to the expressivity of the polynomial functions.

### 3.2 Nonlinearities

We construct nonlinear, polynomial functions of the input by first defining the matrices:

$$ $\displaystyle Z_{1}=Y_{1},$ (2) $\displaystyle Z_{i+1}=(C_{i}\,Z_{i}\,D_{i})\odot Y_{i+1},\quad\text{for }i\in \{1,...,d-1\},$ (3) $$

where $\odot$ indicates the Hadamard (element-wise) product, which is an efficient, parallelizable, and hardware-friendly operation. $C_{i}$ and $D_{i}$ are $N\!\times\!N$ and $D\!\times\!D$ weight matrices that can perform additional token and channel mixing on the $Z_{i}$ tensor as needed, and $Z_{i}\in\mathbb{R}^{N\times D}$. Similarly, we require $C_{i}$ and $D_{i}$ to have structures that allow for efficient multiplications.
In particular, we have the following lemma which is important for the analysis of the scheme,

###### Lemma 1 .

The elements of $Z_{i}$ are homogeneous polynomials of degree $i$ in the input $X$.

###### Proof.

We proceed by induction. In the case $i=1$,

$$ $Y_{1}=A_{1}\,X\,B_{1},$ (4) $$

which shows that each element of $Y_{1}$ is a linear function of the input. Now assume the elements of $Z_{i}$ are homogeneous polynomials of degree $i$ in the input and consider

$$ $Z_{i+1}=(C_{i}\,Z_{i}\,D_{i})\odot(A_{i+1}\,X\,B_{i+1}).$ (5) $$

Observe that the elements of $A_{i+1}\,X\,B_{i+1}$ are homogeneous polynomials of degree $1$ by construction and the elements of $C_{i}\,Z_{i}\,D_{i}$ are homogeneous polynomials of degree $i$ by assumption and construction (linear transformations do not affect this conclusion). The expression is therefore a homogeneous polynomial of degree $i+1$ as claimed.
∎

Next, we compute dense (rather than homogeneous) polynomials of the input to increase diversity, expressivity, and number of degrees of freedom, through the computation of the output tensor $P$,

$$ $[P]_{m,n}=\sum_{i=1}^{d}[W]_{m,n,i}\,[Z_{i}]_{m,n}+[L]_{m,n},$ (6) $$

where $W\in\mathbb{R}^{N\times D\times d}$ is a weight tensor, $1\leq m\leq N$ and $1\leq n\leq D$ are token and embedding dimensions respectively, and $L\in\mathbb{R}^{N\times D}$ is the bias (i.e., $0^{\text{th}}$-order term).

###### Lemma 1 .

###### Proof.

### 3.3 Optional Operations

When an output of size different from to the input is required, we perform a computation of the form,

$$ $O=U\,P\,V,$ (7) $$

where $U\in\mathbb{R}^{F\times N}$, $V\in\mathbb{R}^{D\times G}$, and $O\in\mathbb{R}^{F\times G}$ (required output shape). As discussed, it is necessary that $U$ and $V$ possess structures (e.g., low-rank) to allow efficient application to $P$.

In addition, the $Y_{i}$’s tensors can optionally be normalized. While this is not necessary for the PADRe instances in our experiments, it can be a potentially useful regularization technique.

### 3.4 Overall Framework

Combining all operations together and leveraging Lemma 1, the output takes the following form:(^††\dagger†^††\dagger†$\dagger$We ignore the optional resizing here, which is just a linear transformation.)

$$ $\displaystyle[P]_{m,n}$ $\displaystyle=\sum_{i=1}^{d}[W]_{m,n,i}\,[Z_{i}]_{m,n}+[L]_{m,n}\sim\sum_{k\in \mathbb{N}^{N\times D}:|k|\leq d}\pi_{k}\,x^{k},$ (8) $$

where $k$ is a multi-index, $|k|=\sum_{m=1}^{N}\sum_{n=1}^{D}k_{m,n}$ is the absolute degree, $\{\pi_{k}\}$ are coefficients (which have a complex polynomial relationship with the parameters/weights), and $x^{k}=\prod_{m=1}^{N}\prod_{n=1}^{D}x^{k_{m,n}}_{m,n}$ are monomials. Eq. (8) is a degree-$d$ multi-dimensional polynomial in the entries of the input $X$; that is, PADRe uses hardware-friendly and efficient operations to produce an output tensor whose entries are (potentially normalized) polynomials of the entries of the input tensor. In particular, any proposed scheme that can be written in this form falls within the PADRe framework.

### 3.5 Unifying Framework

Eq. (8) provides a general, quantitative representation of PADRe’s output. This formula in fact encompasses many recently-proposed schemes.
Such observation may not be obvious at first sight but becomes clear following basic algebraic manipulations. For several recently proposed methods, we summarize their PADRe equivalence in Table 1 and provide detailed proofs to show that they are within our proposed PADRe framework in Appendix B. Further, it can be easily verified that the “attention" part of existing models based on the MetaFormer architecture [] and linear token mixing (e.g., convolution [], MLP [, ]) are degree-$1$ PADRe equivalences.

While we focus on leveraging polynomial functions, PADRe can be generalized to rational functions (i.e., ratios of polynomials) as discussed in Appendix A.2. Although the standard self-attention cannot be represented by polynomials, it can be approximated by high-degree rational functions, as we show in Appendix B.6.

**Table 1: These existing attention schemes can be shown to be PADRe equivalences or can be approximated by PADRe. Detailed proofs can be found in Appendix B.**
| Attention Schemes | PADRe Equivalence | Proofs |
| --- | --- | --- |
| SimA [] | Normalized, homogeneous degree-3 polynomial | B.1 |
| Conv2Former [] | Homogeneous degree-2 polynomial | B.2 |
| Hyena [] | Homogeneous degree-$N$ polynomial (order-$N$ operator, commonly $3$) | B.3 |
| Mamba [] | Homogeneous degree-3 polynomial (approx.) | B.4 |
| Castling-ViT [] | Degree-3 polynomial | B.5 |
| Self-attention [] | High-degree rational function/ratio of polynomials (approx.) | B.6 |

### 3.6 Computational Characteristics

We briefly analyze the number of parameters and computational cost of constructing an output $P\in\mathbb{R}^{N\times D}$ given an input $X\in\mathbb{R}^{N\times D}$ using PADRe. First, let us consider the number of parameters. We find $\mathcal{O}(N\cdot d)$ parameters for weight matrices $\{A_{i}\}$ and $\{C_{i}\}$, $\mathcal{O}(D\cdot d)$ parameters for weight matrices $\{B_{i}\}$ and $\{D_{i}\}$,(^††\dagger†^††\dagger†$\dagger$We assume a linear number of parameters is sufficient to characterize the matrices.) and $\mathcal{O}(N\cdot D\cdot d)$ for the tensor $W$. This results in a total of

$$ $\mathcal{O}(N\cdot D\cdot d)$ (9) $$

parameters. For a fixed degree $d$, this scales like $\mathcal{O}(N\cdot D)$ which is linear in the size of the input. Therefore, the *memory cost is linear* in the size of the input.

As for the computational cost, we find that computing all $Y_{i}$’s costs $\mathcal{O}(N\cdot D\cdot d)$ FLOPs using the assumed structure.(^††\dagger†^††\dagger†$\dagger$We assume the structure allows us to perform matrix-vector products in linear time.) Further, each quantity $Y_{i}$ may be computed in parallel (data parallelism).
Then, computing all $Z_{i}$’s carries a cost of $\mathcal{O}(N\cdot D\cdot d)$ FLOPs, which are computed sequentially based on the current scheme. Finally, computing the output $P$ comes at a cost of $\mathcal{O}(N\cdot D\cdot d)$ FLOPs.
In total, the cost of computing the output $P$ from the input $X$ is therefore,

$$ $\mathcal{O}(N\cdot D\cdot d),$ (10) $$

which is also linear in the size of the input for a fixed degree $d$.

We note the scheme only leverages highly-efficient matrix operations and not complex functions (e.g., $\mathrm{exp}(\cdot),\,\mathrm{cos}(\cdot)$), which are often expensive on hardware and can cause quantization complications.

## 4 Experiments

In this section, we comprehensively evaluate PADRe as a drop-in replacement of the standard self-attention, in terms of both accuracy and efficiency. More specifically, we implement a specific instance of PADRe following the design principles laid out in Section 3, and use it to substitute the self-attention in transformer models in several vision applications, such as image classification, single-image object detection, point cloud object detection, and multi-camera depth estimation. We then measure the latency of our PADRe instances on two compute platforms: a server GPU and a mobile Neural Processing Unit (NPU), and compare with the standard self-attention. We show that PADRe provides similar or better accuracy on all the tasks, while being significantly faster, especially when the number of input tokens is large.

### 4.1 Implementation of PADRe

As illustrated in Fig. 1, given an input tensor $X$, we first apply channel mixing and token mixing operations to generate $Y_{i}$, for each $i\in\{1,...,d\}$; these mixing operations correspond to the linear transformations in Eq. 1. More specifically, we implement channel mixing using pointwise convolutions or linear layers. Token mixing is implemented with convolutions applied in the token dimension. When the input possesses an inherent 2D structure (e.g., image features), we use 11x11 2D convolutions to mix the tokens, after reshaping them into a 2D image format. On the other hand, when the input does not present a dense 2D structure, such as sparse 3D voxels from a point cloud, we treat them as a sequence of tokens and use 1D convolutions with kernel size 11 for mixing.

Next, we compute $Z_{i}$ for each degree $i$ based on $Y_{i}$ and $Z_{i-1}$, as described in Eq. 3. Channel and token mixings are further applied to $Z_{i}$ before it is used to produce degree-$(i+1)$ terms. Finally, all the $Z_{i}$’s are added together to generate the output $P$. Note that unlike the standard multi-head attention and existing alternative attention methods, we use a single head in all our experiments.(^††\dagger†^††\dagger†$\dagger$While using multiple heads is possible and straightforward to implement, we found that a single head was sufficient for PADRe to match the accuracy of multi-head attention models.)

Figure: Figure 1: Our implementation of a PADRe-based attention drop-in replacement module. In our experiments, we use this to substitute the standard self-attention parts in existing transformer models. Note that this implementation approach is just a specific instance of our general PADRe framework.
Refer to caption: extracted/5733533/figures/network_1.jpg

In this specific instance, we do not have a degree-$1$ term in the summation that generates $P$. This is because there already exists a skip connection from $X$ to the output of the attention block in transformer, which provides a degree-$1$ term. In addition, we do not require an explicit degree-0 0 term (i.e., $L$ in Eq. 6) in the final summation, since the final convolutional layers provide the bias terms.

### 4.2 Performance Evaluation on Vision Applications

Image Classification:
First, we consider the standard image classification task on ImageNet [].(^††\dagger†^††\dagger†$\dagger$ImageNet terms of use can be found here: [https://www.image-net.org/download](https://www.image-net.org/download).) We use DeiT [] as the baseline, including the tiny, small, and base network configurations, and replace the multi-head attentions with PADRe. We follow the default DeiT training setting (300 epochs without distillation) to train our models. Each training experiment takes $\sim$2 days on 8 NVIDIA A100 GPUs. We use the top-1 accuracy to measure classification performance.(^††\dagger†^††\dagger†$\dagger$We use the mmpretrain repository: [https://github.com/open-mmlab/mmpretrain](https://github.com/open-mmlab/mmpretrain) for model training and testing. Both mmpretrain and DeiT codebases are under Apache License 2.0.)

Table 2 summarizes the results. It can be seen that by replacing self-attention with our PADRe instance, we achieve similar or better classification performance. Further, PADRe can run significantly faster on actual hardware platforms and scales much more efficiently w.r.t. increasing number of tokens (see Section 4.3).

Table 2 also compares PADRe with several recently proposed alternative attention methods, using DeiT baseline and the default DeiT training procedure for a fair comparison. We see that several proposed methods, such as Hydra attention [] and linear-angular attention [], fail to maintain accuracy. Other methods do improve accuracy, but are not as efficient as PADRe on actual compute platforms (Section 4.3).

**Table 2: Results on ImageNet-1K classification. By replacing standard self-attention with our PADRe instance, we achieve similar or better accuracy when comparing to the DeiT baseline and recently proposed alternative attention methods.**
| Baselines | Attention Mechanisms | Top-1 Accuracy$\,\uparrow$ | FLOPs (G) | #Params (M) |
| --- | --- | --- | --- | --- |
| DeiT-Tiny [] | Standard self-attention [] | 72.2 | 1.3 | 5 |
| Hydra attention [] | 68.3 | 1.1 | 6 |  |
| Efficient attention [] | 70.2 | 1.1 | 6 |  |
| Linear-angular attention [] | 70.8 | 1.1 | 6 |  |
| Enhanced linear attention [] | 72.9 | 1.1 | 6 |  |
| k-NN attention [] | 73.0 | 1.3 | 6 |  |
| Focused linear attention [] | 74.1 | 1.1 | 6 |  |
| Vision Mamba [] | 76.1 | 2.6 | 7 |  |
| PADRe-2 (ours) | 74.5 | 1.2 | 5 |  |
|  | PADRe-3 (ours) | 76.5 | 1.7 | 7 |
| DeiT-Small [] | Standard self-attention [] | 79.8 | 4.6 | 22 |
| SimA [] | 79.8 | 4.6 | 22 |  |
| k-NN attention [] | 80.1 | 4.6 | 22 |  |
| Vision Mamba [] | 80.5 | 9.6 | 26 |  |
| PADRe-2 (ours) | 80.4 | 4.8 | 20 |  |
|  | PADRe-3 (ours) | 80.5 | 6.1 | 26 |
| DeiT-Base [] | Standard self-attention [] | 81.8 | 17.6 | 86 |
| Hydra attention [] | 76.4 | 16.9 | - |  |
| Hyena [] | 78.5 | - | 87 |  |
| PADRe-2 (ours) | 81.4 | 18.8 | 80 |  |

Single-Image Object Detection:
In this part, we consider a commonly used, end-to-end transformer-based object detector: DETR [].(^††\dagger†^††\dagger†$\dagger$We use the official DETR repository: [https://github.com/facebookresearch/detr](https://github.com/facebookresearch/detr), which is under Apache License 2.0.) DETR consists of three main components: a ResNet backbone for image feature extraction, a transformer encoder with 6 (self-attention) layers, and a transformer decoder with 6 (self- and cross-attention) layers. In this application, we replace the self-attention in the 6 encoder layers with our PADRe instance. We train and evaluate our DETR with PADRe replacement on the COCO detection benchmark [].(^††\dagger†^††\dagger†$\dagger$COCO is under a CC BY 4.0 license and the terms of use can be found here: [https://cocodataset.org/#termsofuse](https://cocodataset.org/#termsofuse).) We follow the short-schedule [] to train the model for 300 epochs on 8 NVIDIA A100 GPUs, which takes $\sim$4 days. We use the standard metric of Average Precision (AP) to measure detection performance.(^††\dagger†^††\dagger†$\dagger$See [https://cocodataset.org/#detection-eval](https://cocodataset.org/#detection-eval) for more details on COCO object detection metrics.)

The results in Table 3 show that by substituting the standard self-attention with PADRe in the transformer encoder layers, we achieve a similar detection accuracy. In the encoder part, our PADRe instance also significantly reduces the computational cost by 55%.

**Table 3: Results on COCO detection, using DETR as the baseline. We achieve a similar accuracy but with a considerably lower computational cost.**
| Model | Attention Mechanisms | AP$\,\uparrow$ | FLOPs (G) | #Params (M) |
| --- | --- | --- | --- | --- |
| DETR [] | Standard self-attention [] | 40.6 | 10.1 | 6.1 |
| PADRe-2 (ours) | 40.2 | 4.5 | 2.7 |  |

Point Cloud Object Detection:
Next, we evaluate PADRe on a more complex vision application where the input is no longer a single image. In this case, the input is a 3D point cloud. We consider the 3D point cloud object detection task on nuScenes [], which involves detecting vehicles and other traffic elements based on a LiDAR point cloud.(^††\dagger†^††\dagger†$\dagger$nuScenes is under a CC BY-NC-SA 4.0 license and the terms of use can be found here: [https://www.nuscenes.org/terms-of-use](https://www.nuscenes.org/terms-of-use).) We select a transformer-based 3D object detection model, DSVT [], which achieves state-of-the-art performance on nuScenes.(^††\dagger†^††\dagger†$\dagger$We use the official DSVT repository: [https://github.com/Haiyang-W/DSVT](https://github.com/Haiyang-W/DSVT), which is under Apache License 2.0.) In DSVT, a 3D transformer backbone extracts point cloud features, after which the features are projected to the birds-eye-view and further processed by 2D networks and a final detection head. We replace the self-attention in the 3D backbone with PADRe and follow the same training procedure as in []; training takes $\sim$2 days on 4 NVIDIA V100 GPUs. We use standard nuScenes metrics to measure detection performance, including mean average precision (mAP) and nuScenes Detection Score (NDS); NDS is computed as a weighted sum of 6 detection metrics over 10 object classes (see [] for more details on the metrics). As the 3D voxels are sparse and cannot be arranged in a dense 2D format, we treat them as a sequence and use 1D convolutions to perform token mixing.

Table 4 summarizes the 3D object detection results on nuScenes validation set. After replacing the self-attentions with PADRe, we achieve similar detection accuracy (in terms of both NDS and mAP) as compared to the original model.
We then measure the computational costs of the 3D transformer backbones in the baseline and our model. It can be seen that our modified DSVT 3D backbone uses 28% fewer FLOPs while maintaining detection accuracy as compared to the baseline.

**Table 4: Results on nuScenes LiDAR-based 3D object detection. We achieve a similar accuracy with a significantly lower computational cost.**
| Baseline | Attention Mechanisms | NDS$\,\uparrow$ | mAP$\,\uparrow$ | FLOPs (G) | #Params (M) |
| --- | --- | --- | --- | --- | --- |
| DSVT [] | Standard self-attention [] | 71.1 | 66.4 | 14.9 | 1.1 |
| PADRe-2 (ours) | 71.1 | 66.0 | 10.7 | 0.9 |  |

### 4.3 Latency Evaluation on Hardware Platforms

We evaluate and compare the on-device latencies of our PADRe instance and the standard self-attention. Our goal is to provide a practical comparison of computational efficiency across different hardware platforms. Specifically, we measure latencies on two compute platforms:
1) Nvidia A100 GPU, and 2) Qualcomm^\faRegistered Hexagon${}^{\scalebox{0.4}{\text{TM}}}$ NPU (included in the Snapdragon^\faRegistered 8 Gen 3 processor on Samsung S24),(^††\dagger†^††\dagger†$\dagger$Snapdragon and Qualcomm branded products are products of Qualcomm Technologies, Inc. and/or its subsidiaries. Qualcomm patented technologies are licensed by Qualcomm Incorporated.) which is an AI accelerator specialized for neural network workloads. On the NPU, we quantize the network weights and activations to Int8 when measuring latency.

Since our objective is to replace the standard self-attention, we specifically evaluate only the “attention” components on-device. More specifically, for PADRe, we run the instance described in Section 4.1 and shown in Fig. 1, and for the standard transformer, we run a multi-head attention operation. Using the hyperparameters of DeiT-Tiny, we set the embedding dimension to 192 in both PADRe and self-attention, and use 3 heads in the multi-head attention. Note that we only use a single head in our PADRe instance in all our experiments.
We evaluate on different numbers of input tokens, including: 256, 1024, 2304, and 4096, which correspond to input image resolutions of 256$\times$256, 512$\times$512, 768$\times$768, 1024$\times$1024, assuming a 16$\times$16 patchification (as used in ViT).

Fig. 2 shows the on-device latencies w.r.t. number of tokens. We see that the standard attention mechanism exhibits a significant latency growth as the number of tokens increases. In contrast, our proposed PADRe maintains linear latency proportional to the number of tokens. This confirms that PADRe provides a significant efficiency advantage over self-attention when deployed on actual hardware, including large GPUs and resource-constrained mobile NPUs.

Figure: Figure 2: On-device (GPU and NPU) latency comparison for PADRe vs. standard self-attention. It can be observed that the latency of self-attention escalates significantly with an increase in the number of input tokens. In contrast, PADRe demonstrates a linear growth pattern.
Refer to caption: extracted/5733533/figures/profiling_both.jpg

### 4.4 Ablation Study on Polynomial Degree in PADRe

In this study, we scale up the polynomial degree of our PADRe instance following the implementation in Section 4.1, and show its effect on the model performance and computational costs. Table 5 shows that by increasing the polynomial degree, the network obtains stronger modeling capacity, leading to performance improvements. Specifically, when using degree $3$, PADRe has a significantly better accuracy over the baseline model, while being more computationally efficient on device. By further increasing the degree beyond 3, we observe additional small accuracy gains.

**Table 5: ImageNet classification accuracy when scaling up PADRe degree.**
| Baseline | Attention Mechanisms | Top-1 Accuracy$\,\uparrow$ |
| --- | --- | --- |
| DeiT-Tiny [] | Standard self-attention [] | 72.2 |
| PADRe-2 (ours) | 74.5 |  |
| PADRe-3 (ours) | 76.5 |  |
| PADRe-4 (ours) | 76.9 |  |

Table $\dagger$ ‣ 6 shows the computational costs of our PADRe instance with different degrees, and compares those with the standard self-attention and other recently-proposed efficient attention methods such as linear-angular attention [] and focused linear attention []. We follow the same setup in the previous section, where we set the channel to 192, use 3 heads in self-attention and existing efficient attentions, and 1 head in our PADRe instance. The number of input tokens is fixed to be $4096$.

Increasing the degree in PADRe allows for efficiently scaling up model capacity and performance. Even with degree-$4$, PADRe remains significantly faster than self-attention, especially on the mobile NPU. While existing efficient attentions use fewer FLOPs, some can have higher latency on device. It is noteworthy that FLOPs is not an accurate indicator of model efficiency on actual hardware; similar observations have also been made in other works, e.g., []. For instance, operations that are not well optimized or cause memory bottlenecks on device are not reflected by FLOP count.

**Table 6: FLOPs and on-device (both GPU and NPU) latency comparison of standard self-attention, recent efficient attention methods, and PADRe of various degrees.(^††\dagger†^††\dagger†$\dagger$We do not report the latencies of these alternative attention methods on NPU, as it would require significant effort to rewrite the implementation in order to run them on the mobile NPU.) The measurement is based on the attention parts only, with 4096 input tokens.**
| Attention Mechanisms | FLOPs (G) | Latency on GPU (ms) | Latency on NPU (ms) |
| --- | --- | --- | --- |
| Standard self-attention [] | 1.31 | 0.43 | 68.60 |
| Linear-angular attention [] | 1.22 | 0.53 | - |
| Focused linear attention [] | 1.25 | 0.08 | - |
| PADRe-2 (ours) | 1.10 | 0.01 | 1.62 |
| PADRe-3 (ours) | 2.28 | 0.02 | 3.17 |
| PADRe-4 (ours) | 3.27 | 0.04 | 4.22 |

## 5 Discussion

Limitations and Future Work: In this paper, we have extensively studied using PADRe to replace the standard self-attention and have shown that this is more computationally efficient and maintains model performance. However, we did not present empirical results of its application to cross-attention and/or multi-modal operations. In principle, PADRe can be extended for cross-attention, a discussion of which is provided in Appendix A.1; this would be interesting as part of future work.

We observe that model performance begins to saturate when scaling up the degree of our implemented PADRe instance beyond 3 (see Table 5). We have yet to establish the source of this phenomenon. One possibility may have to do with the numerical stability of monomials. In future work, we will investigate the use of different, more stable, polynomial bases, such as orthogonal polynomials.

Broader Impacts: Our proposed PADRe approach provides an efficient replacement to the standard self-attention, which requires extensive computation and memory costs, especially when the input size is large (e.g., high-resolution images and videos). By reducing the computation and memory footprints of large AI models, PADRe can generate positive societal impacts, such as reducing energy consumption and making AI models more accessible to users with limited computational resources.
During the study, several experiments failed to finish and could not make it into the paper (due to server issues, lack of time, etc.), which caused extra energy consumption.

## 6 Conclusions

In this paper, we proposed a unified, polynomial drop-in attention replacement framework, PADRe, which leverages polynomials to replace the expensive standard self-attention in transformers. More specifically, we presented the key components for constructing an expressive, efficient, and hardware-friendly scheme to substitute self-attention. In addition, PADRe is a general framework that encompasses many recently-proposed alternative attention schemes; we provided detailed proofs to show this. In order to demonstrate the efficacy of PADRe on practical applications and actual hardware platforms, we implemented a specific instance of PADRe based on the design principles, using efficient and hardware-friendly operations. Through extensive experiments, we showed that PADRe provides an efficient replacement of self-attention that maintains model performance and incurs much lower computational costs. Specifically, we showed that our PADRe instance runs significantly faster on both a power GPU (NVIDIA A100) and a resource-constrained mobile NPU (Qualcomm Hexagon tensor processor), by as much as 43$\times$ when comparing to the standard self-attention.

## Appendix A PADRe Generalizations

In this section, we present a few generalization of the framework, namely, generalization to the cross-attention mechanism, as well as a rational version of the framework.

### A.1 Multi-Modal PADRe (i.e., Cross-Attention)

PADRe can be generalized to multi-modal input in many ways. We present one here. Imagine that $X^{(a)}\in\mathbb{R}^{B\times N\times D}$ represents input from one mode and $X^{(b)}\in\mathbb{R}^{B\times M\times D}$ represents input from another mode. Then, we can define $\hat{Y}^{(a)}_{i}$ and $\hat{Y}^{(b)}_{i}$ for each mode just as before (linear transformations should be chosen so that they are of equal sizes). To capture nonlinear interactions among modes, the difference lies in the computation of $\hat{Z}_{j}$. In this case, we write,

$$ $\displaystyle Z_{j}=\odot_{i=1}^{j}Y^{(c_{i})}_{i},$ (11) $$

where $c_{i}\in\{a,b\}$ for each $i$. For instance, if the degree is $d=3$ and $c=(a,a,b)$, then one finds that,

$$ $\displaystyle Z_{j}$ $\displaystyle=Y^{(a)}_{0}$ (12) $\displaystyle Z_{j}$ $\displaystyle=Y^{(a)}_{0}\odot Y^{(a)}_{1}$ (13) $\displaystyle Z_{j}$ $\displaystyle=Y^{(a)}_{0}\odot Y^{(a)}_{1}\odot Y^{(b)}_{2}.$ (14) $$

It is important for the sequence is not trivial, i.e., to contain cross-terms, not $(a,a,a)$ nor $(b,b,b)$. This will capture nonlinear interactions between modes. Also, note that more than two modes could be involved, and that multiple sequences may be considered.

### A.2 Rational Approximation

This section describes a generalization of PADRe. Although we found that the standard PADRe based on polynomial approximations offered great performance in the context of ViT, it is expected that certain applications will benefit from this more powerful framework. In particular, the following approach is inspired by the improved effectiveness of PADRe (rational) approximants compared with polynomial (Taylor) approximants.

The main difference is that the drop-in replacement takes the form of a *rational functions* rather than a *polynomial function*. As discussed in , in the one-dimensional case, the former have much higher power than the latter. Further, from an intuitive perspective, adding flexibility to the denominator allows one to better capture "poles" (or areas of high variability) often observed when using softmax. Finally, the introduction of a function at the denominator might be seen as a replacement for the normalization generally used prior to the attention mechanism.

From a technical standpoint, we construct a degree-$[d/e]$ rational PADRe as follows: just as before, we let,

$$ $Y_{i}=A_{i}\,X\,B_{i}$ (15) $$

for $i=1,...,d+e$, and,

$$ $\displaystyle\mathcal{K}_{j}$ $\displaystyle=\odot_{i=1}^{j}\hat{Y}_{i},$ (16) $\displaystyle\mathcal{L}_{k}$ $\displaystyle=\odot_{i=1}^{k}\hat{Y}_{d+i},$ (17) $$

for $j=1,...,d$ and $k=1,...,e$. Then, we define

$$ $\displaystyle\mathcal{N}_{a,b}$ $\displaystyle=\sum_{j=1}^{d}W_{a,b,j}\,[\mathcal{K}_{j}]_{a,b}+V_{a,b}$ (18) $\displaystyle\mathcal{D}_{a,b}$ $\displaystyle=\sum_{k=1}^{e}Q_{a,b,j}\,[\mathcal{L}_{j}]_{a,b}+P_{a,b}.$ (19) $$

Finally, the output (prior to resizing) takes the form (i.e., element-wise "Hadamard division"),

$$ $O_{a,b}=\frac{\mathcal{N}_{a,b}}{\mathcal{D}_{a,b}}$ (20) $$

Finally, note that care should be taken at training time to avoid a vanishing denominator (and ensuing numerical instability). This may be done by, for example, squaring the denominator and adding a small regularizing addition term.

## Appendix B PADRe- Special cases

### B.1 SimA [ 13 ]

SimA (Simple Attention), as its name suggests, is a simple mechanism for replacing the standard attention. The output $O$ of SimA can be expressed as,

$$ $O=\hat{Q}\,\hat{K}^{T}\,V$ (21) $$

where,

$$ $\displaystyle\hat{Q}^{i}=\frac{Q^{i}}{||Q^{i}||_{1}},\,\hat{K}^{i}=\frac{K^{i} }{||K^{i}||_{1}},$ (22) $$

and $\cdot^{i}$ indicates the $i^{th}$ column of a matrix ($\ell_{1}$ column-wise normalization), and where,

$$ $\displaystyle Q=X\,W_{Q},\,K=X\,W_{K},\,V=X\,W_{V},\,$ (23) $$

for some appropriate projection matrices as usual. Substitution implies that,

$$ $\displaystyle O_{m,n}$ $\displaystyle=\sum_{i=1}^{D}\hat{Q}_{m,i}\left(\sum_{j=1}^{N}\hat{K}_{j,i}\,V_ {j,n}\right)$ (24) $\displaystyle=\sum_{i=1}^{D}\frac{[X\,W_{Q}]_{m,i}}{||Q^{i}||_{1}}\left(\sum_{ j=1}^{N}\frac{[X\,W_{K}]_{j,i}}{||K^{i}||_{1}}\,[X\,W_{V}]_{j,n}\right)$ (25) $\displaystyle=\sum_{i=1}^{D}\sum_{j=1}^{N}\,\frac{1}{||Q^{i}||_{1}||K^{i}||_{1 }}\,\left(\sum_{k_{Q}=1}^{D}\sum_{k_{K}=1}^{D}\sum_{k_{V}=1}^{D}\left([W_{Q}]_ {k_{Q},i}[W_{K}]_{k_{K},i}\,[W_{V}]_{k_{V},n}\right)\,\left(X_{m,k_{Q}}\,X_{j, k_{K}}\,X_{j,k_{V}}\right)\,\right)$ (26) $$

This should be recognized as a homogeneous degree-$3$ polynomial with coefficients of the form,

$$ $\frac{\left([W_{Q}]_{k_{Q},i}[W_{K}]_{k_{K},i}\,[W_{V}]_{k_{V},n}\right)}{||Q^ {i}||_{1}||K^{j}||_{1}}$ (27) $$

i.e., exhibiting a (normalized) polynomial relationship with the parameters/weights. Therefore, as claimed, SimA is a special case of the PADRe framework.

### B.2 Conv2Former [ 11 ]

Given an input $X\in\mathbb{R}^{H\times W\times C}$, Conv2Former replaces the self-attention mechanism with a block of the form,

$$ $\displaystyle Z$ $\displaystyle=A\odot V$ (28) $\displaystyle A$ $\displaystyle=\mathrm{DConv}_{k\times k}\left(W_{1}\,X\right)$ (29) $\displaystyle V$ $\displaystyle=W_{2}\,X$ (30) $$

where $Z$ is the output, $\odot$ represents the Hadamard product, $W_{1}$ and $W_{2}$ are weight matrices of two linear layers, and $\mathrm{DConv}_{k\times k}$ denotes a depthwise convolution with kernel size $k\times k$, which is really just an efficient linear transformation applied to each channel.

Letting $C$ be a tensor representing the convolution kernel and expanding formulas, we find that,

$$ $\displaystyle\left[\mathrm{DConv}_{k\times k}\left(W_{1}\,X\right)\right]_{i,j ,k}$ $\displaystyle=\sum_{p,q}C_{i-p,j-q,k}\,\left(\sum_{m,n}[W_{1}]_{(p,q),(m,n)}\, X_{m,n,k}\right)$ (31) $\displaystyle\left[W_{2}\,X\right]_{i,j,k}$ $\displaystyle=\sum_{s,t}[W_{2}]_{(i,j),(s,t)}\,X_{s,t,k}$ (32) $$

Therefore, it follows that,

$$ $\displaystyle[Z]_{i,j,k}=\left[A\odot V\right]_{i,j,k}$ $\displaystyle=\left(\sum_{p,q}C_{i-p,j-q,k}\,\left(\sum_{m,n}[W_{1}]_{(p,q),(m ,n)}\,X_{m,n,k}\right)\right)\cdot\left(\sum_{s,t}[W_{2}]_{(i,j),(s,t)}\,X_{s, t,k}\right)$ (33) $\displaystyle=\sum_{m,n}\sum_{s,t}\,\left([W_{2}]_{(i,j),(s,t)}\,\sum_{p,q}C_{ i-p,j-q,k}[W_{1}]_{(p,q),(m,n)}\right)\,X_{m,n,k}\,X_{s,t,k}$ (34) $$

This should be recognized as a degree-2 (homogeneous) polynomial in the input $X$, and therefore an specific instance of the PADRe framework.

### B.3 Hyena [ 20 ]

The proposed Hyena operator is defined as follows (Definition 3.1, []),

###### Definition B.1 .

*(Order-$N$ Hyena Operator)*
Let $(x^{0},x^{1},...,x^{N})$ be projections of the input and let $h^{1},...,h^{N}$ be a set of learnable filters. The Hyena-$N$ operator is defined by the recurrence:

$$ $\displaystyle z^{1}_{t}$ $\displaystyle=x^{0}_{t}$ (35) $\displaystyle z^{n+1}_{t}$ $\displaystyle=x^{n}_{t}\;(h^{n}\ast z^{n})_{t},\;\;n=1,...,N$ (36) $\displaystyle y_{t}$ $\displaystyle=z^{N+1}_{t}$ (37) $$

That is, Hyena proposes to use compositions of $N$ element-wise multiplications and convolutions as a drop-in replacement for the attention. Such operators possess these two important properties:

- 1.
Homogeneous polynomial functions: the order-$N$ Hyena Operator corresponds to the evaluation of a degree-$(N+1)$ homogeneous polynomial (or equivalently $(N+1)$-tensor) which coefficients are non-linear functions of the parameters.(^††\dagger†^††\dagger†$\dagger$Higher levels of non-linearity ($N$) increase expressivity, while the global nature of convolutions and projections ensures that all correlations among tokens are taken into account at multiple scales.)
- 2.
Fast evaluation: the non-linear dependence of the coefficients (or equivalent the tensor core elements) on the parameters are such that the operator may be evaluated rapidly (i.e., in quasi-linear time).

More precisely, letting $\chi$ be the original input vector and $n:=(n_{0},n_{1},...,n_{N})\in\mathbb{N}^{N+1}$ be multi-indices,(^††\dagger†^††\dagger†$\dagger$We use $\chi$ in order to differentiate between the Hyena notation which uses $x$ as projections rather than input.) we argue that the order-$N$ Hyena operator may be written as,

$$ $y=\sum_{n}\eta_{n}\,\left(\prod_{i=0}^{N}\chi_{n_{i}}\right)$ (38) $$

for some coefficients $\{\eta_{n}\}$ that are polynomial functions of the parameters ($\{h^{i}\}$ and linear projections). This is a a homogeneous polynomial of degree $(N+1)$. In light of these observations, one finds that *Hyena is in fact a special case of PADRe*.

We present analysis of the Order-$N$ Hyena Operator which shows it is a specific instances of PADRe. In particular, we argue that the latter really is a homogeneous polynomial of degree $(N+1)$, which may also be interpreted as an $(N+1)$-tensor (multilinear functional/operator). Following the above definition and Remark 3.1 ([]), an order-$N$ Hyena operator may be expressed as,

$$ $y=x^{N}\cdot(h^{N}\ast(x^{N-1}\cdot(h^{N-1}\ast(...))))$ (39) $$

This expression can be expanded as follows,

$$ $\displaystyle y_{t_{N}}=x^{N}_{t_{N}}\,\sum_{t_{N-1}=0}^{L}h^{N}_{t_{N}-t_{N-1 }}\,\left(x^{N}_{t_{N-1}}\,\sum_{t_{N-2}=0}^{L}h^{N-1}_{t_{N-1}-t_{N-2}}\, \left(...\left(x^{1}_{t_{1}}\,\sum_{t_{0}=0}^{L}h^{1}_{t_{1}-t_{0}}\,x^{0}_{t_ {0}}\right)\right)\right)$ (40) $$

Using commutativity to re-arrange the latter, we find that

$$ $\displaystyle y_{t_{N}}$ $\displaystyle=\sum_{i=1}^{N}\,\sum_{t_{i-1}=0}^{L}\,\left(\prod_{i=1}^{N}\,h^{ i}_{t_{i}-t_{i-1}}\right)\,\left(\prod_{i=0}^{N}x^{i}_{t_{i}}\right),$ (41) $$

where $t_{i}\in\{0,...,L\}$ for every $i$. Now, recall that each vector $x^{i}$ represents a projection of the input $\chi\in\mathbb{R}^{M}$ onto a $L$-dimensional space.(^††\dagger†^††\dagger†$\dagger$$M$ is used to represent the dimension of the input, whatever it is.) Let us write,

$$ $x^{i}_{t_{i}}=[P^{i}\,\chi]_{t_{i}}=\sum_{m_{i}=1}^{M}P^{i}_{t_{i},m_{i}}\, \chi_{m_{i}}.$ (42) $$

Upon substitution, we find that,

$$ $\displaystyle\prod_{i=0}^{N}x^{i}_{t_{i}}$ $\displaystyle=\prod_{i=0}^{N}\left(\sum_{m_{i}=1}^{M}P^{i}_{t_{i},m_{i}}\,\chi _{m_{i}}\right)$ (43) $\displaystyle=\sum_{m_{0}=0}^{M}...\sum_{m_{N}=0}^{M}\,\left(\prod_{j=0}^{N}P^ {j}_{t_{j},m_{j}}\right)\,\left(\prod_{j=0}^{N}\chi_{m_{j}}\right)$ (44) $$

and therefore,

$$ $\displaystyle y_{t_{N}}$ $\displaystyle=\sum_{i=1}^{N}\,\sum_{t_{i-1}=0}^{L}\,\left(\prod_{i=1}^{N}\,h^{ i}_{t_{i}-t_{i-1}}\right)\,\left(\sum_{m_{0}=0}^{M}...\sum_{m_{N}=0}^{M}\, \left(\prod_{j=0}^{N}P^{j}_{t_{j},m_{j}}\right)\,\left(\prod_{j=0}^{N}\chi_{m_ {j}}\right)\right)$ (45) $\displaystyle=\sum_{m_{0}=0}^{M}...\sum_{m_{N}=0}^{M}\,\left(\sum_{i=1}^{N} \sum_{t_{i-1}=0}^{L}\,\left(\prod_{i=1}^{N}\,h^{i}_{t_{i}-t_{i-1}}\,\prod_{j=0 }^{N}P^{j}_{t_{j},m_{j}}\right)\right)\,\left(\prod_{j=0}^{N}\chi_{m_{j}}\right)$ (46) $\displaystyle=\sum_{m_{0}=0}^{M}...\sum_{m_{N}=0}^{M}\,\eta_{(m_{0},m_{1},..., m_{N})}\,\left(\prod_{j=0}^{N}\chi_{m_{j}}\right)$ (47) $$

as claimed. The nonlinear dependence of the coefficients $\eta_{m}$ on the parameters $\{h^{i}\}_{i=1}^{N}$ and $\{P^{j}\}_{j=0}^{N}$ takes the form,

$$ $\displaystyle\eta_{m}$ $\displaystyle=\sum_{t_{0}=0}^{L}...\sum_{t_{N-1}=0}^{L}\,\left(\prod_{i=1}^{N} \,h^{i}_{t_{i}-t_{i-1}}\,\prod_{j=0}^{N}P^{j}_{t_{j},m_{j}}\right)$ (48) $\displaystyle=\sum_{t_{0}=0}^{L}...\sum_{t_{N-1}=0}^{L}\,\left(P^{N}_{t_{N},m_ {N}}\,\left(\prod_{i=1}^{N}\,h^{i}_{t_{i}-t_{i-1}}\,P^{i-1}_{t_{i-1},m_{i-1}} \right)\right)$ (49) $\displaystyle=P^{N}_{t_{N},m_{N}}\,\prod_{i=1}^{N}\left(\sum_{t_{i-1}=0}^{L}h^ {i}_{t_{i}-t_{i-1}}\,P^{i-1}_{t_{i-1},m_{i-1}}\right)$ (50) $$

That is, the coefficients are products of convolutions applied to the projections in each dimension (with the exception of the last dimension which is a mere scalar multiplication). This demonstrates the claim.

###### Definition B.1 .

### B.4 Mamba [ 8 ]

In this section, we argue that the mamba SSM framework is in fact a special case of our PADRe framework. The mamba framework is a state-space model based on the following recursion:

$$ $\displaystyle h_{t}$ $\displaystyle=\bar{A}h_{t-1}+\bar{B}x_{t}$ (51) $\displaystyle y_{t}$ $\displaystyle=Ch_{t}$ (52) $$

where $t$ is the timestep, $h_{t}$ is the hidden state at time $t$, $x_{t}$ is the input at time $t$ and $y_{t}$ is the output at time $t$. The "bars" above the matrices is mamba notation for discretization, which takes the form,

$$ $\displaystyle\bar{A}$ $\displaystyle=e^{\Delta{A}}$ (53) $\displaystyle\bar{B}$ $\displaystyle=(\Delta A)^{-1}\,\left(e^{\Delta A}-I\right)\,\Delta B$ (54) $$

where $\bar{A}$ is an $N\times N$ matrix chosen to be *diagonal*, $\bar{B}$ is an $N\times 1$ vector and $C$ is an $1\times N$ vector.

We now discuss technical details. First, we have the following lemma,

###### Lemma 2 .

If $h_{0}=0$, the discretized system of Equations 51 has solution,

$$ $y_{t}=\sum_{n=0}^{t}\left(C\,\bar{A}^{t-n}\,\bar{B}\right)\,x_{n}$ (55) $$

###### Proof.

We proceed by induction. Clearly for $t=0$, we get,

$$ $y_{0}=C\,\bar{B}\,x_{0}=C\,\bar{A}^{0}\,\bar{B}\,x_{0}$ (56) $$

Now assume the induction hypothesis holds and consider,

$$ $\displaystyle y_{t+1}$ $\displaystyle=C\,\left(\bar{A}h_{t}+\bar{B}x_{t+1}\right)$ (57) $\displaystyle=C\,\left(\bar{A}\sum_{n=0}^{t}\left(\bar{A}^{t-n}\,\bar{B}\right )\,x_{n}+\bar{B}x_{t+1}\right)$ (58) $\displaystyle=C\,\sum_{n=0}^{t+1}\,\left(\bar{A}^{t+1-n}\,\bar{B}\right)\,x_{n}$ (59) $$

which proves the claim.
∎

The main difference with regular SSM is that mamba allows $B$, $C$ and $\Delta$ to be functions of the input (Algorithm 2). Specifically, $B$ and $C$ are *linear* function of the input whereas $\Delta$ is the softplus of a linear function of the input, i.e.,

$$ $\displaystyle B$ $\displaystyle=W_{B}\,x$ (60) $\displaystyle C$ $\displaystyle=x^{T}\,W_{C}$ (61) $\displaystyle\Delta(x)$ $\displaystyle=\frac{1}{\beta}\,\log\left(1+e^{\beta\,(\pi+xW_{\Delta})}\right)$ (62) $$

where $W_{B}$ is a $N\times L$ matrix, $W_{C}$ is a $L\times N$ matrix, $W_{\Delta}$ is a (rank-$1$) $D\times D$ matrix, $B$ and $C$ are size $N\times 1$ and $1\times N$ vectors, respectively, and $\pi$ is some parameter.

For what follows, we make the assumption that,

$$ $0\leq\Delta(x)\ll 1$ (63) $$

uniformly in the input $x$. In this case, we find that,

$$ $\displaystyle\bar{A}$ $\displaystyle=e^{\Delta{A}}=I+\Delta\,A+O(\Delta^{2})$ (64) $\displaystyle\bar{B}$ $\displaystyle=(\Delta A)^{-1}\,\left(e^{\Delta A}-I\right)\,\Delta B$ (65) $\displaystyle=A^{-1}\,\left(I+\Delta\,A+O(\Delta^{2})-I\right)\,B$ (66) $\displaystyle=\Delta\,B+O(\Delta^{2})$ (67) $$

Therefore, under our assumption, it is reasonable to expect that,

$$ $\displaystyle\bar{A}$ $\displaystyle\approx I+\Delta A$ (68) $\displaystyle\bar{B}$ $\displaystyle\approx\Delta\,B=\Delta\,W_{B}\,x$ (69) $\displaystyle C$ $\displaystyle=x^{T}\,W_{C}$ (70) $$

Substituting these expressions in the recursion, we find that,

$$ $\displaystyle y_{t}$ $\displaystyle=C\,\sum_{n=0}^{t}\bar{A}^{t-n}\,\left(\bar{B}\,x_{n}\right)$ (71) $\displaystyle\approx\Delta\,\left(x^{T}\,W_{C}\right)\,\left(\sum_{n=0}^{t} \bar{A}^{t-n}\,x_{n}\right)\,\left(W_{B}\,x\right)$ (72) $\displaystyle=\sum_{n=0}^{t}\left(\sum_{i,j=1}^{L}\left[W_{C}\,\bar{A}^{t-n}\, W_{B}\right]_{i,j}\,\left(x_{i}\cdot x_{j}\cdot x_{n}\right)\right)$ (73) $\displaystyle=\sum_{n=0}^{t}\sum_{i,j=1}^{L}c^{(t,n)}_{i,j}\,\left(x_{i}\cdot x _{j}\cdot x_{n}\right)$ (74) $$

where,

$$ $\displaystyle c^{(t,n)}_{i,j}$ $\displaystyle:=\left[W_{C}\,\bar{A}^{t-n}\,W_{B}\right]_{i,j}$ (75) $$

This should be recognized as a degree-3 homogeneous polynomial in the input $x$. Further, the coefficients $c^{(t,n)}_{i,j}$ exhibit a polynomial relationship in the parameters ($A$, $W_{B}$, $W_{C}$) (assuming a truncated Taylor series for the exponential). Furthermore, the matrix $A$ exhibit such structures that the latter may be computed rapidly. Thus, it follows that Mamba is a special case of our PADRe framework.

###### Lemma 2 .

###### Proof.

### B.5 Castling-ViT [ 36 ]

In this section, we argue that Castling-ViT also fall within the bounds of the PADRe framework. Recall that at inference, the attention replacement used by Castling-ViT can be written as,

$$ $O=\frac{1}{\pi}\,Q\,\left(K^{T}\,V\right)+\left(\frac{1}{2}\,I+M_{DW}\right)\,V$ (76) $$

where,

$$ $\displaystyle Q=X\,W_{Q},\,K=X\,W_{K},\,V=X\,W_{V},$ (77) $$

$I$ is the identity, and $M_{DW}$ represents a learnt depth-wise convolution (i.e., a fast linear transformation). Upon substitution, we find that the left-hand term in Eq.(76) is a degree-3 homogeneous polynomial in the input $X$, whereas the right-hand term is a linear function of the input. All in all, this implies that the output $O$ of the Castling ViT module is a degree-$3$ polynomial that can be computed rapidly (i.e., in linear time). It is therefore an instance of the PADRe framework.

### B.6 Standard Attention [ 30 ]

In this section, we discuss the standard attention and argue that it may also be interpreted as an (approximate) instance of PADRe, albeit an instance of the *rational* framework (Section A.2) with a large degree (larger than practically acceptable).

To see how this is so, we write the attention explicitly as,

$$ $\displaystyle\left[\mathtt{Attention}(Q,K,V)\right]_{m,n}$ $\displaystyle=\left[\mathtt{softmax}\left(\frac{Q\,K^{T}}{\sqrt{d_{k}}}\right) \,V\right]_{m,n}$ (78) $\displaystyle=\sum_{i}\,\frac{e^{\left[\frac{Q\,K^{T}}{\sqrt{d_{k}}}\right]_{m ,i}}}{\sum_{j}e^{\left[\frac{Q\,K^{T}}{\sqrt{d_{k}}}\right]_{m,j}}}\,V_{i,n}$ (79) $$

where,

$$ $\displaystyle Q=X\,W_{Q},\,K=X\,W_{K},\,V=X\,W_{V},\,$ (80) $$

for some weight/projections matrices, as usual. Now comes the main observation: thanks to the normalization $\frac{1}{\sqrt{d_{k}}}$, one may expect that the elements $\frac{Q\,K^{T}}{\sqrt{d_{k}}}$ will generally lie in some canonical interval $[-L,L]$ for some $0<L\in\mathbb{R}$. In this case, the exponential functions may be expanded using the Taylor series as,

$$ $\displaystyle e^{\left[\frac{Q\,K^{T}}{\sqrt{d_{k}}}\right]_{m,i}}=\sum_{l=0}^ {d}\frac{\left[\frac{Q\,K^{T}}{\sqrt{d_{k}}}\right]^{l}_{m,i}}{l!}+O\left( \frac{L^{d+1}}{(d+1)!}\right)$ (81) $$

for some degree $d$ sufficiently large that the error term is on the order of $0<\epsilon\ll 1$. Upon substitution and expansion, we therefore find that,

$$ $\displaystyle\left[\mathtt{Attention}(Q,K,V)\right]_{m,n}$ $\displaystyle\approx\frac{\sum_{i}\,\sum_{l=0}^{d}\frac{\left[Q\,K^{T}\right]^ {l}_{m,i}}{d_{k}^{l/2}\,l!}\,V_{i,n}}{\sum_{j}\sum_{p=0}^{d}\frac{\left[\frac{ Q\,K^{T}}{\sqrt{d_{k}}}\right]^{p}_{m,j}}{d_{k}^{p/2}\,p!}}$ (82) $$

Finally, since $Q$, $K$ and $V$ are all linear transformations of the input data, we conclude that the attention can be arbitrarily approximated by a rationale function of degree $(2d+1)/(2d)$. The degree $d$, however, may have to be unpractically large in order reach proper accuracy. Nonetheless, this shows that the standard attention falls within the PADRe framework as well.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: See weakness.
soundness: 2
strengths: The unified formulation proposed in this paper is intriguing, as it establishes mathematical connections across different attention variants. I believe this polynomial-based approach has great potential for impact, offering a foundation for designing more efficient attention mechanisms in the future.
summary: The paper proposes a novel approach to replace full-attention with a unified polynomial formulation that achieves linear complexity. The authors show that various efficient attention mechanisms can be seen as special cases within the PADRe framework. They also implement a specific instance of PADRe and observe a favorable trade-off between efficiency and accuracy, even with a polynomial degree as low as 3.
weaknesses: 1): Figure 1 illustrates a specific design within the PADRe framework, but a straightforward baseline is missing. For example, for degree-
$d$ approximation, directly ensembling $d$ MLP-Mixers can be a strong baseline, and this approach can also be efficiently parallelized.

2): While the paper introduces a unified formulation, it remains unclear how readers can directly leverage it to enhance existing linear attention mechanisms. In Eq. (8), the coefficients $\pi_k$ are noted to have a complex dependency on parameters, specifically on the properties of the transformation matrices $A_i$, $B_i$, $C_i$ and $D_i$.  For example, in Eq. (82), a practical approximation formulation for the softmax operation can be more concrete. 

3): The experimental results demonstrate comparable performance with previous methods, but the paper defines the transformation matrices as pointwise and 2D convolutions to mix channels and tokens, respectively. However, it is well-established that adding locality can boost ViT performance (e.g., LocalViT). To fully evaluate the PADRe's effectiveness, testing on NLP benchmarks and models—without adding any locality inductive bias—would provide a more compelling validation.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See the weakness.
soundness: 3
strengths: 1. Good insight and strong theoretical proofs. Approximating self-attention with polynomial approximants sounds interesting. The authors have provided solid proofs that the proposed PADRe can cover other special cases in recent works.

2. Experiments have well demonstrated the effectiveness of their approach, especially considering the FLOPs reduction in Table 3 and Table 4.

3. The paper is in a good structure. Figures are clear as well.
summary: This paper introduces Polynomial Attention Drop-in Replacement (PADRe), a new method which aims to replace the conventional self-attention in Transformers. Specifically, PADRe replaces self-attention with polynomial approximants. Furthermore, it adopts mobile-friendly operations such as  Hadamard products to achieve a good balance between efficiency and performance.  The authors have provided rigorous proofs that the proposed PADRe is a general approach that covers previous Hyena, Mamba, SimA, etc. Extensive experiments have also demonstrated the effectiveness of their approach.
weaknesses: 1. The paper acknowledges that model performance begins to saturate (Lines 462-463) when the polynomial degree exceeds 3, suggesting potential numerical stability issues with higher-degree polynomials. Investigating more stable polynomial bases could be beneficial.

2. As the author already states, this paper lacks experiments on LLMs. It would be great if we can know the performance of PADRe on LLMs. This could be more attractive to the community than experimenting on ViTs.

3. In Table 1, the improvement over DeiT-Base is minor. Could it be that PADRe works more effectively for smaller Transformers?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Please see the weaknesses above.
soundness: 3
strengths: 1. This paper addresses an important problem in vision transformer design.

2. This paper provides a theory-driven solution with a thorough mathematical justification.
summary: This paper proposes PADRe, a unifying framework for replacing the attention mechanism based on
polynomial approximants to provide an efficient, linear-complexity alternative to the
standard self-attention without sacrificing accuracy.
weaknesses: 1. How would the method work for large language models (LLMs), with much longer sequence lengths?

2. How to decide the optimal degree of PARDe?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - How does replacing self-attention with PADRe change the number of model parameters? Is there any increase or decrease? This will be helpful in understanding PADRe's impact on mobile devices since parameter count is directly correlated to storage, DRAM usage etc.
soundness: 3
strengths: - Framework that helps building alternative self-attention based transformer architecture
- advantageous over Flash attention
summary: This paper proposes a framework for drop in self attention replacement. This framework allows for building any polynomial time self-attention module, relies on polynomial approximates at the input, uses Hadamard product to introduce non-linearities. The framework at its current stage does not look into cross-attention and/or multi-modal attention although in principal it can be extended to cross-attention
weaknesses: - The authors evaluated PADRe on DeiT. It might be a good idea to evaluate PADRe on ViT (Dosovitskiy et al.). This will demonstrate the impact of PADRe even more since ViT is widely used as an image encoder for many other models such as CLIP, Vision LLMs etc.
- Since you are evaluating on device, it would be good to show peak memory consumption, anything that shows energy efficiency improvement when running PADRe over standard self-attention

Meta-review: The PADRe paper introduces a polynomial attention mechanism as an efficient drop-in replacement for self-attention in vision transformers, offering hardware-friendly operations and maintaining accuracy while significantly reducing computational costs.

After rebuttal and discussions, this paper receives three positive ratings and one negative rating. After carefully  reviewing the paper and all the reviewers' comments and discussions, The AC agrees to accept the paper and strongly recommends incorporating the content from the rebuttal into the final version.

Decision: Accept (Poster)