S4ND: Modeling Images and Videos as Multidimensional Signals Using State Spaces
Abstract
Visual data such as images and videos are typically modeled as discretizations of inherently continuous, multidimensional signals.
Existing continuous-signal models attempt to exploit this fact by modeling the underlying signals of visual (e.g., image) data directly.
However, these models have not yet been able to achieve competitive performance on practical vision tasks such as large-scale image and video classification.
Building on a recent line of work on deep state space models (SSMs), we propose S4ND, a new multidimensional SSM layer that extends the continuous-signal modeling ability of SSMs to multidimensional data including images and videos.
We show that S4ND can model large-scale visual data in 1 D, 2 D, and 3 D as continuous multidimensional signals and demonstrates strong performance by simply swapping Conv2D and self-attention layers with S4ND layers in existing state-of-the-art models.
On ImageNet-1k, S4ND exceeds the performance of a Vision Transformer baseline by 1 . 5% when training with a 1 D sequence of patches, and matches ConvNeXt when modeling images in 2 D. For videos, S4ND improves on an inflated 3 D ConvNeXt in activity classification on HMDB-51 by 4% .
S4ND implicitly learns global, continuous convolutional kernels that are resolution invariant by construction, providing an inductive bias that enables generalization across multiple resolutions.
By developing a simple bandlimiting modification to S4 to overcome aliasing, S4ND achieves strong zeroshot (unseen at training time) resolution performance, outperforming a
baseline Conv2D by 40% on CIFAR-10 when trained on 8 × 8 and tested on 32 ×images.
When trained with progressive resizing, S4ND comes within ∼ 1% of a high-resolution model while training 22% faster
Introduction
Modeling visual data such as images and videos is a canonical problem in deep learning.
In the last few years, many modern deep learning backbones that achieve strong performance on benchmarks like ImageNet [53] have been proposed.
These backbones are diverse, and include 1 D sequence models such as the Vision Transformer (ViT) [13], which treats images as sequences of patches, and 2 D and 3 D models that use local convolutions over images and videos (ConvNets) [37, 25, 56, 58, 60, 42, 24, 32, 49, 64, 15].
A commonality among modern vision models capable of achieving state-of-the-art (SotA) performance is that they treat visual data as discrete pixels rather than continuous-signals. However, images ∗ Equal contribution.
[Figure] [Figure] ( 1/∆ 002. Discretize $ and videos are discretizations of multidimensional and naturally continuous signals, sampled at a fixed rate in the spatial and temporal dimensions.
Ideally, we would want approaches that are capable of recognizing this distinction between data and signal, and directly model the underlying continuous-signals.
This would give them capabilities like the ability to adapt the model to data sampled at different resolutions.
A natural approach to building such models is to parameterize and learn continuous convolutional kernels, which can then be sampled differently for data at different resolutions [16, 50, 54, 20, 21].
Among these, deep state space models (SSM) [20], in particular S4 [21], have achieved SotA results in modeling sequence data derived from continuous-signals, such as audio [17].
However, a key limitation of SSMs is that they were developed for 1 D signals, and cannot directly be applied to visual data derived from multidimensional “ND” signals.
Given that 1 D SSMs outperform other continuous modeling solutions for sequence data [21], and have had preliminary success on image [21] and video classification [31], we hypothesize that they may be well suited to modeling visual data when appropriately generalized to the setting of multidimensional signals.
Our main contribution is S4ND, a new deep learning layer that extends S4 to multidimensional signals.
The key idea is to turn the standard SSM (a 1 D ODE) into a multidimensional PDE governed by an independent SSM per dimension.
By adding additional structure to this ND SSM, we show that it is equivalent to an ND continuous convolution that can be factored into a separate 1 D SSM convolution per dimension.
This results in a model that is efficient and easy to implement, using the standard 1 D S4 layer as a black box.
Furthermore, it can be controlled by S4’s parameterization, allowing it to model both long-range dependencies, or finite windows with a learnable window size that generalize conventional local convolutions [22].
We show that S4ND can be used as a drop-in replacement in strong modern vision architectures while matching or improving performance in 1 D, 2 D, and 3 D. With minimal change to the training procedure, replacing the self-attention in ViT with S4-1D improves top- 1 accuracy by 1 . 5% , and replacing the convolution layers in a 2 D ConvNeXt backbone [42] with S4-2D preserves its performance on ImageNet-1k [10].
Simply inflating (temporally) this pretrained S4-2D-ConvNeXt backbone to 3 D improves video activity classification results on HMDB-51 [38] by 4 points over the pretrained ConvNeXt baseline.
Notably, we use S4ND as global kernels that span the entire input shape, which enable it to have global context (both spatially and temporally) in every layer of a network.
Additionally, we propose a low-pass bandlimiting modification to S4 that encourages the learned convolutional kernels to be smooth.
While S4ND can be used at any resolution, performance suffers when moving between resolutions due to aliasing artifacts in the kernel, an issue also noted by prior work on continuous models [50].
While S4 was capable of transferring between different resolutions on audio data [21], visual data presents a greater challenge due to the scale-invariant properties of images in space and time [52], as sampled images with more distant objects are more likely to contain power at frequencies above the Nyquist cutoff frequency.
Motivated by this, we propose a simple criteria that masks out frequencies in the S4ND kernel that lie above the Nyquist cutoff frequency.
The continuous-signal modeling capabilities of S4ND open the door to new training recipes, such as the ability to train and test at different resolutions.
On the standard CIFAR-10 [36] and Celeb-A [43] datasets, S4ND degrades by as little as 1 . 3% when upsampling from low- to high-resolution data (e.g. 128 × 128 → 160 × 160 ), and can be used to facilitate progressive resizing to speed up training by 22% with ∼ 1% drop in final accuracy compared to training at the high resolution alone.
We also validate that our new bandlimiting method is critical to these capabilities, with ablations showing absolute performance degradation of up to 20%+ without it.
Related Work
Image Classification.
There is a long line of work in image classification, with much of the 2010s dominated by ConvNet backbones [37, 25, 56, 58, 60].
Recently, Transformer backbones, such as
ViT [13], have achieved SotA performance on images using self-attention over a sequence of 1 D patches [40, 39, 62, 71, 12].
Their scaling behavior in both model and dataset training size is believed to give them an inherent advantage over ConvNets [13], even with minimal inductive bias.
Liu et al. [42] introduce ConvNeXt, which modernizes the standard ResNet architecture [25] using modern training techniques, matching the performance of Transformers on image classification.
We select a backbone in the 1 D and 2 D settings, ViT and ConvNeXt, to convert into continuous-signal models by replacing the multi-headed self-attention layers in ViT and the standard Conv2D layers in ConvNeXt with S4ND layers, maintaining their top-1 accuracy on large-scale image classification.
S4 & Video Classification.
To handle the long-range dependancies inherent in videos, [31] used 1D S4 for video classification on the Long-form Video Understanding dataset [67].
They first applied a Transformer to each frame to obtain a sequence of patch embeddings for each video frame independently, followed by a standard 1D S4 to model across the concatenated sequence of patches.
This is akin to previous methods that learned spatial and temporal information separately [33], for example using ConvNets on single frames, followed by an LSTM [27] to aggregate temporal information.
In contrast, modern video architectures such as 3 D ConvNets and Transformers [24, 32, 49, 64, 15, 35, 41, 67, 2, 1] show stronger results when learning spatiotemporal features simultaneously, which the generalization of S4ND into multidimensions now enables us to do.
Continuous-signal Models.
Visual data are discretizations of naturally continuous signals that possess extensive structure in the joint distribution of spatial frequencies, including the properties of scale and translation invariance.
For example, an object in an image generates correlations between lower and higher frequencies that arises in part from phase alignment at edges [47].
As an object changes distances in the image, these correlations remain the same but the frequencies shift.
This relationship can potentially be learned from a coarsely sampled image and then applied at higher frequency at higher resolution.
A number of continuous-signal models have been proposed for the visual domain to learn these inductive biases, and have led to additional desirable properties and capabilities.
A classic example of continuous-signal driven processing is the fast Fourier transform, which is routinely used for filtering and data consistency in computational and medical imaging [11].
NeRF represents a static scene as a continuous function, allowing them to render scenes smoothly from multiple viewpoints [45].
CKConv [51] learns a continuous representation to create kernels of arbitrary size for several data types including images, with additional benefits such as the ability to handle irregularly sampled data.
FlexConv [50] extends this work with a learned kernel size, and show that images can be trained at low resolution and tested at high resolution if the aliasing problem is addressed.
S4 [21] increased abilities to model long-range dependancies using continuous kernels, allowing SSMs to achieve SotA on sequential CIFAR [36].
However, these methods including 1D S4 have been applied to relatively low dimensional data, e.g., time series, and small image datasets.
S4ND is the first continuous-signal model applied to high dimensional visual data with the ability to maintain SotA performance on large-scale image and video classification.
Progressive Resizing.
Training times for large-scale image classification can be quite long, a trend that is exacerbated by the emergence of foundation models [4].
A number of strategies have emerged for reducing overall training time.
Fix-Res [61] trains entirely at a lower resolution, and then fine-tunes at the higher test resolution to speed up training in a two-stage process.
Mix-andMatch [28] randomly samples low and high resolutions during training in an interleaved manner.
An effective method to reduce training time on images is to utilize progressive resizing.
This involves training at a lower resolution and gradually upsampling in stages.
For instance, fastai [14] utilized progressive resizing to train an ImageNet in under 4 hours.
EfficientNetV2 [60] coupled resizing with a progressively regularization schedule, increasing the regularization as well to maintain accuracy.
In
EfficientNetV2 and other described approaches, the models eventually train on the final test resolution.
As a continuous-signal model, we demonstrate that S4ND is naturally suited to progressive resizing, while being able to generalize to unseen resolutions at test time.
Preliminaries
State space models.
S4 investigated state space models, which are linear time-invariant systems that map signals u ( t ) (cid:55)→ y ( t ) and can be represented either as a linear ODE (equation (1)) or convolution (equation (2)).
Its parameters are A ∈ C N × N and B , C ∈ C N for a state size N .
[Equation]
Basis functions.
For the clearest intuition, we think of the convolution kernel as a linear combination (controlled by C ) of basis kernels K n ( t ) (controlled by A , B )
[Equation]
Discretization.
The SSM (1) is defined over a continuous-time axis and produces continuous-time convolution kernels (2)(3).
Given a discrete input sequence u 0 , u 1 , . . .
sampled uniformly from an underlying signal u ( t ) at a step size ∆ (i.e. u k = u ( k ∆) ), the kernel can be sampled to match the rate of the input.
Note that instead of directly sampling the kernel, standard discretization rules should be applied to minimize the error from the discrete to the continuous-time kernel [21].
For inputs given at different resolutions, the model can then simply change its ∆ value to compute the kernel at different resolutions.
We note that the step size ∆ does not have to be exactly equal to a “true sampling rate” of the underlying signal, but only the relative rate matters.
Concretely, the discrete-time kernel depends only on the product ∆ A and ∆ B , and S4 learns separate parameters ∆ , A , B .
S4.
S4 is a special SSM with prescribed ( A , B ) matrices that define well-behaved basis functions, and an algorithm that allows the convolution kernel to be computed efficiently.
Variants of S exist that define different basis functions, such as simple diagonal SSMs [23], or one that defines truncated Fourier functions K n ( t ) =
sin(2 πnt ) I ([0 , 1]) [22] (Fig. 1).
These versions of S have easy-to-interpret basis functions that will allow us to control the frequencies in the kernel (Section 4.2).
Method
We describe the proposed S4ND model for the 2D case only, for ease of notation and presentation.
The results extend readily to general dimensions; full statements and proofs for the general case are in
Appendix A. Section 4.1 describes the multidimensional S4ND layer, and Section 4.2 describes our simple modification to restrict frequencies in the kernels.
Fig.
1 illustrates the complete S4ND layer.
4.1 S4ND
We begin by generalizing the (linear time-invariant) SSM (1) to higher dimensions.
Notationally, we denote the individual time axes with superscripts in parentheses.
Let u =
u
( t (1) , t (2) ) and y =
y ( t (1) , t (2) ) be the input and output which are signals R 2 → C , and x = ( x (1) ( t (1) , t (2) ) , x (2) ( t (1) , t (2) ))
∈ C N (1) × N (2) be the SSM state of dimension N (1) × N (2) , where x ( τ ) : R 2 → C N ( τ ) .
Definition 1 (Multidimensional SSM) .
Given parameters A ( τ ) ∈ C N ( τ ) × N ( τ ) , B ( τ ) ∈
C N ( τ ) × 1 ,
C ∈ C N (1) × N (2) , the 2D SSM is the map u
(cid:55)→ y defined by the linear PDE with initial condition x (0 , 0) = 0 :
[Equation]
Note that Definition 1 differs from the usual notion of multidimensional SSM, which is simply a map from u ( t ) ∈
C n (cid:55)→ y ( t ) ∈
C m for higher-dimensional n, m > 1 but still with 1 time axis.
However, Definition 1 is a map from u ( t 1 , t 2 ) ∈ C 1 (cid:55)→ y ( t 1 , t 2 ) ∈ C 1 for scalar input/outputs but over multiple time axes.
When thinking of the input u ( t (1) , t (2) ) as a function over a 2D grid,
Definition 1 can be thought of as a simple linear PDE that just runs a standard 1D SSM over each axis independently.
Analogous to equation (2), the 2D SSM can also be viewed as a multidimensional convolution.
Theorem 1.
(4) is a time-invariant system that is equivalent to a 2D convolution y
=
K ∗ u by the kernel
[Equation]
This kernel is a linear combination of the N (1) × N (2) basis kernels { K (1) n (1) ( t (1) ) ⊗ K (1) n (2) ( t (2) ) :
n (1) ∈
[ N (1) ] , n (2) ∈
[ N (2) ] } where K ( τ ) are the standard 1D SSM kernels (3) for each axis .
However, a limitation of this general form is that the number of basis functions N (1) × N (2) × . . .
grows exponentially in the dimension, increasing the parameter count (of C ) and overall computation dramatically.
This can be mitigated by factoring C as a low-rank tensor.
Corollary 4.1.
Suppose that C ∈ C N (1) × N (2) is a low-rank tensor C
= (cid:80) ri =1 C (1) i ⊗ C (2) i where each C ( τ ) i ∈ C N ( τ ) .
Then the kernel (5) also factors as a tensor product of 1D kernels
K
[Equation]
t [Equation] In our experiments, we choose C as a rank-1 tensor, but the rank can be freely adjusted to tradeoff parameters and computation for expressivity.
Using the equivalence between (1) and (2), Corollary 4.1 also has the simple interpretion as defining an independent 1D SSM along each axis of the multidimensional input.
4.2 Resolution Change and Bandlimiting
SSMs in 1D have shown strong performance in the audio domain, and can nearly preserve full accuracy when tested zero-shot on inputs sampled at very different frequencies [21].
This capability relies simply on scaling ∆ by the relative change in frequencies (i.e., if the input resolution is doubled, halve the SSM’s ∆ parameter).
However, sampling rates in the spatial domain are often much lower than temporally, leading to potential aliasing when changing resolutions.
A standard technique to avoid aliasing is to apply a low-pass filter to remove frequencies above the Nyquist cutoff frequency.
For example, when A is diagonal with n -th element a n , each basis function has simple form
K n ( t ) =
e t a n B n .
Note that the frequencies are mainly controlled by the imaginary part of a n .
We propose the following simple method: for any n such that a n · ∆ < 12 α , mask out the corresponding coefficient of the linear combination C n (equation (3)).
Here α is a hyperparameter that controls the cutoff; theoretically, α = 1 . 0 corresponds to the Nyquist cutoff if the basis functions are pure sinusoids.
However, due to the decay e ℜ ( a n ) arising from the real part as well as approximations arising from using finite-state SSMs, α often has to be set lower empirically.
Experiments
We evaluate S4ND on large-scale image classification in Section 5.1 in the 1 D and 2 D settings, followed by activity classification in videos in Section 5.2 in the 3 D setting, where using S4ND as a drop-in replacement for standard deep learning layers matches or improves performance in all settings.
In Section 5.3, we performed controlled ablations to highlight the benefits of S4ND as a continuous-signal model in images [Table]
5.1 S4ND in 1D & 2D: Large-scale Image Classification
First
[Table]
we show that S4ND is a drop-in replacement for existing visual modeling layers such as 1 D self-attention and 2 D local convolutions, with no degradation in top- 1 performance when used in modern backbones such as ViT [13] and ConvNeXt [42] on ImageNet-1k [10] [Table] [Table] [Table]
[Table] [Table] [Table] Baselines and Methodology.
We consider large-scale image classification on the ImageNet-1k dataset, which consists of 1000 classes and 1 .
3 M images.
We start with two strong baselines: ViT-B ( base , 88 M parameters) for processing images in the 1 D setting and ConvNeXt-T ( tiny , 28 . 4 M) in the 2 D setting.
(We omit the postfix “B” and “T” for brevity).
More recent works using Transformers on images have surpassed ViT, but we focus on the original ViT model to highlight specifically the drop-in capability and performance difference in self-attention vs.
S4ND layers. We first swap the selfattention layers in ViT with S4ND layers, and call this model S4ND-ViT.
Notably, we simplify ViT by removing the positional encodings, as S4ND does not require injecting this inductive bias.
Similarly, we swap the local Conv2D layers in the ConvNeXt blocks with S4ND layers, which we call S4ND-ConvNeXt, a model with global context at each layer.
Both S4ND variants result in similar parameter counts compared to their baseline models.
Training.
For all ImageNet models, we train from scratch with no outside data and adopt the training procedure from [62, 69], which uses the AdamW optimizer [44] for 300 epochs, cosine decay learning rate, weight decay 0.05, and aggressive data augmentations including RandAugment [7],
Mixup [72], and AugMix [26].
We add RepeatAug [29] for ConvNeXt and S4ND-ConvNeXt.
The initial learning rate for ViT (and S4ND-ViT) is 0 .
001 , while for ConvNeXt (and S4ND-ConvNeXt) it is 0 . 004 .
See Appendix B.1 for additional training procedure details.
Results.
Table 1 shows top- 1 accuracy results for each model on ImageNet.
After reproducing the baselines, S4ND-ViT was able to moderately boost performance by + 1 . 5 % over ViT, while S4ND-ConvNeXt matched the original ConvNeXt’s performance.
This indicates that S4ND is a strong primitive that can replace self-attention and standard 2 D convolutions in practical image settings with large-scale data.
5.2 S4ND in 3D: Video Classification
Next, we demonstrate the flexible capabilities of S4ND in settings involving pretraining and even higher-dimensional signals.
We use the activity recognition dataset HMDB-51 [38] which involves classifying videos in 51 activity classes.
Baselines and Methodology.
Prior work demonstrated that 2 D CNNs (e.g. pretrained on ImageNet) can be adapted to 3 D models by 2 D to 3 D kernel inflation (I3D [5]), in which the 2D kernels are repeated temporally N times and normalized by 1 /N .
Our baseline, which we call ConvNeXt-I3D, uses the 2 D ConvNeXt pretrained on ImageNet (Section 5.1) with I3D inflation. Notably, utilizing
S4ND in 3 D enables global context temporally as well. We additionally test the more modern spatial-temporal separated 3 D convolution used by S3D [68] and R(2+1)D [65], which factor the 3 D convolution kernel as the outer product of a 2 D (spatial) by 1 D (temporal) kernel.
Because of its flexible factored form (Section 4.1), S4ND automatically has these inflation capabilities.
We inflate the pretrained S4ND-ConvNeXt simply by loading the pretrained 2 D model weights for the spatial dimensions, and initializing the temporal kernel parameters A (3) , B (3) , C (3) from scratch.
We note that this model is essentially identical to the baseline ConvNeXt-S3D except that each component of the factored kernels use standard 1 D S4 layers instead of 1D local convolutions.
Finally, by varying the initialization of these parameters, we can investigate additional factors affecting model training; in particular, we also run an ablation on the kernel timescales ∆ .
Training.
Our training procedure is minimal, using only RGB frames (no optical flow).
We sample clips of 2 seconds with 30 total frames at 224 × 224 , followed by applying RandAugment; we performed a small sweep of the RandAugment magnitude for each model.
All models are trained with learning rate 0 . 0001 and weight decay 0 .
Additional details are included in Appendix B.2. [Table]
Results.
Results are presented in Table 2.
Our baselines are much stronger than prior work in this setting, 8% top- 1 accuracy higher than the original I3D model in the RGB frames only setting, and confirming that separable kernels (S3D) perform at least as well as standard inflation ( +0 . 53% ).
S4ND-ConvNeXt-3D improves over the baseline ConvNeXt-I3D by +4 .
0% with no difference in models other than using a temporal S4 kernel.
This even exceeds the performance of I3D when trained on optical flow.
Finally, we show how S4ND’s parameters can control for factors such as the kernel length (Table 2).
Note that our temporal kernels K 3 ( t 3 ) are always full length (30 frames in this case), while standard convolution kernels are shorter temporally and require setting the width of each layer manually as a hyperparameter [68].
S4 layers have a parameter ∆ that can be interpreted such that 1∆ is the expected length of the kernel.
By simply adjusting this hyperparameter, S4ND can be essentially initialized with length- 1 temporal kernels that can automatically learn to cover the whole temporal length if needed.
We hypothesize that this contributes to S4ND’s improved performance over baselines. [Table] [Table]
5.3 Continuous-signal Capabilities for Images
Images are often collected at varied resolutions, even with the same hardware, so it is desirable that models generalize to data sampled at different resolutions.
We show that S4ND inherits this capability as a continuous-signal model, with strong zero-shot performance when changing resolutions, and the ability to train with progressively resized multi-resolution data.
We perform an ablation to show our proposed bandlimiting modification is critical to achieving strong performance when changing resolutions.
Setup.
We focus on image classification on 2 D benchmark datasets: a dataset with lowresolution images (CIFAR-10) and one with higher-resolution images (Celeb-A).
For each dataset, we specify a base image resolution ( base ), and two lower resolutions ( mid , low ), summarized in Table 3. To highlight that
S4ND’s continuous capabilities are independent of backbone, we experiment with [-
 [Table]-]2 different 2 D backbones: an isotropic, fixed-width model backbone on CIFAR-10 [36] and a small ConvNeXt backbone on Celeb-A [43].
For each backbone, we compare S4ND’s performance to Conv2D layers as a standard, widely used baseline.
Additional details can be found in Appendix B.3.
We first verify that S4ND achieves comparable test classification performance with baseline Conv2D models.
Table 4 and Fig. 3 show that we exceed the performance of Conv2D on both tasks, with S4ND improving over Conv2D models by 0 .
4% on CIFAR-10 and 0 .
3% on Celeb-A.
Zero-Shot Resolution Change.
We train S4ND and Conv2D models at either low or mid resolution, and test them at the base resolution for each dataset.
We also compare to FlexConv [50] on CIFAR10, which is the current SotA for zero-shot resolution change.
Compared to training and testing at the base resolution, we expect that Conv2D should degrade more strongly than S4ND, since it cannot adapt its kernel appropriately to the changed resolution.
Table 4 and Fig.
3 show that S4ND outperforms Conv2D on mid → base by 15+ points and low → base by 40+ points on CIFAR-10, and 9+ points and 3+ points on Celeb-A.
In fact, S4ND yields better performance on low → base (a more difficult task) than Conv2D does from mid → base (an easier task), improving by 1 . 1% on
CIFAR-10 and 3% on Celeb-A. Compared to FlexConv on CIFAR-10 mid → base , S4ND improves zero-shot performance by 5+ points, setting a new SotA.
Progressively Resized Training.
We provide an exploration of training with progressive resizing [14, 60] i.e. training in multiple stages at different resolutions.
The only change we make from standard training is to reset the learning rate scheduler at the beginning of each stage (details in Appendix B.3).
We compare S4ND and the Conv2D baseline with progressive resizing in Table 5.
For CIFAR-10, we train with a low → base , 80 − 20 epoch schedule, and perform within ∼ 1% of an S4ND model trained with base resolution data while speeding up training by 21 . 8% .
We note that
Conv2D attains much higher speedups as a consequence of highly optimized implementations, which we discuss in more detail in Section 6.
For Celeb-A, we explore flexibly combining the benefits of both progressive resizing and zero-shot testing, training with a low → mid , 16 − 4 epoch schedule that uses no base data.
We outperform Conv2D by 7 . 5%+ , and attain large speedups of 50%+ over training at the base resolution. D [Table] M ODEL
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] Effect of Bandlimiting. Bandlimiting in S4ND is critical to generalization at different resolutions.
We analyze the effect of the bandlimiting parameter α on CIFAR-10 performance when doing zero- shot resolution change.
We additionally vary the choice of basis function K n ( t ) used in the SSM.
Fig.
4 shows that zero-shot performance on base degrades for larger values of α , i.e. for cutoffs that do not remove high frequencies that violate the Nyquist cutoff.
As we would expect, this holds regardless of the choice of basis function.
In Fig.
5, we visualize learned kernels with and without bandlimiting, showing that bandlimiting improves smoothness.
Appendix B.3 includes additional experiments that analyze α on Celeb-A, and when doing progressively resized training.
Discussion
Summary.
We introduced S4ND, a multidimensional extension of S4 that models visual data as continuous valued signals.
S4ND is the first continuous model that matches SotA baselines on large-scale 1 D and 2 D image classification on ImageNet-1k, as well as outperforming a strong pretrained model in a 3 D video classification setting.
As a continuous-signal model, S4ND inherits useful properties that are absent from standard visual modeling approaches, such as zero-shot testing on unseen resolutions without a significant performance drop.
Limitations.
A limitation of S4ND is its training speed in high dimensions.
In the 1 D image setting,
S4ND-ViT has similar training speed to ViT; however, in the 2 D setting, the S4ND-ConvNeXt was 2 × slower than the baseline ConvNeXt.
We remark that vanilla local convolutions have been heavily optimized for years, and we expect that layers such as S4ND can be substantially sped up with more optimized implementations.
Our core computational primitives accounting for 65% of our runtime (FFT, pointwise operations, inverse FFT) are all bottlenecked by reading from and writing to GPU memory [57].
With a more optimized implementation that fuses these operations [9] (i.e., loading the input once from GPU memory, perform all operations, then write result back to GPU memory), we expect to speed these up by 2-3 × .
Further discussion can be found in Appendix C.
Future work.
We presented a first step in using continuous-signal models in images and videos, and believe this opens the door to new capabilities and directions.
For example, recent video benchmark datasets are significantly larger than the HMDB-51 dataset used in our experiments [34, 18, 8, 46, 55], and can be explored.
In addition, we demonstrated capabilities in zero-shot resolution spatially , but even less work has been done on zero-shot testing temporally .
This opens up an exciting new direction of work that could allow models to be agnostic to different video sampling rates as well, capable of testing on higher unseen sampling rates, or irregular (non-uniform) sampling rates.
As models become larger and combine multiple modalities, S4ND has shown strong promise in being able to better model underlying continuous-signals and create new capabilities across visual data, audio, time-series (e.g., wearable sensors, health data), and beyond. Acknowledgments
We thank Arjun Desai, Gautam Machiraju, Khaled Saab and Vishnu Sarukkai for helpful feedback on earlier drafts.
This work was done with the support from the HAI-Google Cloud Credits Grant
{+Program. We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF
+}
(Beyond Sparsity), CCF1563078 (Volume to Velocity), and(RTML);
{+ONR under No. N
+}
(Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Salesforce, Total, the HAI-AWS Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare.
The Mobilize Center is a Biomedical Technology Resource Center, funded by the NIH National Institute of Biomedical Imaging and Bioengineering through Grant P41EB027060. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] The main method can be implemented based on the released S4 repository, and we will publically release our code.
All hyperparameters and experimental details are reported in the Appendix.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Yes (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] Full experimental results including standard deviations can be found in the Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Reported in the Appendix. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
A Theory Details
We formalize and generalize the notation of Section 4.1 and prove the results.
For the remainder of this section, we fix a dimension D (e.g. D = 2 for a 2-dimensional SSM).
The
D -dimensional SSM will be a map from a function u : R D → C to y : R D → C .
Definition 2 (Indexing notation) .
Let [ d ] denote the set { 1 , 2 , . . .
, d } .
Let [0] denote the empty set {} .
Given a subset I ⊆
[ D ] , let −
I denote its complement [ D ] \ I .
Let [ − d ] = −
[ d ] = { d + 1 , . . .
, D } .
Definition 3 (Tensor products and contractions) .
Given a ∈ C N 1 and b ∈ C N 2 , let a ⊗ b ∈ C N 1 × Nbe defined as ( a ⊗ b ) n 1 ,n 2 = a n 1 b n 2 .
Given a tensor x ∈ C N 1 ×···× N D and matrix A ∈ C N 1 × N 1 , define A · (1) x ∈ C N 1 ×···× N D as ( A · (1) x ) n 1 ,...,n D = (cid:88)
[Equation]
which is simply a matrix multiplication over the first dimension.
Let · ( τ ) similarly denote matrix multiplication over any other axis.
It will be easier to work directly from the convolutional definition of SSM (equation (5)).
We will then show equivalence to a PDE formulation (equation (4)).
Definition 4 (Multidimensional SSM) .
Let N 1 , . .
. , N D be state sizes for each of the D dimensions.
The D -dimensional SSM has parameters
A ( τ ) ∈ C N ( τ ) × N ( τ ) , B ( τ ) ∈
C N ( τ ) , and
C ∈ C N (1) ×···× N ( D ) for each dimension τ ∈
[ D ] . and is defined by the map u (cid:55)→ y given by
[Equation]
Note that at all times t
=
( t (1) , . . .
, t ( D ) ) , the dimension of the state is x ( t )
∈ C N 1 ×···× N D .
Finally, it will be convenient for us to define “partial bindings” of u, x, y where one or more of the coordinates are fixed.
Definition 5 (Partial binding of u ) .
Let I ⊆
[ D ] and let t I =
( t ( I 1 ) , . . .
, t ( I | I | ) ) be an partial assignment of the time variables.
Define u
I ( t I ) to be the function R D
−| I | (cid:55)→ C where the I indices are fixed to t I .
For example, if D = 3 and I = [1] , then [Equation]
is a function from R [-
-]2 → C mapping ( t (2) , t (3) ) (cid:55)→ u ( t (1) , t (2) , t (3) ) .
Definition 6 (Partial binding of x ) .
[Equation]
Note that x [0] () : R D (cid:55)→ C N 1 ×···× N D coincides with the full state x .
The following more formal theorem shows the equivalence of this convolutional LTI system with a standard 1D SSM differential equation in each dimension.
Theorem 2.
Given a partial state x
I ( t I ) and a time variable t ( τ ) for τ
̸∈ I , let I ′
=
I ∪ { τ } .
Then the partial derivatives satisfy
[Equation]
I
[Equation]
Proof.
WLOG assume I =
[ d ] , since all notions are permutation-independent.
We will consider differentiating the state x
I with respect to the time variable t ( d +1) .
The key fact is that differentiating a convolution ddt ( f ∗ g ) is equivalent to differentiating one of the operands f ∗ ( ddt g ) .
[Equation]
The following corollary follows immediately from linearity of the convolution operator, allowing the order of convolution and inner product by C to be switched.
Corollary A.1.
The output y is equivalent to y =
K ∗ u where
[Equation]
This completes the proof of Theorem 1.
Finally, Corollary 4.1 is an immediate consequence of the Kronecker mixed-product identity
[Equation]
B Experimental Details
We use PyTorch for all experiments, and build on the publicly available S4 code.
B.1 Image Classification
ImageNet training: all models were trained from scratch with no outside data using 8 Nvidia [-A
-]A
GPUs.
For both ViT and ConvNeXt experiments, we follow the procedure from T2T-ViT [69] and the original ConvNeXt [42], respectively, with minor adjustments.
Preprocessing and dataloading was done using the TIMM [66] library.
In S4ND-ViT, we turn off weight decay and remove the class token prepending of the input sequence.
For the ConvNeXt models, we add RepeatAug [29], as well as reduce the batch size to 3840 for S4ND-ConvNeXt.
S4ND specific settings include a bidirectional S4 kernel followed by Goel et al [17], and a state dimension of 64 for the SSMs.
B.2 Video Classification
HMDB-51 training: all models were trained from scratch on a single Nvidia A100 GPU.
We use the
Pytorchvideo library for data loading and minimal preprocessing of RGB frames only (no optical
[Table]
flow).
During training
[Table]
we randomly sample 2 second clips from each video.
At validation and test time, 2 second clips are sampled uniformly to ensure that an entire video is seen.
For each ConvNeXt 3D video model (baselines and S4ND versions), we fix the hyperparameters in Table
[Table]
and do a sweep over the RandAugment[7] settings num layers = { 1 , 2 } and magnitude = { 3 , 5 , 7 , 9 } .
S4ND specific settings were similar to image classification, a bidirectional S4 kernel and a state dimension of 64 for the SSMs.
[Table]
B.3 Continuous Time Capabilities Experiments
We use the CIFAR-10 and Celeb-A datasets for all our experiments.
Below, we include all experimental details including data processing, model training and hyperparameters, and evaluation details.
As noted in Table 3, we consider 3 resolutions for each dataset, a base resolution that is considered the standard resolution for that dataset, as well as two lower resolutions low and mid .
Images at the lower resolutions are generated by taking an image at the base resolution, and then downsampling using the resize function in torchvision , with bilinear interpolation and antialiasing turned on.
CIFAR-10.
The base resolution is chosen to be 32 × 32 , which is the resolution at which models are generally trained on this dataset.
The lower resolutions are 16 × 16 and 8 × 8 .
We use no data augmentation for either zero-shot resolution change or progressive resizing experiments.
We train with standard cross-entropy loss on the task of 10 -way classification.
For reporting results, we use standard accuracy over the 10 classes.
We use a simple isotropic model backbone, identical to the one used by Gu et al.
[21] for their (1D)
S4 model.
The only difference is that we use either the S4ND or Conv2D layer in place of the S block, which ensures that the model can accept a batch of 2 D spatial inputs with channels.
For results reported in the main body, we use a 6 × 256 architecture consisting of 6 layers and a model dimension of 256 .
For all CIFAR-10 experiments, we train for 100 epochs, use a base learning rate of 0 .
01 and a weight decay of 0 .
Celeb-A. On Celeb-A, we consider 160 × 160 as the base resolution, and 128 × 128 and 64 × 64 as the lower resolutions.
Images on Celeb-A are generally 218 × 178 , and we run a CenterCrop(178) →
Resize(160) transform in torchvision to resize the image.
We use no data augmentation for either zero-shot resolution change or progressive resizing experiments.
We train with standard binary cross-entropy loss on the task of 40 -way multilabel attribute classification on Celeb-A. For reporting results, we use the multilabel binary accuracy, i.e. the average binary accuracy across all 40 tasks.
Similar to our ImageNet experiments, we use a ConvNeXt model as the basic backbone for experiments on Celeb-A.
However, use a particularly small model here, consisting of 4 stages with depths (3 , 3 , 3 , 3) and corresponding model dimensions (64 , 128 , 256 , 512) .
For S4ND, we simply follow the same process as ImageNet, replacing all depthwise Conv2D layers with S4ND.
A minor difference is that we use a global convolution in the stem downsampling for S4ND, rather than the smaller kernel size that we use for ImageNet.
We use a drop path rate of 0 .
1 for both S4ND and the
Conv2D baseline in all experiments.
For all Celeb-A experiments, we train for 20 epochs, use a base learning rate of 0 . 004 , and use automatic mixed precision for training.
Other fixed hyperparameters.
For all experiments, we use bidirectional S4 kernels following Goel et al.
[17], and use a state dimension of 64 for the S4 SSMs.
We use AdamW as the optimizer.
B.3.1 Zero-Shot Resolution Change
For Table 4, we simply train on a single resolution among base , mid and low , and then directly test at all higher resolutions.
For model selection at a particular test resolution, we select the best performing model for that test resolution i.e. we checkpoint the model at the epoch where it performs best at the test resolution in question.
Note that we refer to “test resolutions” but report validation metrics in Table 4, as is standard practice for additional experiments and ablations.
CIFAR-10.
We use a batch size of 50 , and for both methods (and all resolutions), we use a cosine decay schedule for the learning rate with no restarts and a length of 100000 , and a linear warmup of 500 steps.
For baseline Conv2D hyperparameters, we compare the performance of Conv2D with and without depthwise convolutions.
Otherwise, all hyperparameters are fixed to the common values laid out in the previous section (and are identical to those used for S4ND).
For S4ND hyperparameters, we sweep the choice of initialization for the state space parameters A, B in the S4 model among { legs, fourier, random-linear, random-inv } 2 and sweep bandlimit ( α ) values among { 0 . 05 , 0 . 10 , 0 . 20 , 0 .
50 , ∞} .
Celeb-A.
We use a batch size of 256 when training on the lower resolutions, and 128 when training at the base resolution.
For both methods (and all resolutions), we use a cosine decay schedule for the learning rate with no restarts and a length of 13000 (batch size 256 ) or 26000 (batch size 128 ), and a linear warmup of 500 steps.
For baseline ConvNeXt hyperparameters, we use a weight decay of 0 .
1 after an initial sweep over weight decay values { 0 . 05 , 0 . 1 , 0 . 2 , 0 . 5 } .
For S4ND hyperparameters, we sweep the choice of initialization for the state space parameters
A, B in the S4 model among { legs, fourier, random-linear, random-inv } , sweep bandlimit ( α ) values among { 0 . 05 , 0 . 10 , 0 . 20 , 0 .
50 , ∞} , and use a weight decay of 1 . 0 .
The value of the weight decay was chosen based on an initial exploratory sweep over weight decay values { 0 . 1 , 0 . 2 , 0 . 5 , 1 . 0 , 5 . 0 } .
B .3.2 Progressive Resizing
For Table 5, we train with a resizing schedule that progressively increases the resolution of the data.
We report all metrics on validation sets, similar to our zero-shot experiments.
We only use 2 resizing stages in our experiments, as we found in preliminary experiments that 2+ stages had little to no benefit.
All hyperparameters and training details are identical to the zero-shot resolution change experiments, except for those we describe next.
When training with multiple stages, we reset the cosine learning rate scheduler at the beginning of each stage (along with the linear warmup).
The length of the scheduler is varied in proportion to the length of the stage, e.g. on CIFAR-10, for a stage 50 epochs long, we change the length of the decay schedule from 100000 steps to 50000 steps (and similarly for
Celeb-A).
We always use a linear warmup of 500 steps for each stage.
Additionally, Table 5 denotes the length (in epochs) of each stage, as well as their training resolutions.
Another important detail is that we set the bandlimit parameter carefully for each stage.
We found that bandlimiting in the first stage is critical to final performance, and without bandlimiting the performance on the base resolution is substantially worse.
Bandlimiting is generally not useful for a stage if it is training at the base resolution i.e. at the resolution that we will be testing at (in our experiments, this is the case only for second stage training on CIFAR-10).
On CIFAR-10, we use a bandlimit α of 0 .
10 or 0 .
20 for the first stage, depending on whether the first stage was training at 8 × 8 ( low ) or 16 ×
( mid ) respectively.
In the second stage, we always train at the base resolution and use no bandlimiting, since as stated earlier, it has a negligible effect on performance.
For Celeb-A, we set the bandlimit parameter to 0 .
1 for both stages.
Note that on Celeb-A, we do not train at the base resolution at all, only training at low in the first stage and mid in the second stage.
C Discussion
Runtime profiling and potential optimization.
We profile the runtime of an S4ND layer on an A100 GPU, and plot the time taken by each operation in Fig.
We see that the majority of the time (between 65% and 80%) are taken by the FFT, pointwise operation, and inverse FFT.
The higher-dimensional FFT is a very standard scientific primitive that should be substantially optimizable.
These operations are memory-bound, that is, the runtime is dominated the time to read/write to GPU memory.
There is currently no library support for fusing these steps, so for each of those steps the data has to be loaded from GPU memory, arithmetic operations are performed, then the result is written back to GPU memory.
One potential optimization is kernel fusion : the input could be loaded once, all the arithmetic operations for all three steps are performed, then the final result is written back to GPU memory.
We expect that with library support for such optimization, the S4ND runtime can be reduced by 2-3 × .
[Figure]