ModeRNN: Harnessing Spatiotemporal Mode Collapse in Unsupervised Predictive Learning

1 INTRODUCTION . Predictive learning is an unsupervised learning paradigm that has shown the ability to discover the spatiotemporal modes of visual dynamics ( Xu et al. , 2019 ; Goyal et al. , 2021 ) . However , for largescale and real-world datasets ( see Figure 1 ) , the modes in visual dynamics can be highly entangled and difficult to learn due to the richness of data environments , the diversity of object interactions , and the complexity of motion patterns . For clarity , in the following discussion , spatiotemporal modes are considered to have the following properties : 1 . A spatiotemporal mode refers to a representation subspace that corresponds to a family of similar , but not predefined , visual dynamics . 2 . Multiple spatiotemporal modes naturally exist in real-world data , even in a single frame . 3 . We assume the i.i.d . setup to allow all videos to share the same set of spatiotemporal modes in a dataset . Different data may have different compositional structures over the modes . Under these assumptions , video prediction models are required to ( i ) decouple the potentially mixed spatiotemporal modes from raw video frames , ( ii ) understand the compositional structures on top of the learned modes , and ( iii ) learn the state transitions based on the compositional structures . Otherwise , since the learned dynamics with respect to different modes may interfere and compete during training , it remains challenging for the prior art in video prediction to generate less blurry future frames based on an ambiguous understanding of mixed physical processes . We refer to this empirical phenomenon as spatiotemporal mode collapse ( STMC ) , which is mainly caused by the collapse of learned representations into invalid subspaces when compromising to multiple spatiotemporal modes in the training set . Unlike the widely concerned mode collapse problem in generative adversarial networks , STMC has not drawn much attention because predictive learning is supposed to be well constrained by the image reconstruction loss . However , due to the limitation of model size , STMC occurs when the model can not effectively decouple mixed spatiotemporal modes and infer their underlying structures . As a result , its responses to different modes tend to lose diversity and may collapse to a meaningless average of multiple representation subspaces of valid modes . In Figure 1 ( left ) , we can observe the existence of STMC on a large-scale video dataset named RoboNet ( Dasari et al. , 2019 ) , in which potential spatiotemporal modes may come from seven different robot platforms ( e.g. , Baxter and WidowX ) , four data collection environments ( e.g. , Berkeley and Stanford ) , and a variety of unlabeled robot control tasks ( e.g. , pushing and grasping ) . An additional outcome of STMC is that we can achieve a performance gain when training individual models in separate subsets with remarkably different visual dynamics , as shown in Figure 1 ( right ) . However , such a dilemma prevents the model from growing into big ones that allow scalable training on large-scale , natively multimodal spatiotemporal sequences . We explore STMC for the first time in unsupervised predictive learning . The core idea is to provide a strong inductive bias for the predictive model to discover the compositional structures of latent modes . To this end , we propose ModeRNN , a new modular recurrent architecture that learns structured hidden representations through a set of mode slots1 , where each of them responds to the representation subspace of a single spatiotemporal mode . ModeRNN also introduces a decoupling-aggregation framework to process the slot features in three stages , which is completely different from existing predictive models with modular architectures ( Xu et al. , 2019 ; Goyal et al. , 2021 ) . The first stage is recurrent state interaction and slot binding , in which we use the multi-head attention mechanism ( Vaswani et al. , 2017 ) to enable the memory state to interact with the input state and previous hidden state of RNNs . We name the memory state “ slot bus ” , because for each sequence , it is initialized from a multi-variate Gaussian distribution with learnable parameters , and thereafter refined using the slot features at each time step . By using the slot bus as the queries , multi-head attention can naturally decouple modular components from hidden representations and bind them to particular mode slots . Features in each slot are then independently modeled using per-slot convolutional parameters . The second stage in each ModeRNN unit is slot fusion , motivated by the assumption that , there can be multiple spatiotemporal modes in a single video and similar videos can be represented by similar compositional structures over the mode slots . Therefore , we assign slot features with learnable importance weights and aggregate them into a unified hidden representation , which is then used in the third stage to update the slot bus and generate the output state of the ModeRNN unit . We empirically show the existence of STMC on five datasets , and include the results on three realworld datasets in the manuscript , including the large-scale RoboNet dataset that has various data collection environments and multiple robot control tasks , the KTH dataset with six types of human actions that has been widely used by previous literature , and the radar echo dataset for precipitation forecasting that contains time-varying modes of seasonal climates . In addition , we include results on a Mixed Moving MNIST dataset and the Human3.6M dataset in the appendix . In a series of quantitative and visualization results , we demonstrate the effectiveness of ModeRNN in mitigating STMC and learning from highly entangled visual dynamics . 1The concept of “ slot ” was initially introduced by Locatello et al . ( 2020 ) to denote the object-centric features in static scene understanding . We borrow this term here for the subspaces of spatiotemporal representations . 2 RELATED WORK . RNN-based predictive models . Many deep learning models based on RNNs have been proposed for spatiotemporal prediction ( Ranzato et al. , 2014 ; Srivastava et al. , 2015 ; Shi et al. , 2015 ; Oh et al. , 2015 ; De Brabandere et al. , 2016 ; Villegas et al. , 2018 ) . Shi et al . ( 2015 ) integrated 2D convolutions into the recurrent state transitions of standard LSTM and proposed the convolutional LSTM ( ConvLSTM ) network , which can model the spatial correlations and temporal dynamics in a unified recurrent unit . More recent approaches have extended the prediction ability of ConvLSTM in different aspects ( Wang et al. , 2017 ; Oliu et al. , 2018 ; Wang et al. , 2019b ; a ; Yao et al. , 2020 ; Guen & Thome , 2020 ; Yu et al. , 2019 ; Su et al. , 2020 ; Lin et al. , 2020 ; Lee et al. , 2021 ) . For example , as an important compared model of our approach , SA-ConvLSTM ( Lin et al. , 2020 ) incorporates selfattention in the recurrent state transitions in ConvLSTM to obtain more global context information across time . However , unlike our approach , it does not learn decoupled representations to understand individual components in complex visual dynamics . Besides deterministic models , probabilistic models were proposed to explicitly consider the uncertainty in future prediction ( Mathieu et al. , 2016 ; Vondrick et al. , 2016 ; Tulyakov et al. , 2018 ; Xu et al. , 2018 ; Wang et al. , 2018 ; Denton & Fergus , 2018 ; Castrejon et al. , 2019 ; Kwon & Park , 2019 ; Bhagat et al. , 2020 ) . We use a typical stochastic video generation approach ( Denton & Fergus , 2018 ) based on conditional VAE as a compared model . Unsupervised predictive learning for spatiotemporal disentanglement . Previous work has focused on learning to disentangle the spatial and temporal features from visual dynamics ( Denton et al. , 2017 ; Guen & Thome , 2020 ; Hsieh et al. , 2018 ; Wu et al. , 2021 ) . These methods factorize spatiotemporal data into feature subspaces with strong priors , e.g. , assuming that the spatial information is temporally invariant . Another line of work is to learn predictive models for unsupervised scene decomposition such as ( Xu et al. , 2019 ; Hsieh et al. , 2018 ) . Unlike the above models , our approach uses a set of modular architectures in the recurrent unit to represent the mixed spatiotemporal dynamics . The most relevant work to our method is the Recurrent Independent Mechanism ( RIM ) ( Goyal et al. , 2021 ) , which consists of largely independent recurrent modules that are sparsely activated and interact via soft attention . ModeRNN is different from RIM in three aspects . First , it is specifically designed to tackle STMC in real-world environments . Second , it learns modular features by incorporating multi-head attention in the recurrent unit , and performs state transitions on compositional features with learnable importance weights . Third , the modular structures in ModeRNN are frequently activated responding to the mixed visual dynamics . ModeRNN is compared with the state of the art in Section 4 , including SA-ConvLSTM ( Lin et al. , 2020 ) , PhyDNet ( Guen & Thome , 2020 ) , CrevNet ( Yu et al. , 2019 ) , RIM ( Goyal et al. , 2021 ) , and LMC ( Lee et al. , 2021 ) . 3 MODERNN . We propose ModeRNN to reduce spatiotemporal mode collapse ( STMC ) in unsupervised predictive learning . The key idea is to build a decoupling-aggregation framework to model the recurrent state transitions of mixed spatiotemporal modes . In this section , we first discuss the basic network components in ModeRNN and then describe the details in the decoupling-aggregation recurrent unit . 3.1 MODE SLOTS & SLOT BUS . Mode slots . The decoupling-aggregation framework is built upon a set of hidden representations named mode slots . The term slot is in part borrowed from previous work for unsupervised scene decomposition ( Locatello et al. , 2020 ) . We use it here to respond to a family of similar visual dynamics , that is , we aim to bind each mode slot to the representation subspace of each spatiotemporal mode one-to-one . Slot features can be viewed as latent factors that can explicitly improve the unsupervised decoupling of mixed dynamics across the dataset . Slot bus . Assuming that multiple spatiotemporal modes naturally co-exist in real-world videos , all slots dynamically respond with different importance weights to form compositional representations , which are then used to update a long-term memory state , termed slot bus . The hierarchical structure of mode slots and the slot bus leads to a better understanding of the complex and highly mixed dynamic patterns without mode annotations . From similar data samples , the model is allowed to learn similar compositional structures over the slots . On the contrary , for distinct visual dynamics , it shows significant differences in the learned importance weights to update the slot bus features . Therefore , it provides a solution to STMC . Specifically , the slot bus is initialized from a learnable , multi-variate Gaussian distribution , whose mean and variance encode the global priors for the entire dataset . 3.2 MODECELL . To learn and leverage the mode slots , we introduce a novel recurrent unit named ModeCell , which follows a decoupling-aggregation framework with three modules , i.e. , the state interaction and slot binding module , the adaptive slot fusion module , and the slot bus transition module .

Summary: The manuscript proposes a novel architecture with a slot-based decoupling-aggregation framework for unsupervised sequence prediction. The model is motivated by preventing spatio-temporal mode collapse, which affects many existing methods. The experiments clearly show that the model addresses this issue and performance comparisons across three commonly used datasets are presented.