通过广义空间传播网络实现平行序列建模

我们提出了广义空间传播网络（GSPN），这是一种针对视觉任务优化的新型注意力机制，能够本质上捕获 2D 空间结构。现有注意力模型，包括 transformer、线性注意力和类似 Mamba 的状态空间模型，都将多维数据处理为 1D 序列，导致空间一致性和效率受损。GSPN 通过直接作用于具有空间一致性的图像数据，并采用线扫描方法形成稠密的成对连接来克服这些限制。GSPN 的核心是稳定-上下文条件（Stability-Context Condition），它确保 2D 序列中稳定、上下文感知的传播，并将平方地图上 N 个元素的有效序列长度为缩减至 $\sqrt{N}$，显著提升计算效率。凭借可学习的、输入依赖的权重且不依赖位置嵌入，GSPN 在保持空间保真性方面取得优势，并在视觉任务中实现了最先进的性能，包括 ImageNet 分类、受类指导的图像生成和文本到图像生成。值得注意的是，GSPN 在生成 16K 图像时，将 SD-XL 使用 softmax 注意力加速超过 $84\times$。