TPAMI26 | 全面升级PolaFormer++:极性感知计算流 +通道级尖峰,实现线性注意力的新天花板

作者:机器之心 发布:2026-09-18 07:43 收录:2026-09-18 17:40 1 次阅读 约 2920 字
摘要:一种增强的极性感知线性注意力机制。
推荐理由:本文涵盖「TPAMI」、「基础模型」、「多模态」等多个主题,重点关注 TPAMI。


作者信息:本文一作孟维康是哈尔滨工业大学(深圳)与鹏城实验室联合培养的博士生,本科毕业于哈尔滨工业大学,主要研究方向是大规模基础模型的高效训练和推理算法。通讯作者张正教授是哈尔滨工业大学(深圳)长聘教授、博士生导师,长期从事高效能多模态机器学习研究。


继 ICLR 2025 的 PolaFormer 之后,哈工大(深圳)与鹏城实验室合作的扩展版本 PolaFormer++ 近日被 IEEE TPAMI 正式接收。新版本在理论框架、特征映射设计和实验覆盖面上全面升级:首次给出判定特征映射是否具备「降熵尖锐性」的理论判据,并提出极性感知的通道级尖锐(PaCS)特征映射,在六大类视觉核心任务上全面验证了线性注意力的潜力。



    • 论文题目:PolaFormer++: Polarity-Aware Linear Attention with Channel-wise Spikiness
    • 作者:Weikang Meng, Yadan Luo, Liangyu Huo, Yingjian Li, Yaowei Wang, Zheng Zhang, Heng Tao Shen
    • 单位:哈尔滨工业大学(深圳),鹏城实验室,昆士兰大学,同济大学
    • 代码开源:https://github.com/ZacharyMeng/PolaFormerPP
    • 论文链接:https://ieeexplore.ieee.org/document/11657470
    • 会议版论文(ICLR 2025):https://arxiv.org/abs/2501.15061


    一、引入:线性注意力的两大「先天不足」


    Transformer 的核心——Softmax 自注意力,复杂度是序列长度的平方 O(N2d)。处理长视频、高分辨率图像时,计算和显存开销急剧上升,严重限制了实际部署。


    线性注意力用核化特征映射 ϕ(⋅) 替换 Softmax,把计算改写为 ϕ(Q)(ϕ(K)⊤V),借助矩阵乘法结合律将复杂度降为线性的 O(Nd′²)。但效率的提升是以表达能力为代价的,差距主要来自两方面:


    1. 负值丢失。为保证注意力权重非负,现有方法(如基于 ReLU、1+ELU 的特征映射)不得不把 query 和 key 中的负值全部置零。将 q、k 按正负部分解后可以看到,内积实际上包含四项:



    前两项是同号维度间的交互,后两项是异号维度间的交互。基于 ReLU 的映射直接把负成分抹掉,后两项(neglected negatives)完全丢失 —— 模型再也无法完整地重建原始的 q・k 相似度。


    2. 注意力分布信息熵过高。没有了 Softmax 的指数缩放,线性注意力的权重分布趋于均匀,熵更高、不「尖锐」。模型难以区分强弱 q-k 对,对关键 token 的关注力被稀释。此前的补救方案(如 FLatten Transformer 的固定幂函数)对所有通道施加统一的尖锐程度,忽略了不同通道本就需要不同程度的尖锐化这一事实。


    二、PolaFormer 回顾:极性感知线性注意力


    PolaFormer(ICLR 2025)的核心思想是极性分解:把 q 和 k 逐元素拆成正部与负部(q=q+−q−,k=k+−k−),让正负成分平等地参与相似度计算。同号交互与异号交互分别在同号流(Same-polar Flow)异号流(Opposite-polar Flow)中独立处理;value 向量沿通道维一分为二,分别承载两路的响应;最后用两个可学习的极性门控矩阵 Gs、Go 对两路输出做 Hadamard 加权,以可学习的方式近似「减法」操作、补偿松弛减法带来的影响,同时保证注意力权重非负。


    三、PolaFormer++:通道级尖锐度的全面进化


    图 1:PolaFormer++ 整体框架。query-key 对按极性被显式拆分为同号流与异号流,两路输出分别由可学习符号感知矩阵 Gs、Go 调控;极性感知的通道级尖锐特征映射(PaCS)ϕs、ϕo 分别作用于两路,在保持尖锐性的同时区分不同通道对不同极性流的贡献。


    在 PolaFormer 的基础上,PolaFormer++ 迈出了关键一步:不同通道对相似度的贡献不同,所需的尖锐程度也不同。为此,作者提出了 PaCS(Polarity-aware Channel-wise Spiky)特征映射。

    3.1 PaCS:给每个通道配一个可学习的「温度」


    图 2:PaCS 示意。左:极性分解后的向量由不同的极性流驱动,正负成分获得独立的尖锐化映射;右:基函数与逐通道温度参数组合,构建通道级差异化的特征映射,使输入向量呈现多尺度、通道级的尖锐性。


    具体设计上,作者以指数函数为基函数,为每个通道引入一个可学习的温度系数 p(由 sigmoid 归一化):


      • ϕ(x;p)=(exp(p1x1),…,exp(pdxd))−1

      • 同号流与异号流各配一套温度向量 ps 和 po,分别得到特征映射 gs(⋅;ps) 与 go(⋅;po)


      这样一来,每个通道、每条极性流都有自己专属的「尖锐度旋钮」,模型可以自适应地放大判别性强的响应、压平次要响应,把强弱注意力信号有效分离。


      3.2 网络真的学到了「逐通道差异化」吗?


      作者将训练后同号流两个温度参数(s1,s2)按通道画成散点图。可以看到,点云大范围弥散、显著偏离 y = x 对角线——说明网络为不同通道学到了高度独立的尖锐化行为,而不是趋同于一个固定幂次。


      图 3:每个点代表一个特征通道的一对温度参数(s1,s2)(s1,s2)。显著偏离 y=x 对角线表明网络为每种极性学习到了高度独立的分布。


      类似的,对极性门控矩阵 Gs 与 Go 在多个层上做 PCA 可视化,两类矩阵在特征空间中形成明显分离的簇,证实可学习混合策略确实捕捉到了同号与异号值之间互补的关系:


      图 4:不同层中 Gs(红)与 Go(黄)的 PCA 可视化。两类门控矩阵分布明显分离,说明二者学到了互补的极性调控模式。


      四、实验:六大视觉核心任务


      PolaFormer++ 构建了 b0–b5 共 6 个规格的层级式视觉骨干家族(7M–114M 参数),并在六大类任务上做了系统验证。


      4.1 ImageNet-1K 图像分类


      PolaFormer++ 在各规格上均取得同级别最优或极具竞争力的精度:



      精度 - 算力曲线上,PolaFormer++(红色实线)整体位于最上方,在各个计算档位都压过近三年的高效模型:


      图 6:ImageNet-1K 上 Top-1 精度 vs. FLOPs 曲线,PolaFormer++(红色)在各档位全面领先。


      4.2 效率:长序列下最高 5.25× 提速


      在 4K–200K 序列长度上实测吞吐与峰值显存:普通 Softmax 注意力(Vanilla-Torch)在 65K 之后直接 OOM;Flash Attention 虽不爆显存,但速度被 PolaFormer++ 越甩越远 ——在 200K 序列长度下提速达 5.25×,且显存占用始终保持最低。


      图 7:吞吐(折线)与峰值显存(柱状)随序列长度的变化。PolaFormer++ 在长序列下显著优于 Flash Attention,最高提速 5.25×。


      4.3 COCO 目标检测与实例分割


      以 RetinaNet、Mask R-CNN(1× 与 3× schedule)为框架:


      • RetinaNet 1×:PolaFormer++-b2 达 46.0 APb,超过 RMT-T(45.1)与 RAVLT-T(44.1);b3 达 47.3 APb
      • Mask R-CNN 1×:b3 达 49.7 APb / 44.6 APm,以更少参数量超过 InternImage-S(69M,47.8 / 43.3);b4 达 51.5 / 45.9,超过 RMT-B(73M,51.1 / 45.5);
      • Mask R-CNN 3×:b4 达 52.9 APb / 46.2 APm,超过 RMT-B(52.2 / 46.1)与 VMamba-S 等对手;
      • 最大规格 b5(114M,522G)在 Mask R-CNN 1× 下达 51.9 APb / 46.1 APm,算力显著低于同精度模型。


      4.4 ADE20K / Cityscapes 语义分割



      PolaFormer++-b3 在 ADE20K 上以 28G FLOPs 达到 50.5% mIoU,分别比 EfficientViT-L1 和 SegNeXt-B 高 1.3% 和 2.0%,且算力显著更低;Cityscapes 上 83.4% mIoU 同样领先。可视化结果中,PolaFormer++ 对行人、人行道、车辆的边界刻画更精细,复杂街景下也能完整分割远处人群:


      图 8:Cityscapes 语义分割可视化对比。红框标出关键差异区域,PolaFormer++(第一行)边界更细、目标区域更完整。


      4.5 扩散生成(PolaDiT)


      将 DiT 中的注意力替换为 PolaFormer++ 得到 PolaDiT,在 ImageNet-1K 类条件生成上:


      • DiT-S/2 规模:PolaDiT 取得最优 FID 58.30 / IS 25.29,优于 DiG(62.06)、MHLA(59.80)等近期线性注意力变体;
      • DiT-B/2 规模:FID 36.00 / IS 41.56,继续保持线性注意力变体中的最优成绩,展现出优秀的可扩展性。


      4.6 3D 新视角合成(Pola-GNT)


      在 GNT 框架下替换点积注意力,LLFF 数据集 8 个场景平均:PSNR 从 27.25 dB 提升至 27.34 dB,LPIPS 从 0.1019 降至 0.1009,SSIM 从 0.8868 提升至 0.8885,三项指标全面改善,证明其在 3D 空间推理与几何建模中的适用性。


      4.7 轻量级图像超分


      将 PolaFormer++ 接入 ESRT 与 MambaIRv2:


      • 精度:MambaIRv2-Pola++ 在全部 5 个 benchmark 上超过此前最优的 ESC(如 ×4 下 Set5 32.59 dB / Manga109 31.38 dB);
      • 效率:Urban100 ×3 下 ESRT-Pola++ 延迟仅 85.7ms,相比 ESRT(327.4ms)降低 73.8%,显存节省 90.0%;Manga109 上显存节省高达 92.5%。MambaIRv2 换成 Softmax 版本后在 Urban100、Manga109 等大分辨率测试集上直接 OOM,在小分辨率集上延迟也飙升近 20 倍(Set5:1504.8ms vs 78.6ms),而 Pola++ 版本几乎不增加开销。

      图 9:×4 超分可视化对比。PolaFormer++ 在建筑边缘、斑马条纹、金属网格等高频细节上重建出更锐利、更连续的纹理。


      4.8 消融与稳定性


      • 逐模块消融(ImageNet-1K,b0):极性机制基线 76.3%,加入门控矩阵 Gs/Go、可学习尖锐化、PaCS 等组件后逐步提升,完整模型达 76.9%
      • 与线性注意力同台对比(约 29M 参数、4.5G 量级 FLOPs 的统一设置):PolaFormer++ 以 82.8% 超过 Efficient Attn(81.0)、FLatten Attn(82.1)、Agent Attn(82.6)、InLine Attn(82.4)、LinDiff(82.3)等所有基线,且保持线性复杂度;
      • 训练稳定性:BF16 混合精度、全局 batch 2048/4096/8192 三档设置下训练 300 epoch,loss 与梯度范数全程稳定,无爆炸、无 NaN、无发散。

      图 10:不同全局 batch size 下 PolaFormer++-b0 的 BF16 训练稳定性,loss(实线)与梯度范数(虚线)均平稳收敛。


      五、结语


      PolaFormer++ 从两个根本性问题出发 ——负值信息丢失注意力分布不够尖锐—— 给出了统一且有理论支撑的解法:


      1. 极性感知:极性分解 + 双流交互 + 可学习门控矩阵,让 q、k 中的每个元素都平等参与注意力计算,完整重建相似度信息;
      2. 通道级尖锐度:PaCS 特征映射用两组可学习温度向量,为每个通道、每条极性流自适应调节尖锐程度;
      3. 通用理论判据:基于 Schur - 凹性首次给出判定特征映射是否具备降熵性质的严格标准,将 PolaFormer 的特例证明推广为通用框架。


      从分类、检测、分割,到超分、扩散生成、3D 合成,PolaFormer++ 在六大类任务上实现了精度与效率的更优平衡,为视觉 Transformer 中的线性注意力建立了坚实的理论与实践基础。代码已开源,欢迎试用与拓展!


      图片


      © THE END 

      转载请联系本公众号获得授权

      投稿或寻求报道:liyazhou@jiqizhixin.com

      转载声明:本文转载自原发布平台 (作者:机器之心), 原文标题《TPAMI26 | 全面升级PolaFormer++:极性感知计算流 +通道级尖峰,实现线性注意力的新天花板》, 查看原文。 版权归原作者及原发布平台所有,本站仅作收录与展示,未对正文内容作实质性修改; 若涉及侵权请联系本站处理。