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
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Efficient Vision Transformer, Attention Approximation, Computer Vision, Deep Learning

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.

Decision: accept (poster)
Reviewer scores: 5, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730233607957

thread_type:
official_review

<thread_object>
creator: reviewer_jnag

<content>
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

strengths:
- Framework that helps building alternative self-attention based transformer architecture
- 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.
- 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

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730333995376

thread_type:
official_review

<thread_object>
creator: reviewer_mgyz

<content>
summary:
This paper proposes PADRe, a unifying framework for replacing the attention mechanism based on
polynomial approximants to provide an efficient, linear-complexity alternative to the
standard self-attention without sacrificing accuracy.

strengths:
1. This paper addresses an important problem in vision transformer design.

2. This paper provides a theory-driven solution with a thorough mathematical justification.

weaknesses:
1. How would the method work for large language models (LLMs), with much longer sequence lengths?

2. How to decide the optimal degree of PARDe?

questions:
Please see the weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730466234661

thread_type:
official_review

<thread_object>
creator: reviewer_o4xt

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

strengths:
1. Good insight and strong theoretical proofs. Approximating self-attention with polynomial approximants sounds interesting. The authors have provided solid proofs that the proposed PADRe can cover other special cases in recent works.

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

3. The paper is in a good structure. Figures are clear as well.

weaknesses:
1. The paper acknowledges that model performance begins to saturate (Lines 462-463) when the polynomial degree exceeds 3, suggesting potential numerical stability issues with higher-degree polynomials. Investigating more stable polynomial bases could be beneficial.

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

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

questions:
See the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730824200310

thread_type:
official_review

<thread_object>
creator: reviewer_lp3x

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

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-
$d$ approximation, directly ensembling $d$ MLP-Mixers can be a strong baseline, and this approach can also be efficiently parallelized.

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

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

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732304820484

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We would like to thank the reviewers and ACs for their efforts in this review process. And we thank the reviewers for their constructive comments. We are encouraged that the reviewers find our work to have a great potential for impact, with strong mathematical foundation and extensive experiments demonstrating effectiveness of PADRe, and being more advantageous than FlashAttention. 

We have addressed all the reviewers’ comments. Specifically, some reviewers asked about applying PADRe to LLMs. We note that existing models like Mamba and Hyena have already been used for LLMs and that they fall within PADRe framework, showing the successes of our PADRe framework for LLMs. We have also performed an experiment applying our specific PADRe instance to LLM, which further confirms the effectiveness of PADRe and indicates its great potential. More information is provided in our detailed responses below.
</content>
</thread_object>

</thread>


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

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