PaperlayerView exact revision on arXiv

Source-linked AI summary

PolaFormer: Polarity-aware Linear Attention for Vision Transformers

Weikang Meng, Yadan Luo, Xin Li, Dongmei Jiang, Zheng Zhang

arXiv:2501.15061v2cs.CVcs.AI

TL;DR

Softmax attention is computationally expensive, and existing linear attention can lose negative query-key information while producing less sharp attention distributions. PolaFormer models same-signed and opposite-signed interactions separately and uses learnable power rescaling; experiments report improvements of up to 4.6% across vision tasks and benchmarks. The approach retains linear sequence-length complexity but remains subject to low-rank and feature-map limitations.

  • Problem

    Existing linear attention improves efficiency but can discard negative query-key interactions and produce higher-entropy, less discriminative attention than softmax.

  • Method

    PolaFormer uses polarity-aware streams for same-signed and opposite-signed interactions, learnable power rescaling, convolution, and sign-aware coefficient matrices.

  • Results

    PolaFormer improves performance by up to 4.6% across various vision tasks and the Long Range Arena benchmark while maintaining linear complexity.

  • Takeaways & Limitations

    The method provides a linear-complexity attention design that more comprehensively represents query-key relationships and restores sharper attention responses.

  • Takeaways & Limitations

    The method addresses low-rank degeneracy with convolution, and prior fixed-norm rescaling may be suboptimal across datasets.

Abstract

from arXiv · show

Linear attention has emerged as a promising alternative to softmax-based attention, leveraging kernelized feature maps to reduce complexity from quadratic to linear in sequence length. However, the non-negative constraint on feature maps and the relaxed exponential function used in approximation lead to significant information loss compared to the original query-key dot products, resulting in less discriminative attention maps with higher entropy. To address the missing interactions driven by negative values in query-key pairs, we propose a polarity-aware linear attention mechanism that explicitly models both same-signed and opposite-signed query-key interactions, ensuring comprehensive coverage of relational information. Furthermore, to restore the spiky properties of attention maps, we provide a theoretical analysis proving the existence of a class of element-wise functions (with positive first and second derivatives) that can reduce entropy in the attention distribution. For simplicity, and recognizing the distinct contributions of each dimension, we employ a learnable power function for rescaling, allowing strong and weak attention signals to be effectively separated. Extensive experiments demonstrate that the proposed PolaFormer improves performance on various vision tasks, enhancing both expressiveness and efficiency by up to 4.6%.

1 INTRODUCTION

Standard softmax attention captures long-range dependencies but has quadratic cost, while linear attention improves efficiency yet loses negative interactions and sharp attention structure. PolaFormer addresses both issues with polarity-aware streams and learnable rescaling.

  • Standard self-attention has quadratic complexity O(N^2), creating computational overhead for long sequences and high-resolution images.
  • Linear attention reduces complexity by replacing softmax with kernelized feature maps, but its attention weights are often more uniform and less specific.
  • Non-negative feature maps discard negative-negative and positive-negative query-key interactions, limiting relational coverage and discriminative power.
  • Without exponential scaling, linear attention produces lower-entropy distributions that distinguish less effectively between strong and weak query-key pairs.
  • PolaFormer separates same-signed and opposite-signed interactions into two streams, then combines them using learnable sign-aware scaling.
  • A channel-wise learnable power function restores sharper attention responses, and experiments report performance improvements of up to 4.6%.

2 RELATED WORK

Related work improves transformer efficiency by restricting attention or approximating softmax with kernel-based linear attention. These approaches differ in how they preserve contextual information and attention-weight properties.

  • Localized and sparse attention reduce computation by restricting attention to smaller windows or sparser token patterns.
  • Restricting attention can sacrifice contextual information, while linear-attention methods motivate further work on attention weights and low-entropy kernel designs.
  • Kernel-based linear attention replaces softmax with feature-map dot products to reduce complexity from O(N^2) to O(N).
  • Existing linear-attention variants use feature maps such as ReLU and 1 + ELU, while Cosformer adds cosine-based re-weighting and FLatten Transformer uses a power operation.

3 PRELIMINARY

The preliminary section contrasts quadratic self-attention with kernel-based linear attention. Linearization reduces sequence-length scaling, but feature-map choices introduce non-negativity and entropy-related limitations.

  • Self-attention projects token inputs into query, key, and value vectors and has per-head complexity O(N^2d).
  • Kernel-based linear attention approximates the softmax kernel exp(q_i k_j^T) using feature-map products and reassociates matrix multiplication.
  • Associative matrix multiplication reduces per-head complexity to O(Nd′^2), which scales linearly with sequence length.
  • Feature maps are designed to preserve non-negative attention values and low entropy, using activations such as ReLU or 1 + ELU.
  • Non-negative feature maps lose information from original negative dot-product values, while fixed rescaling norms may be suboptimal across datasets.

4 PROPOSED APPROACH

PolaFormer captures both same-signed and opposite-signed query-key interactions, then uses learnable rescaling to reduce entropy while retaining linear complexity. Its formulation separates polarity streams, mixes their responses, and theoretically motivates power-function feature maps for sharper attention.

  • Polarity-aware attention: PolaFormer separates same-signed and opposite-signed query-key interactions to recover information discarded by prior linear attention methods.The mechanism independently computes interactions across positive-positive, negative-negative, positive-negative, and negative-positive components.
  • Polarity-aware attention: Learnable mixing combines same-signed and opposite-signed similarities without directly subtracting them, avoiding non-negativity violations and unstable training.Value vectors are split into same-signed and opposite-signed halves for separate response handling.
  • Entropy reduction: The method analyzes positive sequence entropy to restore the sharper attention distributions typically associated with softmax-based attention.The analysis treats each attention row as a generalized unnormalized positive sequence and defines positive sequence entropy for it.
  • Entropy reduction: A differentiable element-wise function with positive first and second derivatives can strictly reduce the transformed sequence’s positive sequence entropy.The theorem states PSE(⟨g(x), g(y1)⟩, ..., ⟨g(x), g(yN)⟩) < PSE(⟨x, y1⟩, ..., ⟨x, yN⟩).
  • Entropy reduction: ReLU and ELU +1 do not satisfy the theorem’s derivative conditions across their entire domains, explaining why they fail to reduce entropy effectively.The paper uses this observation to motivate an alternative rescaling function.
  • Learnable power functions: PolaFormer uses channel-wise learnable power functions with exponents greater than 1 to rescale positive and negative components while preserving linear complexity.The feature maps apply the learnable power function to ReLU(x) and ReLU(−x) for queries or keys; the complexity analysis demonstrates linear scaling in sequence length.

5 EXPERIMENTS

Experiments evaluate PolaFormer across image classification, detection, segmentation, efficiency, ablations, and long-range sequence tasks. Across these settings, the method generally improves baseline performance while retaining favorable efficiency.

  • Experimental Setup: PolaFormer is evaluated on ImageNet-1K classification, COCO detection and instance segmentation, ADE20K semantic segmentation, and Long Range Arena tasks.ImageNet-1K variants are trained from scratch, while downstream segmentation models use ImageNet-1K pretrained weights.
  • ImageNet-1K Classification: PolaFormer consistently outperforms baseline models on ImageNet-1K, with DeiT-T-PolaFormer surpassing other DeiT variants by 0.5% to 6.3%.PVT-T/S-PolaFormer also improves over corresponding baselines by 3.7% and 2.1% at comparable FLOPs.
  • Efficiency Analysis: 1.15× and 1.12×: PVT-PolaFormer inference is faster on RTX3090 and RTXA6000, while Swin-PolaFormer reaches 1.32× on the reported platforms.Efficiency curves compare accuracy against FLOPs and runtime on ImageNet-1K.
  • Object Detection and Instance Segmentation: 2.3% to 4.6%: PVT-T-PolaFormer surpasses baselines with RetinaNet and Mask R-CNN on COCO, while Swin-T-PolaFormer reaches 49.1% APb 75.The Swin-T result is a 1.4% improvement over the original Swin-T with Mask R-CNN.
  • Object Detection and Instance Segmentation: PolaFormer backbones consistently outperform original backbones across the reported COCO detection, instance-segmentation, and ADE20K semantic-segmentation settings.The paper uses RetinaNet, Mask R-CNN, Cascade Mask R-CNN, Semantic FPN, and UperNet configurations.
  • Ablation Study: 1.8%: adding polarity coefficients Gs and Go improves ablation performance, indicating learned complementarity between same-signed and opposite-signed values.Depth-wise convolution achieves 74.6% in the reported component comparison.
  • Ablation Study: On Long Range Arena, the learnable scaling factor is evaluated with α = 3, 5, and 7; classification results differ by no more than 2% across α variations.PolaFormerα=3 achieves the best reported performance, while α depends primarily on model size and context length.

6 CONCLUSION

The conclusion presents PolaFormer as a linear-complexity transformer that restores non-negative, spiky attention through polarity-aware similarity and learnable rescaling. Experiments support compatibility with attention-based models and a favorable performance–efficiency balance.

  • Conclusion: PolaFormer has linear complexity and is designed to preserve non-negative and spiky attention weights.Its similarity computation models omitted negative interactions, while a learnable power function lowers entropy through rescaling.
  • Conclusion: PolaFormer combines polarity-aware similarity, learnable power rescaling, convolution, and polarity coefficients to address negative interactions, entropy, low rank, and complementary signed relationships.The conclusion describes these components as the basis of the proposed efficient transformer.
  • Conclusion: Experiments across vision tasks and Long Range Arena show compatibility with most attention-based models and a better balance between performance and efficiency.The conclusion reports this as the overall experimental outcome.

A APPENDIX

The appendix contains the proof and supporting lemmas for Theorem 1, implementation details, and additional efficiency, initialization, entropy, and attention-map analyses.

  • Appendix Contents: The appendix provides the mathematical proof and supporting lemmas for Theorem 1.It also includes implementation details for the vision-task experiments.
  • Appendix Contents: Additional appendix materials cover training settings, long-sequence efficiency, G initialization comparisons, attention-distribution entropy, and attention-map visualizations.These items are listed as supplementary analyses and implementation details.

A.1 PROOF OF THEOREM 1

The proof shows that element-wise functions with positive first and second derivatives can strictly reduce positive-sequence entropy in transformed linear-attention scores. This supports replacing softmax with element-wise scaling while retaining linear complexity and lower entropy.

  • Lemma 1: Lemma 1 shows that the function induced by element-wise g preserves positivity together with positive first and second derivatives.The construction applies g independently to the dimensions of x and y.
  • Lemma 2: For two positive values, applying any function with positive first and second derivatives yields PSE(f(a), f(b)) ≤ PSE(a, b).The proof compares the original ratio c = a/b with the transformed ratio d = f(a)/f(b), where d > c > 1.
  • Entropy comparison: The entropy comparison gives H2 < H1, so functions satisfying the derivative conditions decrease entropy.The argument defines h(x) and compares the original and transformed two-value distributions.
  • Conclusion: The resulting element-wise scaling removes the need for softmax while providing linear complexity and lower entropy.The conclusion connects the theorem’s scaling effect to the attention mechanism.

A.2 IMPLEMENTATION DETAILS

The implementation details specify task-specific training configurations for classification, object detection, semantic segmentation, and Long Range Arena experiments. These settings vary optimizers, learning rates, schedules, batch sizes, and training duration by task.

  • Classification: Classification models use AdamW for 400 epochs, including 20 warm-up epochs, with learning rate 1e−3 at batch size 1024.The weight decay is 5e−2, using the official Swin Transformer training framework.
  • Object Detection: Object detection uses pretrained PVT or Swin backbones with RetinaNet, Mask R-CNN, or Cascade Mask R-CNN under task-specific schedules.The stated optimizer is AdamW with learning rate 1e−4 and weight decay 1e−4.
  • Semantic Segmentation: Semantic segmentation uses pretrained PVT and Swin models with SemanticFPN or UperNet and AdamW optimization.Training runs for 40000 iterations for PVT-SFPN and 160000 for Swin-UperNet models.
  • Long Range Arena: Long Range Arena models are trained from scratch with AdamW and task-specific batch sizes and learning rates.The configurations cover ListOps, text classification, Pathfinder, image classification, and retrieval.

A.3 LONG SEQUENCE EFFICIENCY

The Long-Range Arena benchmark is used to assess PolaFormer’s scalability in long-sequence and high-resolution settings. The reported experiments demonstrate efficiency and scalability across vision and NLP applications.

  • Scalability evaluation: Long-Range Arena experiments evaluate PolaFormer’s efficiency and scalability for high-resolution vision tasks and long-sequence NLP applications.The benchmark is used specifically to assess scalability in these settings.
  • Reported metrics: Table 6 reports accuracy, throughput, and peak memory for the compared models.A denotes accuracy, T throughput, and M peak memory cost.

A.4 COMPARISON OF THE RESULTS WITH DIFFERENT INITIALIZATIONS OF COEFFICIENTS MATRICES

The initialization study compares five coefficient-matrix initialization strategies on a 4k-token Long Range Arena text-classification task. Attention visualizations further examine entropy and similarity characteristics across attention mechanisms.

  • Initialization comparison: The initialization comparison tests Kaiming uniform, zeros, Normal(0,1), Uniform(0,1), and ones.All strategies are evaluated on the Long Range Arena text-classification task with sequence length 4k.
  • Initialization comparison: The study maintains the experimental setup used in Table 4 while varying only the coefficient-matrix initialization strategy.The comparison is conducted on the TEXT task.
  • Attention distributions: Figure 5 compares entropy and one-row attention-score distributions for standard self-attention, linear attention, and PolaFormer.The reported comparison states that PolaFormer has lower entropy than linear attention.
  • Attention maps: Figure 6 visualizes additional attention-map examples to examine similarity calculation and focus on relevant locations.The passage attributes this behavior to the designed kernel function.
Loading 2501.15061v2…