面向扩散转换器的谱状态空间方法：用于扩散 Transformer 中的空间混合
机器学习
2026-03-17 v1
摘要
视觉转换器（尤其是用于生成模型）的成功受到自注意力二次计算成本和弱空间归纳偏置的限制。我们提出了 PDE-SSM，一种用于空间混合的状态空间模块，以可学习的扩散-对流-反应偏微分方程取代注意力机制。该算子通过建模基于物理的动态信息流，编码强大的空间先验，取代所有到所有 token 交互。在傅里叶域中求解 PDE 可实现全局耦合，时间复杂度近似线性为 ，为注意力提供了一个原则且可扩展的替代方案。我们将 PDE-SSM 集成到基于流匹配的生成模型中，得到基于 PDE 的扩散转换器 PDE-SSM-DiT。实验结果表明，PDE-SSM-DiT 在性能上与最先进的扩散转换器相当，同时大幅降低计算开销。我们的結果表明，正如 1D 环境中 SSM 取代注意力一样，多维 PDE 算子为下一代视觉模型提供了高效且富有归纳偏置的基础。
引用
@article{arxiv.2603.13663,
title = {PDE-SSM: A Spectral State Space Approach to Spatial Mixing in Diffusion Transformers},
author = {Eshed Gal and Moshe Eliasof and Siddharth Rout and Eldad Haber},
journal= {arXiv preprint arXiv:2603.13663},
year = {2026}
}