Conditional Object-Centric Learning from Video

1 INTRODUCTION . Humans understand the world in terms of separate objects ( Kahneman et al. , 1992 ; Spelke & Kinzler , 2007 ) , which serve as compositional building blocks that can be processed independently and recombined . Such a compositional model of the world forms the foundation for high-level cognitive abilities such as language , causal reasoning , mathematics , planning , etc . and is crucial for generalizing in predictable and systematic ways . Object-centric representations have the potential to greatly improve sample efficiency , robustness , generalization to new tasks , and interpretability of machine learning algorithms ( Greff et al. , 2020 ) . In this work , we focus on the aspect of modeling motion of objects from video , because of its synergistic relationship with object-centric representations : On the one hand , objects support learning an efficient dynamics model by factorizing the scene into approximately independent parts with only sparse interactions . Conversely , motion provides a powerful cue for which inputs should be grouped together , and is thus an important tool for learning about objects . Unsupervised multi-object representation learning has recently made significant progress both on images ( e.g . Burgess et al. , 2019 ; Greff et al. , 2019 ; Lin et al. , 2020 ; Locatello et al. , 2020 ) and on video ( e.g . Veerapaneni et al. , 2020 ; Weis et al. , 2020 ; Jiang et al. , 2020 ) . By incorporating objectcentric inductive biases , these methods learn to segment and represent objects from the statistical structure of the data alone without the need for supervision . Despite promising results these methods are currently limited by two important problems : Firstly , they are restricted to toy data like moving 2D sprites or very simple 3D scenes and generally fail at more realistic data with complex textures ( Greff et al. , 2019 ; Harley et al. , 2021 ; Karazija et al. , 2021 ) . And secondly , it is not entirely clear how to interface with these models both during training and inference . The notion of an object is ambiguous and task-dependent , and the segmentation learned by these models does not necessarily align with the tasks of interest . The model might , for example , over-segment a desired object into distinct parts , or alternatively fail to segment into the desired parts . Ideally , we would like to be able to provide the model with hints as to the desired level of granularity during training , and flexibly query the model during inference to request that the model detects and tracks a particular object or a part thereof . In this paper we introduce a sequential extension of Slot Attention ( Locatello et al. , 2020 ) that we call Slot Attention for Video ( SAVi ) to tackle the problem of unsupervised / weakly-supervised multiobject segmentation and tracking in video data . We demonstrate that 1 ) using optical flow prediction as a self-supervised objective and 2 ) providing a small set of abstract hints such as the center of mass position for objects as conditional inputs in the first frame suffices to direct the decomposition process in complex video scenes without otherwise requiring any priors on the size of objects or on the information content of their representations . We show successful segmentation and tracking for synthetic video data with significantly higher realism and visual complexity than the datasets used in prior work on unsupervised object representation learning . These results are robust with respect to noise on both the optical flow signal and the conditioning , and that they generalize almost perfectly to longer sequences , novel objects , and novel backgrounds . 2 SLOT ATTENTION FOR VIDEO ( SAVI ) . We introduce a sequential extension of the Slot Attention ( Locatello et al. , 2020 ) architecture to video data , which we call SAVi . Inspired by predictor-corrector methods for integration of ordinary differential equations , SAVi performs two steps for each observed video frame : a prediction and a correction step . The correction step uses Slot Attention to update ( or correct ) the set of slot representations based on slot-normalized cross-attention with the inputs . The prediction step uses self-attention among the slots to allow for modeling of temporal dynamics and object interactions . The output of the predictor is then used to initialize the corrector at the next time step , thus allowing the model to consistently track objects over time . Importantly , both of these steps are permutation equivariant and thus are able to preserve slot symmetry . See Figure 1 for a schematic overview of the SAVi architecture . Encoder For each time-step t ∈ { 1 . . . T } the corresponding video frame xt is first passed through a small convolutional neural network ( CNN ) encoder ( here , a stack of five convolutional layers with ReLUs ) , where we concatenate a linear positional encoding at the second-to-last layer . The resulting grid of visual features is flattened into a set of vectors ht = fenc ( xt ) ∈ RN×Denc , where N is the size of the flattened grid ( i.e. , width * height ) and Denc is the dimensionality of the CNN feature maps . Afterwards , each vector is independently passed through a multi-layer perceptron ( MLP ) . Slot Initialization SAVi maintains K slots , each of which can represent a part of the input such as objects or parts thereof . We denote the set of slot representations at time t as St = [ s1t , . . . , sKt ] ∈ RK×D , where we use the calligraphic font to indicate that any operation on these sets is equivariant ( or invariant ) w.r.t . permutation of their elements . In other words , the ordering of the slots carries no information and they can be freely permuted ( in a consistent manner across all time steps ) without changing the model output . We consider two types of initializers : conditional and unconditional . In the conditional case , we encode the conditional input either via a simple MLP ( in the case of bounding boxes or center of mass coordinates ) or via a CNN ( in the case of segmentation masks ) . For slots for which there is no conditioning information available ( e.g . if K is larger than the number of objects ) , we set the conditional input to a fixed value ( e.g. , ‘ −1 ’ for bounding box coordinates ) . The encoded conditional input forms the initial slot representation for SAVi . In the unconditional case , we either randomly initialize slots by sampling from a Gaussian distribution independently for each video ( both at training and at test time ) , or by learning a set of initial slot vectors . Corrector The task of the corrector is to update the slot representations based on the visual features from the encoder . In SAVi this is done using the iterative attention mechanism introduced in Slot Attention ( Locatello et al. , 2020 ) . Different from a regular cross-attention mechanism ( e.g . Vaswani et al. , 2017 ) which is normalized over the inputs , Slot Attention encourages decomposition of the input into multiple slots via softmax-normalization over the output ( i.e . the slots ) , which makes it an appealing choice for our video decomposition architecture . When using a single iteration of Slot Attention , the corrector update takes the following form : Ut = 1 Zt N∑ n=1 At , n v ( ht , n ) ∈ RK×D , At = softmax K ( 1√ D k ( ht ) · q ( St ) T ) ∈ RN×K , ( 1 ) where Zt = ∑N n=1At , n and denotes the Hadamard product . k , q , v are learned linear projections that map to a common dimension D. We apply LayerNorm ( Ba et al. , 2016 ) before each projection . The slot representations are then individually updated using Gated Recurrent Units ( Cho et al. , 2014 ) as ŝkt = GRU ( u k t , s k t ) . Alternatively , the attention step ( followed by the GRU update ) can be iterated multiple times with shared parameters per frame of the input video . For added expressiveness , we apply an MLP with residual connection ŝkt ← ŝkt +MLP ( LN ( ŝkt ) ) and LayerNorm ( LN ) after the GRU when using multiple Slot Attention iterations , following Locatello et al . ( 2020 ) . Predictor The predictor takes the role of a transition function to model temporal dynamics , including interactions between slots . To preserve permutation equivariance , we use a Transformer encoder ( Vaswani et al. , 2017 ) . It allows for modeling of independent object dynamics as well as information exchange between slots via self-attention , while being more memory efficient than GNN-based models such as the Interaction Network ( Battaglia et al. , 2016 ) . Slots are updated as follows : St+1 = LN ( MLP ( S̃t ) + S̃t ) , S̃t = LN ( MultiHeadSelfAttn ( Ŝt ) + Ŝt ) . ( 2 ) For MultiHeadSelfAttn we use the default multi-head dot-product attention mechanism from Vaswani et al . ( 2017 ) . We apply LayerNorm ( LN ) after each residual connection . Decoder The network output should be permutation equivariant ( for per-slot outputs ) or invariant ( for global outputs ) with respect to the slots . Slots can be read out either after application of the corrector or after the predictor ( transition model ) . We decode slot representations after application of the corrector using a slot-wise Spatial Broadcast Decoder ( Watters et al. , 2019 ) to produce per-slot RGB predictions of the optical flow ( or reconstructed frame ) and an alpha mask . The alpha mask is normalized across slots via a softmax and used to perform a weighted sum over the slot-wise RGB reconstruction to arrive at a combined reconstructed frame : yt = K∑ k=1 mkt ykt , mt = softmax K ( m̂kt ) , m̂kt , y k t = fdec ( ŝkt ) . ( 3 ) Training Our sole prediction target is optical flow for each individual video frame , which we represent as RGB images using the default conversion in the literature ( Sun et al. , 2018 ) . Alternatively , our framework also supports prediction of other image-shaped targets , such as reconstruction of the original input frame . We minimize the pixel-wise squared reconstruction error ( averaged over the batch ) , summed over both the temporal and spatial dimensions : Lrec = T∑ t=1 ‖yt − ytruet ‖2 . ( 4 ) 3 RELATED WORK . Object-centric representation learning There is a rich literature on learning object representations from static scenes ( Greff et al. , 2016 ; Eslami et al. , 2016 ; Greff et al. , 2017 ; 2019 ; Burgess et al. , 2019 ; Engelcke et al. , 2020 ; Crawford & Pineau , 2019 ; Lin et al. , 2020 ; Locatello et al. , 2020 ) or videos ( van Steenkiste et al. , 2018 ; Kosiorek et al. , 2018 ; Stelzner et al. , 2019 ; Kipf et al. , 2020 ; Crawford & Pineau , 2020 ; Creswell et al. , 2021 ) without explicit supervision . PSGNet ( Bear et al. , 2020 ) learns to decompose static images or individual frames from a video into hierarchical scene graphs using motion information estimated from neighboring video frames . Most closely related to our work are sequential object-centric models for videos and dynamic environments , such as OP3 ( Veerapaneni et al. , 2020 ) , R-SQAIR ( Stanić & Schmidhuber , 2019 ) , ViMON ( Weis et al. , 2020 ) , and SCALOR ( Jiang et al. , 2020 ) , which learn an internal motion model for each object . SIMONe ( Kabra et al. , 2021 ) auto-encodes an entire video in parallel and learns temporally-abstracted representations of objects . OP3 ( Veerapaneni et al. , 2020 ) uses the same decoder as SAVi and a related dynamics model , but a less efficient inference process compared to Slot Attention . In an attempt to bridge the gap to visually richer and more realistic environments , recent works in object-centric representation learning have explored integration of inductive biases related to 3D scene geometry , both for static scenes ( Chen et al. , 2020 ; Stelzner et al. , 2021 ) and for videos ( Du et al. , 2021 ; Henderson & Lampert , 2020 ; Harley et al. , 2021 ) . This is largely orthogonal to our approach of utilizing conditioning and optical flow . A recent related method , FlowCaps ( Sabour et al. , 2020 ) , similarly proposed to use optical flow in a multi-object model . FlowCaps uses capsules ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) instead of a slot-based representation and assumes specialization of individual capsules to objects or parts of a certain appearance , making it unsuitable for environments that contain a large variety of object types . Using a slot-based , exchangeable representation of objects allows SAVi to represent a diverse range of objects and generalize to novel objects at test time . We discuss further recent related works on attention-based modular architectures ( Santoro et al. , 2018 ; Goyal et al. , 2021a ; b ; c ) , object-centric models for dynamic visual reasoning ( Yi et al. , 2020 ; Ding et al. , 2021a ; Bar et al. , 2021 ) , and supervised attention-based object-centric models ( Fuchs et al. , 2019 ; Carion et al. , 2020 ; Kamath et al. , 2021 ; Meinhardt et al. , 2021 ) in Appendix A.1 . Video object segmentation and tracking The conditional tasks we consider in our work are closely related to the computer vision task of semi-supervised video object segmentation ( VOS ) , where segmentation masks are provided for the first video frame during evaluation . Different from the typical setting , which is addressed by supervised learning on fully annotated videos or related datasets ( e.g . Caelles et al. , 2017 ; Luiten et al. , 2018 ) , we consider the problem where models do not have access to any supervised information beyond the conditioning information on the first frame ( e.g . a bounding box for each object ) . Several recent works have explored pre-training using self-supervision ( Li et al. , 2019 ; Jabri et al. , 2020 ; Caron et al. , 2021 ) or image classification ( Zhang et al. , 2020 ) for the semi-supervised VOS task . These models rely on having access to segmentation masks in the first frame at evaluation time . We demonstrate that multi-object segmentation and tracking can emerge even when segmentation labels are absent at both training and test time . There is a rich literature on using motion cues to segment objects in the computer vision community . These methods use motion information at test time to , for example , cluster trajectories in order to segment independently moving objects ( Faktor & Irani , 2014 ) or estimate multiple fundamental matrices between two views ( Isack & Boykov , 2012 ) , to name a few . Closest to our work is a contemporary method ( Yang et al. , 2021 ) that trains a Slot Attention model on isolated optical flow data for foreground-background segmentation of a single object , independently for individual video frames and without using visual observations . Our method , on the other hand , supports multi-object environments and only relies on motion information as a training signal but otherwise operates directly on textured visual information , which allows it to segment static scenes at test time where optical flow is unavailable and consistently represent and track multiple objects throughout a video .

Summary: This paper considers the problem of learning object-centric representations from videos. Different from existing methods that mostly model this problem under a pure unsupervised setting by reconstructing video frames, this work introduces two improvements. The first is offering weak "hints" (can be seen as additional inputs) during training, which can avoid falling into a trivial solution like previous unsupervised approaches always did. The "hints" here can be in form of pixel-wise masks, bounding boxes, or even as simple as centers of mass.  The second improvement regards the supervision signals. Instead of reconstructing raw pixels, they propose to predict optical flows. In easy scenes the optical flows are actually similar to pixel-wise masks: values inside an object are prone to be consistent. Therefore the learning is eased, and then making it possible to train on more realistic datasets. Extensive experiments are conducted to validate the effectiveness of these two improvements.