arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2603.13663cs.LG

PDE-SSM:一种用于扩散变换器中空间混合的谱状态空间方法

PDE-SSM: A Spectral State Space Approach to Spatial Mixing in Diffusion Transformers

University of British Columbia, Vancouver, BC Canada(不列颠哥伦比亚大学,加拿大温哥华,BC省);University of Cambridge, Cambridge, United Kingdom(剑桥大学,英国剑桥)

Eshed Gal, Moshe Eliasof, Eldad Haber

AI总结:

本文提出PDE-SSM,通过学习可扩散-反应偏微分方程替代自注意力机制,解决视觉变换器在生成建模中的二次成本和弱空间归纳偏置问题,实现高效且可扩展的替代方案。

摘要:

视觉变换器的成功尤其在生成建模中受到二次成本和自注意力弱空间归纳偏置的限制。我们提出PDE-SSM,一种空间状态空间块,用可学习的对流-扩散-反应偏微分方程替代注意力。该操作符通过物理基础动态建模信息流,而非所有到所有token交互,编码强空间先验。在傅里叶域求解PDE获得全局耦合,复杂度为O(N log N),提供原理上和可扩展的注意力替代方案。我们将PDE-SSM整合到流匹配生成模型中,获得基于PDE的扩散变换器PDE-SSM-DiT。实证结果表明,PDE-SSM-DiT在性能上匹配或超过最先进的扩散变换器,同时大幅减少计算量。我们的结果表明,类似于一维情况下SSM取代注意力,多维PDE运算提供高效且具有丰富归纳偏置的基础,为下一代视觉模型奠定基础。

英文摘要:

The success of vision transformers-especially for generative modeling-is limited by the quadratic cost and weak spatial inductive bias of self-attention. We propose PDE-SSM, a spatial state-space block that replaces attention with a learnable convection-diffusion-reaction partial differential equation. This operator encodes a strong spatial prior by modeling information flow via physically grounded dynamics rather than all-to-all token interactions. Solving the PDE in the Fourier domain yields global coupling with near-linear complexity of $O(N \log N)$, delivering a principled and scalable alternative to attention. We integrate PDE-SSM into a flow-matching generative model to obtain the PDE-based Diffusion Transformer PDE-SSM-DiT. Empirically, PDE-SSM-DiT matches or exceeds the performance of state-of-the-art Diffusion Transformers while substantially reducing compute. Our results show that, analogous to 1D settings where SSMs supplant attention, multi-dimensional PDE operators provide an efficient, inductive-bias-rich foundation for next-generation vision models.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑