NeurIPs Paper Reviews 2023 #3

Convolutional State Space Models for Long-Range Spatiotemporal Modeling

Jimmy T.H. Smith, Shalini De Mello, Jan Kautz, Scott W. Linderman, Wonmin Byeon

Linear State Space Models are a classical tool in Signal Processing and Control Theory. They also became in recent years an object of interest for Machine Learning researchers, when it was discovered that composing properly structured and initialized models of this type with nonlinear activations results in a new type of scalable and efficient sequential architecture. Such architectures, which can be thought of as sophisticated linear RNNs, have shown impressive performance on a number of classification tasks requiring processing long sequences, audio processing, and recently state of the art results in language modelling at a scale of few billion parameters. The model proposed in this paper is a type of State Space Model architecture, that has been designed for processing spatiotemporal data.

The key idea is to use the convolution operation in defining the transition, input and output operations to the state space model. This is analogous to using the convoluiton operation to define the step of ConvRNNs like ConvLSTM and ConvGRU, which are commonly used in spatiotemporal modelling. Since convolution is a linear operation, such definition results in a linear state space model with a distinct structure. The authors opt to restrict the state transition convolution to be pointwise, and integrate this design variant into previously proposed S5 layer. Furthermore, they leverage initialization schemes from previous SSM models and implement an efficient parallel scan for the proposed layer. The resulting architecture, which is named ConvS5, combines the fast stateful inference and training time linear in the sequence length of RNN models with parallelizability of transformers. The proposed model either matches or outperforms state of the art models on a number of spatiotemporal benchmarks, while also comparing favourably to them in terms of required computational resources.

Given that video is the next frontier of generative AI, I am looking forward to seeing how well generative video models with a ConvSSM backbone are going to perform.

How to Scale Your EMA

Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko, Eeshan Gunesh Dhekane, Xavier Suau, Russ Webb

Model EMA is a copy of a given machine learning model, with parameters equal to the exponential moving average along the optimization trajectory of the original model. This object is employed in Machine Learning for a variety of purposes. In Supervised Learning, using model EMA instead of last iterate of training process often improves stability and generalization. It is also commonly used in semi-supervised and self-supervised learning as the teacher model.

A performance-critical hyperparameter of model EMA is the momentum, which is a number between zero and one chosen by the practitioner. In the era of training very large models, hyperparameter tuning is performed on smaller models and then hyperparameters for the larger model are chosen based on the result of this tuning, using appropriate scaling rules. The authors of this article derive a practical and theoretically grounded guideline on how the EMA momentum hyperparameter should be chosen when the batch size is scaled. The key result is that as the batch size is multiplied by K, the EMA momentum should be raised to the power K.

The authors derive this scaling rule by analysing the limiting SDE corresponding to dynamics of SGD with various batch sizes, and validate it by throughout empirical testing on several tasks. The empirical results are especially impressive, as the authors are able to recover almost perfectly matching training curves for most of the supervised learning experiments, even when the batch size is scaled by a factor of 256. For the pseudo-labelling problem and self-supervised learning problem recovering the training dynamics turns out to be more difficult, as it is hard to replicate the training dynamics of the base model in the early phases, especially at large batch size. The authors manage to recover a very close replication anyway, with simple interventions.

I really appreciated this paper for providing an elegant, easy to understand and implement result with broad applicability.