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: This paper proposes a mechanism to reduce spatiotemporal mode collapse in unsupervised predictive learning. To achieve this goal, the proposed method is built upon the idea that different latent modes in the same data domain should share a set of hidden representation subspaces, which can be represented with various compositional structures based on the features in each subspace. Experimental results show improvement over simpler baselines such as PredRNN or ConvLSTM.