GSPN-2: Efficient Parallel Sequence Modeling

We thank the reviewer for the thorough and insightful feedback. Below we respond to each concern and clarify the current scope, design choices, and planned extensions.

Machine Learning Tasks. Our experiments in the main paper focus on image classification (ImageNet-1K) and text-to-image generation (Stable Diffusion XL) to evaluate both discriminative and generative capabilities. Besides, as GSPN-2 is designed for high-resolution spatial modeling, we have extended experiments to dense prediction tasks across various benchmarks in Appx E. We also provide preliminary integration of GSPN-2 into SigLIP2 in the response below to assess its capacity to scale with large and downstream generalization.

Fig. 1 Should be Expanded to the Downstream Tasks. Yes, we did – the corresponding comparison has been fully explored in the paper: in Tab.2 of the main paper, we show an extensive comparison with both VIT (all with Flash Attn) based (orange block), and mamba-based (green block) architectures. In Fig. 1 of Appx B, we provided a comprehensive analysis of accuracy, model size, and throughput for GSPN-2 compared to SoTA architectures (e.g., ConvNet, Transformer-based, Mamba-based). These results highlight the competitiveness of GSPN-2—beyond speed gains alone—and establish it as a general-purpose architecture.

Detection and Segmentation. We did provide both in Appx E: we have extended GSPN-2 with high-resolution inputs (e.g., 1024x1024 resolution) to evaluate on segmentation, depth estimation, and 3D probing benchmarks such as ADE20K, VOC2012, Probe3D, NYUDv2 and PASCAL in Appx E. Preliminary results at Tab. 2 indicate comparable or superior performance to transformer-based backbones with much better inference latency. We will move this experiment into the main paper in the final version.

SigLip. Thank you for the valuable suggestion. We are currently integrating GSPN-2 into vision encoder backbones (e.g., SigLIP2-Base) using a subsample of 1 million samples from the DataComp1B dataset for rapid evaluation, with testing on ImageNet Zero-shot. The table shows a significant improvement in inference speed, e.g., 180 times faster compared with flash-attention and 5 times faster, including overheads like MLP.

IN-1K is Relatively Old. We agree that ImageNet-1K is limited in scale. However, it remains a universally adopted and effective benchmark in attention-based literature [15-29], especially for comparing multiple model sizes (e.g., tiny, base, large) within practical GPU resource constraints. For instance, state-of-the-art methods like SigLIP and Radio require substantial GPU resources (multiple weeks on 16×8 GPUs), making them impractical for extensive variant analysis. Additionally, beyond IN-1K, we have successfully validated GSPN-2’s scalability to higher-resolution tasks (e.g., 1024×1024) in segmentation (ADE20K, VOC2012, NYUDv2, PASCAL) and text-to-image generation (SDXL) in Appx E, confirming its effectiveness across diverse and challenging dense-prediction benchmarks.

Only small models in Tab. 2. To enforce fair comparison, we follow the common setting used in previous relevant works (e.g., VMamba, VisionMamba), where models are evaluated from tiny to large size on a 224×224 resolution. This choice aligns with the standard benchmarking practices in the field for efficient models. However, we acknowledge the importance of validating performance on high-resolution inputs, which are more practical for real-world applications. To address this, we have extended our experiments to both high-resolution text-to-image generation and dense prediction tasks using 1024×1024 resolution, as presented in Appx E. These experiments better reflect the capabilities of GSPN-2 in handling high-resolution visual data

Ablation on $C_{procy}$. We selected $C_{proxy}=2$ as a minimal proxy to demonstrate how aggressively the input can be compressed while still preserving accuracy. We provide an ablation on $C_{procy} \in \{2, 4, 8, 16, 32\}$ in the table below, analyzing the trade-off between accuracy and throughput. To make a fair comparison, we maintain almost the same model size for different $C_{procy}$ by increasing the number of blocks if necessary. We found this value yielded.

As shown, using $C_{proxy} = 2$ achieves the best latency–accuracy trade-off, with no measurable loss in accuracy and the highest throughput. This suggests that the proxy dimension serves primarily as a low-rank bottleneck for efficient propagation and does not need to match the full input dimensionality to be effective. We believe that the global context captured by the spatial recurrence appears to be preserved even through such a compressed intermediate representation.

Quantitative metrics for T2I generation. We have included CLIP-T scores and FID metrics in Tab 1 in Appx D, along with the quality-runtime trade-off at Fig. 3.

Choice of 128-channel setting. We chose 128 channels to match prior system benchmarks (e.g., VMamba). However, we agree that large channel counts would be more realistic, so we update Table 1 with results for $C \in \{192, 768,1152\}$ with appropriate resolutions below.

Backward Benchmark. We agree, and will highlight backward pass performance more explicitly in the main text. Note that the backward function is very similar to the forward function in terms of kernel launch design, and thus the speedup of backward is similar to the forward behaviour. We will introduce it in the revised paper.

Run on Different Accelerators. Yes, our CUDA acceleration design is fully compatible with all CUDA-enabled GPUs, relying only on standard libraries such as ATen and THC. Moreover, our algorithmic components—line-scan recurrence and channel-compressed formulations—can be readily ported to other accelerator backends (e.g., Triton or TileLang). The core idea of kernel fusion can also be adapted to accelerators that provide similar memory hierarchies. However, to achieve optimal performance, hardware-level control over thread blocks and shared memory is beneficial.

Dear Reviewer Da5V,

We sincerely appreciate your valuable review, which has greatly helped us clarify the scope, design choices, and broader applicability of our work.

Following your valuable suggestions, we have revised our paper and submitted a detailed rebuttal. In summary, we have:

• Highlighted experiments beyond ImageNet-1K already included in Appx E like segmentation, depth estimation, and 3D probing benchmarks (ADE20K, VOC2012, Probe3D, NYUDv2, PASCAL) with high-resolution inputs (e.g., 1024×1024), showing competitive or superior performance with much better inference latency.

• Integrated GSPN-2 into SigLIP2-Base, demonstrating up to 180× module-level and 5× block-level speedups over Flash Attention, while maintaining comparable zero-shot accuracy.

• Highlighted comprehensive comparisons with both Transformer- and Mamba-based SoTA architectures, highlighting GSPN-2’s competitiveness in accuracy, model size, and throughput already included in Fig. 1 in Appx B.

• Added a detailed ablation study on the compressive proxy dimension ($C_{\text{proxy}} \in {2,4,8,16,32}$), showing that $C_{\text{proxy}}=2$ achieves the best latency–accuracy trade-off with no loss in accuracy.

• Highlighted quantitative metrics (CLIP-T, FID) for text-to-image generation tasks already included in Appx D.

• Updated throughput benchmarks for larger channel counts ($C \in {192, 768, 1152}$) and clarified backward-pass speedups.

• Discussed portability to other accelerators (e.g., Triton, TileLang) and broader algorithmic generality.

We hope these updates have addressed your concerns. As the discussion period draws to a close, we would greatly appreciate it if you could let us know whether our response has resolved your questions.

Thank you again for your time and consideration.

Sincerely,

Authors of Paper 943

We thank the reviewer for the constructive feedback. We respect your concerns and have taken them seriously. In the following, we provide clarifications and additional context to better address the issues raised and to more clearly communicate the value of our work.

Coverage of Flash Attention Baselines. Flash Attention is already present in every ViT-based baseline we report: all orange blocks in Table 1 and Fig. 1 (Appx B), the SDXL row in Fig. 5 (main paper), and the ViT-L RADIO model in Table 2 (Appx E) are Flash Attention implementations. In the strongest head-to-head test—SDXL vs. GSPN-2 SDXL—we replace each soft-max/Flash layer with a single GSPN layer and fine-tune by distillation: image quality is maintained while inference latency drops by $\approx 92\times$. These results show that our gains are measured against the latest optimized attention baselines.

Comparison with other SOTA algorithm-system codesign work. In addition, we have compared with other algorithm–system co-designs such as VMamba and MambaVision (Fig 1 in Appx B), to highlight differences in computation patterns and practical GPU efficiency. Additionally, Appx B provides a detailed trade-off analysis among accuracy, model size, and throughput for the entire system. We will further clarify them in the revised paper.

Limited Algorithmic Contribution, a System Implementation. We address the question from the following aspect:

• The importance of CUDA optimization. CUDA optimizations have been widely explored for softmax attention (e.g., FlashAttention, TensorCore-based solutions). However, efficient CUDA support for recurrent or linear-propagation methods—despite their promising scalability for unlimited-length contexts (e.g., SSM, Mamba, Linear Attention, and GSPN)—is severely lacking. Our work explicitly targets this important gap.

• Just a System Implementation? No, in contrast, we found that CUDA-level optimizations alone were insufficient. Even aggressive kernel tuning left us with a runtime that still grew linearly with channel and batch dimensions—a fundamental algorithmic bottleneck. To solve this, we introduced a novel GSPN micro-design (Sec. 3.3) that significantly reduces this dependency without sacrificing accuracy, clearly an intrinsic algorithmic advancement.

• Algorithm-system Co-design: Unlike typical algorithm-only papers, we tightly fused this new micro-design into a single optimized CUDA kernel—exactly demonstrating an algorithm-system co-design. This integrated solution unlocks the full potential of GSPN in practical scenarios, underscoring that our contribution is both algorithmically substantial and systemically necessary. We appreciate your feedback and will revise the paper to better highlight these algorithmic contributions and explicitly clarify their integration with our system-level optimizations.

Presentation Clarity. Thanks for pointing out the issue. We will clarify the connection between each bottleneck and its corresponding solution, and restructure Section 3 for better flow in the revised paper.

How Channel-shared Weight Mitigates the Concurrency Bottleneck? They mitigate the concurrency bottlenecks by reducing the number of independent CUDA blocks needed during kernel launch. In GSPN-1, per-channel weights ($W_i$,$c$) required independent computation for each (N, C) slice, quickly saturating the GPU's grid dimension and memory queues. In contrast, channel-sharing allows one weight matrix per row/column ($W_i$), enabling:

• Shared memory reuse: once W_i is loaded into shared/L1 cache, it’s reused across all C channels within the block, improving cache hit rates and lowering register pressure. See +2D thread block ($1.4\times$), +shared memory cache ($1.5\times$), and +channel-shared gates ($1.3\times$) at Fig.3.
• Fewer parameter fetches per step, accelerating global memory throughput (confirmed by Nsight in Table 1).

This significantly improves parallel efficiency on modern GPUs—especially when combined with 2D blocks.

Relate Compressive Proxy Dimension to Low-rank Approximation. Yes—conceptually, the proxy dimension ($C_{proxy}$) is a low-rank projection of the original $C$-dimensional input. Specifically, we project $x\in R^{N\times C\times H\times W}$ to $x_{proxy} \in R^{N\times C_{proxy}\times H\times W}$ via a lightweight linear layer (Sec. 3.2), apply the GSPN recurrence in this reduced space, and then project back to $C$ channels. This reduces CUDA workload from $N\times C\times H$ slices to $N\times C_{proxy}\times H$, mitigating the concurrency bottleneck under high $C$ or $N$. We empirically found no measurable loss in accuracy across classification and generation tasks. For example, GSPN-2-S (with compression ratio 2:1) improves accuracy over GSPN-S while using less compute (9.2G vs. 9.0G MACs). We also provide an ablation on $C_{procy} \in \{2, 4, 8, 16, 64\}$ in the below table, analyzing the trade-off between accuracy and throughput.

This suggests that the proxy dimension serves primarily as a low-rank bottleneck for efficient propagation and does not need to match the full input dimensionality to be effective. We believe that the global context captured by the spatial recurrence appears to be preserved even through such a compressed intermediate representation.

We sincerely thank the reviewer for the encouraging feedback and thoughtful comments. We truly appreciate your recognition of the contributions and are glad that the strengths of our work came through. Below, we address your questions and suggestions in more detail.

Code Availability. Yes, we will release it upon acceptance. In the meantime, we ensure reproducibility by (1) providing extensive implementation details of our CUDA kernel design (Section 3.1–3.4). (2) Detailing block/grid configuration, memory layout, shared memory strategy, and concurrency management (Figure 2, Figure 3, and Table 1). (3) Describing all architectural modifications, including the proxy compression dimension and channel-shared weights (Section 3.2–3.3). We believe these elements provide a clear path to reproducing the kernel and model behaviors even before code release.

Significance outside GSPN. First, improving GSPN is non-trivial. GSPN [1], because it is already a top-performing attention architecture that breaks the $O(N^2)$ softmax wall with $O(\sqrt{N})$ propagation, yet its open-source code runs at < 8 % GPU utilisation. Our work directly addresses this bottleneck, making GSPN fast enough to be a practical alternative to softmax attention. Similar to recent studies on Mamba, we believe improving GSPN’s efficiency is both necessary and impactful. Moreover, the compressed-channel architecture can reduce the channel-linear cost inside any 2D recurrence. Because the design only assumes a line-scan recurrence and a shared (scalar) affinity map, it can be plugged into other propagation-based models—including SSM, Mamba/MambaVision, 2D RNNs, and LSTMs—to cut memory-bandwidth pressure without changing their update equations. Likewise, the single-kernel fusion pattern —or its Triton/TileLang analogue—can accelerate those models on all CUDA-class GPUs. In short, our contributions generalise: the compressed-channel idea is an algorithmic building block for any linear propagation method, and the fused-kernel blueprint shows how to reach hardware limits, making the work broadly relevant rather than GSPN-specific.

Algorithmic Contribution. Our paper does not merely optimize implementation. We identify that GSPN’s main efficiency bottleneck is the large channel dimension (often thousands) typically used in vision models. To fundamentally address this, we propose a novel channel sharing and compression approach (Sec. 3.3) that significantly reduces memory and computational load—this is a novel micro design, not merely an engineering optimization. Additionally, we integrate this compressed-channel structure directly into a new, highly efficient CUDA kernel. Thus, our method tightly couples algorithmic and system-level improvements to unlock GSPN’s full practical potential.

Comparison with the SOTAs other than Speed. Our experiments already show that GSPN-2 matches or outperforms the strongest non-GSPN backbones in accuracy and parameter count, not just latency. On ImageNet-1K, GSPN-2-T delivers 83.0 % top-1 with 4.2G MACs, beating LocalVMamba-T (82.7 %, 5.7G MACs) and other light ViTs at the same scale. For larger models, GSPN-2-B reaches 84.7 % with fewer MACs than ViT-B and VMamba-B. Beyond classification, we directly replace Flash-Attention with our layer in SDXL and obtain matching text-to-image quality (same FID and CLIP-score) while preserving the 92 × speed-up reported in Fig. 5. On dense prediction, our Appx E shows GSPN-2 backbones achieve 47.9 mIoU on ADE20K and 82.5 mIoU on VOC2012, on par with or better than transformer baselines yet at much lower inference cost. These results demonstrate that the proposed channel-compressed design is competitive with state-of-the-art models in accuracy, compute, and memory, so its value is not confined to speeding up GSPN—it offers a strong alternative backbone in its own right.

Accuracy vs. Speed Trade-off. Thank you for this suggestion. We agree that clearly showing the trade-off between accuracy and speed is important. Besides Fig. 3, we also show a more comprehensive analysis of accuracy, model size, and throughput in Fig. 1 (Appx B). To improve clarity, we will move key results from Appendices B and D into the main paper. We note that in Fig. 1, the trade-off isn’t obvious, as most of our models simultaneously achieve top results in both accuracy and throughput.

Light-weighted Models. While Fig. 1 (Appx B) focuses on isolating our core design improvements over various baselines, we have also compared GSPN-2 against several SOTA lightweight backbones in both accuracy and speed in the table below:

Despite having a similar compute budget, GSPN-2-T outperforms all listed lightweight models by a wide margin in accuracy—by +4.4% over MobileViT-S and +5.1% over ConvNeXt-XT—while maintaining a comparable number of MACs. Furthermore, GSPN-2-T achieves much higher throughput than transformer and CNN-based models of similar scale, or even lighter scale. This demonstrates that GSPN-2 offers an excellent trade-off between model latency and performance.

We appreciate the reviewer’s positive evaluation and insightful observations. Your detailed technical feedback is highly valuable, and we have carefully addressed each point below to further clarify and strengthen our contributions.

Justification for Channel-Agnostic Propagation. As we have briefly stated at line-193, Softmax attention already uses one scalar affinity per token pair, shared across all channels; the extra $C$ per-channel weight matrices in the original GSPN likely duplicate that information. By replacing them with a single shared $W$ plus per-channel gains $\Lambda$, we (i) align GSPN’s affinity operator with standard attention, (ii) drop the hidden channel-dependent cost, and (iii) keep channel diversity through $\Lambda$. Formally, Eq. 5 shows the full $N\times N$ spatial kernel is preserved; sharing $W$ just factorises the original block matrix into $W\otimes I(C)$, a rank-1 compression that removes redundant parameters while retaining dense global interactions. This is the same low-rank idea behind attention and grouped convolutions, where weight sharing regularises the model and cuts over-fitting risk. We will add this analysis to the revised paper.

SM Utilization for Small Workloads. This trade-off is a known limitation of large-kernel fused CUDA designs under light parallel workloads. We are actively exploring fine-grained tiling, micro-batching, and block fusion across directions as potential directions to address this in future work. Despite this limitation, GSPN-2 still outperforms GSPN-1 and state-space alternatives in both accuracy and runtime, even under small-scale settings, as shown in Fig. 4 and Tab. 1.

Font size. We will increase the font size and improve the figure in the revised paper.

Generalization to Other Architectures. We view GSPN-2 as a blueprint for applying GPU-aware design to 2D structured attention mechanisms, and anticipate future extensions to SSM, 2D RNN, LSTM, linear attention, RWKV variants, etc. The core optimization strategies in GSPN-2 are highly generalizable to any operation with 2D directional recurrence (e.g., causal convolutions, 2D scan layers / structured RNNs). For channel-sharing and proxy compression, these are modular strategies that can be transplanted to transformer blocks, gated MLPs, and SSMs like Mamba for reducing memory and increasing kernel fusion efficiency. Memory optimization tactics (e.g., shared memory reuse, coalesced access) can also benefit any scan-like operation where partial results are reused over spatial dimensions.

Impact on Downstream Tasks (see Appx E). We agree that evaluating downstream dense prediction tasks is important. In Appx E, we already extend GSPN-2 to evaluate on segmentation, depth estimation, and 3D probing benchmarks such as ADE20K, VOC2012, Probe3D, NYUDv2, and PASCAL, as shown in Tab.2. These experiments confirm that GSPN-2 maintains strong dense prediction performance—comparable or superior to transformer-based backbones—while significantly reducing attention latency, especially for high-resolution inputs (e.g., 1024x1024). We will move these segmentation results into the main paper to clearly highlight GSPN-2’s advantages on dense vision tasks.