√N并行+84倍计算加速！英伟达港大全新图像注意力：空间结构都保留

GSPN是一种新型视觉注意力机制，通过线性扫描和稳定性-上下文条件，高效处理图像空间结构，显著降低计算复杂度。通过线性扫描方法建立像素间的密集连接，并利用稳定性-上下文条件确保稳定的长距离上下文传播，将计算复杂度显著降低至√N量级。

Transformer 及其核心的注意力机制在自然语言处理和计算机视觉等领域带来了革命性进展，展现出强大的深度上下文建模和数据间复杂依赖关系捕捉能力。

然而，其在处理视觉数据时面临两大核心挑战：

1. 二次计算复杂度使其难以高效处理高分辨率图像等长上下文数据；
2. 忽略空间结构，将多维图像视为无结构的一维标记序列，破坏了图像固有的空间连贯性，而这种信息对于依赖空间关系的视觉任务至关重要。

为克服效率瓶颈，近期研究如线性注意力和状态空间模型（如 Mamba） 致力于将复杂度降低至线性。

然而，这些方法在提升效率的同时，依然未能有效保留和利用图像的关键二维空间结构信息，本质上仍是序列化处理。

尝试将一维光栅扫描（raster scan）扩展至二维的线扫描方法（line scan）是增强空间连贯性的一种思路。

但二维线性传播面临严峻挑战：标量权重变为连接像素与前序邻居的矩阵权重。在传播过程中累积的矩阵乘法极易导致稳定性问题——矩阵特征值过大引发指数增长（不稳定），过小则导致信号迅速衰减（信息消失）。

因此，在二维空间中同时实现稳定性和维持长距离上下文的有效传播，是一个亟待解决的难题。

针对上述挑战，来自英伟达、香港大学和UCSD的研究人员提出广义空间传播网络（GSPN），一种专为视觉任务优化的新型注意力机制，其核心优势在于直接操作空间连贯的图像数据，通过高效的线扫描方法建立密集的像素间连接。

论文地址：https://arxiv.org/abs/2501.12381

项目主页：https://whj363636.github.io/GSPN/

代码：https://github.com/NVlabs/GSPN

GSPN成功的关键在于其提出的稳定性-上下文条件（Stability-Context Condition），该条件确保了跨二维序列的稳定长上下文传播，并将具有N个元素的图像的复杂度显著降低至√N 量级。

因此，GSPN能够在保持卓越空间保真度的同时，实现极高的计算效率，并在ImageNet分类、类引导图像生成及文本到图像生成等任务中达到先进性能。例如，在生成16K图像时，GSPN相比基于softmax注意力的SD-XL加速超过84倍。

论文第一作者为王弘焌，香港大学统计系博士三年级学生，目前为NVIDIA research intern，研究方向包括高效基础模型、开放世界理解。

GSPN方法

二维线性传播

二维线性传播通过逐行或逐列的顺序处理进行。对于二维图像，其遵循线性循环过程，隐藏层通过前一行的隐藏状态和当前输入计算得出。

将隐藏状态和输入的行向量连接成序列后，可表示为输入与一个下三角矩阵的乘积，输出则为输入的加权和，该公式可类比为带因果掩码的非归一化线性注意力机制，其中额外的传播矩阵调制注意力强度。

稳定性-上下文条件

在传播过程中上述累积的矩阵乘法极易导致稳定性问题。

为实现稳定且有效的长距离传播，研究人员引入定理1和定理2（统称为稳定性-上下文条件）。

定理1指出，若所有矩阵均为行随机矩阵，则满足各元素加权和为1

定理2表明，行随机矩阵可确保传播过程的稳定性。行随机矩阵的定义为元素非负且每行元素之和为1，乘积仍为行随机矩阵，这为稳定传播提供了数学基础。

传播层的关键实现

对于二维线性循环过程，研究人员对前序状态的三邻居连接来计算当前时刻的隐藏层（每个像素连接前一行的三个相邻像素）以提高参数效率。

文中同时还提出GSPN的两种变种，全局GSPN和局部GSPN：

全局GSPN捕捉整个序列的长距离依赖，局部GSPN通过将空间维度划分为非重叠组来限制传播序列长度，提高效率。

最后，通过四方向集成确保全像素连接，形成密集成对连接。

对每个传播方向的矩阵元素应用 sigmoid 函数并归一化，以保证行随机性。

通过定制的CUDA内核实现线性传播层，采用并行化结构，在批量、通道和与传播方向正交的行/列上实现全并行化，有效减少内核循环长度，实现高效可扩展的线性传播。

GSPN架构

GSPN是一个通用序列传播模块，可无缝集成到各种视觉任务的神经网络中。针对判别任务和生成任务设计了不同的GSPN块，均基于核心GSPN模块构建：

• GSPN模块：通过共享1×1卷积进行降维，再通过三个独立的1×1卷积生成依赖于输入的参数，用于二维线性传播，这些投影和传播封装在模块化的GSPN单元中。

• 图像分类架构：采用Swin-Transformer的四级分层架构，通过堆叠设计良好的GSPN块，在相邻层级间进行下采样操作，平衡计算效率和表示能力。

• 类条件图像生成架构：重新设计生成架构，通过向量嵌入加法集成时间步和条件信息，包含跳跃连接和线性投影，去除位置嵌入并引入FFN进行通道混合。

• 文本到图像生成架构：将GSPN模块直接集成到Stable Diffusion架构中，替换所有自注意力层，利用预训练权重初始化参数，加速训练。

实验结果

图像分类

在ImageNet-1K分类任务中，GSPN在参数数量相当的情况下优于现有序列模型，GSPN在从小型到基础配置的模型规模上表现出一致的性能提升，证明了其可扩展性。

类条件图像生成

与多种基线方法相比，GSPN-XL/2在ImageNet 256×256类条件生成任务中建立了新的最先进性能，GSPN-L/2仅使用先前模型65.6%的参数就获得了更优的FID和IS分数，GSPN-B/2在收敛时仅使用DiT-XL/2 20.3%的参数就实现了有竞争力的性能，验证了GSPN的效率和可扩展性。

文本到图像生成

GSPN由于其归一化权重满足稳定性-上下文条件，无需额外归一化即可适应任意分辨率，在不使用任何预训练权重且在相同训练轮数内达到了与baseline相当的性能。

此外，GSPN在单块A100 GPU上生成16K×8K分辨率图像可实现约84倍的加速。

总结

研究人员提出了广义空间传播网络（GSPN），这是一种用于视觉任务中并行序列建模的新型注意力机制。

通过稳定性-上下文条件确保稳定且上下文感知的传播，GSPN在保持效率的同时将序列复杂度减少到√N

实验表明，GSPN在多个视觉任务中实现了最先进的结果和显著的加速，展示了其在视觉任务中的效率和潜力。

未来，GSPN有望在更多视觉领域及视觉多模态模型中发挥重要作用，推动下一代视觉理解和生成基础结构的发展。

参考资料：

https://arxiv.org/abs/2501.12381

文章来自于“新智元”，作者“LRST”。