A 2-Dimensional State Space Layer for Spatial Inductive Bias

Thank you for the comprehensive review and the constructive feedback.

As mentioned by the authors, one major limitation of the method is its high training time cost. It can double the training time compared with the baseline, limiting its application to the training of large models on large benchmarks.

We refer you to the main response, section 2 “Latency and Memory Analysis”. In summary, first, we optimized our code, which resulted in a 30% decrease in inference time and a 25% reduction in training time. Second, for the ConvNeXt variants, our training time is only 1.26 times longer than the standard CUDA-dedicated and highly-optimized implementation of the Conv2D-based ConvNeXt (note that our code is not optimized for CUDA). Lastly, we wish to highlight that, from a theoretical perspective, our layer is quite efficient, as analyzed in Section 3.3 'Computation and Complexity,' and empirically in Appendix H.

The experiments limit to image classification problems and 2-D inductive bias can be very important for dense prediction. It would be better to also evaluate the proposed method for dense prediction problems like segmentation, depth estimation etc. In response to your suggestion regarding the evaluation of our method for dense prediction tasks such as segmentation and depth estimation, we acknowledge the significance of such assessments. However, it is important to note that our current resource constraints limit our ability to conduct extensive experiments in these areas, particularly due to the typical requirements for pretraining on large datasets like ImageNet-22\1K or the utilization of Self-Supervised Learning (SSL) techniques.

Instead, we have chosen to demonstrate the effectiveness of our method by incorporating our layer into ConvNeXt (main response, section 1) and conducting additional experiments on multiple datasets. We believe that these experiments offer a different perspective on the efficacy of our approach and provide valuable insights.

The complex (C) variant of the proposed method may exhibit instability, obtaining very bad results, e.g. Table 1, Table 6. I would recommend the authors to include a detailed explanation of this situation and any potential ways of avoiding the instability.

In all our experiments, we have maintained consistency by using the same hyperparameters that the backbone was originally trained with, without engaging in hyperparameter tuning. In the instances where instability was observed, such as in Table 1 and Table 6, we found that decreasing the learning rate to a smaller value (1e-4 instead of 1e-3) helped stabilize the training process.

It is important to note that various architectures may indeed require different hyperparameters when applied to different datasets. However, we decided to retain the original hyperparameters in order to adhere to the framework's original setup without extensive hyperparameter adjustments.

I would suggest the authors to clearly explain each model in each table. For example, it would be much clear if the authors can explain what is 'ViT w/ MixFFN' and so on before describing the results. Also, are the results of the proposed method shown in Table 3 obtained by using SSM-r or SSM-c?

Thank you for your feedback. We will remove the reference to the MixFFN experiment from our tables, as it was included erroneously. We will review and ensure all models in each table are clearly explained. Additionally, for any further clarification needed on specific experiments or tables, your guidance would be invaluable. Regarding Table 3, the table showcases the SSM-r version.

Thank you for the comprehensive review and the constructive feedback.

“Weaknesses: Extended Training Time”

Please refer to 2. 'Latency and Memory Analysis” in the shared response for more details. In summary, first, we optimized our code, which resulted in a 30% decrease in inference time and a 25% reduction in training time. Second, for the ConvNeXt variants, our training time is only 1.26 times longer than the standard CUDA-dedicated and highly-optimized implementation of the Conv2D-based ConvNeXt (note that our code is not optimized for CUDA). Lastly, we wish to highlight that, from a theoretical perspective, our layer is quite efficient, as analyzed in Section 3.3 'Computation and Complexity,' and empirically in Appendix H.

"Weaknesses: Simplicity of Tasks: .. A more critical evaluation would involve assessing the performance of the baseline methods after large-scale pre-training ... can obtain better performance in downstream tasks after finetuning (etc., prompt tuning) ..”

Such experiments rely on large pre-trained models (trained on ImageNet-22\1K or through SSL schemes), and therefore they are very costly and are out of our budget’s reach. Please refer to 'Nature of our Benchmarks' in the main response for more details. We acknowledge the value of such experiments, however, we believe that our extensive evaluation across six datasets, with several backbones (ConvNeXt, ViT, Swin, Mega, DeiT), and a variety of model sizes—from DeiT/Mega-Base (85M params) and Swin-T (27M params) to ConvNeXt-Micro (9M params)—provides a sufficient assessment to justify the empirical contribution of our 2-D SSM layer.

“Limited to Transformer-based Architectures: .. It is uncertain how well the method would perform on convolutional networks (ConvNets) and how it compares to similar inductive bias methods, such as ConvNeXt. A comprehensive evaluation across different network architectures and inductive bias strategies is needed to fully understand ..”

Following the review, we have added results using the ConvNeXt backbone across four datasets and three model sizes. Please refer to the section '2-D SSM as an Alternative to Standard Convolution Layers (Experiments with ConvNeXt)' in the main response, and to Tables 6 and 7 in the revised manuscript for detailed information. As the results show, our 2-D SSM layer outperforms the Conv2D layers in these backbones for all tested experiments. We consider these results very promising and extend our gratitude to the reviewer for this constructive suggestion.

Thank you for the comprehensive review and the constructive feedback.

In the experiments, the proposed method is primarily benchmarked against SSM-based methods. However, various approaches employ convolution for embedding, for example in [1], and the use of convolutional layers or simple components for positional encoding is a common practice [2, 3, 4, 5]. It would be instructive to compare the proposed method with Vision Transformers that incorporate the methods to exploit spatial equivariance, which could serve as additional baselines. Furthermore, while it may not be a critical flaw, the implementation of the proposed method appears to be somewhat more complicated and less straightforward than simply employing convolutional layers.

Your suggestion to compare our method with Vision Transformers that utilize convolution for embedding and positional encoding is indeed valuable. We refer to the main response, part “2-D SSM as an alternative to standard Convolution layers (Experiments with ConvNeXt)”, where we outline the results of replacing the convolutional layer in ConvNext with our 2-D SSM layer. We believe this result showcases the effectiveness of our layer when compared to standard convolutional layers (which can be considered as representative cases of methods that exploit spatial equivariance).

Addressing the perceived complexity of our proposed method, we understand it may initially seem complex relative to conventional convolutional layers. Yet, this complexity is offset by notable advantages. Our method employs a global 2-D kernel, ensuring a consistent parameter count regardless of kernel size. Alongside this, our parameterization technique results in kernels formed within a sub-vector space, leading to reduced noise interference. Our approach leads to enhanced performance over ViT & Convnext backbones, across all datasets, underscoring the efficacy of our more complex design.

I believe that a larger model and dataset size would more effectively leverage strong spatial equivariance [6]. This implies that the proposed method may not be effective in environments with substantial data and model scales. Considering that training on datasets like IN1K with 'Base' or larger models has become a norm in the era of Vision Transformers, the applicability of the proposed method in such standard real-world scenarios could be limited. Additionally, the focus of the experiments on smaller datasets and models further suggests potential constraints in its utility for larger-scale tasks.

Please refer to 'Nature of our Benchmarks' in the main response for more details. We acknowledge the value of such experiments; however, we believe that our extensive evaluation across six datasets, with several backbones (ConvNeXt, ViT, Swin, Mega, DeiT), and a variety of model sizes—from DeiT/Mega-Base (85M params) and Swin-T (27M params) to ConvNeXt-Micro (9M params)—provides a sufficient assessment to justify the empirical contribution of our 2-D SSM layer.

Why is the D omitted in Equation 2? Does its inclusion empirically reduce the performance?

This was done solely for brevity. We interpret 'D' as a parameter-based skip-connection in our model. Noting that key papers, like [2], introduce the State-Space Layer with 'D' and later omit it, we agree it could be clearer to present the full equation initially. We are open to making this change for enhanced clarity in our paper.

In Figure 4, it is observed that the 2-D SSM markedly enhances performance when 100% of the dataset is utilized, as opposed to 20% or less. This result is counterintuitive. Providing explanations for this phenomenon would enhance the comprehensiveness of the paper.

It seems that there may be a misinterpretation. The dotted line in the test accuracy graph represents the 2-D SSM and is consistently above the straight line, which indicates the baseline, at all data points. This shows that the 2-D SSM performs better across all dataset utilization levels. Furthermore, we use the same hyperparameters for all experiments. Those hyperparameters were optimized for the baseline (attention without our layer) for the full-dataset regime, which could explain such phenomena. [1]: SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers. Enze Xie, Wenhai Wang, Zhiding Yu, Anima Anandkumar, Jose M. Alvarez, Ping Luo [2]: Efficiently Modeling Long Sequences with Structured State Spaces. Albert Gu, Karan Goel, Christopher Ré

In the experiments, the proposed method is primarily benchmarked against SSM-based methods. However, various approaches employ convolution for embedding, for example in [1], and the use of convolutional layers or simple components for positional encoding is a common practice [2, 3, 4, 5]. It would be instructive to compare the proposed method with Vision Transformers that incorporate the methods to exploit spatial equivariance, which could serve as additional baselines. Furthermore, while it may not be a critical flaw, the implementation of the proposed method appears to be somewhat more complicated and less straightforward than simply employing convolutional layers.

In our earlier response, we added a comparison to a layer with a spatial inductive bias by comparing it with Conv2D within the state-of-the-art ConvNeXt. We would like to introduce another set of experiments using a ViT variant called MixFFN from SegFormer [1]. MixFFN omits positional encoding and integrates a 3x3 depth-convolution within the MLP, specifically before the hidden layer. In the table below, our modified approach—where the 2-D SSM is placed within the MLP in a manner akin to MixFFN, rather than preceding the attention layer and excluding the positional encoding—demonstrates superior performance over MixFFN. We tested it on CIFAR-100 and Tiny-ImageNet.

[1]: SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers. Enze Xie, Wenhai Wang, Zhiding Yu, Anima Anandkumar, Jose M. Alvarez, Ping Luo

Thank you for the comprehensive review and the constructive feedback.

“.. the improvement is tiny and it's hard to see the real benefits of the proposed idea. Actually this tiny improvement may disappear when hyper parameters vary a little bit”

We believe that demonstrating performance improvements across six datasets (ImageNet, ImageNet-100, Celeb-A, Tiny ImageNet, CIFAR-100, and CIFAR-10), multiple backbones (ViT, Swin, Mega, DeiT, and isotropic designs), and various model sizes (such as DeiT-T, DeiT-S, DeiT-B) provides compelling evidence of our method's consistent value. Furthermore, our improvements achieved SOTA results on some of the benchmarks (Imagenet-100 and Celeb-A) and remained consistent during our additional experiments (see ablations in Figure 4 and Table 8). Moreover, we wish to reemphasize that for most benchmarks, we did not perform any hyper-parameter tuning, instead, we used configurations optimized for the baseline method. The only exception was in specific cases where we increased the stochastic depth, which we consider as minimal tuning. We think that it is reasonable to assume that the gap will increase by optimizing the hyper-parameters

"More experiments will be needed to justify"

We have added results using the ConvNeXt backbone across four datasets and three model sizes. Please refer to the section '2-D SSM as an Alternative to Standard Convolution Layers (Experiments with ConvNeXt)' in the shared response, and Tables 6 and 7 in the revised manuscript for detailed information.

“Besides image level classification, can you run experiments on pixel level tasks such as semantic segmentation or instance segmentation by plug in this new layer into, e.g., Mask2Former or Pyramid Vision Transformer?”

Traditionally, such experiments rely on large pre-trained models (trained on ImageNet-22\1K or through SSL schemes), and therefore are very costly and out of our budget’s reach. Please refer to 'Nature of our Benchmarks' in the main response for more details. We acknowledge the value of such experiments; however, we believe that our extensive evaluation across six datasets, with several backbones (ConvNeXt, ViT, Swin, Mega, DeiT), and a variety of model sizes—from DeiT/Mega-Base (85M params) and Swin-T (27M params) to ConvNeXt-Micro (9M params)—provides a sufficient assessment to justify the empirical contribution of our 2-D SSM layer.