PolaFormer++: Polarity-Aware Linear Attention with Channel-Wise Spikiness Sets New Benchmark

Transformer's Softmax self-attention has quadratic complexity O(N²d), limiting deployment on long sequences like high-resolution images and videos. Linear attention replaces Softmax with a kernel feature map φ(·) to rewrite attention as φ(Q)(φ(K)ᵀV), reducing complexity to linear O(Nd'²). However, this efficiency comes at the cost of expressiveness due to two fundamental issues:

Negative value loss : To keep attention weights non-negative, existing feature maps (e.g., ReLU, 1+ELU) zero out negative components in queries and keys. Decomposing q·k into four interaction terms shows that same-sign interactions (positive-positive, negative-negative) are preserved, but opposite-sign interactions (positive-negative, negative-positive) are completely discarded, preventing faithful reconstruction of the original similarity.

High attention entropy : Without Softmax's exponential scaling, linear attention weights become near-uniform, yielding high entropy and blunt distributions. Previous fixes like FLatten Transformer apply a fixed power function uniformly across all channels, ignoring that different channels require different sharpness levels.

PolaFormer Recap: Polarity-Aware Linear Attention

PolaFormer (ICLR 2025) introduced polarity decomposition : each element of q and k is split into positive and negative parts (q = q⁺ − q⁻, k = k⁺ − k⁻). Same-sign and opposite-sign interactions are processed independently in a Same-polar Flow and an Opposite-polar Flow . The value vector is split along channels to carry responses for each flow. Two learnable polarity gating matrices Gₛ, Gₒ perform Hadamard-weighted combination to approximate subtraction and compensate for relaxation effects while keeping weights non-negative.

PolaFormer++: Channel-Wise Spikiness Evolution

PolaFormer++ adds PaCS (Polarity-aware Channel-wise Spiky) feature mapping . Recognizing that different channels contribute differently to similarity and need distinct sharpness, the authors design a per-channel learnable temperature for each polarity flow.

3.1 PaCS: A Learnable Temperature per Channel

Using exponential base functions, each channel gets a learnable temperature coefficient p (normalized via sigmoid):

φ(x; p) = (exp(p₁x₁), …, exp(p_d x_d)) − 1

Separate temperature vectors pₛ and pₒ are learned for same-polar and opposite-polar flows, yielding feature maps gₛ(·; pₛ) and gₒ(·; pₒ). Thus every channel and every polarity flow has its own "sharpness knob", allowing adaptive amplification of discriminative responses and suppression of minor ones.

3.2 Does the Network Learn Channel-Wise Differentiation?

Scatter plots of the two temperature parameters (s₁, s₂) for same-polar flow across channels show wide dispersion significantly deviating from the y = x diagonal, confirming highly independent sharpness behaviors per channel rather than convergence to a fixed power. PCA visualization of gating matrices Gₛ and Gₒ across layers reveals clearly separated clusters, proving the learnable mixing strategy captures complementary relationships between same-sign and opposite-sign values.

Experiments: Six Core Vision Tasks

PolaFormer++ builds a hierarchical backbone family (b0–b5, 7M–114M parameters) and validates on six task categories.

4.1 ImageNet-1K Classification

PolaFormer++ achieves top-tier accuracy at every scale. On the accuracy-FLOPs curve, it consistently outperforms efficient models from the past three years across all compute budgets.

4.2 Efficiency: Up to 5.25× Speedup on Long Sequences

Measured throughput and peak memory on 4K–200K sequence lengths: vanilla Softmax attention OOMs after 65K; Flash Attention avoids OOM but falls further behind in speed — PolaFormer++ achieves 5.25× speedup at 200K length while maintaining the lowest memory footprint.

4.3 COCO Detection & Instance Segmentation

RetinaNet 1×: PolaFormer++-b2 reaches 46.0 APᵇ (vs. RMT-T 45.1, RAVLT-T 44.1); b3 reaches 47.3 APᵇ .

Mask R-CNN 1×: b3 achieves 49.7 APᵇ / 44.6 APᵐ with fewer parameters than InternImage-S (69M, 47.8/43.3); b4 reaches 51.5 / 45.9 , surpassing RMT-B (73M, 51.1/45.5).

Mask R-CNN 3×: b4 hits 52.9 APᵇ / 46.2 APᵐ , beating RMT-B (52.2/46.1) and VMamba-S.

Largest b5 (114M, 522G FLOPs) attains 51.9 APᵇ / 46.1 APᵐ at significantly lower compute than peers.

4.4 ADE20K / Cityscapes Semantic Segmentation

PolaFormer++-b3 scores 50.5% mIoU on ADE20K with 28G FLOPs, +1.3% over EfficientViT-L1 and +2.0% over SegNeXt-B at lower compute. On Cityscapes it reaches 83.4% mIoU . Visualizations show sharper boundaries for pedestrians, sidewalks, vehicles, and complete segmentation of distant crowds in complex street scenes.

4.5 Diffusion Generation (PolaDiT)

Replacing attention in DiT with PolaFormer++ yields PolaDiT. On ImageNet-1K class-conditional generation:

DiT-S/2 scale: FID 58.30 / IS 25.29 , beating DiG (62.06), MHLA (59.80), and other recent linear attention variants.

DiT-B/2 scale: FID 36.00 / IS 41.56 , maintaining best-in-class among linear attention variants and demonstrating strong scalability.

4.6 3D Novel View Synthesis (Pola-GNT)

Swapping dot-product attention in GNT with PolaFormer++ on LLFF (8 scenes): PSNR improves from 27.25 to 27.34 dB , LPIPS drops from 0.1019 to 0.1009 , SSIM rises from 0.8868 to 0.8885 — all three metrics improve, confirming applicability to 3D spatial reasoning and geometric modeling.

4.7 Lightweight Image Super-Resolution

Integrating PolaFormer++ into ESRT and MambaIRv2:

Accuracy : MambaIRv2-Pola++ surpasses previous best ESC on all 5 benchmarks (e.g., ×4 Set5 32.59 dB, Manga109 31.38 dB).

Efficiency : ESRT-Pola++ on Urban100 ×3 reduces latency to 85.7 ms (vs. 327.4 ms, 73.8% reduction ) and saves 90.0% memory ; on Manga109 memory saving reaches 92.5% . Softmax-based MambaIRv2 OOMs on large-resolution sets and slows ~20× on small sets (Set5: 1504.8 ms vs 78.6 ms), while Pola++ adds negligible overhead.

4.8 Ablation & Stability

Module ablation (ImageNet-1K, b0): Polarity baseline 76.3%; adding gating matrices, learnable sharpness, and PaCS stepwise improves to 76.9% .

Unified linear attention comparison (~29M params, 4.5G FLOPs): PolaFormer++ scores 82.8% , exceeding Efficient Attn (81.0), FLatten Attn (82.1), Agent Attn (82.6), InLine Attn (82.4), LinDiff (82.3) while retaining linear complexity.

Training stability : BF16 mixed precision with global batch sizes 2048/4096/8192 for 300 epochs shows stable loss and gradient norms — no explosions, NaNs, or divergence.

Conclusion

PolaFormer++ addresses two root causes — negative value loss and insufficient attention sharpness — with a unified, theoretically grounded solution:

Polarity awareness : Polarity decomposition + dual-flow interaction + learnable gating matrices let every element of q and k participate equally, fully reconstructing similarity information.

Channel-wise sharpness : PaCS feature mapping uses two learnable temperature vectors to adaptively tune sharpness per channel per polarity flow.

General theoretical criterion : Based on Schur-concavity, the authors provide the first rigorous criterion to judge whether a feature map possesses entropy-reducing (sharpness) properties, generalizing PolaFormer's specific proof into a universal framework.

Across classification, detection, segmentation, super-resolution, diffusion generation, and 3D synthesis, PolaFormer++ establishes a stronger accuracy-efficiency trade-off for linear attention in vision Transformers. Code is open-sourced at https://github.com/ZacharyMeng/PolaFormerPP.