PADRe: A Unifying Polynomial Attention Drop-in Replacement for Efficient Vision Transformer
Authors: Pierre-David Letourneau
Manish Kumar Singh
Hsin-Pai Cheng
Shizhong Han
Yunxiao Shi
Dalton Jones
Matthew Harper Langston
Hong Cai
Fatih Porikli
Venue: ICLR 2025 Poster (2025)
Keywords: Efficient Vision Transformer
Attention Approximation
Computer Vision
Deep Learning
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: Polynomial-based replacement of standard self-attention for efficient vision transformer that is hardware friendly and preserves accuracy

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

Published as a conference paper at ICLR 2025
A UNIFYING POLYNOMIAL ATTENTION DROP-IN REPLACEMENT FOR EFFICIENT VISION TRANSFORMER
PADRE:
Pierre-Davie Letourneau'
Manish Kumar Singh*
Hsin-Pai Cheng
Shizhong Han
Yunxiao Shi Dalton Jones
Matthew Harper Langston
Hong Cai Fatih Porikli
Qualcomm Al Researcht
(pletcurn,masi) masi, hsinpaic, shi zhan, yunxshi, daltjone, blangsto, hongcai, www.fppiiieppppppppp /@qti qua com.mccoommmm
ABSTRACT
We present Polynomial Attention Drop-in Replacement (PADRe), 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 specific instances of our PADRe framework. PADRe leverages polynomial func- tions and draws upon established results from approximation theory. enhancing computational efficiency without compromising accuracy. PADRe's key compo- nents include multiplicative nonlincarities, which we implement using straightfor- ward, hardware tardware-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 accu- racy compared to traditional self-attention. We assess the effectiveness of PADRe as a drop-in replacement for self-attention across diverse computer vision tasks. These tasks include image classification, image-based 2D object detection, and 3D point cloud object detection. Empirical results demonstrate that PADRe runs significantly faster than the conventional self-attention (11x ~ 43x faster server GPU and mobile NPU) while maintaining similar accuracy when substituting self-attention in the transformer models.
INTRODUCTION
Transformers have been pivotal in recent advancements across natural language processing, computer vision, and speech processing. In computer vision, Vision Transformers (ViTs) [9] particularly significant, driving innovations in areas such as open-world recognition [15], image and video generation [25; 3], and large multi-modal models [19; (9].
However, the computational and memory costs associated with self-attention operations in trans- formers can increase quadratically relative to input size. This demanding challenge 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, unifying Polynomial Attention Drop-in Replacement scheme. PADRe is a novel approach that can effectively substitute standard self-atrention at low (linear computational and memory) costs without sacrificing accuracy. PADRe's unifying framework provides guide for designing alternative transformer architectures. In particular, observe and demonstrate that many recently-prope proposed attention replacement mechanisms (c.g.. Hyena [24], Mamba [11], Cony2Former [14], SimA [16], Castling-Vis ViT [41]), as well as the standard
"qqual contribution
QQalcoom AI Research initiative of Qualcomm Technologies, Inc.

---

Published as a conference paper at ICLR 2025
attention itself [35], may in fact be interpreted as specific instances of the PADRe framework (See Table 1).
Our framework relies 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 [33]. We use these existing results, together with the aforementioned observations, as guiding principles. We also leverage simple, mobile-friendly operations, such 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 follows:
We propose PADRe, a unifying framework for replacing the attention mechanism based on polynomial approximants, which provides an efficient, linear-complexity alternative the standard self-attention without sacrificing accuracy.
We show that PADRe provides a unifying formulation for many of the eecently-proposee attention replacement schemes, including Hyena [24], Mamba [11], Conv2Former [14], SimA [16], and Castling- ViT [41].
We implement 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 progressively downsampled [20]. This strategy allows the deeper transformer layers 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 [17; 13: 34; 38). While this approach considerably improves efficiency, it still suffers from the inherent quadratic complexity when input size becomes large. Building upon hierarchical and/or manvolltioo-transfommen
hybrid architectures, some studies propose more efficient, alternative attention schemes, e.g., ReLU-based attention [5], transposed attention [22], convolutional modulation [14], additive attention [27], shift-add attention [40], linear- angular attention [41]. However, all these models are designed to process 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 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 [12]. Another line of works study how to use fewer tokens to reduce computation while still utilizing the standard self-attention operation, e.g., [1; 29].
Efficient Attention Drop-in Replacements: To enhance the computational efficiency of vision
transformers, various alternative attention mechanisms have been introduced. For instance, Swin [20] proposes windowed attention, where self-attention computation is limited local neighborhoods of an image. Although this theoretically reduces computation. the intricate reshaping and indexing operations pose practical challenges resource-constrained hardware platforms. Several works
2

---

Published as a conference paper at ICLR 2025
propose linear-cost attentions [16; 41; 12; 27]. Some of them (e.g., [16: 12]) need to compute normalization factors at inference, which requires summing over all the tokens. This is inefficient on memory-commee constrriind mobile platforms, especially when the number of tokens is large. Others (e.g., [41; 27]) 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 such that the coefficients and the output itself may be computed efficiently (e.g., in linear time).
For N N tokens and embedding feature channel dimension D E N, consider input X E RNDD,
XD We denote as the degree of our framework; this parameter determines the strengths of the nonlinearities discussed below. PADRe operations are performed on each batch element identically.
We omit batch dimension for brevity and consider a single head; more than one head is used, each 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 discussed below in the following sections.
3.1 LINEAR TRANSFORMATIONS
Linear transformations consist of the following operations: for = I,... compute
Y A X B;
(I)
where A; and B; are N x N and Dx D weight matrices, and Y; NxD,, We impose structure on the matrices A; and B; 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 form of mixing among tokens (in the spatial dimension in the case of visual input) and in the embedding/chammi dimension, respectively. Coupled with nonlinearities (Section 3.2) this leads 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:
Zi Y1, Zi+1 (C, D.) O Y141, for E {l,.., d- 1}.
where indicates the Hadamard (element-wise) product, which is an efficient, parallelizable, and mendware-friendyy operation. C₁ and D; are N x N and D D weight matrices that can perform additional token and channel mixing on the Z; tensor as needed, and Z, E NxD. Similarly, we require C and D, 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 Z4 are homogeneous polynomials of degree i in the input X.
Proof. We proceed by induction. In the case i = 1,
Yi A, X Bi
which shows that each element of Y1 is a linear function of the input. Now assume the elements of Z, are homogeneous polynomials of degree i in the input and consider
Zi+1 (C, Z, D:) @ (A:+1 X B141).
(5)
*This is technically not the only possible choice; other polynomial bases, e.g., Chebyshev, Legendre, could he used for instance.
3

---

Published as a conference paper at ICLR 2025
Observe that the elements of A,+1 X B,+1 are homogeneous polynomials of degree 1 by construc- tion and the elements of C, Z, D, are homogeneous polynomials of degree by assumption and construction (linear transformations do not affect this conclusion). The expression is therefore homogeneous polynomial of degree i 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] mm Σ Wimmn [Zelm.n + [4]m
(6)
where W E NxDd is a weight tensor, 1 S m S N and S S D are token and embedding dimensions respectively, and L NxD is the bias (i.e., 0thorder 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 == DUPPV
(7)
where E RPN, V& E RxG and FXG (required output shape). As discussed, it is necessary that U and V possess structures (e.g., low-rank) allow efficient application to P.
In addition, the Yi's tensors can optionally be normalized. While this is not necessary for the PADRe instances in experiments, it can be potentially useful regularization technique.
3.4 OVERALL FRAMEWORK
Combining all operations together and leveraging Lemma 1, the output takes the following form:
PP Pmm
Σ 1=1
[mm,,, mm [Z)mm, + [T]mm
Σ Th x
(8)
KENIY
Em=L EP km. is the absolute degree, { are coeffi- cients (which have complex polynomial relationship with the parameters/weights), and x² IT= Im,n are monomials of the input variables. Equation is degree-d multi-dimensional polynomial in the entries of the input X; that is, PADRe uses hardware-friends and efficient op- crations 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.
where ls is multi-index [k]
3.5 UNIFYING FRAMEWORK
Equation 8 provides general, quantitative representation of PADRe's output. This formula in fact 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 I and provide detailed proofs to show that they are within our proposed PADRe framework in Appendix B. Further, can be easily verified that the "attention" part of existing models based on the MetaFormer architecture [42] and linear token mixing (e.g., convolution [21], MLP [31; 30)) are also 1 PADRe equivalences.
encompasses many
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. In addition, 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.
We ignore the optional resizing here, which is just a linear transformation.
4

---

Published as a conference paper at ICLR 2025
Attention Schemes
PADRe Equivalence
Proots
SimA [16]
Normalized, homogeneous degree-3 polynomial
B.1
Conv2Former [14]
Homogeneous degree-2 polynomial
B.2
Hyena [24]
Homogeneous degree-N polynomial (order-N operator, commonly 3)
B.3
Mamba [11]
Homogeneous degree-3 polynomial (approx.)
B.4
Castling- ViT [41]
Degree-3 polynomial
B.5
Self-attention [35]
High-degree rational function/ratio of polynomials (approx.)
B.6
Table These existing attention schemes can be shown be PADRe equivalences or can be approximated by PADRe. Detailed proofs can be found Appendix
3.6 COMPUTATIONAL CHARACTERISTICS
We briefly analyze the number of parameters and computational cost of constructing an output P E I NDD given an input X E RNND using PADRe. First, let us consider the number of parameters. Under our structural assumptions, we find O(N - d) parameters for weight matrices {4.} and {Ct}. ( d) parameters for weight matrices {B.} and (D:). and O(N D d) for the tensor W. Specifically, by construction, the PADRe framework requires that matrices have sufficient structure for an application in linear time. This implies no more than linear scaling of the parameters. This results in a total of
O(N D dd)
parameters. For a fixed degree d, this scales like ONN D) which is linear in the size of the input. Therefore, the memory cost linear in the size of the input.
As for the computational cost, we find that computing all Y,'s costs O(N D. d) FLOPs using the assumed structure. Further, each quantity Y, may be computed in parallel (data parallelism). Then, computing all Z:'s ss carries a cost of O(N D d) FLOPs, which are computed sequentially based on the current scheme. Finally, computing the output comes at a cost of O(N D-d) FLOPs. In total, the cost of computing the output P from the input X is therefore,
O(N D-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.. exp(-), cos(-)), which are often expensive on hardware and can cause quantization complications.
EXPERIMENTS
In this section, we comprehensively evaluate PADRe as a drop-in replacement of the standard self- attention for computer vision tasks. in terms of both accuracy and efficiency. More specifically, we implement specific instance of PADRe following the design principles laid out Section 3, and use it to substitute the self- self-attention in transformer models in several vision applications, such as image classification, single-image object detection, and point cloud object detection. We then measure the latency of our PADRe instances on two compute platforms: 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 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,, for each i E (1....d) these mixing operations correspond to the linear transformations in Eq. More specifically, we implement channel mixing using pointwise con- volutions or linear layers. Token mixing is implemented with convolutions applied in the token dimension. When the input possesses an inherent 2D structure (c.g., image features), we use 11x11 2D convolutions mix the tokens, after reshaping them into 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, treat them as sequence of tokens and use ID convolutions with kernel size 11 for mixing.
Next, we compute Z; for each degree i based on Y; and Z;-1, as described in Eq. 3. Channel and token mixings are further applied to Z, before it is used to produce degree-(i + 1) terms. Finally,
5

---

Published as a conference paper at ICLR 2025
all the Z 's are added together to generate the output P. Note that unlike the standard attention and existing alternative attention methods, we use a single head in all our experiments.
Y
Z/2
Degree-2 terme
Channel Mixinng
Token Mixing
Channel Mixinng
Y,
Channel Mixingg
Token Mixinn
PADRe-2
Tolion Mixing
Y,
Degree-3 tarms
Channel Mixing
Token Mixing
Channal Moxing
PADRe-3
Token Mixing
Ya
Degree-d terms
Channel Mixing
Taken Mixing
Charnel Mixing
PADRe-d
Our implementation of PADRe based attention drop-in in replacement module. In our experiments, we this substitute the standard attentttentoon parts in existing transformer models. Note that this imple- mentation approach is just specifie instance of general PADRe framework.
Figure 1:
In this specific instance, we do not have 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 degree-1 term. In addition, do not require an explicit degree-O 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
First, consider the standard image classification task ImageNet [26]. We use DeiT [32] 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 ~~2 days on 8 NVIDIA A 100 GPUs. We use the top-1 accuracy to measure classification performance.
Image Classification:
Table 2 summarizes the results. It can be seen that by replacing self elff-atentioon with our PADRe instance, we achieve similar 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 fair comparison. We see that several proposed methods, such as Hydra attention [2] and linear-angular attention [41], fail to maintain accuracy. Other methods do improve accuracy, are not as efficient as PADRe on actual compute platforms (Section 4.3).
Single-Image Object Detection: In this part, we consider a commonly used, end-to-end transformer- based object detector: DETR [6]. DETR consists of three main components: ResNet backbone for image feature extraction, a transformer encoder with (self-attention) layers, and a transformer decoder with 6 (self- and cross-attentinn
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 the COCO detection benchmark [18]. We follow the short-schedule [6] train the model for 300 epochs on 8 NVIDIA A100 GPUs, which takes 44 days. We use the standard metric of Average Precision (AP) measure detection performance.'
The results in Table show that by substituting the standard self-attention with PADRe in the transformer encoder layers, we achieve a similar detection accuracy. the encoder part, our PADRe instance also significantly reduces the computational cost by 55%.
While using multiple heads is possible and straightforward implement, found that single head was sufficient for PADRe match the accuracy multi-head attention models. See https: / /ocoootoosseeeeeo org/ #detect i con-eval for more details COCO object detection metrics.

---

Published as a conference paper at ICLR 2025
Baselines
Attention Mechanisms
Top-1 Accuracy
FLOPS (G)
#Params (M)
Standard self attention [9]
72.2
1.3
5
Hydra attention [2]
68.3
1.1
6
Efficient attention [28]
70.2
1.1
6
Linear-angular attention [41]
70.8
1.1
6
DeiT-Tiny [32]
Enhanced linear attention [5]
72.9
1.1
6
k-NN attention [37]
73.0
1.3
6
Focused linear attention [12]
74.1
1.1
6
Vision Mamba [43]
76.1
2.6
7
PADRe-2 (ours)
74.5
1.2
5
PADRe-3 (ours)
76.5
1.7
7
Standard self-attention [9]
79.8
4.6
22
SimA [16]
79.8
4.6
22
DeiT-Small [32]
k-NN attention [37]
80.1
4.6
22
Vision Mamba [43]
80.5
9.6
26
PADRe-2 (ours)
80.4
4,8
20
PADRe-3 (ours)
80.5
6.1
26
Standard self-attention [9]
81.8
17.6
86
Hydra attention [2]
76.4
16.9
DeiT-Base [32]
Hyena [24]
78.5
87
PADRe-2 (ours)
81.4
18.8
80
Table Results on ImageNet-Iee 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.
Model
Attention Mechanisms
AP
FLOPs (G)
#Params (M)
Standard self offf-attenionn [9]
40.6
10.1
6.1
DETR [6]
PADRe-2 (ours)
40.2
4.5
2.7
Table 3: Results on COCO detection, using DETR as the baseline. We achieve similar accuracy but with
considerably lower computational cost.
Bascline
Attention Mechanisms
NDS
mAP
FLOPs (G)
#aarams (M)
Standard self attention [9]
71.1
66.4
14.9
1.1
DSVT [36]
PADRe-2 (ours)
71.1
66.0
10.7
0.9
Table 4: Results on nuScenes LiDAR-bases 3D object detection. We achieve similar accuracy
significantly lower computational cost.
Point Cloud Object Detection: Next, we evaluate PADRe on a more complex vision application where the input is no longer single image. In this case, the input is 3D point cloud. We consider the 3D point cloud object detection task on nuScenes [4], which involves detecting vehicles and other traffic elements based LiDAR point cloud. We select transformer-based 3D object detection model, DSVT [36], which achieves state-of-the-an performance on nuScenes. DSVT, 3D transformer backbone extracts point cloud features, after which the features are projected to the birds-eyc-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 in [36]; training takes ~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 weighted sum of 6 detection metrics over 10 object classes (see [4] for more details on the metrics). As the 3D voxels are sparse and cannot be arranged in a dense format, treat them as sequence and use ID convolutions to perform token mixing.
Table 4 summarizes the 3D object detection results on nuScenes validation set. After replacing the self-attentions with PADRe, achieve similar detection accuracy (in terms of both NDS and mAP) as compared to the original model. then measure the computational costs of the 3D transformer backbones in the baseline and model. can be seen that modified DSVT 3D backbone uses 28% fewer FLOPs while maintaining detection accuracy as compared to the baseline.
4.3 LATENCY AND MEMORY EVALUATION HARDWARE PLATFORMS
We evaluate and compare the on-device latencies of our PADRe instance and the standard self- attention. Our goal is to provide practical comparison of computational efficiency across different

---

Published as a conference paper at ICLR 2025
Latency Nvidia A100 GPU
Latency Qualcomm Hexagon
Standard SA
PADRu-2
PADRu-3
Standard SA
PADRe-2
PADRe-3
0.5
80
70
Latency (ms)
0.3
Latency
02
10
1024
2304
Number of tokens
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 increase in the number of input tokens. contrast, PADRe demonstrates linear growth pattern.
hardware platforms. Specifically, we measure latencies on two compute platforms: Nvidia A100 GPU, and 2) Qualcomm "exagon NPU (included in the Snapdragon" 8 Gen 3 processor on Samsung S24), which is an 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 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 multi-head attention operation. Using the hyperparameters of DeiT-Tiny, 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 experiments. We evaluate different numbers of input tokens, including: 256, 1024, 2304, and 4096, which correspond to input image resolutions of 256x256, 512x512, 768 x76, 1024x 1024, assuming 16x 16 patchification (as used in ViT).
set the embedding dimension to 192 in both PADRe and
Fig. shows the on-device latencies w.r.t. number of tokens. We see that the standard attention mechanism exhibits significant latency growth the number of tokens increases. In contrast, our proposed PADRe maintains linear latency proportional to the number of tokens. This confirms that PADRe provides significant efficiency advantage over self-attention when deployed on actual hardware, including large GPUs and resource-comstrainee mobile NPUs.
Furthermore, we evaluate the memory consumption. Considering one layer, given 4096 input tokens, PADRe's 's peak memory consumption is 106MB while that of self-attentio is 145MB, 37% higher than PADRe. Such advantages are expected become even more prominent as one increases the number of tokens, given the linear scaling of PADRe.
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 inereasing the polynomial degree, the network obtains stronger modeling capacity, leading to performance improvements. Specifically, when using degree 3, PADRe has 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
Standard self-attention [9]
72.2
PADRe-2 (ours)
74.5
DeiT-Tiny [32]
PADRe-3 (ours)
76.5
-4 (ours)
76.9
Table 5: ImageNet classification accuracy when scaling up PADRe degree.
Table 6 shows the computational costs of our PADRe instance with different degrees, and compares those with the standard self-attention and other recently-propose efficient attention methods such
SSapdragonn and Qualcomm branded products are products of Qualcomm Technologies, Inc. and/or its subsidiaries. Qualcomm patented technologies are licensed by Qualcomm Incorporated.
8

---

Published as a conference paper at ICLR 2025
Attention Mechanisms
FLOPss
Latency on GPU (ms)
Latency NPU (ms)
Standard self attention [9]
1.31
0.43
68.60
Linear angular attention [41]
1.22
0.53
Focused linear attention [12]
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. The measurement based on the attention parts only, with 4096 input tokens.
as linear-angular attention [41] and focused linear attention [12]. 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 head our PADRe instance. The number of input tokens is fixed 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 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., [7]. For instance, operations that are not well optimized or cause memory bottlenecks on device are not reflected by FLOP count.
5 DISCUSSION
One may consider directly approximating sufficiently regular nonlineari- ties, including the self-attention, using Taylor expansion. We have performed such attempt, and have observed the accuracy to be greatly affected by small errors in the approximation, thus requiring an impractically large degree. Our PADRe framework, by contrast, learns a different nonlinearity without affecting the output accuracy in any significant way. This more flexible framework, although inspired by Taylor's 's theorem, is not equivalent, and has been shown to be beneficial.
Relation to Taylor Series:
Comparing with FlashAttention [8]: the softmax computation, but still requires O(N2) computation. In addition, it is only designed for specific hardware, e.g., Nvidia A100 GPU. Our proposed PADRe, the other hand, requires only ((N) memory and computation costs, and is generally applicable. Furthermore, PADRe amenable low-level optimization (e.g., tiling) for on-device deployment it consists of hardware-friendly operations.
reduces memory cost ( by decomposing
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 -otss-attention, a discussion of which provided in Appendix A. 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. In addition, depending the task and data, it is possible to find the optimal degree via existing neural architecture search methods, e.g., [10].
While we mainly apply PADRe computer vision tasks, it can be applied for any use case with transformer models. Specifically, PADRe can be easily extended to sequential applications like natural language processing and time series analysis. In preliminary study, we used a reduced Llama2 with 300M parameters trained on SlimPajama as the baseline. Then, we replaced the self-attention with PADRe and retrained the model. On the test set, PADRe achieves similar cross-entropy. We will consider further investigations of LLMs as part of future work.
We do not report the latencies of these alternative attention methods on NPU, it would require significant effort rewrite the implementation in order run them on the mobile NPU.
9

---

Published as a conference paper at ICLR 2025
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 specificallly, we presented the key components for constructing an expressive, efficient, and hardware- friendly scheme to substitute self-attention. passes 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 the design principles, using efficient and hardware-frienee 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 power GPU (NVIDIA A A100) and coource-constrainee mobile NPU (Qualcomm Hexagon tensor processor), by as much as 43 x when comparing to the standard self-attention.
addition, PADRe a general framework that encom-
operations. Through extensive experiments, we showed that PADRe
REFERENCES
[1] Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffmman. Token merging: Your ViT but faster. In International Conference on Learning Representations, 2023.
[2] 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.
[3] 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.
[4] 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 IEEF/CVF Conference Computer Vision and Pattern Recognition, pages 1621111631,, 2020.
[5] Han Cai, Junyan Li, Muyan Hu, Chuang Gan, and Song Han. Efficientvit: Multi-scale linear attention for high-resolution dense prediction, 2024.
[6] 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.
[7] 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 the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
[8] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory--fficiient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 1334411339,, 2022.
[9] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x16 words: Transformers for image recognition at scale. International Conference Learning Representations, 2020.
[10] Chengyue Gong and Dilin Wang. Nasvit: Neural architecture search for efficient vision transformers with gradient conflict-aware supernet training. Proceedings of the International Conference on Learning Representations, 2022.
[11] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0077 2023.
10

---

Published as a conference paper at ICLR 2025
[12] Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang. Flatten transformer: Vision transformer using focused linear attention. Proceedings of the IEEE/CVF International Conference Computer Vision, 2023.
[13] 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.
[14] Qibin Hou, Cheng-Ze Lu, Ming-Ming Cheng, and Jiashi Feng. Conv2former: A simple transformer-style convnet for visual recognition. arXiv preprint arXi::221.11944, 2022.
[15] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander Berg, Wan- Wan Lo, et al. Segment anything. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4015 4026, 2023.
[16] 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.
[17] 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, 15934444-1 2022.
[18) Tsung-Y 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.
[19] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in neural information processing systems, 36, 2023.
[20] 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.
[21] Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. Proceedings of the IEEE/CVF conference on Computer Vision and Pattern Recognition, 2022.
[22] Muhammad Maaz, Abdelrahman Shaker, Hisham Cholakkal, Salman Khan, Syed Waqas Zamir, Rao Muhammad Anwer, and Fahad Shahbaz Khan. Edgencxt: Efficiently amalgamated enn- transformer architecture for mobile vision applications, 2022.
[23] Donald J Newman. Rational approximation to |z/. Michigan Mathematical Journal, 11(1):11,
1964.
[24] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel 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.
[25] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ 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.
[26] 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, 1 115:211-252, 2015.
[27] AbdeIrahman Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang, and Fahad Shahbaz Khan. Swiftformer: Efficient additive attention for transformer-baut real-time mobile vision applications, 2023.
[28] 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-3533 2021.
11

---

Published as a conference paper at ICLR 2025
[29] Manish Kumar Singh, Rajeev Yasarla, Hong Cai, Mingu Lee, and Fatih Porikli. Tosa: Token selective attention for efficient vision transformers. arXiv preprint arXiv:2406.0888. 08816, 2024.
[30] Ily O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, al. Mlp- mixer: An all-mlp architecture for vision. Advances in neural information processing systems, 34:24261-24272, 2021.
[31] 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 Pattern Analysis and Machine Intelligence, 45(4):5314-5322 2022.
[32] 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.
[33] Lloyd N Trefethen. Approximation theory and approximation practice, extended edition. SIAM, 2019.
[34] 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 Computer Vision, pages 5785-5795, 2023.
[35] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[36] 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.
[37] 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.
[38] Chenglin Yang, Siyuan Qiao, Qihang Yu, Xiaoding Yuan, Yukun Zhu, Alan Yuille, Hartwig Adam, and Liang-Chit Chich Chen. Moat: Alternating mobile convolution and attention brings strong vision models. In The Eleventh International Conference on Learning Representations, 2022.
[39] Zhengyuan Yang, Linjie Li, Kevin Lin, Jianfeng Wang. Chung- Lin, Zicheng Liu, and Lijuan Wang. The dawn of Imms: Preliminary explorations with gpt-4v(ision). arXiv preprint arXiv:2309. /742/, 2023.
[40] 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.
[41] 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.
[42] 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 paltern recognition, pages 10819-10829, 2022.
[43] 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 arXin:2401.09411. 2024.
12

---

Published as a conference paper at ICLR 2025
A PADRE GENERALIZATIONS
In this section, we present a few generalization of the framework, namely, generalization to the cross-attention mechanism, well rational version of the framework.
A.1 MULTI-MODAL PADRE (L.E.,
PADRe can be generalized to multi-modal input in many ways. We present one here. Imagine that X(a) E RN x D represents input from one mode and (b) R M represents input from another mode. Then, we can define ((()) and ((()) 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 2,. this case, we write,
2, =
(11)
where & E (a, b) for each i. For instance, if the degree is = 3 and c = (a, a, b), then one finds that,
Zj = Yy(a)
(12)
2, = Y/(a) c v(())
(13)
Zj Youh о Y y
(14)
It important for the sequence is not trivial, i.e. contain cross-terms, not (a, a, a) nor (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 generalization of PADRe. Although we found that the standard PADRe based on polynomial approximations offered great performance in the context of computer vision, is expected that certain applications will benefit from this more powerful framework. In particular, the following approach is inspired by the improved effectiveness 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 polynomial function. As discussed in [23] for instance, in the one-dimensional case, the former have much higher approximation power than the latter. Further, from an intuitive perspective, adding flexibility to the denominator allows one better capture "poles" (or areas of high variability) often observed when using so f tmax. Finally, the introduction of a function at the denominator might be seen 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 = A X B
(15)
for 1,.., e, and,
K2 = o Vi,
(16)
C o
(17)
for j = 1,.., d and ls 1,.... Then, define
Na.b = Σ
Wa,bis (Kylabb + Va,b
(18)
j=1
Dab = Σ
Qa.ii [Cyla,b + Pa.D.
(19)
k=1
Finally, the output (prior to resizing) takes the form (i.e., wlement-wise "Hadamard division") "),
No.
Oa.b = Dab
(20)
13

---

Published as a conference paper at ICLR 2025
The cost of computing such approximation amounts to that of computing the numerator and denominator. Since both are PADRe-like, we expect a cost of roughly 2x that of the polynomial-only scheme, and thus linear scaling in the size of the input N.
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 small regularizing positive term.
B
PADRE- SPECIAL CASES
B.1
SIMA [16]
SimA (Simple Attention). as its name suggests, is simple mechanism for replacing the standard attention. The output of SimA can expressed as,
O == Q=QKV
(21)
where,
Q' Q' 110'11
K
K
(22)
||
and ! indicates the 7th column of a matrix (Ei column-wise normalization), and where,
= X Wo. K X WK. V x Wv.
(23)
for appropriate projection matrices as usual. Substitution implies that,
Om.
mmi
Σ Kys Vyn
(24)
X Wolma
[ Wklaa
[x Wvljn
(25)
12
||<l
ΣΣ 1 ky=1
(WvJavn) (Xm,ko Nj Xi
|Q
(26)
This should be recognized as homogeneous degree-3 polynomial with coefficients of the form,
[Wolag.ilV (W. Ja [Wv]ev.)
(27)
1|Q [
i.e., exhibiting a (normalized) polynomial relationship with the parameters/weights. Therefore, as claimed, SimA is special case of the PADRe framework.
B.2 CONV2FORMER [14]
Given an input e "xxxC Conv2Formez replaces the self-attention mechanism with a block of the form,
2 = AoV V
(28)
= DConvk (W1X)
(29)
V W2 X
(30)
where Z is the output, @ represents the Hadamard product, W; and W2 are weight matrices of two linear layers, and DConvex& denotes a depthwise convolution with kernel size & x k, which is really just efficient linear transformation applied each channel.
14

---

Published as a conference paper at ICLR 2025
Letting C' be tensor representing the convolution kernel and expanding formulas, find that,
[DConvex (Wix)) = Cc-pp-oo EWiilpss)
Xmmnk
(31)
P-1
W2 x 11,,A
Σ
W2 (el(0.))))))) Xa,ek
(32)
Therefore, it follows that,
(m) (imms)) m,n
[Z] = [A O Vlesas
Σ
Ci-p.s-..
Σ
W IWWllass.o X.,ck
P.q
s.t
(33)
ΣΣ
(Wllos).....
ECips--olWWlooooot... Xm,uk Xo,t,
P.q
(34)
This should recognized as a degree-2 (homogeneous) polynomial in the input X, and therefore an specific instance of the PADRe framework.
B.3 HYENA [24]
The proposed Hyena operator is defined follows (Definition 3.1, (24)),
Definition B.1. (Order-N Hyena Operator) Let (x+ x² ...,..,,%%) be projections of the input and let h1 h,,.., N be a set of learnable filters. The Hyena-N operator is defined by the recurrence:
(35)
2
n+1
xtt (h * =="),, Je, I...,, N
(36)
V1
zt N+1
(37)
That is, Hyena proposes 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 N Hyena Operator corresponds to the evaluation of degree-(M 1) homogeneous polynomial (or equivalently (N + )t-tessrr) which coefficients are non-linear functions of the parameters.
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 X be the original input vector and n := (no, Ri...., RN) E NN+1 be multi- indices, we argue that the order-N Hyena operator may written as,
y Σ The
II ==0
X
(38)
for some coefficients } that are polynomial functions of the parameters ((h") and linear projec- tions). This is a homogeneous polynomial of degree (N 1). In light of these observations, one finds that Hyena in fact a special case of PADRe.
We present analysis of the Order-N N Hyena Operator which shows it is specific instances of PADRe. In particular, we argue that the latter really is a homogeneous polynomial degree (N + 1). which
1 Higher levels of non- non-lincarity (N) increase expressivity, while glohal nature convolutions and projections ensures that all correlations among tokens are taken into account at multiple scales. We use x in order differentiate between the Hyena notation which projections rather than input.
15

---

Published as a conference paper at ICLR 2025
may also be interpreted an (N 1)-tensor (multilinear functional/operrttri Following the above definition and Remark 3.1 ([24)), an order-N Hyena operator may be expressed as,
y = x (h (xN-11 (N--1 (..)))))
(39)
This expression can be expanded as follows,
(-(sth T
1) (40)
Vty
PNN
N
N PNN-tN-1
N
Σ
1N-21-
N-1
Σ
h -t Fto
tN
1=
fx
2
=0=0
Using commutativity re-arrange the latter, find that
L
N
MiN
Σ
II
(41)
to 0
f-- m0
where t2 E (O,... L} for every Now, recall that each vector æf represents projection of the input RM onto L-dimensional space. Let us write,
M Σ m:=1
o = [P* xle,
Ph
Xm,-
(42)
Upon substitution, we find that,
N
N
II
Σ Am;=1
Ph
ta Xen.
(43)
E=
II
Pd
tymms
II
Xms
(44)
mi =
1=0
and therefore,
L
M
M
N
Uix
Σ
Σ
hi
-
Σ
Σ
ps
my
II j=0
Xms
(45)
===
IN-1=0
mp=0
Te N =
j=0
M
M
L
N
N II 1=0
N II /=0
II
pd
sy.m;
Xms
(46)
mg
m
f==
by
1
Xm,
(47)
mgg
M
claimed. The nonlinear dependence of the coefficients 7m the parameters {N' and (Pi) 1 takes the form,
N
7m
P
tyymn
(48)
fo=0
fy-1=1 =0
i=1
j=0
N
1.me 1)
PN
Pixmn
II
pi-
(49)
10=0
1==1
N
P tx
II
Σ
hy
pi-1
(50)
mn
fi-1m
(-1=0
That is, the coefficients are products of convolutions applied to the projections in each dimension (with the exception of the last dimension which is mere scalar multiplication.. This demonstrates
the claim.
M used represent the dimension of the input, whatever is.
16

---

Published as a conference paper at ICLR 2025
B.4 MAMBA [11]
In this section, we argue that the mamba SSM framework is in fact a special case of our PADRe framework. The mamba framework is state-space model based on the following recursion:
he Ah4- Br,
(51)
Yt = Chh,
(52)
where is the timestep, h4 the hidden state time H, is the input time and y/ the output at time t. The "bars" above the matrices is mamba notation for discretization, which takes the form,
Ā AA
(53)
B = ( (cAA- 1) AB
(54)
where À is N N matrix chosen to be diagonal, B is N vector and C is an 1 N vector. We now discuss technical details. First, we have the following lemma, Lemma 2. If ho = 0, the discretized system of Equations 51 has solution,
y1
(CA' A' A B) In
(55)
n=0
Proof. We proceed by induction. Clearly for = 0, get,
yo C B x C. A 2
(56)
Now assume the induction hypothesis holds and consider.
Ut+1 = C (Ahe + BI(+1)
(57)
= C
A
(A'-a B) In + Bzits)
(58)
n=0
H+1
C
(Āt+l-n B)
(59)
n=0
which proves the claim.
The main difference with regular SSM is that mamba allows B, C and A to be functions of the input (Algorithm 2). Specifically, B and C' are linear function of the input whereas A is the softllus of a linear function of the input, i.c.
B WB x
(60)
C T Wc
(61)
A(r) log (1 d (T+zWa))
1
(62)
3
where W3 is N x L matrix, We is a L x N matrix, WA is a (rank-1) D x D matrix, B and C' are size N x 1 and N vectors, respectively, and is some parameter.
For what follows, we make the assumption that,
5 A()) <
(63)
uniformly in the input æ. this case, we find that,
- - (
(64)
B = (AA) (c
- (6AA- AB
(65)
= = (1+AA+0(23) 1) B
(66)
= =AB+O(A))
(67)
17

---

Published as a conference paper at ICLR 2025
Therefore, under our assumption, it is reasonable expect that,
A R + AA
(68)
B x A = A ABT
(69)
= x² T' Wc
(70)
Substituting these expressions in the recursion, we find that,
yt =
A (B
(71)
Wc)
Σ A
2n
(WBx)
(72)
M==
We A 1t-n WBles A"B"Wllss (os Iy xnn
(73)
n=0
1===1
(I.n) ij
(x) + Tzj Fa)
(74)
where,
(t.n) A WWB i.j
(75)
This should be recognized as 3 homogeneous polynomial in the input x. Further, the coefficients Cij
(t.x) exhibit a polynomial relationship in the parameters (A, WB, Wc) (assuming truncated Taylor series for the exponential). Furthermore, the matrix exhibit such structures that the latter may be computed rapidly. Thus, it follows that Mamba is special case of our PADRe framework.
B.5 CASTLING- VIT [41]
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,
0 -= - (KT v) + (2 Mpw) V
(76)
π
where,
Q = X = V= X
(77)
is the identity, and MDw represents a learnt depth-wise convolution (i.e., fast linear transforma- tion). Upon substitution, we find that the left-hand term in Eq.equation 76 is degree-3 homogeneous polynomial in the input X, whereas the right-hand term is linear function of the input. All in all, this implies that the output O of the Castling VIT module is 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 [35]
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,
QT
[Attention(Q.k K, V) = sof tmax
(78)
Vdle
QTT
e Vdk
Vin
(79)
QKT
E, e Vic
m.j
18

---

Published as a conference paper at ICLR 2025
where,
= X K = X V = =XWV
(80)
for some weight/projections matrices, as usual. Now comes the main observation: thanks to the normalization va. may expect that the elements QK will generally lie in canonical
vdd interval |-1,, I. for some < L. E IR. In this case, the exponential functions may be expanded using the Taylor series as,
QKTI
QKT V
vdx
L (d + 1)!
mi
m.
O
(81)
1!
1=0
for some degree d sufficiently large that the enror term is on the order of 0 < < truncated Taylor expansion may be assumed be accurate over this fixed interval only.
1. Note that such
Upon substitution and expansion, we therefore find that,
Q QQKTI
E: Σ
V
d
|Ateettion(, K, K)V)
(82)
R
Σ; Ep=
m.
d' pt
and since Q, K and V are all linear transformations of the input data, we conclude that the attention can be arbitrarily approximated by rationale function degree (2d + 1) 1)/(..). The degree however, may have to be impractically large in order reach proper accuracy, and therefore the approx- imation attention itself, albeit a member of PADRe, may not offer competitive performance in practice. Nonetheless, this shows that the standard attention falls within the PADRe framework as well.
Finally, let highlight that such analysis is not limited to the softmax and the attention mechanism in particular, and may in fact be applied to any sufficiently regular (e.g., smooth) nonlinear functions (such activations) substitution, one finds that the resulting approximation is in fact a polynomial itself (given that all its building block such). As part of some future work, intend generalize PADRe this context as well.
long the input remains bounded in some interval. Upon making such
19


Reviews:
[{"review_id": "1oe0gSqbyN", "reviewer": "Reviewer_Lp3X", "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.", "questions": "See weakness.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "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.", "weaknesses": "1): Figure 1 illustrates a specific design within the PADRe framework, but a straightforward baseline is missing. For example, for degree-\n$d$ approximation, directly ensembling $d$ MLP-Mixers can be a strong baseline, and this approach can also be efficiently parallelized.\n\n2): 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. \n\n3): 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1I2Z25sJbx", "reviewer": "Reviewer_o4xT", "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.", "questions": "See the weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n\n2. Experiments have well demonstrated the effectiveness of their approach, especially considering the FLOPs reduction in Table 3 and Table 4.\n\n3. The paper is in a good structure. Figures are clear as well.", "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.\n\n2. 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.\n\n3. In Table 1, the improvement over DeiT-Base is minor. Could it be that PADRe works more effectively for smaller Transformers?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4dqD9ryIsw", "reviewer": "Reviewer_MgyZ", "summary": "This paper proposes PADRe, a unifying framework for replacing the attention mechanism based on\npolynomial approximants to provide an efficient, linear-complexity alternative to the\nstandard self-attention without sacrificing accuracy.", "questions": "Please see the weaknesses above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper addresses an important problem in vision transformer design.\n\n2. This paper provides a theory-driven solution with a thorough mathematical justification.", "weaknesses": "1. How would the method work for large language models (LLMs), with much longer sequence lengths?\n\n2. How to decide the optimal degree of PARDe?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uOzCkS88se", "reviewer": "Reviewer_JNAG", "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", "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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Framework that helps building alternative self-attention based transformer architecture\n- advantageous over Flash 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.\n- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)