A 2-Dimensional State Space Layer for Spatial Inductive Bias
Authors: Ethan Baron
Itamar Zimerman
Lior Wolf
Venue: ICLR 2024 poster (2024)
Keywords: SSM
Dimensional State Spaces
Spatial Representation
Primary area: general machine learning (i.e., none of the above)
TL;DR: A 2-Dimensional State Space Layer for Spatial Inductive Bias

Abstract: A central objective in computer vision is to design models with appropriate 2-D inductive bias. Desiderata for 2-D inductive bias include two-dimensional position awareness, dynamic spatial locality, and translation and permutation invariance. To address these goals, we leverage an expressive variation of the multidimensional State Space Model (SSM). Our approach introduces efficient parameterization, accelerated computation, and a suitable normalization scheme. Empirically, we observe that incorporating our layer at the beginning of each transformer block of Vision Transformers (ViT), as well as when replacing the Conv2D filters of ConvNeXT with our proposed layers significantly  enhances performance for multiple backbones and across multiple datasets. The new layer is effective even with a negligible amount of additional parameters and inference time. Ablation studies and visualizations demonstrate that the layer has a strong 2-D inductive bias. For example, vision transformers equipped with our layer exhibit effective performance even without positional encoding. Our code is attached as supplementary.

Published as a conference paper at ICLR 2024
A 2-D 2-DIMENSIONAL STATE SPACE LAYER FOR SPATIAL INDUCTIVE BIAS
Ethan Baron", Itamar Zimerman' & Lior Wolf The Blavatnik School of Computer Science, Tel Aviv University (barone, mmerman1 @emmmm...t..........e....e....ee at..aa..i.............. wol fecsfcc....
ABSTRACT
A central objective in computer vision is to design models with appropriate 2- D inductive bias. Desiderata for 2-D inductive bias include two-dimensional position awareness, dynamic spatial locality, and translation and permutation invariance. To address these goals, we leverage an expressive variation of the multidimensional State Space Model (SSM). Our approach introduces efficient parameterization, accelerated computation, and a suitable normalization scheme. Empirically, we observe that incorporating our layer at the beginning of each transformer block of Vision Transformers (ViT), as well as when replacing the Conv2D filters of ConvNeXT with our proposed layers significantly enhances performance for multiple backbones and across multiple datasets. The new layer is effective even with negligible amount of additional parameters and inference time. Ablation studies and visualizations demonstrate that the layer has a strong 2-D inductive bias. For example, vision transformers equipped with our layer exhibit effective performance even without positional encoding. Our code is available at this git https URL
INTRODUCTION
Incorporating image-specific inductive bias into computer vision networks could play crucial role in their success, by shaping the hypothesis space way that fits image data and improves generalization. Common ingredients of image-specific inductive bias include two-dimensional neighborhood structure, locality, translation equivariance and invariance, and extraction of hierarchical features. Traditionally, it was injected into the model through the backbone architecture. However, more recently, has been modeled as part of the data. For example, two-dimensional neighborhood structures are typically expressed in one of two ways: (i) Vision Transformers (Dosovitskiy et al. (02T) use 1-D positional encoding (Vaswani et af.] 2017), which considered weak inductive bias. (ii) ConvNets employ 2-D kernels, which provide strong priors on the underlying image structure (Ulyanov et al. 2018).
Most ConvNets employ relatively small filters in the convolution layers, and the balance between local and global features is handled by increasing the receptive field with depth. However, other kernel sizes be beneficial. For example, ConvNeXt improved ResNet by 0.7% on Imagenet, by only increasing its kernel size from 3 x 3 to 7 x 7 (Liu et al. (202). More generally, using fixed-size filters limits the type of dependencies the layer can capture.
The objective of this study is to develop a new layer that is adept at integrating both local and global spatial features, with a particular focus on a doodimmnniioall neighborhood structure. We accomplish this by building on recent developments in 1-D SSM-based layers, which are renowned for capturing various types of dependencies. By extending this 1-D concept to 2-D, our layer deeply rooted in control theory, and much like its predecessors, maintains a strong bias towards position awareness and parameter efficiency.
Our main contribution is the 2D-SSM layer, which is new spatial layer based on Roesser's s model for multidimensional state space (Kung et al. 1977). We show that simple design choices,
'These authors contributed equally this work.

---

Published as a conference paper at ICLR 2024
such as diagonalization and normalization, make the layer numerically stable and efficiently computed without recurrence using a 2-D convolution (left panel of Fig. D. Our layer has some unique properties, including: (i) A strong inductive bias towards two-dimensional neighborhood and locality, which stems from the multi-dimensional recurrent rule. As far as we know, this novel concept does not appear in other layers, (ii) The new layer can capture unrestricted controllable context. The SSM parameters of the layer can be focused on short or long, horizontal, vertical, or diagonal dependencies (middle panel of Fig. []. (iii) The layer parameter-efficien and can express kernels of any length via 8 scalars. Visualization and ablation studies demonstrate these key aspects of our layers. Finally, the layer is well-grounded in control theory, and further theoretical analysis shows that it generalizes S4ND (Nguyen et al. 2022) and proves its greater expressiveness.
Empirically, we show that our layer can be used as a general-purpose booster for vision transformers (the schematic architecture is illustrated in the right panel of Fig. D. with negligible additional parameters and computation at inference. Furthermore, it appears that our 2D-SSM surpasses standard methods, such as incorporating positional encoding, in effectively integrating positional bias into Vision Transformers. Finally, we demonstrate that incorporating the new layer instead of 2D convolutions, within the ofate-of-the-arr CNN architecture ConvNeXt Liu et al. (2022). improves its accuracy, without significantly changing the number of parameters.
2 BACKGROUND AND NOTATIONS
Framing Our research delves into two emerging research domains. The first domain focuses on the development of multii-aees global convolution techniques. Although 1-D (long) global convolution has shown promise in 1-D sequence modeling, leveraging methods such as SSM (Dao al. 2022 Gu et al.) 2021a by Gupta 2022; Mehta et al. (2022) or other recent approaches (Fu al. 2023) PPoli et al.) 2023). (Li et al. 2022b). its applicability and effectiveness in modern computer vision tasks remain uncertain. Our work aims to explore and highlight the potential of these techniques in this domain, by extending them into 2-D.
The second domain investigates the synergistic combination of attention and SSM in I-D modeling
across various domains Ma et al. 2022; [aon et al. 2023 Islam et al. 2022] Zuo et al. 2022 Dao
et al. 2022 Mehta al., (022). For example, the based H3 (Dao et al. 2022) outperforms
GPT-Nco NPTT-No-2..7B 2.7B (Black et 2021) (as well as other transformers of the same size) with only
attention layers. However, the question of whether these components are complementarity in 2-D modeling remains unanswered. We provide empirical evidence supporting the complementary nature of these components.
State Space Model (SSM) The state space model maps an input scalar function u(t): : R + IR to N-D latent state xt)) E - before projecting an output signal y(t) R R:
a(l) = Ax(t)= + Bu((), y(t) = Cx(()) + Du(l)
(1)
The use of SSMs is widespread across numerous scientific disciplines and is closely associated with latent state models, such as Hidden Markov Models. There is a well-known connection between linear time-invariant SSMs, such as and continuous convolution, thus allowing efficient training using the aforementioned equation as a discrete convolution. The S4 (Gu al. 202 and LSSL (Gu et al., (0211) layers leveraged the SSM as a black-box representation in deep sequence modeling, with learned parameters A, B, C, and D, and achieved strong results on several tasks, especially ones that require handling long-range dependencies, such as the Long Range Arena (LRA) (Tay et al.)
2020), audio generation (Goel et al. 2022), and long-text processing (Mehta et al. 2022 Golub & He 2016 Dao et al. 2022).
The underlying reasons for the suitability of S4 and SSMs for modeling long sequences were recently analyzed (Li et al. 2022b Fu al.) 2023). It was found that (i) employing global kernels with decaying structures, and (ii) using regularized kernels, are both critical design choices in this area.
Multi-dimennional State Space Layers As far as we know, S4ND (Nguyen al. 2022 is the only previous SSM- layer that can naturally handle multidimensional data. S4ND is built on top of S4, and contains M separate instances of S4, where M is the number of axes. On the forward path, each S4 layer, which we denote as SSM, factors a one-dimensional local kernel ky, and a global kernel K is computed as the outer products of those local kernels.
2

---

Published as a conference paper at ICLR 2024
Equivalent Views the 2-O SSM: Convelutional and Recumrett
Strong Image-Specific Inductive Blas
Drop-in 2-0 Filter Booster for VIT
Recument Audu
2-0 SSM
Patch-Based Modals
ConvNeet with 2-0 SSM
SSM
Bay
C3
Unear
GELU
GELU
2-0 Global Convolutional Nareel
Drop-in Bocster for VIT
kanal
tapa
Outaa
Our block
2.0 9BM
VIT Block
(Left) The 2-D SSM layer is parameterized by A, B, C, and D. is built on top of two-axis linear recurrent and can be efficiently computed using 2-D convolution. (Center) Since the layer based on two-dimensional recurrence, exhibits a strong bias toward positional awareness. The recurrent is unrestricted, allowing the layer to operate on 2-D sequences of any length. The values of A2, 12, Ag. and A₃ control the layer's 's focus, enabling it to capture short or long spatial dependencies in horizontal, vertical, diagonal directions, can be easily integrated into ViT by applying it to the two-dimensional sequence of patches at the beginning of each transformer block, by replacing the standard Conv2D) in ConvNeXt.
Figure 1:
opposed to patch- models. (Right) The layer
Roesser's 2D-State Space Model The attempt to extend the classical SSM for 2-D multi-axes
systems was thoroughly studied the past. (Kung et al. 1977. Fising 1978 Fornasini & Marchesini 1978) Kurek 1985 Givone & Roesser 1972. Hinamoto 1980) notes few different formulations of the problem. We employ Roesser's s SSM model (Kung et al. 1977) as our discrete multi-axial model, which is the most general form of 2-axial and N-axial state-space models. As opposed other SSMs already in use in machine learning, this SSM uses M states, one per axis, rather than just one. The model in 2-D form is presented here:
A2 A3
2h 1
B B2
Jij
Mi. = [G1
C2)
ut
(2)
where the state Zi, R2N is the concatenation of the horizontal i E RV and vertical x²
R
states, the system matrices are A2, As, At E NNN and the input and output matrices are B1, B2, C1 RVV There is also learned parameter that behaves as a skip-connection. omitted from now on for brevity.
Additional related work is presented in Appendix
2.1 NOTATION
The notation follows closely as possible the notation used in the state-space layer literature Guu al. 2021a Gupta 2022; Gu et al. (022). Specifically, we use H as the number of channels, N as the state's hidden dimension, L as the sequence length, Rsam the number of non-shared channels, and treat A, B, C, D R as the system matrices. Note that for brevity the system matrices are treated vealevalued, even though also test complex version of them. The number of axes is set to M.
Although N-D SSM can be used in an N N-Dimensional manner, since our paper focuses on using this model regularization method for VIT backbones and for simplifying the reading experience, will treat as 2-D SSM.
L; IR is the sequence length along the axes, Liot L1 La is the total signal size, and Lynaz = max(ax L1, L2).
The kernel notation introduces K, which is the computed 2-D kernel such that Y U K, where denotes discrete convolution.
Vi.j = C, is + C2 = Σ Σ (G,K,, ij + C2b",3))
(3)
osisio5isi
3

---

Published as a conference paper at ICLR 2024
3 METHOD
In this section, we present the core of the 2-D SSM layer, which is our main technical contribution. This layer maps 2-dimensional sequence U Yi.j where Wij Mi,j R for all S i S Li and 0) - & L2. Similarly to previous SSM-based layers, we extend the core of our layer to multi-directional and multidimensionall layer, detailed in Appendix
3.1 2-D RECURRENT STATE SPACE MODEL AS CONVOLUTION
Eq. 2 is defined in recursive manner. However, for reasons of efficiency, it is best to define operators in closed form and in way that is concurrent with parallel processing. Inspired by previous work (Gu al.
2021a; Gupta 2022 Ma al. 2022 Gu al. 2022), we exploit the fact that the SSM is linear and can therefore be expressed as a convolution with a kernel K. To do so, we first unroll the recurrent rule and then describe it in closed- form formulation.
For simplicity, we assume that the initial states are zeros Vj 2 0 -11 O, and vi 2 0: 000r,, The horizontal and vertical states
0,j are:
x² 0,0
00,0
x 000.0 B2u0,0
By applying the recurrent rule once at each axis:
x² 1,0 Auooo + A2 + BiUiO:
x¹0,1 Asli + A4B20o.0 + B2u0,1 (5)
x" 1.0 B2
2¹ 1
(6)
Next, we compute x² 1,1 given x² "01, and x" 0'0,1 0,1.
xt 1.1 AlAaBiuo. A Ai B2uo, + A2Biuu, + BIUI.I.
(7)
= 40,0 kj + k0,0
(8)
and in general
Σ Σ
k
Σ
Σ
k2 154133'
(9)
osisiosisi
osigiosisi
where as explained in the notation, each element k1, k"t is an aggregation of A1, A2, Ag, A1. B1, B2 multiplications (c.g Eq. 7 and associated with a single path from co- ordinate (0 to (i,j), as presented in Fig 2
By plugging Eq. 9 in Eq. 2 one obtains:
yi C1ah is + Cax" by
Σ Σ
(Cik" + C2k"
(10)
osisiosisi
and the convolutional kernel K is
vi.j: C,ke 2, + C2k"
(11)
3.2 EFFICIENT AND STABLE PARAMETERIZATION
Parameter diagonalization Computing K is difficult for two reasons. Firstly, the number of elements in each k is exponential in and j. since it is equivalent to the number of paths from (0, 0) to (i, Secondly, calculating the powers non-diagonal matrices A,, A2, A3, A4 becomes costly when L1, L2 are large. To overcome these challenges, we parameterized the system matrices A1, A2, Ag. A4 as diagonal. This allows for efficient summation at most 2Lmaz elements in k.,- Although diagonalization limits the expressiveness of our SSM, previous works have shown its
effectiveness in one-dimensional cases (Gu et al. 2022 Gupta 2022 Mehta et al. 2002) Gupta et al.
2022).
4

---

Published as a conference paper at ICLR 2024
Figure 3: The kernels before and after the modifica- tions of Sec. 3.2 Each column is created by the same Aq...A4, B1 B2, C1, C2 E R parameters. The first row is the normalized 2-D SSM formulation explained in 2 the second the outcome of E and and performing Eq. [3] which is the kernel for- mulation we use. The bottom left corner of each heatmap is Ko.0- The figures demonstrate that before the relaxation, the kernels displayed a diagonal tendency while afterward, they exhibited more diverse and versatile pattern.
Figure 2: Examples of paths
from coordinate (i, (0, 0)
to (i.j) (4, 4). Each path
represents a sequence of recur- sive calls for Eq. 2
Limiting the parameters A, E NNN is now a diagonal matrix. We'll denote the eigenvalues of A; by (a1, .22.....x)) == diag(A, A.). Each a value behaves separately until the computation of (Eq. Thus, for a; 1, limg-roo a = co. Therefore, we limit a; E [0, 1] by parameterized it by ag sigmoid(á:)) and optimizing d; instead.
Normalization There is scaling problem that arises from Eq. 7 Since kt,1 A
even if A1, ArA4, , can be greater than 1. The same behavior makes the kernel explode.
We would like to keep Vi,j 0 S ki, S 1 and thus we employ straightforward normalization mechanism, in which we divide by two each time we compute Eq.
2
This Eq. is thus replaced by:
h [11 x i.j+1
A1A2
xh io x s.j.
B B2
0.5
A3
+
Wi
(12)
Relaxation of the kernel impose a kernel formulation that much more biased towards modeling diagonal kernels, i.c., 1- 5l is small, kf is much larger than when j jl is large.
When imposing a "divide by two" normalization for every kig
Thus, we relax the normalization in the first row and column as follows. When calculating k kb.t we use Eq. 2 Additionally, for Ke Ko,j, we use G 2C1, C2 2C₂ in the following manner:
Ko, Ciko, + C2Kb,
(13)
Figure illustrates examples of different kernels before and after the relaxation.
We note that these modifications are straightforward and probably not optimal. The development of optimal normalization according to the kernel formulation is an interesting direction for future work.
3.3 COMPUTATION AND COMPLEXITY
Training Complexity The training process has two parts. First, calculate the kernel K. The calculation itself is explained thoroughly in Appendix(Bland and has time complexity of (Ltor Liman N)) which is not dependent on B. and is much faster than naive computation thanks to several design choices, such as parameters diagonalization, pre preprocesssig, and a sophisticated caching procedure.
Next, we apply convolution between K and U, using the classical procedure of FFT, element- wise multiplication, and inverse FFT. The complexity of this step (B log(L))))) where the batch size. Hence, the total complexity is O [Toma N B log( Trot) Trtt)
5

---

Published as a conference paper at ICLR 2024
Inference Complexity During inference, the 2-D SSM can pre-compute the convolution kernels, resulting in an additional computation cost of only the 2-dimensional convolution of the layer input and the signal, which takes Ltot Ltoe) operations. Therefore, the additional computation overhead relative to the vanilla ViT is minimal, as the quadratic complexity dominates the overall complexity. For empirical analysis of the time and memory complexities, refer the reader Appendix
3.4 COMPLEX AND REAL SSMS
While most SSM-based deep learning models, e.g S4 (Gu et al. 2021 or DLR (Gupta et al. 2022 use complex SSM, MEGA used EMA, which can be interpreted as a restriction of the diagonal SSM to real numbers. Our 2-D SSM layer can be built over real complex parametrization. The complex-Se -SMM based model that we examined is described in detail in Appendix E.
4
MODEL ANALYSIS
We study the expressiveness of the 2-D SSM layer, as well as its unique spatial inductive bias.
4.1 EXPRESSIVENESS
We compare the expressiveness of our 2-D SSM layer with S4ND (Nguyen et al. 2022) a very recent layer that is also based on multidimensional multi-axes SSM. We first introduce the key differences between our layer and S4ND, and then demonstrate the expressiveness gap.
The relationship between 2-D SSM and S4ND The main difference between S4ND) and 2-D) SSM is that S4ND runs a standard I-D SSM over each axis independently, and those functions are combined to form global kernel. In contrast, our model learns multi-axes data directly. This difference arises from the fact that the 2-D SSM has additional system matrices, A2, Ag, which aggregate and process information from different axes.
functions over
2D-SSM is a generalization of S4ND When restricted to 2-dimensional space, given 2-D SSM model, S4ND is obtained by restricting A2, Ag to be zeros, and setting A1, A4 the system matrices of S4ND. Additionally, to replicate S4ND in 2D-SSM, one should initialize the states the 2D-SSM by the kernels factorized from S4ND.
Tensor rank as a criterion for expressiveness Over the years, several criteria were proposed for measuring the expressiveness neural and classical models, such as VC-dimension (Vapnik & Chervonenkis 2015), norms, and Rademacher complexity (Bartlett & Mendelson 2002). Inspired by (Cohen et al.) 2016). we employ tensor rank as our measure, and prove the followings theorems:
Theorem 4.1. The 8 parameters of the 2-D SSM can express full-rank kernels
Theorem 4.2. S4ND can express kernels of rank 1 only.
Assumptions For simplicity, we assume that both the 2-D SSM and S4ND layers contain one channel and one hidden dimension. In this case, the SSM matrices are scalars. When the assumption about the number of channels is omitted, the rank of kernels that S4ND express increases from 1 to N. However, this comes with the cost of M Nr additional parameters, where M the number of axes and is the required rank for S4ND. It should be noted that the authors of S4ND did not evaluate the performance of models with > 1.
Under these assumptions, the proof of Theorem is specified in Appendix C.I The proof of 4.2 trivial and derives from the fact that to compute a global multi-axi axis kernel K. S4ND takes the outer product operation the per-axis kernels km C m for all m [M]. Since each kernel is vector, it clear that:
rank () rank(ki @ 62 ....kM. =
(14)
4.2 SPECIFC INDUCTIVE BIAS
Our layer is grounded by two-dimensional linear recurrent (Eq. Fig. [left), as a result, positional information is taken into account by design when the kernel K is computed from the parameters A, B, C, and D. This is unique property without counterpart in other modern layers. Furthermore, as can be seen Sec. 5] our method is highly effective in inserting positional bias into the transformer, even outperforming positional encoding in some cases.
Two-dimensional Position-Awarenes:

---

Published as a conference paper at ICLR 2024
Table 1: Results using ViT, MEGA and Swin backbones the
Table 2: ViT & Celeb-A resutls.
Tiny ImageNet (T-IN) and CIFAR- 100 (C100) datasets. No
Model
Top
#Param
hyperparameter tuning except for stochastic depth.
DeiT-T
88.43
5.532M
Model
C100 T-IN Train Time # Of Parameters
DeiT-T w. SSM-r
89.76
5,537M
DeiT-T w. SSM-c
89.84
5.54IM
VIT
73.81 57.07
1x
2.71M (Ix)
VIT w/ S4ND.
72.60 56.10
2.22x
2.72M (1.003x)
DeiT-S
89.66
21.681M
DeiT-S w. SSM-r
90.24
21 21.688M
VIT w/ SSM-r.
77.21 60.01
2.66x
2.72M (1.003)
DeÎT-S W. SSM-c
90.38
21.691M
VIT w/ SSM-c.
76.46 59.08
2.89x
2.73M (1.01x)
DeÏT-B
90.13
85.829M
Mega-ablate
56.44
1.1x
2.75M (1x)
DeÎT-B W. SSM-r
90.45
85.84IM
Mega
72.27 54.49
1.28x
2.98M (L.0xx)
DeiT-B w. SSM-c
90.73
85.845M
Mega w/ S4ND
74.9 56.65
1.46x
2.80M (1.0Zx)
Swin-T
91.48
27 27.550M
Mega 2-D SSM-r
76,02 57.95
2.03x
2.80M (1.0x)
Swin-T SSM-r
91.68
27.556M
Mega 2-D SSM-c 75.09 56.51
2.03x
2.84M (1.03x)
Swin-T SSM-c
91.78
27.558M
Swin
76.87 60.87
1x
7.15M (Ix)
Swin-reprod.
77.98 61.29
1x
7.15M (1x)
Table 3: VIT & ImageNet-100 100 runs
Swin w/ EMA
77.01 60.13
1.39x
7.52M (1.05x)
with and without 2D-SSM
Swin w/ S4ND
79.26 64.6
1.29x
7.18M
(1.004x)
Models
Baseline
w/ 2D-SSM
Swin w/ SSM-r
80.12
65.77
2.16x
7.25M (1.01x)
Swin w/ SSM-c
3.28 12.76
2.18x
7.26M (1.02x)
DeiT-T 78.21
81.16
DeiT-S 82.27
84.82
Swin-T 81.06
82.29
Controllable Unrestricted Spatial Context A significant limitation of both CNNs and transformers the need to choose an appropriate patch size, which can result loss of global or local context. In contrast, our layer implements a controllable global convolution, which benefits from global context, and can control the effective receptive field and efficiently capture spatial local dependencies. Moreover, our layer is not confined to patch-like dependencies and can effectively capture diverse local and global features in horizontal, vertical, or diagonal directions (See I middle).
Modeling Symmetries Similar to other CNNs, our layer exhibits translation equivariance the kernels slide across the input during computation. Additionally, as detailed Appendix [ our layer's core extends over multiple directions, which allows to accommodate rotations and reflections naturally.
Parameter Efficiency In contrast CNNs, which parameterize filters of size H W with at least HW parameters, our layer has a fixed and small number of parameters (9 parameters per channel, A1, A2, A3, A4, B1, B2, Ci, C22 D R), however, those parameters can be expanded into unbounded two-dimensional kernels. Furthermore, we use parameter sharing for the SSM parameterization across channels, similarly CNNs or state-space layers and donate neam as the number of non-shared channels.
EXPERIMENTS
We assess our 2-D SSM layer as an inductive bias within various ViT-based and CNN backbones. We demonstrate the universality our method by incorporating it into various backbones. such as ViT. DeiT, Swin, ConvNeXt and report improved results over the baselines on ImageNet- Ik, ImageNet- 100, Celeb-A, Tiny-ImageNet and CIFAR-100, with negligible additional parameters and without hyperparameter tuning, except for stochastic depth. For comprehensive overview of the experimental setup. please see Appendix
Swin, DeiT and ViT We adopted a straightforward approach to incorporate our 2-D SSM layer into the aforementioned backbone structures. Prior to the Transformer Block, we apply the 2-D SSM to the input signal
RL x D illustrated in Fig. | We highlight that in the case of Swin Transformer, the 2-D SSM operates on the patch level and not the window level and therefore injects additional 2-D bias between windows.
We tested Swin and VIT over the small datasets reported by (Lee et al. 202I) baseline. As shown in Tab. [ with the ViT backbone, we improve 0.8% on CIFAR- 100 and 0.5% on ImageNe,, and with the Swin backbone, we improve 3.25% on CIFAR-100 100 and 4.25% Tiny ImageNet.
and CIFAR- 100, using the results
7

---

Published as a conference paper at ICLR 2024
Table 4: Accuracy the CIFAR-10
Table 5: ImageNet- 1K accuracy of MEGA variants.
grayscale classification task, which is part of the Long Range Arena.
Model
Top 1
Top
# of Parameters
MEGA-ablate
66.97
88.17
5.91M
Models
Image LRA
EMA
6.21M 69.73
89.76
Transformer
42.94
2D-SSM-R
70.11
90.19
5.96M
S4-v1
87.26
S4-v2
88.65
CCNN-ID
88.90
CCNN-2D
91.12
S4ND
89,90
Hyena
91.20
That
More CIPAA
Mega Ablate
81,00
Soenn Sast Tiny-imageNet
Tiny Imergeeet
Mega
90.44
Dell 2 Siz Calecc
MEGA 2-D SSM
91.31
Training Set Size (%6
Figure 4: The effect of the training set size.
The DeiT and Swin backbones were tested on the large-scale Celeb-A dataset (Liu et al. 2015) and ImageNet-100. Celeb-A involves a 40-way multi-label attribute classification. We report aggregate accuracy across all 40 tasks. As can be seen in Tab 2 the complex version outperforms the real in all experiments and achieves 1.41% %, 0.72%. 0.6%. and 0.3% improvements over the baseline for DeiT sizes Tiny, Small, and Base, and Swin-Tiny respectively. The ImageNet-100 benchmark a known subset of ImageNet-1K K.. The results of these experiments are presented in Tab. 3 The hyperparameters for each model are taken from the corresponding paper and git repository. For DeiT, our 2-D SSM layer resulted in a performance enhancement of 2. 2.95% (from 78.21% % to 81.16%) on DeiT-T and 2.55% (from 82.27% to 84.82% on DeiT-S. On the Swin-T, the performance improved by 1.23%, elevating the accuracy from 81.06% to 82.29%.
Mega In Mega, we replace the EMA mechanism with 2-D SSM, which means that only perform our layer on Q, K. We compare original Mega (with EMA) vs Mega-ablate (without EMA) and Mega 2-D SSM. We examined our model on CIFAR-I0 Grayscale in an isotropic manner (without decreasing the image size along the architecture, and without patches), which is part of the Long Range Arena benchmark (Tay al. 2020). As shown in Tab we improved the result by almost 1% over MEGA, obtaining statc-of-the-an results, including very recent approaches such as Hyena (Poli et al. 2023) and CCNN-2D (Knigge et al. 2023). We also conduct an experiment on the ImageNet-IK dataset (Deng et al., 2009), and as shown in Tab. [ we improve over MEGA's ViT-T results by 0.4% in both Top accuracy and Top 5 accuracy. Finally, we check other small datasets (Tab. and find superior results for combining Mega with 2-D SSM over baseline or other methods.
ConvNeXt Non we replace the x 7 Conv2D filters with 2-D SSM layers. We eval- uate our ConvNeXt variants on four datasets, in- cluding high-resolution (Imagenet- 100 and Celeb-A) and low-resolution (Tiny-Imagenet and CIFAR- 100) datasets, across three model sizes: Con-NeXt-Ting NomMNeXtMMir and ConyNeXI-Nano, These back- bones were introduced in previous work, the differ- ences between the backbones, and more details about the experimental setup are provided in Appendix G As highlighted in Tab. [ and Tab. 7 both our com- plex and real variants consistently outperform the original ConvNeXt across all datasets and model sizes, without changing the parameter amount sig- nificantly. It is important to note that we do not perform hyper-panc parameter tuning, instead, we use the exact same hyper-paramene parameters optimized for the CNN baseline, yet our results achieve SOTA performance on ImageNet- 100 and CelebA datasets.
Model
IM-100
Celeb
# Param
T-Conv2D
89.72
91.70
27.900 27.90M
T-2D SSM.R
90.81
91.81
27.68M
T-2D SSM.C
90.47
91.96
27.79M
M-Conv2D
88.51
91.00
9.27M
M-S4ND
N.A
91.30
9.60M
M-2D SSMR
89.81
91.61
9.23M
M-2D SSM.C
89.16
91.75
9.33M
Table 6: Micro (M) and Tiny (T) ConvNeXt
results Imagenet- 100 and Celeb-A.
Model
C100
T-IM
Param
N-Conv2D
79.10
67.09
6.48M
N-2D SSMR
80.16
69.10
6.51M
N-2D SSM.C
79.92
67.46
6.63M
Table 7: Results of Conv on Tiny-
Imagenet (T-IM) and CIFAR- 100 (C100).
8

---

Published as a conference paper at ICLR 2024
Table 8: Ablations. For each model and dataset, examine the effect of using original positional encoding and complex (C) vs. real (R) SSM. The column ó represents the average difference for models with and without PE. As can be seen, our models much more resistant PE removal.
Dataset:
Tiny-INet (Swin)
CIFARION (ViT)
CelebA (DeiT-T)
CIFARIO (Mega-ISO)
Avg.
Model
with PE
w/o
with w/o
with PE w/o PE
with PE
w/o PE
Baseline
61.29
58.97
73.26
64.09
88.43
87.99
90.44
75.21
-6.79
+Ours (R) 65.77
65.44
74.07
74.89
89.76
89.63
91.31
90.68
-0.07
+Ours (C) 3.28
2.16
73.91
74.67
89.84
89.83
90.46
90.79
-0.01
Comparisons against S4ND S4ND) is the only N-Dimensional SSM-based layer known to us. Nguyen al. (2022) reported results on Celeb-A with ConvNeXI-Micro. As shown Tab. our real and complex variants show improvements over S4ND by 0.31% % and 0.45%, respectively, both with a smaller number of parameters. We further compare our results against by substituting SSM in VIT. Mega, and Swin. We conducted experiments on CIFAR 100 and Tiny Imagenet. As indicated in Tab. S4ND performs very poorly when integrated into the original ViT backbone on CIFAR- (lower than the baseline by 1.21%) and sub-optimally when integrated into Swin and MEGA (achieves 1% or more lower aceuracy on CIFAR-100 and Tiny-ImageNet for both backbones).
Sample Complexity We examine the behavior of our model with different backbones on different datasets over the baseline. As can be seen in Fig. [ 2-D SSM maintains improved results over the baseline for all backbones, which shows the data-cfficient quality of our model.
Removing the positional encoding (PE) We compare the empirical results obtained with real complex kernels, with and without PE in Tab. g Evidently, complex- can be superior or inferior to real- based SSM, depending on the scenario. Additionally, we find that complex-SSM, has tendency exhibit instability during training, which result in poor performance. Stabilizing these models is an important direction for future research.
Running ViT backbones without PE decreases performance dramatically. In contrast, when 2D-SSM layer is inserted into these backbones, they benefit from PE, and even without PE they outperform the original backbones with PE. These findings support an innovative approach to introducing positional bias in ViT: rather than encoding positional information directly into the operators.
representation, it integratede into the computation by incorporating
6 LIMITATIONS
Despite the promising results presented, the current implementation of our proposed layer has relatively slow training times. One possible approach mitigating this challenge is to use multi- dimensional parallel scanners, which could potentially reduce the training time of our layer. The main idea is extending the work of S5 (Smith et al.) 2002), which leverages 1-D parallel scanners apply SSM on 1-D sequences to multi-dimensiona parallel scanners and multi-dimensional sequences.
CONCLUSIONS
We present novel spatial SSM-based layer that is more general than existing ones, encodes positional information by design, and is able to model spatial relations more expressively than other SSMs, including S4ND. When added to various backbones, it is able to improve classification results on the various benchmarks without optimizing any aspect or other parts of the architecture. In future work, would like to study the behavior of the layer in the context of spatial vision tasks, such as video processing, image segmentation, phrase grounding, and image inpainting. the last task, the recursive view of the layer could be applied directly impute missing pixels efficiently.

---

Published as a conference paper at ICLR 2024
8 ACKNOWLEDGMENTS
This work was supported by grant from the Tel Aviv University Center for Al and Data Science (TAD), and the Blavatnik Family Foundation. This research was also supported by the Ministry of Innovation, Science & Technology .Israel (1001576154) and the Michael J. Fox Foundation (MJFT-002407).. The contribution of IZ is part Ph.D. thesis research conducted at Tel Aviv University.
REFERENCES
Training convnext on cifar-10.
1o/ml
corvnext-cifar10 Accessed: 2023- 11-19.
Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463 482, 2002.
Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. Gpt-neo: Large scale autore- gressive language modeling with mesh-tensorflow. If you use this software, please cite using these metadata, 58, 2021.
Nadav Cohen, Or Sharir, and Amnon Shashua. On the expressive power of deep learning: A tensor analysis. Conference on learning theory, pp. 698-728. PMLR, 2016.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoe V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
Zihang Dai, Hanxiao Liu, Quoc V Le, and Mingxing Tan. Coatnet: Marrying convolution and attention for all data sizes, arXiv preprint arXiv:2 006.00660000 2021.
Tri Dao, Daniel Y Fu, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprins arXiv:2212. 14052, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Ieee, 2009.
Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo. Cswin transformer: A general vision transformer backbone with cross-shaped windows, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2021. URL httptttss///peeeeiiiii net/fooum: id=YicbFdNTTT
Rikus Eising. Realization and stabilization of 2-d systems. IEEE Transactions on Automatic Control, 23(5):793-799, 1978.
Ettore Fornasini and Giovanni Marchesini. Doubly-indexed dynamical systems: State-space models and structural properties. Mathematical systems theory, 12(1):59-72, 1978.
Daniel Y Fu, Elliot L Epstein, Eric Nguyen, Armin W Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficien long convolutions for sequence modeling. arXiv preprint 2302::2302.06641 2023.
Donald Givone and Robert Roesser. Multidimensional linear iterative circuits general properties. IEEE Transactions Computers, 106-10):06771173 1972.
10

---

Published as a conference paper at ICLR 2024
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. In International Conference on Machine Learning, pp. 7616-7633, PMLR, 2022.
David Golub and Xiaodong He. Character-level question answering with attention. arXiv preprint arXiv:1604.00727, 2016.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00396, 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in Neural Information Processing Systems, 34, 2021b.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the paramcterization and initialization of diagonal state space models. arXiv preprint arXiv:2206. /193,, 2022.
Ankit Gupta. Diagonal state spaces are as effective as structured state spaces. arXiv preprint arXiv:2203.14343 2022.
Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear rnns. arXiv preprint urXiv:2212.00766 2022.
Ts Hinamoto. Realizations of a state- space model from two-dimensional onput-output map. IEEE Transactions on Circuits and Systems, 27(1):36-44, 1980.
Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, and Daniel Soudry. Augment your batch: Improving generalization through instance repetition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8129-81328 8138, 2020.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. European conference on computer vision, pp. 646 6666661.66 Springer, 2016.
Md Mohaiminul Islam, Mahmudul Hasan, Kishan Shamsundar Athrey, Tony Braskich, and Gedas Bertasius. Efficient movie scene detection using state-space transformers. arXiv preprint arXiv:2212.14427 2022.
David M Knigge, David W Romero, Albert Gu, Efstratios Gavves, Erik J Bekkers, Jakub M Tomczak, Mark Hoogendoorn, and Jan-Jakob Sonke. Modelling long range dependencies in nd: From task-specific to general purpose enn. arXiv preprint arXiv:2301.105444 2023.
Sun-Yuan Kung, B.C. Levy, M. Morf, and T. Kailath. New results in 2-d systems theory, part i:: 2-d state-space models realization and the notions of controllability, observability, and minimality. Proceedings of the IEEE, 65(6):945-63)
1977. doi: 10. 00.119//R0C...77.00522
J Kurek. The general state-space model for two-dimensional linear digital system. IEEE Transac- fions on Automatic Control, 3006)600-- 1985.
Seung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets. arXiv preprint 12.1922 2021.
Kunchang Li, Yali Wang, Junhao Zhang, Peng Gao, Guanglu Song, Yu Liu, Hongsheng Li, and Yu Qiao. Uniformer: Unifying convolution and self-attention for visual recognition, 2022a.
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? arXiv preprint arXiv:2210.09298,, 2022b.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEFF/CVF international conference on computer vision, pp. 10012-10022, 2021.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11976 1986, 2022.

---

Published as a conference paper at ICLR 2024
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv: 10655, 2022.
Harsh Mchta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. 13947, 2022.
Eric Nguyen, Karan Goel, Albert Gu, Gordon W Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos as multidimensional signals using state spaces. arXiv preprint 10.i::221000688,, 2022.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:230. 10866, 2023.
George Saon, Ankit Gupta, and Xiaodong Cui. Diagonal state space augmented transformers for speech recognition. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023,
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprini wrXiv:2208.049933 2022.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:201 1.04006, 2020.
Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sublayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Deep image prior. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp.
2018.
Vladimir N Vapnik and A Ya Chervonenkis. On the uniform convergence of relative frequencies of events to their probabilities. In Measures of complexity, pp. 11-30. Springer, 2015.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. Proceedings of the IEEE/CVF international conference computer vision, pp. 6023-6032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint 1rXiv:1710.09412. 2017.
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yang. Random erasing data augmenta- tion. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 13001-13008, 2020.
Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao and Jianfeng Gao. Efficient long sequence modeling via state space augmented transformer. arXiv preprint arXiv:2212.08116, 2022.
12

---

Published as a conference paper at ICLR 2024
A COMPUTING THE KERNEL
We discuss kk. The same calculations hold for Ki.5h
J can be written as:
2+ Zomas
Vi.j j : KNj
Σ
(15)
For brevity and since it is not material for the method, we limit our exposition of the different power combinations of the system matrices A1.A2. A3, A1 and the input matrices B1, Bg. As noted above, for each ki, there are most 2Lmax elements.
The problem of finding cs for each element in the summation is generalization of Pascal's 's triangle. order to calculate the kernel, we calculate all the coefficients and the powers of A1...A4 up to the size of L1. L2 and cache them before the training process.
Pre-Processing
During training, we employ a matrix multiplication process to compute ki, with the learned parame- ters AI.... A4,
B1, B2.
Thus, for each cell there are at most 2Immaz elements, and for each element, we save a constant number of values (the values of z and cz). As a result, the size of the cached matrix is bounded by mooo))
It should be noted that currently in our method we cache the coefficients as One Hot Encoding and not the coefficient itself, and thus our specific implementation we need to multiply the time complexity and memory complexity by Lmaz.
B TIME AND MEMORY COMPLEXITY OF OUR METHOD
To understand the complexity of our method, we will outline each step of the computation, starting from calculating the cache, creating the kernel during training, and computing the output Y afterward.
As before, for brevity, we will refer only to the xa 1.j> ki matrices caching, but the same holds for the vertical matrices. For simplicity, we assume H = (number of channels).
Caching As noted in Section 3.3 for each cell K., in the horizontal kernel there are at most 2Lmaz elements. Also as noted in Section 3.3] Z1, 22, ø3; 24 <= 2Lmaz. Thus, for each element in Eq. [5] we save 21, 22, 23, 24, 25 and Or2 values. In total, for each cell, we save XI Limas values, where XI is
small constant. We have Leoe cells and thus the total coefficient cached tensor for calculating K₁. is sized
X1 Liar
(16)
From now on we will denote the tensors of horizontal coefficients as CACHE E Rxi X lemux, Notice that there is CACHE tensor for each parameter, meaning CACHER CACHE A2i CACHE CACHE, CACH
Creating the Kernel For brevity, we use real-valued diagonal A; E 10, (after the sigmoid). First, we calculate the Vandermonde Matrix for each A1, Ag, As, A1 eigenvalues up the highest power that exists in the kernel, which is 2L and denote VAN, = V anddermonde(4, E 2 X.N
Again, we have Luat cells in the horizontal kernel, each cell having T-maz elements. We take CACH F., which holds for each element its A; power, 2; and creates matrix that holds for each element its corresponding A value.
O = V AANICCCHEE. E I hest x X2rmas x
(17)
Now we multiply the matrices element-wise to obtain the final value of each element:
Opra Opre-additn addition = 0%, @ O%, o O% 2 O%, c O c O E R Tuo T.z N
(18)
where denotes element-wise multiplication. Now for each cell, we sum all the elements in the summation ke,, meaning summing over the second dimension:
h Opoat-addition summ0'. Oht additionn d 1) X
(19)
13

---

Published as a conference paper at ICLR 2024
Again, all the above steps are employed for the vertical axis well, thus are finally able to compute the kernel by using C1, C2 E 1
K = Opost adstiadittion
C1 + Og O aost-nadiiion C2 E RLixLa
(20)
Remembering that we actually used Nsam channels, the kernel size is E E 1LL,xL288am It should be noted here that the calculation of the kernel is not dependent the batch size B.
Forward Pass Let B denote the batch size, would like to convert input signal e BBx L2xH to output signal Y E R BBxLx LaxH After calculating the kernel, convert U, to the frequency domain through FFT:
U = FFT(U), ; FFT(K)
(21)
IFIFFTU O Kÿ)
(22)
This whole process costs us ( log(L(o:))
Thus, total forward pass time complexity is: Lox N + LTOT log LTOT)
(23)
Implementation detail
We implemented the caching and matrix multiplication process with One Hot Encoding Vector of the powers and not by using floats representing the powers themselves. Thus, the size of each COEFF in our implementation is multiplied by Lmaas is the time complexity of Eq. and
C EXPRESSIVENESS
Theorem C.I. One channel of 2-D SSM can express full-rank kernels
Proof. We start by restricting the system, output and input matrices:
A1 A2 = A3 1, A4 0, C1 = C2 = 0,
Bi = 12 B2==0
=
(24)
For simplicity we assume that the values of the initial states are O:
- i m. = 0, =
(25)
suffices to show that (i) K triangular matrix, and (ii) the diagonal of contains non-zero elements. First, by plugging 24
25 into the recurrent rule can be simplified:
yh. = x 1.j
x s.j
1- i- Usis
x" = x 1-1
(26)
Given this simplification 26 both (i) and (ii) can be easily proven by induction on the diagonals of K.
To provide more insight into the proof, Eq. 27 illustrates the values of K.
17
6
(27)
4
1
And in general, since its clear from 26 that
Vi. d 1.5 yi
91-1,1-1 Wh
Vj
ko. =1
(28)
easy to understand that the upper triangular of K can obtained from rotation of Pascal's triangle.
14

---

Published as a conference paper at ICLR 2024
D MODEL EXTENSION
D.1 BIDIRECTIONAL SSM
Using the State Space model, when calculating oy one only considers z1. where S i, - j. To benefit from bidirectionality, we employ version where we transpose the kernel in two or four directions (shown in the ablation experiments) and then sum the results. A similar mechanism for the 1-D case used in S4, MEGA, and elsewhere
D.2 MULTI-AXIS MULTIDIMENSSONAL SSM
To enrich the dependencies that can be captured by kernels, weighted linear combination of kernels. Under the assumption that the system matrices are diagonal, the N coordinates of the states are not dependent on each other, and can thus be calculated independently. Specifically, Eq. can be re-written separately per coordinate Vgg [N]
Aواي = Σ Σ
(29)
osiciusisi
Therefore, increasing N will increase the number kernels that make up K. By using this structure, the kernel capture variety of dependencies, where each coordinate focuses different type of dependency. Furthermore, this extension adds relatively negligible runtime when working with large batches, since the batch dimension does not affect the complexity when the kernel is computed. Therefore, increasing N will increase the number of kernels that compose K.
E
JUSTIFY DESIGN CHOICES
E.1 OUR COMPLEX 2D-SSM
As explained in Sec. 3.3 our models employ real rather than complex SSMs. For reproducibility.
here provide detailed description of our complex SSM variant: The complex variant of our 2-D SSM model, still assumes Vt, U4.j E I. Vt, Vi.j E RR and employs:
A1, Ag, A1 E NNN, BB, B2 Nxl ,C,, C2 ClaN
(30)
(diagonal matrices as above), and therefore
2 E CN 22 e CN
(31)
The output remains real number Big E R2, and thus the real part of Eq. 2 is used, namely y Re(M.
For complex SSM, we save Aiangter Almasiss E NNN and we calculate A E NNN in the following manner:
A Atante 2Tsiggooid AAionate) Abradems = sigmoid( A
(32)
Airadius (cos (Alangte) i sin siinaaa NNN
(33)
The same goes for B1, B2. As for C1, C2. perform the same operation without limiting the radius size (not applying a sigmoid to Crradio...
E.2 No WEIGHT DECAY ON THE SSM CORE
While vision transformers and MEGA regularize the model via weight decay, SSM-based layers typically do not apply this (Guptal 2022 Gu et al. 2021a), since drastically lowers the models' ability to learn, especially in the context of long-range dependencies. In general, higher values of the A,, B, C parameters do not seem to correspond with overfitting. Therefore, our method does not employ weight decay those parameters.
15

---

Published as a conference paper at ICLR 2024
F EXPERIMENTAL SETUP
We use PyTorch for all experiments. As a deliberate decision we choose to not perform hyper- parameter tuning of the backbone and training procedure, apart from stochastic depth. All experiment results were averaged over seeds = [0. 1, 2]. For all datasets and backbones, we set Nlssm = 8, = 16 for all -SM-Based variants (SSM-2D real & complex and S4ND).
Cifar-100 and Tiny imagenet For both datasets, we use baseline the experiments performed by (Lee et al., 2021). This means we follow DeiT's (Touvron al. 2021 application of a long list of data augmentation and regularization methods, including Cutmix Yunn et al. 2019), Mixup (Zhang et al. 2017), stochastic depth (Huang et al. 2016), repeated augmentation (Hoffer et al.) 2020), Rand-Auge Augment (Cubuk et al, 2020), and random erasing (Zhong et al. 2020). Adam W was used the optimizer. Weight decay was set to 0.05 (apart from SSM layer where it was set to 0), batch size to 128, and warm-up to 10. All models were trained for 100 epochs, and cosine learning rate 0.003. In certain scenarios, models converge faster compared to the baseline approach. We discovered that slight modification, specifically doubling the stochastic depth, proved to be instrumental in maximizing the model's performance.
decay used. The initial learning rate
noticed that
When comparing S4ND, we used the same parameter scheme being used in 2-D SSM perform valid comparison, by making C E C" instead of C' E "H as in the original paper.
CelebA For Celeb-A, the original image size 178x218. it is resized to 224x224 match DeiT (Tou vron et al. 2021) backbone and patch size. The dataset includes a 40-way multi-label attribute classi- fication. We are reporting an average accuracy of all 40 tasks. We use the same data augmentation, and hyperparameters as DeiT, and train the models for 20 epochs, similar to the training procedure of S4ND (gguyen al. 2022) on this datasets.
Imagenet and CIFAR-10 Grayscale We use the exact same training procedure including hyper- parameters, data augmentation and training environment as used in the git repository of the base- line (Ma et al. 2022) for those datasets.
G CONVNEXT VARIANTS
The ConvNeXt model marks a significant advancement in CNN architectures, building upon the foundation of classic CNNs such ResNets and inspired from the progress in ViT. incorporates three main areas of enhancement: (i) micro design changes, including the adoption of improved activation functions and normalization methods. (ii) macro design adjustments, such as implementing depthwise convolutions with larger kernels and modern stage compute ratio, and (iii) modern training regime that extends the training process, employs the Adam optimizer, and incorporates advanced augmentation and regularization strategies like Stochastic Depth and Label Smoothing.
ConvNeXt Variants In the original ConvNeXt paper Liu et al. (2022). five model sizes were proposed:
1. ConvNeXt-T (Tiny): é= (96, 192, 384, 768), B (3, (,,3,
2. ConvNeXt-S (Small): C = (96, 192, 384, 768), B 3,, 3, 27,
3. ConvNexi-B (Base): Ö (128, 256, 512, 1024), B = (3, 3, 27, 3)
4. ConvNeXt-L (Large): Ö (192, 384, 768, 1536), B = (3, 2727,3)
5. ConvNeXI-XL (Extra Large): Ö (256, 512, 1024, 2048), á 3,,73)
where Ò represents the number of channels and B the number of blocks in each stage. Later, for low-resource environments and small dataset applications, two additional variants were proposed:
1. Co (Micro): Č (64, 128, 256, 512), 3 = (3, 3,
2. ConvNeXt-N (Nano): Ć (664, 128, 256, 512)), B 2,
16

---

Published as a conference paper at ICLR 2024
The ConyNeXt-Miero variant was proposed in S4NDNguwe al. (2022) for the Celeb-A A dataset, and the Nano variant, which introduced in this paper, is taken from this link and was designed for the CIFAR-10 dataset.
Experimental Setup for ConyNeXt For experiments with ConvNeXt-M we employ the exact hyperparameters used in Nguyen al. (2022) In the case of ConyNeXt-T, we followed the hyperpa- rameter specifications from the original ConvNeXt paper, with an exception for the number of training epochs in the Celeb-A experiments, which were adjusted to align with the baselines reported in the S4ND paper (20 epochs), ensuring a balanced comparison. Regarding Con-NeXI-N, our approach mirrored the training protocol found at at which involves 100-epoch training period. We applied this scheme to CIFAR-100 and of CIFAR-101 10, but maintained the same hyperparameters.
datasets, diverging from the original experiment' ' use
H INFERENCE PERFORMANCE ANALYSIS
To provide a comprehensive analysis of the 2-D SSM layer, we empirically examined both its latency and memory footprint.
Inference Profiling
To assess the efficiency integrating our model into a Vision Transformer (ViT) backbone, we analyzed the inference latency. Fig. 6 displays pie chart profiling the latency within single block, encompassing elements like the attention mechanism, 2-D SSM layer, normalization layers, residuals, and feed-forward networks (FFNs). Notably, our 2-1) SSM layer accounts for an additional 20.5%.
We further compared the execution time and memory usage of our layer during inference with the Self-Attention mechanism. As illustrated Fig. 7 we evaluated different 2-D input dimensions. Our layer introduces less overhead overall, and notably, as the 2-D sequence length extends (especially beyond 100), the performance of the SSM-2D layer appears to converge more favorably compared to self-attention.
Fig. 8 presents the memory consumption data in logarithmic scale. Here, the same trend is observed: the attention mechanism's memory usage increases quadratically relative the 2-D SSM layer the sequence length (L) grows. For example, for image sized 130x 130x,00 the attention mechanism uses 2c4MB while 2-D SSM only needs 5e2MB.
I OTHER RELATED WORK
Vision Transformer Backbones To demonstrate the versatility and efficacy of our 2-D SSM layer plug-and-play component compatible with various ViTs, evaluate its performance when integrated into the following backbone architectures: (i) VIT The original ViT that employs self- attention 1-D sequence of patches; it used learnable 1-D position encoding. (fi) Swin The Swin Transformer (Liu et al.] 2021) refines VIT by incorporating hierarchical structure and local connections within windows. employs a shifted windowing scheme and stage-wise processing to efficiently capture global context, which enhances its performance across various vision tasks. (iii) Mega The Mega (Ma al. 2022) model introduced single-head gated attention mechanism enhanced with an exponential moving average, which imparts position- -aware local dependencies to the attention mechanism. This approach addresses the limitations of the Transformer's attention mechanism, such as weak inductive bias. Mega demonstrates superior performance across a variety of tasks, including the Long Range Arena, neural machine translation, language modeling, and image and speech classification. (iv) DeiT (Touvron et al. 2021 is adaptation of ViT, incorporating class-token designed for distillation purposes.
The ConvNeXT Liu et al. (2022) Backbone is the state-of-the-art CNN architecture, which offers a viable alternative to the Transformer-ban based models while maintaining competitive accuracy and computational efficiency. In experiments, convolutional layers with our 2-D SSM layer to evaluate our method in a cutting-edgi setting. Additional details about the Nent model and the variants use are provided in Appendix G
maintain the architecture and replace the
17

---

Published as a conference paper at ICLR 2024
Residual
FFN
2.6%%12.5%
2D-SSM-Wrapper
Attention 49.8%
20.5%
3.9%
2D-SSM-Core
LayerNorm
Figure 6: Pie chart depicting the inference time distribution for each component in VIT backbone with an integrated 2-D SSM layer. 2-D-SSM-Core represents operations directly associated with the layer, while AD-SSM--rapppe encompasses surrounding operations like normalization and skip connections.
Average Execution Time vs. Sequence Length (5000 Experiments)
2-0 SSM
Self-Attention
fvrraae Time (ms)
0.29
60
100
120
Sequence Length
Figure 7: Comparison of Inference Time: Self-Attentinn mechanism vs. 2-D SSM, analyzed across various image sizes (L denotes one axis of the image, meaning image size is LxL).
Adding positional bias into transformers By design, Vision Transformers are permutation invari- ant, and thus a lot of work was put into injecting bias into them. Besides the standard positional encoding, the following methods are proposed:
Exponential Moving Average (EMA) The EMA is common technique for smoothing time-series data and prioritizing recent data points. It is computed using EMA, (1 -a) EMAA-1 + 4, where F.M. A, is the current EMA value, E.M F.M. A is the previous value and is the smoothing
18

---

Published as a conference paper at ICLR 2024
Peak Memory Usage 2-D Transformer Layers
10
-x- Self-Attention
10"
Pesk Memory (MB)
10a
10a
40
60
80
100
120
Figure 8: Comparison of memory consumption: Self-Attention mechanism vs. 2-D SSM, analyzed across various image sizes (I. denotes axis of the image, meaning image size is Lxl.), with memory usages (in mega bytes) presented log-scale.
factor: It is being used in MEGA Ma et al. 2022 to incorporate positional awareness bias into the attention.
Other 2-D bias contributions By design, Vision Transformers are permutation invariant, and thus a lot work was put into injecting 2-D bias into them. A particular research direction emphasizes the introduction of positional bias via various positional encoding methods. For instance, the Swin Transformer (Liu et al.) 2021) employs a learnable bias term referred to as relative positional bias. In alternative avenues of research, efforts have been made modify the attention window through diverse techniques, such as incorporating a two-dimensional local bias by cropping (Dong et al. 2021) the attention window or integrating convolutional neural networks (CNNs) with attention
mechanisms (Dai et al. 2021) Li et al. 2022a).
19


Reviews:
[{"review_id": "1b5PSD4jnM", "reviewer": "Reviewer_CN5B", "summary": "The paper found that a 2D recurrent state space model (SSM) can be computed as convolution and proposed a SSM based layer which can be seamlessly plug into ViT.  Experiments show this new layer can improve ViT classification accuracy slightly.", "questions": "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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1/ A new SSM based layer which has sound theoretical justification (I did not check the math carefully), and can be calculated as Convolutions.\n2/ This SSM based layer can be easily plugged into ViT and good results are achieved.", "weaknesses": "Although it seems the new layer is based on sound theoretical justification, and experiment results show that it indeed works, 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. Thus it is really a stretch to claim the benefits of the new layer. More experiments will be needed to justify.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PwwrzPCADs", "reviewer": "Reviewer_ZC1E", "summary": "The submission introduces a novel layer based on a variation of the multidimensional State Space Model (SSM), aimed at enhancing 2-D inductive bias in computer vision models. The 2D-SSM layer is designed to be integrated into Vision Transformers (ViT), contributing to improved model performance across various ViT backbones and datasets, without adding substantial parameters. The authors underscore the layer’s ability to bring about strong 2-D inductive bias, highlighting its performance even in the absence of positional encoding, and showcasing its robustness through ablation studies and visualizations.", "questions": "Please refer to the weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The central innovation of this work is the 2D-SSM layer, which is grounded in Roesser’s model for multidimensional state space, and benefits from efficient parameterization, accelerated computation, and suitable normalization. The layer introduces a strong inductive bias toward 2-D neighborhood and locality, captures unrestricted controllable context, and is highly parameter-efficient, being able to express kernels of any length via just eight scalars. Through empirical evaluation, the authors demonstrate that their layer acts as a general-purpose booster for vision transformers, surpassing standard methods like positional encoding in effectively integrating positional bias, all the while maintaining efficiency in terms of parameters and computation at inference.\n\nThe work is well-grounded in control theory, with the authors providing theoretical analysis to show that their layer generalizes S4ND and exhibits greater expressiveness. The submission includes supplementary code, facilitating reproducibility and practical application of the proposed method. Overall, this work presents a significant contribution to the field of computer vision, introducing a novel layer that addresses key challenges in 2-D inductive bias and demonstrates notable performance enhancements for Vision Transformers.", "weaknesses": "1. **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.\n\n2. **Simplicity of Tasks**: The tasks used to evaluate the method are considered too simple, raising questions about whether the inductive bias introduced is specifically beneficial for such tasks. A more critical evaluation would involve assessing the performance of the baseline methods after large-scale pre-training to assess if the learned inductive bias during pretraining can obtain better performance in downstream tasks after finetuning (etc., prompt tuning). If such an approach yields better results, the practical significance of the proposed 2D-SSM becomes unclear.\n\n3. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D5jSCExaSc", "reviewer": "Reviewer_mKTe", "summary": "This paper aims to integrate spatial inductive biases into neural network architectures, such as Vision Transformers, through the application of a two-dimensional state space model (SSM). By incorporating certain assumptions, this method ensure that the computational complexity remains tractable. Experimental results demonstrate that the proposed method surpasses previous SSM baselines across a range of scenarios.", "questions": "1. Why is the `D` omitted in Equation 2? Does its inclusion empirically reduce the performance? \n2. 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The introduction of a two-dimensional State Space Model (SSM) is an intuitive approach for incorporating spatial inductive biases within neural networks.\n2. The suggested technique can be easily integrated into various neural network models.\n3. Particularly in scenarios with small size of data, the proposed method demonstrates improved performance compared to baselines that are based on SSM.\n4. I appreciate the comparison of the proposed method with S4ND in Section 4.1, as it effectively articulates the proposed method's strengths. Although a more generalized approach does not necessarily guarantee enhanced real-world expressiveness—occasionally it may even compromise the stability of the training process—the conducted experiments effectively demonstrate the proposed method's practicality.", "weaknesses": "1. One limitation of the proposed method, as highlighted in Section 6, is its computational complexity. The method approximately doubles the training time, imposing a considerable computational load. Furthermore, while the added complexity during inference may not substantially contribute to the overall computational demand, a detailed report of the actual inference times would be beneficial for a comprehensive understanding of the method's characteristics.\n2. 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.\n3. 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.\n\n.\n\n[1] Xiao, Tete, et al. \"Early convolutions help transformers see better.\" *Advances in neural information processing systems* 34 (2021): 30392-30400.\n\n[2] Chu, Xiangxiang, et al. \"Conditional positional encodings for vision transformers.\" arXiv preprint arXiv:2102.10882 (2021).\n\n[3] Wu, Kan, et al. \"Rethinking and improving relative position encoding for vision transformer.\" Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021.\n\n[4] Liu, Ze, et al. \"Swin transformer v2: Scaling up capacity and resolution.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n\n[5] Chu, Xiangxiang, et al. \"Twins: Revisiting the design of spatial attention in vision transformers.\" Advances in Neural Information Processing Systems 34 (2021): 9355-9366.\n\n[6] Gruver, Nate, et al. \"The lie derivative for measuring learned equivariance.\" arXiv preprint arXiv:2210.02984 (2022).\n\n---\n\n**Post rebuttal**\n\nThank you for your efforts in conducting additional experiments. These further results have significantly improved the manuscript and mitigated many of its technical drawbacks. However, my concerns about the limited impact of this research persist. Therefore, I will retain my rating of 'Weak Accept'.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tu5xNhiWqJ", "reviewer": "Reviewer_SZ7R", "summary": "The paper focuses on developing new method for injecting 2-D inductive bias into Vision Transformer for computer vision problems. To achieve this, the authors propose to leverage an expressive variation of the multidimensional State Space Model (SSM) with the proposed efficient parameterization, accelerated computation and suitable normalization scheme. The paper show that by incorporating the proposed layer at the beginning of each transformer block of ViT improves the performance of various ViT backbones, such as Mega, for various datasets for image classification. It is also shown that the method achieves effective results without positional encoding.", "questions": "* 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? \n\n* It is great to see the analysis of the inference cost. I would recommend to also give comparisons of the inference time in Table 1. It is also suggested to include comparisons of the memory cost in Table 1 and the details about the platform for experiments.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The paper proposes a new method for encoding image-specific inductive bias for ViT in Computer vision problems.\n\n* The proposed method is shown to be effective for Image Classification in various datasets with various ViT backbones with negligible amount of additional parameters and inference time. \n\n* The method can achieve good performance without positional encoding.", "weaknesses": "* 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.\n\n* 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.\n\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)