Neural spatio-temporal reasoning with object-centric self-supervised learning

1 INTRODUCTION . Artificial intelligence research has long been divided into rule-based approaches and statistical models . Neural networks , a classic example of the statistical approach , certainly have limitations despite their massive popularity and success . For example , experiments with two recently released video question-answering datasets , CLEVRER ( Yi et al. , 2020 ) and CATER ( Girdhar & Ramanan , 2020 ) , demonstrate that neural networks fail to adequately reason about spatio-temporal and compositional structure in visual scenes . While the networks perform adequately when asked to describe their inputs , they tend to fail when asked to predict , explain , or consider counterfactual possibilities . By contrast , a neuro-symbolic model called NS-DR ( Yi et al. , 2020 ) appears to be much better suited to predicting , explaining , and considering counterfactual possibilities with this data . The model leverages independent neural networks to detect objects , infer dynamics , and syntactically parse the question . A hand-coded symbolic executor interprets the questions grounded on the outputs of the networks . The fact that hybrid models employing both distributed ( neural ) representations and symbolic logic can sometimes perform better has led some to consider neuro-symbolic hybrids to be a more promising model class compared to end-to-end neural networks ( Andreas et al. , 2016 ; Yi et al. , 2018 ; Marcus , 2020 ) . There is evidence from other domains , however , that neural networks can indeed adequately model higher-level cognitive processes . For example , in some symbolic domains ( such as language ) , neural networks outperform hybrid neuro-symbolic approaches when tasked to classify or predict ( Devlin et al. , 2018 ) . Neural models have also had some success in mathematics , a domain that , intuitively , would seem to require the execution of formal rules and manipulation of symbols ( Lample & Charton , 2020 ) . Somewhat surprisingly , large-scale neural language models such as GPT-3 ( Brown et al. , 2020 ) can acquire a propensity for arithmetic reasoning and analogy-making without being trained explicitly for such tasks , suggesting that current neural network limitations are ameliorated when scaling to more data and using larger , more efficient architectures ( Brown et al. , 2020 ; Mitchell , 2020 ) . A key motivation of our work , therefore , is to reconcile existing neural network limitations in video domains with their ( perhaps surprising ) successes in symbolic domains . One common element of these latter results is the repeated application of self-attention processes ( Vaswani et al. , 2017 ) to sequences of discrete ‘ entities ’ . Here , we apply this insight to videos of physical interactions between sets of objects , where the input data to models are continuously-valued pixel arrays at multiple timesteps ( together with symbolic questions in certain cases ) . A key design decision is the appropriate level of granularity for the discrete units underlying the selfattention computation . What is the visual analogue to a word in language , or a symbol in mathematics ? We hypothesize that the discrete entities acted upon by self-attention should correspond to semantic entities relevant to the task . For tasks based on visual data derived from physical interactions , these entities are often times objects ( van Steenkiste et al. , 2019 ; Battaglia et al. , 2018 ) . To extract representations of these entities , we use MONet , an unsupervised object segmentation model ( Burgess et al. , 2019 ) , but we leave open the possibility that other object-estimation algorithms might work better . We propose that a sufficiently expressive self-attention model acting on entities corresponding to physical objects will exhibit , on video datasets , a similar level of higher-level cognition and ‘ reasoning ’ seen when these models are applied to language or mathematics . Altogether , our results demonstrate that self-attention-based neural nets can outperform hybrid neurosymbolic models on visual tasks that require high-level cognitive processes , such as causal reasoning and physical understanding . We show that choosing the right level of discretization is critical for successfully learning these higher-order capabilities : pixels and local features are too fine , and entire scenes are too coarse . Moreover , we identify the value of self-supervised tasks , especially in low data regimes . These tasks ask the model to infer future arrangements of objects given the past , or to infer what must have happened for objects to look as they do in the present . We verify these conclusions in two video datasets , one in which the input is exclusively visual ( CATER ) and one that requires the combination of language ( questions ) and vision ( CLEVRER ) . 2 METHODS . Our principal motivation is the converging evidence for the value of self-attention mechanisms operating on a finite sequences of discrete entities . Written language is inherently discrete and hence is well-suited to self-attention-based approaches . In other domains , such as raw audio or vision , it is less clear how to leverage self-attention . We hypothesize that the application of self-attention-based models to visual tasks could benefit from an approximate ’ discretization ’ process analogous to the segmentation of speech into words or morphemes , and that determining the appropriate level of discretization is an important choice that can significantly affect model performance . At the finest level , data could simply be discretized into pixels ( as is already the case for most machine-processed visual data ) . But since pixels are too-fine grained , some work considers the downsampled “ hyper-pixel ” outputs of a convolutional network to comprise the set of discrete units ( e.g . Zambaldi et al . ( 2019 ) ; Lu et al . ( 2019 ) ) . In the case of videos , an even courser discretization scheme is often used : representations of frames or subclips ( Sun et al. , 2019b ) . The neuroscience literature , however , suggests that biological visual systems infer and exploit the existence of objects , rather than use spatial or temporal blocks with artificial boundaries ( Roelfsema et al. , 1998 ; Spelke , 2000 ; Chen , 2012 ) . Because objects are the atomic units that tasks we consider here focus on , it makes sense to discretize on the level of objects . Numerous object segmentation algorithms have been proposed ( Ren et al. , 2015 ; He et al. , 2017 ; Greff et al. , 2019 ) . We chose to use MONet , an unsupervised object segmentation algorithm that produces object representations with disentangled features ( Burgess et al. , 2019 ) . Because MONet is unsupervised , we can train it directly in our domain of interest without the need for object segmentation labels . To segment each frame into object representations , MONet first uses a recurrent attention network to obtain a set of No “ object attention masks ” ( No is a fixed parameter ) . Each attention mask represents the probability that any given pixel belongs to that mask ’ s object . The pixels assigned to the mask are encoded into latent variables with means µti ∈ Rd , where i indexes the object and t the frame . These means are used as the object representations in our model . More details are provided in Appendix A.1 . The self-attention component is a transformer model ( Vaswani et al. , 2017 ) over the sequence µti . In addition to this sequence of vectors , we include a trainable vector CLS ∈ Rd that is used to generate classification results ; this plays a similar role to the CLS token in BERT ( Devlin et al. , 2018 ) . Finally , for our CLEVRER experiments , where the inputs include a question and potentially several choices , we embed each question word wi ( and choice for multiple choice questions ) in Rd and include these in the sequence of inputs to the transformer . We also append a two-dimensional one-hot vector to µti and wi to indicate whether the input is a word or an object . We pass the sequence consisting of the object latent means µti , the classification token , and the word embeddings ( for CLEVRER ) through a transformer with NT layers . We add a relative positional encoding at each layer of the transformer to give the model knowledge of the word and frame order ( Dai et al. , 2019 ) . The transformed value of the classification token CLS is passed through an MLP ( with one hidden layer of size NH ) to generate the final answer . This general approach is adapted to each of our datasets according to the format of expected answers . The final layer , which operates on the transformed value of CLS , is a softmax over answers for CLEVRER descriptive questions , a logit for each choice for CLEVRER multiple-choice questions , or a softmax over the grid-index of the final location of the snitch for CATER . A schema of our architecture is shown in Figure 1 . Note that in this model , an object in one frame can attend to every object in every frame . We also consider an alternative model with hierarchical attention , which consists of two stages of selfattention with two different transformers . The first transformer acts independently on the objects within each frame along with the word embeddings . The outputs of the first transformer for each frame are concatenated into a single feature vector , one for each frame . The second transformer acts on these feature vectors , treating each frame as an atomic entity . We study the importance of global attention ( objects as the atomic entities ) vs hierarchical attention ( objects , and subsequently frames as the atomic entities ) . The comparison is shown in Table 1 . 2.1 SELF-SUPERVISED LEARNING . Self-supervised learning—unsupervised learning where the data provides its own supervision—has a long history in language processing , where it allows models to acquire useful representations of words , phrases and sentences ( Mikolov et al. , 2013 ; Kiros et al. , 2015 ; Devlin et al. , 2018 ) . Such techniques are also effective in visual domains , for image classification ( Chen et al. , 2020 ) , video analysis ( Han et al. , 2019 ) , and RL agents ( Gregor et al. , 2019 ) . In these cases , it is proposed that the learned representations are more informationally dense , allowing models to satisfy their ultimate objectives more effectively . MONet , the technique we use to estimate object-centric visual representations , can also be considered a self-supervised learning method . We explored whether self-supervised learning could improve the performance of our model beyond the benefits conveyed by object-level representation , i.e . in ways that support the model ’ s interpretation Estimation steps . of the dynamics of the scenes rather than just via improved perception of static observations . Our approach is inspired by the loss used in BERT ( Devlin et al. , 2018 ) , where a bidirectional transformer model is trained to predict certain words that are masked from the input . In our case , we mask object representations , and train the model to infer the content of the masked object representations using its knowledge from all unmasked objects . More formally , we set transformer input = 〈 CLS ; mtiµti|t , i ; wi|i 〉 , where mti ∈ { 0 , 1 } is a masking indicator . We write the output of the transformer as transformer output = 〈 CLS′ ; µ′ti|t , i ; w ′ i|i 〉 . We expect the transformer to understand the underlying dynamics of the video , so that the masked out slot µti could be predicted from µ′ti . We add an auxiliary loss to guide the transformer in learning effective representations capable of this type of dynamics prediction : auxiliary loss = ∑ t , i τtil ( f ( µ ′ ti ) , µ ) , where f is a learned linear mapping to Rd , l a loss function , and τti ∈ { 0 , 1 } are one-hot indicator variables identifying the prediction targets . We propagate gradients only to the parameters of f and the transformer . This auxiliary loss is added to the main classification loss with weighting λ , and both losses are minimized simultaneously by the optimizer . In particular , we do not pretrain the model with only the auxiliary loss . We tested two different loss functions , an L2 loss and a contrastive loss ( formulas given in Appendix A.2 ) , and six different masking schemes , as illustrated in Figure 2 . This exploration was motivated by the observation that video inputs at adjacent timesteps are highly correlated in a way that adjacent words are not . We thus hypothesized that BERT-style prediction of adjacent words might not be optimal . A different masking strategy , in which prediction targets are separated from the context by more than a single timestep , may stimulate capacity in the network to acquire the environment knowledge that permits context-based unrolls and better long-horizon predictions . First , we set mti = 1 ( uniformly ) at random across t and i and τti = 1−mti , generated by fixing the expected proportion of the mti set to 1 ( schema b in Figure 2 ) . While simple , this has the downside of masking out multiple objects per frame , which is potentially problematic since MONet does not assign objects to slot indices in a well defined way . MONet usually switches object-to-slot assignments multiple times in a single video , and these switches occur unpredictably . If multiple slots are masked out , the transformer can not determine with certainty which missing object to assign to each slot , and so the auxiliary loss could penalize the model even if it predicted all the objects correctly . To avoid this problem , we also try constraining the mask such that exactly one mti = 0 for each t ( schema a ) ; this ensures only one slot per frame is masked out , eliminating the ambiguity . To pose harder prediction challenges , we add a buffer between the context ( where mti = 1 ) and the infilling targets ( where τti = 1 ) . For t in this buffer zone , both mti = 0 and τti = 0 ( schemas c–f ) . In the presence of this buffer , we compared prediction ( where the context is strictly before the targets ; schema c , d ) versus infilling ( where the context surrounds the targets ; schema e , f ) . We also compared setting the targets as individual objects ( schema c , e ) versus targets as entire scenes ( schema d , f ) . We visually inspect the efficacy of this self-supervised loss in encouraging better representations ( beyond improvements of scores on tasks ) in Appendix C .

Summary: the paper propose to tackle visual reasoning problem in videos. The proposed solution is to combine MONET (Burgess et al., 2019) with self-attention mechanism (Vaswani et al., 2017) to first encode images into object-centric encodings and aggregate the encodings using self-attention to make the final prediction. The method is shown to outperform neural-symbolic reasoning approaches such as MAC (Hudson & Manning, 2018) and NS-DR (Yi et al 2020.) on image QA and R3D (Girdhar & Ramanan, 2020) on CARTER, which is a video reasoning task / benchmark.