GENESIS: Generative Scene Inference and Sampling with Object-Centric Latent Representations

1 INTRODUCTION . Task execution in robotics and reinforcement learning ( RL ) requires accurate perception of and reasoning about discrete elements in an environment . While supervised methods can be used to identify pertinent objects , it is intractable to collect labels for every scenario and task . Discovering structure in data—such as objects—and learning to represent data in a compact fashion without supervision are long-standing problems in machine learning ( Comon , 1992 ; Tishby et al. , 2000 ) , often formulated as generative latent-variable modelling ( e.g . Kingma & Welling , 2014 ; Rezende et al. , 2014 ) . Such methods have been leveraged to increase sample efficiency in RL ( Gregor et al. , 2019 ) and other supervised tasks ( van Steenkiste et al. , 2019 ) . They also offer the ability to imagine environments for training ( Ha & Schmidhuber , 2018 ) . Given the compositional nature of visual scenes , separating latent representations into object-centric ones can facilitate fast and robust learning ( Watters et al. , 2019a ) , while also being amenable to relational reasoning ( Santoro et al. , 2017 ) . Interestingly , however , state-of-the-art methods for generating realistic images do not account for this discrete structure ( Brock et al. , 2018 ; Parmar et al. , 2018 ) . As in the approach proposed in this work , human visual perception is not passive . Rather it involves a creative interplay between external stimulation and an active , internal generative model of the world ( Rao & Ballard , 1999 ; Friston , 2005 ) . That this is necessary can be seen from the physiology of the eye , where the small portion of the visual field that can produce sharp images ( fovea centralis ) motivates the need for rapid eye movements ( saccades ) to build up a crisp and holistic percept of a scene ( Wandell , 1995 ) . In other words , what we perceive is largely a mental simulation of the external world . Meanwhile , work in computational neuroscience tells us that visual features ( see , e.g. , Hubel & Wiesel , 1968 ) can be inferred from the statistics of static images using unsupervised learning ( Olshausen & Field , 1996 ) . Experimental investigations further show that specific brain areas ( e.g . LO ) appear specialised for objects , for example responding more strongly to common objects than to scenes or textures , while responding only weakly to movement ( cf . MT ) ( e.g. , GrillSpector & Malach , 2004 ) . ∗Corresponding author : martin @ robots.ox.ac.uk In this work , we are interested in probabilistic generative models that can explain visual scenes compositionally via several latent variables . This corresponds to fitting a probability distribution pθ ( x ) with parameters θ to the data . The compositional structure is captured by K latent variables so that pθ ( x ) = ∫ pθ ( x | z1 : K ) pθ ( z1 : K ) dz1 : K . Models from this family can be optimised using the variational auto-encoder ( VAE ) framework ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) , by maximising a variational lower bound on the model evidence ( Jordan et al. , 1999 ) . Burgess et al . ( 2019 ) and Greff et al . ( 2019 ) recently proposed two such models , MONet and IODINE , to decompose visual scenes into meaningful objects . Both works leverage an analysis-by-synthesis approach through the machinery of VAEs ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) to train these models without labelled supervision , e.g . in the form of ground truth segmentation masks . However , the models have a factorised prior that treats scene components as independent . Thus , neither provides an object-centric generation mechanism that accounts for relationships between constituent parts of a scene , e.g . two physical objects can not occupy the same location , prohibiting the component-wise generation of novel scenes and restricting the utility of these approaches . Moreover , MONet embeds a convolutional neural network ( CNN ) inside of an recurrent neural network ( RNN ) that is unrolled for each scene component , which does not scale well to more complex scenes . Similarly , IODINE utilises a CNN within an expensive , gradient-based iterative refinement mechanism . Therefore , we introduce GENErative Scene Inference and Sampling ( GENESIS ) which is , to the best of our knowledge , the first object-centric generative model of rendered 3D scenes capable of both decomposing and generating scenes1 . Compared to previous work , this renders GENESIS significantly more suitable for a wide range of applications in robotics and reinforcement learning . GENESIS achieves this by modelling relationships between scene components with an expressive , autoregressive prior that is learned alongside a sequential , amortised inference network . Importantly , sequential inference is performed in low-dimensional latent space , allowing all convolutional encoders and decoders to be run in parallel to fully exploit modern graphics processing hardware . We conduct experiments on three canonical and publicly available datasets : coloured Multi-dSprites ( Burgess et al. , 2019 ) , the GQN dataset ( Eslami et al. , 2018 ) , and ShapeStacks ( Groth et al. , 2018 ) . The latter two are simulated 3D environments which serve as testing grounds for navigation and object manipulation tasks , respectively . We show both qualitatively and quantitatively that in contrast to prior art , GENESIS is able to generate coherent scenes while also performing well on scene decomposition . Furthermore , we use the scene annotations available for ShapeStacks to show the benefit of utilising general purpose , object-centric latent representations from GENESIS for tasks such as predicting whether a block tower is stable or not . Code and models are available at https : //github.com/applied-ai-lab/genesis . 2 RELATED WORK . Structured Models Several methods leverage structured latent variables to discover objects in images without direct supervision . CST-VAE ( Huang & Murphy , 2015 ) , AIR ( Eslami et al. , 2016 ) , SQAIR ( Kosiorek et al. , 2018 ) , and SPAIR ( Crawford & Pineau , 2019 ) use spatial attention to partition scenes into objects . TAGGER ( Greff et al. , 2016 ) , NEM ( Greff et al. , 2017 ) , and R-NEM ( van Steenkiste et al. , 2018a ) perform unsupervised segmentation by modelling images as spatial mixture models . SCAE ( Kosiorek et al. , 2019 ) discovers geometric relationships between objects and their parts by using an affine-aware decoder . Yet , these approaches have not been shown to work on more complex images , for example visual scenes with 3D spatial structure , occlusion , perspective distortion , and multiple foreground and background components as considered in this work . Moreover , none of them demonstrate the ability to generate novel scenes with relational structure . While Xu et al . ( 2018 ) present an extension of Eslami et al . ( 2016 ) to generate images , their method only works on binary images with a uniform black background and assumes that object bounding boxes do not overlap . In contrast , we train GENESIS on rendered 3D scenes from Eslami et al . ( 2018 ) and Groth et al . ( 2018 ) which feature complex backgrounds and considerable occlusion to perform both decomposition and generation . Lastly , Xu et al . ( 2019 ) use ground truth pixel-wise flow fields as a cue for segmenting objects or object parts . Similarly , GENESIS could be adapted to also leverage temporal information which is a promising avenue for future research . 1We use the terms “ object ” and “ scene component ” synonymously in this work . MONet & IODINE While this work is most directly related to MONet ( Burgess et al. , 2019 ) and IODINE ( Greff et al. , 2019 ) , it sets itself apart by introducing a generative model that captures relations between scene components with an autoregressive prior , enabling the unconditional generation of coherent , novel scenes . Moreover , MONet relies on a deterministic attention mechanism rather than utilising a proper probabilistic inference procedure . This implies that the training objective is not a valid lower bound on the marginal likelihood and that the model can not perform density estimation without modification . Furthermore , this attention mechanism embeds a CNN in a RNN , posing an issue in terms of scalability . These two considerations do not apply to IODINE , but IODINE employs a gradient-based , iterative refinement mechanism which expensive both in terms of computation and memory , limiting its practicality and utility . Architecturally , GENESIS is more similar to MONet and does not require expensive iterative refinement as IODINE . Unlike MONet , though , the convolutional encoders and decoders in GENESIS can be run in parallel , rendering the model computationally more scalable to inputs with a larger number of scene components . Adversarial Methods A few recent works have proposed to use an adversary for scene segmentation and generation . Chen et al . ( 2019 ) and Bielski & Favaro ( 2019 ) segment a single foreground object per image and Arandjelović & Zisserman ( 2019 ) segment several synthetic objects superimposed on natural images . Azadi et al . ( 2019 ) combine two objects or an object and a background scene in a sensible fashion and van Steenkiste et al . ( 2018b ) can generate scenes with a potentially arbitrary number of components . In comparison , GENESIS performs both inference and generation , does not exhibit the instabilities of adversarial training , and offers a probabilistic formulation which captures uncertainty , e.g . during scene decomposition . Furthermore , the complexity of GENESIS increases with O ( K ) , where K is the number of components , as opposed to the O ( K2 ) complexity of the relational stage in van Steenkiste et al . ( 2018b ) . Inverse Graphics A range of works formulate scene understanding as an inverse graphics problem . These well-engineered methods , however , rely on scene annotations for training and lack probabilistic formulations . For example , Wu et al . ( 2017b ) leverage a graphics renderer to decode a structured scene description which is inferred by a neural network . Romaszko et al . ( 2017 ) pursue a similar approach but instead make use of a differentiable graphics render . Wu et al . ( 2017a ) further employ different physics engines to predict the movement of billiard balls and block towers . 3 GENESIS : GENERATIVE SCENE INFERENCE AND SAMPLING . In this section , we first describe the generative model of GENESIS and a simplified variant called GENESIS-S . This is followed by the associated inference procedures and two possible learning objectives . GENESIS is illustrated in Figure 1 and Figure 2 shows the graphical model in comparison to alternative methods . An illustration of GENESIS-S is included Appendix B.1 , Figure 5 . Generative model Let x ∈ RH×W×C be an image . We formulate the problem of image generation as a spatial Gaussian mixture model ( GMM ) . That is , every Gaussian component k = 1 , . . . , K represents an image-sized scene component xk ∈ RH×W×C . K ∈ N+ is the maximum number of scene components . The corresponding mixing probabilities πk ∈ [ 0 , 1 ] H×W indicate whether the component is present at a location in the image . The mixing probabilities are normalised across scene components , i.e . ∀i , j ∑ k πi , j , k = 1 , and can be regarded as spatial attention masks . Since there are strong spatial dependencies between components , we formulate an autoregressive prior distribution over mask variables zmk ∈ RDm which encode the mixing probabilities πk , as pθ ( z m 1 : K ) = K∏ k=1 pθ ( zmk | zm1 : k−1 ) = K∏ k=1 pθ ( z m k | uk ) |uk=Rθ ( zmk−1 , uk−1 ) . ( 1 ) The dependence on previous latents zm1 : k−1 is implemented via an RNN Rθ with hidden state uk . Next , we assume that the scene components xk are conditionally independent given their spatial allocation in the scene . The corresponding conditional distribution over component variables zck ∈ RDc which encode the scene components xk factorises as follows , pθ ( z c 1 : K | zm1 : K ) = K∏ k=1 pθ ( z c k | zmk ) . ( 2 ) Now , the image likelihood is given by a mixture model , p ( x | zm1 : K , zc1 : K ) = K∑ k=1 πk pθ ( xk | zck ) , ( 3 ) where the mixing probabilities πk = πθ ( zm1 : k ) are created via a stick-breaking process ( SBP ) adapted from Burgess et al . ( 2019 ) as follows , slightly overloading the π notation , π1 = πθ ( z m 1 ) , πk = 1− k−1∑ j=1 πj πθ ( zmk ) , πK = 1− K−1∑ j=1 πj . ( 4 ) Note that this step is not necessary for our model and instead one could use a softmax to normalise masks as in Greff et al . ( 2019 ) . Finally , omitting subscripts , the full generative model can be written as pθ ( x ) = ∫∫ pθ ( x | zc , zm ) pθ ( zc | zm ) pθ ( zm ) dzm dzc , ( 5 ) where we assume that all conditional distributions are Gaussian . The Gaussian components of the image likelihood have a fixed scalar standard deviation σ2x . We refer to this model as GENESIS . To investigate whether separate latents for masks and component appearances are necessary for decomposition , we consider a simplified model , GENESIS-S , with a single latent variable per component , pθ ( z1 : K ) = K∏ k=1 pθ ( zk | z1 : k−1 ) . ( 6 ) In this case , zk takes the role of zck in Equation ( 3 ) and of z m k in Equation ( 4 ) , while Equation ( 2 ) is no longer necessary . Approximate posterior We amortise inference by using an approximate posterior distribution with parameters φ and a structure similar to the generative model . The full approximate posterior reads as follows , qφ ( z c 1 : K , z m 1 : K | x ) = qφ ( zm1 : K | x ) qφ ( zc1 : K | x , zm1 : K ) , where qφ ( z m 1 : K | x ) = K∏ k=1 qφ ( zmk | x , zm1 : k−1 ) , and qφ ( zc1 : K | x , zm1 : K ) = K∏ k=1 qφ ( z c k | x , zm1 : k ) , ( 7 ) with the dependence on zm1 : k−1 realised by an RNN Rφ . The RNN could , in principle , be shared with the prior , but we have not investigated this option . All conditional distributions are Gaussian . For GENESIS-S , the approximate posterior takes the form qφ ( z1 : K | x ) = ∏K k=1 qφ ( zk | x , z1 : k−1 ) . Learning GENESIS can be trained by maximising the evidence lower bound ( ELBO ) on the logmarginal likelihood log pθ ( x ) , given by LELBO ( x ) = Eqφ ( zc , zm|x ) [ log pθ ( x | zc , zm ) pθ ( zc | zm ) pθ ( zm ) qφ ( zc | zm , x ) qφ ( zm | x ) ] ( 8 ) = Eqφ ( zc , zm|x ) [ log pθ ( x | z c , zm ) ] − KL ( qφ ( zc , zm | x ) || pθ ( zc , zm ) ) . ( 9 ) However , this often leads to a strong emphasis on the likelihood term , while allowing the marginal approximate posterior qφ ( z ) = Epdata ( x ) [ qφ ( z | x ) ] to drift away from the prior distribution , hence increasing the KL-divergence . This also decreases the quality of samples drawn from the model . To prevent this behaviour , we use the Generalised ELBO with Constrained Optimisation ( GECO ) objective from Rezende & Viola ( 2018 ) instead , which changes the learning problem to minimising the KL-divergence subject to a reconstruction constraint . Let C ∈ R be the minimum allowed reconstruction log-likelihood , GECO then uses Lagrange multipliers to solve the following problem , θ ? , φ ? = argmin θ , φ KL ( qφ ( z c , zm | x ) || pθ ( zc , zm ) ) such that Eqφ ( zc , zm|x ) [ log pθ ( x | z c , zm ) ] ≥ C . ( 10 )

Summary: The authors propose a probabilistic generative latent variable model representing a 2D image as a mixture of latent components. It formulates the scene generation problem as a spatial Gaussian mixture model where each Gaussian component comes from the decoding of an object-centric latent variable. The contribution of the proposed method from previous works is the introduction of an autoregressive prior on the component latents. This allows the model to capture autoregressive dependencies among different components and thus help generate coherent scenes, which has not been shown in the previous works. In the experiments, the authors compare GENESIS with MONet and VAEs qualitatively and quantitatively and show that the model outperforms the baseline in terms of both scene decomposition and generation.