**footnotetext: Equal contribution

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

Pierre-David Letourneau∗ &Manish Kumar Singh∗ &Hsin-Pai Cheng &Shizhong Han
&Yunxiao Shi &Dalton Jones &Matthew Harper Langston &Hong Cai &Fatih Porikli
Qualcomm AI Research†
{ pletourn, masi, hsinpaic, shizhan, yunxshi,
daltjone, hlangsto, hongcai, fporikli } @ qti.qualcomm.com

###### 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$$\,\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) [7] are particularly significant, driving innovations in
areas such as open-world recognition [12], image and video generation [21, 2],
and large multi-modal models [16, 34].

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 [20], Mamba [8],
Conv2Former [11], SimA [13], Castling-ViT [36]), as well as the standard
attention itself [30], 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 [28]. 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 [20], Mamba [8],
Conv2Former [11], SimA [13], and Castling-ViT [36].

  * •

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 [17]. 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 [14, 10,
29, 33]. 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 [4], transposed attention [19],
convolutional modulation [11], additive attention [23], shift-add attention
[35], linear-angular attention [36]. 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 [9].

Efficient Attention Drop-in Replacements: To enhance the computational
efficiency of vision transformers, various alternative attention mechanisms
have been introduced. For instance, Swin [17] 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
[13, 36, 9, 23]. Some of them (e.g., [13, 9]) 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., [36, 23]) 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).

###  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 [37] and linear token mixing (e.g., convolution
[18], MLP [26, 25]) 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.

Attention Schemes | PADRe Equivalence | Proofs
---|---|---
SimA [13] | Normalized, homogeneous degree-3 polynomial | B.1
Conv2Former [11] | Homogeneous degree-2 polynomial | B.2
Hyena [20] |  Homogeneous degree-$N$ polynomial (order-$N$ operator, commonly $3$) | B.3
Mamba [8] | Homogeneous degree-3 polynomial (approx.) | B.4
Castling-ViT [36] | Degree-3 polynomial | B.5
Self-attention [30] | High-degree rational function/ratio of polynomials (approx.) | 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.

###  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 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.

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$ 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 [22].\dagger\dagger$\dagger$ImageNet terms of use can be
found here: https://www.image-net.org/download. We use DeiT [27] 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 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 [1] and linear-angular attention [36], fail to maintain accuracy.
Other methods do improve accuracy, but are not as efficient as PADRe on actual
compute platforms (Section 4.3).

Baselines | Attention Mechanisms |  Top-1 Accuracy$\,\uparrow$ | FLOPs (G) | #Params (M)
---|---|---|---|---
DeiT-Tiny [27] | Standard self-attention [7] | 72.2 | 1.3 | 5
Hydra attention [1] | 68.3 | 1.1 | 6
Efficient attention [24] | 70.2 | 1.1 | 6
Linear-angular attention [36] | 70.8 | 1.1 | 6
Enhanced linear attention [4] | 72.9 | 1.1 | 6
k-NN attention [32] | 73.0 | 1.3 | 6
Focused linear attention [9] | 74.1 | 1.1 | 6
Vision Mamba [38] | 76.1 | 2.6 | 7
PADRe-2 (ours) | 74.5 | 1.2 | 5
| PADRe-3 (ours) | 76.5 | 1.7 | 7
DeiT-Small [27] | Standard self-attention [7] | 79.8 | 4.6 | 22
SimA [13] | 79.8 | 4.6 | 22
k-NN attention [32] | 80.1 | 4.6 | 22
Vision Mamba [38] | 80.5 | 9.6 | 26
PADRe-2 (ours) | 80.4 | 4.8 | 20
| PADRe-3 (ours) | 80.5 | 6.1 | 26
DeiT-Base [27] | Standard self-attention [7] | 81.8 | 17.6 | 86
Hydra attention [1] | 76.4 | 16.9 | -
Hyena [20] | 78.5 | - | 87
PADRe-2 (ours) | 81.4 | 18.8 | 80
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.

Single-Image Object Detection: In this part, we consider a commonly used, end-
to-end transformer-based object detector: DETR [5].\dagger\dagger$\dagger$We
use the official DETR repository: 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
[15].\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. We follow the
short-schedule [5] 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 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%.

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

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 [3], 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. We select a transformer-based 3D object
detection model, DSVT [31], which achieves state-of-the-art performance on
nuScenes.\dagger\dagger$\dagger$We use the official DSVT repository:
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 [31];
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 [3] 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.

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

###  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 Hexagon${}^{\scalebox{0.4}{\text{TM}}}$ NPU (included in the
Snapdragon 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 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.

###  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.

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

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 [36] and focused linear attention [9]. 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., [6]. For instance, operations that are not well optimized or cause
memory bottlenecks on device are not reflected by FLOP count.

Attention Mechanisms | FLOPs (G) | Latency on GPU (ms) | Latency on NPU (ms)
---|---|---|---
Standard self-attention [7] | 1.31 | 0.43 | 68.60
Linear-angular attention [36] | 1.22 | 0.53 | -
Focused linear attention [9] | 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
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.

##  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.

## References

  * [1] Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, and Judy Hoffman.  Hydra attention: Efficient attention with many heads.  In Proceedings of the European Conference on Computer Vision, 2022.
  * [2] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh.  Video generation models as world simulators, 2024.
  * [3] Holger Caesar, Varun Bankiti, Alex H Lang, Sourabh Vora, Venice Erin Liong, Qiang Xu, Anush Krishnan, Yu Pan, Giancarlo Baldan, and Oscar Beijbom.  nuscenes: A multimodal dataset for autonomous driving.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11621–11631, 2020.
  * [4] Han Cai, Junyan Li, Muyan Hu, Chuang Gan, and Song Han.  Efficientvit: Multi-scale linear attention for high-resolution dense prediction, 2024.
  * [5] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko.  End-to-end object detection with transformers.  In Proceedings of the European Conference on Computer Vision, 2020.
  * [6] Jierun Chen, Shiu-hong Kao, Hao He, Weipeng Zhuo, Song Wen, Chul-Ho Lee, and S-H Gary Chan.  Run, don’t walk: Chasing higher flops for faster neural networks.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
  * [7] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  In International Conference on Learning Representations, 2020.
  * [8] Albert Gu and Tri Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  arXiv preprint arXiv:2312.00752, 2023.
  * [9] Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang.  Flatten transformer: Vision transformer using focused linear attention.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2023.
  * [10] Ali Hatamizadeh, Greg Heinrich, Hongxu Yin, Andrew Tao, Jose M Alvarez, Jan Kautz, and Pavlo Molchanov.  Fastervit: Fast vision transformers with hierarchical attention.  In The Twelfth International Conference on Learning Representations, 2023.
  * [11] Qibin Hou, Cheng-Ze Lu, Ming-Ming Cheng, and Jiashi Feng.  Conv2former: A simple transformer-style convnet for visual recognition.  arXiv preprint arXiv:2211.11943, 2022.
  * [12] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al.  Segment anything.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4015–4026, 2023.
  * [13] Soroush Abbasi Koohpayegani and Hamed Pirsiavash.  Sima: Simple softmax-free attention for vision transformers.  In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 2607–2617, 2024.
  * [14] Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, and Jian Ren.  Efficientformer: Vision transformers at mobilenet speed.  Advances in Neural Information Processing Systems, 35:12934–12949, 2022.
  * [15] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.  Microsoft coco: Common objects in context.  In Proceedings of the European Computer Vision Conference, 2014.
  * [16] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee.  Visual instruction tuning.  Advances in neural information processing systems, 36, 2023.
  * [17] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows, 2021.
  * [18] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie.  A convnet for the 2020s.  In Proceedings of the IEEE/CVF conference on Computer Vision and Pattern Recognition, 2022.
  * [19] Muhammad Maaz, Abdelrahman Shaker, Hisham Cholakkal, Salman Khan, Syed Waqas Zamir, Rao Muhammad Anwer, and Fahad Shahbaz Khan.  Edgenext: Efficiently amalgamated cnn-transformer architecture for mobile vision applications, 2022.
  * [20] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré.  Hyena hierarchy: Towards larger convolutional language models.  arXiv preprint arXiv:2302.10866, 2023.
  * [21] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684–10695, 2022.
  * [22] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  International journal of computer vision, 115:211–252, 2015.
  * [23] Abdelrahman Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang, and Fahad Shahbaz Khan.  Swiftformer: Efficient additive attention for transformer-based real-time mobile vision applications, 2023.
  * [24] Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li.  Efficient attention: Attention with linear complexities.  In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pages 3531–3539, 2021.
  * [25] Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al.  Mlp-mixer: An all-mlp architecture for vision.  Advances in neural information processing systems, 34:24261–24272, 2021.
  * [26] Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Gautier Izacard, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, et al.  Resmlp: Feedforward networks for image classification with data-efficient training.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(4):5314–5321, 2022.
  * [27] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  In International Conference on Machine Learning, 2021.
  * [28] Lloyd N Trefethen.  Approximation theory and approximation practice, extended edition.  SIAM, 2019.
  * [29] Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, and Anurag Ranjan.  Fastvit: A fast hybrid vision transformer using structural reparameterization.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 5785–5795, 2023.
  * [30] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [31] Haiyang Wang, Chen Shi, Shaoshuai Shi, Meng Lei, Sen Wang, Di He, Bernt Schiele, and Liwei Wang.  Dsvt: Dynamic sparse voxel transformer with rotated sets.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
  * [32] Pichao Wang, Xue Wang, Fan Wang, Ming Lin, Shuning Chang, Hao Li, and Rong Jin.  Kvt: k-nn attention for boosting vision transformers.  In European Conference on Computer Vision, pages 285–302, 2022.
  * [33] Chenglin Yang, Siyuan Qiao, Qihang Yu, Xiaoding Yuan, Yukun Zhu, Alan Yuille, Hartwig Adam, and Liang-Chieh Chen.  Moat: Alternating mobile convolution and attention brings strong vision models.  In The Eleventh International Conference on Learning Representations, 2022.
  * [34] Zhengyuan Yang, Linjie Li, Kevin Lin, Jianfeng Wang, Chung-Ching Lin, Zicheng Liu, and Lijuan Wang.  The dawn of lmms: Preliminary explorations with gpt-4v(ision).  arXiv preprint arXiv:2309.17421, 2023.
  * [35] Haoran You, Huihong Shi, Yipin Guo, and Yingyan Lin.  Shiftaddvit: Mixture of multiplication primitives towards efficient vision transformer.  Advances in Neural Information Processing Systems, 36, 2023.
  * [36] Haoran You, Yunyang Xiong, Xiaoliang Dai, Bichen Wu, Peizhao Zhang, Haoqi Fan, Peter Vajda, and Yingyan Celine Lin.  Castling-vit: Compressing self-attention via switching towards linear-angular attention at vision transformer inference.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14431–14442, 2023.
  * [37] Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan.  Metaformer is actually what you need for vision.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10819–10829, 2022.
  * [38] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang.  Vision mamba: Efficient visual representation learning with bidirectional state space model.  arXiv preprint arXiv:2401.09417, 2024.

##  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, [20]),

######  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. 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. 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 ([20]), 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.

###  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.

###  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.