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 paper proposes a generative model for images. There's a probability mask per-pixel per-component (which yields mixing probabilities), and then a set of latents per-component that yield an image. The system is tested on a set of scenes like the GQN dataset, stacks of blocks, and the multi-dsprites dataset. The system is better than MONet, although there are a few lingering questions.