Parallel Sequence Modeling via Generalized Spatial Propagation Network

TL;DR

GSPN introduces a spatially coherent, sub-quadratic attention mechanism for vision by performing 2D linear propagation with a Stability-Context Condition that guarantees stable, dense long-range interactions. By line-scanning across rows and columns and using learnable, input-dependent weights, GSPN reduces effective sequence length to $\sqrt{N}$ and eliminates positional embeddings, enabling efficient high-resolution processing. The framework supports global and local variants, a learnable directional merge, and achieves competitive image classification alongside state-of-the-art results in class-conditional generation and text-to-image synthesis, with dramatic speedups (e.g., up to $\sim84\times$) on Diffusion models for very large outputs. Collectively, GSPN provides a robust alternative to transformers for 2D vision tasks, balancing spatial fidelity, stability, and scalability across discriminative and generative settings.

Abstract

We present the Generalized Spatial Propagation Network (GSPN), a new attention mechanism optimized for vision tasks that inherently captures 2D spatial structures. Existing attention models, including transformers, linear attention, and state-space models like Mamba, process multi-dimensional data as 1D sequences, compromising spatial coherence and efficiency. GSPN overcomes these limitations by directly operating on spatially coherent image data and forming dense pairwise connections through a line-scan approach. Central to GSPN is the Stability-Context Condition, which ensures stable, context-aware propagation across 2D sequences and reduces the effective sequence length to $\sqrt{N}$ for a square map with N elements, significantly enhancing computational efficiency. With learnable, input-dependent weights and no reliance on positional embeddings, GSPN achieves superior spatial fidelity and state-of-the-art performance in vision tasks, including ImageNet classification, class-guided image generation, and text-to-image generation. Notably, GSPN accelerates SD-XL with softmax-attention by over $84\times$ when generating 16K images.

Figures (11)

• Figure 1: Comparison of attention mechanisms and computational efficiency.Top: Softmax attention (left), mamba (middle), and our GSPN (right). Bottom: Inference speed comparison across different input sizes on A100 GPU, demonstrating GSPN's superior efficiency for high-resolution input. Dashed lines indicate quadratic extrapolation due to memory explosion. Global and local GSPN are detailed in \ref{['para:glob_locl']}.
• Figure 2: 2D Linear Propagation. For the $i$-th row, each item in the hidden state $h_i$ is computed by: (1) a weighted sum of three neighboring values from the hidden layer $h_{i-1}$, where weights form a normalized tridiagonal matrix $w_i$, and (2) the element-wise product of the current input $x_i$ with $\lambda$. $w_i$ and $\lambda$ are both learnable and input-dependent parameters. The weights in $w_i$ are obtained by applying sigmoid activation followed by row-wise normalization.
• Figure 3: Architecture of GSPN. We design different GSPN blocks for discriminative and generative task. Both blocks share the same GSPN module with our core 2D Linear Propagation mechanism illustrated in \ref{['fig:2dprog']}.
• Figure 4: Design choice of GSPN with benchmark methods in classification accuracy and computational efficiency for both classification and generation tasks. A higher Top-1 accuracy and a lower FID indicate better model capability. Our GSPN can outperform other methods under similar FLOPs.
• Figure 5: Qualitative text-to-image results generated from our SD-XL-GSPN model. We enable generation up to 16K resolution on a single A100 GPU while reducing inference time by up to $84\times$ on the SD-XL model.

Theorems & Definitions (4)

• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4