Paper: A 2-Dimensional State Space Layer for Spatial Inductive Bias (ICLR-2024)

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Presentation):
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?

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.



| Dataset            | Mix-FFN | w/2-D SSM | Delta | #Params Ratio |
|--------------------|---------|-----------|-------|---------------|
| Tiny-ImageNet | 58.94   | 60.01     | 1.07  | 1.0035x       |
| CIFAR-100   | 76.32   | 77.21     | 0.89  | 1.0035x       |

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

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
**Extended Training Time**: The submission raises concerns about the extensive training time required for the proposed method. There is ambiguity regarding whether other methods, if given a comparable increase in computational resources, would yield indistinguishable results. The current results lack persuasive power as they do not strictly control factors, making it difficult to definitively attribute performance gains to the proposed method.

Author response:
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.

Reviewer comment (Experiments):
**Limited to Transformer-based Architectures**: The method primarily targets transformer structures that lack explicit inductive bias design. 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 the method's applicability and effectiveness.

Author response:
> “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.