LAVAE: Disentangling Location and Appearance

1 INTRODUCTION . Many hallmarks of human intelligence rely on the capability to perceive the world as a layout of distinct physical objects that endure through time—a skill that infants acquire in early childhood ( Spelke , 1990 ; 2013 ; Spelke and Kinzler , 2007 ) . Learning compositional , object-based representations of visual scenes , however , is still regarded as an open challenge for artificial systems ( Bengio et al. , 2013 ; Garnelo and Shanahan , 2019 ) . Recently , there has been a growing interest in unsupervised learning of disentangled representations ( Locatello et al. , 2018 ) , which should separate the distinct , informative factors of variations in the data , and contain all the information on the data in a compact , interpretable structure ( Bengio et al. , 2013 ) . This notion is highly relevant in the context of visual scene representation learning , where distinct objects should arguably be represented in a disentangled fashion . However , despite recent breakthroughs ( Chen et al. , 2016 ; Higgins et al. , 2017 ; Kim and Mnih , 2018 ) , multi-object scenarios are rarely considered ( Eslami et al. , 2016 ; Van Steenkiste et al. , 2018 ; Burgess et al. , 2019 ) . We propose the Location-Appearance Variational AutoEncoder ( LAVAE ) , a probabilistic generative model that , without supervision , learns structured , compositional , object-based representations of visual scenes . We explicitly model an object ’ s location and appearance with distinct latent variables , unlike in most previous works , thus providing a highly beneficial inductive bias . Following the framework of variational autoencoders ( VAEs ) ( Kingma and Welling , 2013 ; Rezende et al. , 2014 ) , we parameterize the approximate variational posterior of the latent variables with inference networks that are trained end-to-end with the generative model . Our model learns to correctly count objects and compute a compositional , object-wise , interpretable representation of the scene . Objects are represented independently of each other , and each object ’ s location and appearance are disentangled . Unlike previous approaches that disentangle location and appearance , LAVAE generalizes seamlessly to scenes with many more objects than in the training regime . We demonstrate these capabilities on multi-MNIST and multi-dSprites data sets similar to those by Eslami et al . ( 2016 ) and Greff et al . ( 2019 ) . 2 METHOD . Generative model . We propose a latent variable model for images in which the latent space is factored into location and appearance of a variable number of objects . For each image x with D pixels , the number of objects is modeled by a latent variable n , their locations by n latent variables { z ( i ) loc } ni=1 , and their appearance by { z ( i ) app } ni=1 . We assume the number of objects in every image to be bounded by M . The joint distribution of the observed and latent variables for each data point is : p ( x , zloc , zapp , n ) = p ( x | zloc , zapp ) p ( zloc |n ) p ( zapp |n ) p ( n ) = p ( x | zloc , zapp ) ( n∏ i=1 p ( z ( i ) loc | { z ( j ) loc } j < i ) p ( z ( i ) app ) ) p ( n ) , ( 1 ) where we use the shorthand zapp = { z ( i ) app } n i=1 and similarly for zloc . The generative process can be described as follows . First , the number of objects n is sampled from a categorical distribution p ( n ) = Cat ( M + 1 , a ) , n ∈ { 0 , . . . , M } , ( 2 ) where a is a learned probability vector of size M + 1 . The n location variables are sequentially sampled without replacement from a categorical distribution with D classes : p ( z ( i ) loc | { z ( j ) loc } j < i ) = Cat ( D , b ( i ) ) , b ( i ) = 1 D − i+ 1 ( 1D − ∑i−1 j=1 z ( j ) loc ) , ( 3 ) where i = 1 , . . . , n , and 1D is a vector of ones of length D. To each z ( i ) loc , which is a one-hot representation of an object ’ s location , corresponds a continuous appearance vector z ( i ) app ∼ N ( 0L , IL ) of size L that describes the object . The likelihood function p ( x | zloc , zapp ) is parameterized in a compositional manner . For each image , the visual representation , or sprite , of the ith object is generated from z ( i ) app by a shared function fsprite . Each sprite is then convolved with a 2-dimensional Kronecker delta that is the one-hot representation of the object ’ s location . Finally , the resulting n tensors are added together to give the pixel-wise parameters λ = ( λ1 , . . . , λD ) of the distribution : λ = n∑ i=1 fsprite ( z ( i ) app ) ∗ z ( i ) loc , ( 4 ) where ∗ denotes 2-dimensional discrete convolution . Inference model . The approximate posterior for a data point x has the following form : q ( zapp , zloc , n |x ) = q ( zapp | zloc , n , x ) q ( zloc |n , x ) q ( n |x ) = ( n∏ i=1 q ( z ( i ) app | z ( i ) loc , x ) q ( z ( i ) loc ∣∣ { z ( j ) loc } j < i , x ) ) q ( n |x ) . ( 5 ) Since each data point is assumed independent , the lower bound ( ELBO ) to the marginal log likelihood is the sum of the ELBO for each data point , which is : log p ( x ) ≥ Eq [ log p ( x | zloc , zapp ) ] + Eq [ log p ( zloc , zapp , n ) q ( zloc , zapp , n |x ) ] , ( 6 ) where the first term is the likelihood and the second is the negative Kullback-Leibler ( KL ) divergence between q ( zloc , zapp , n |x ) and p ( zloc , zapp , n ) . The different terms of the KL divergence can be derived as in Appendix A and estimated by Monte Carlo sampling . Two inference networks compute appearance and location feature maps , happ and hloc , both having size D like the input . The inference model for the number of objects n is a categorical distribution parameterized by a function of the location features fcount ( hloc ) . Object locations follow n categorical distributions without replacement parameterized by logits hloc . The vector at location z ( i ) loc in the feature map happ represents the appearance parameters for object i , i.e . the mean and log variance of q ( z ( i ) app | z ( i ) loc , x ) . The overall inference process can be summarized as follows : q ( n |x ) = Cat ( n ; M + 1 , fcount ( hloc ) ) ( 7 ) q ( z ( i ) loc ∣∣∣ { z ( j ) loc } j < i , x ) = Cat ( z ( i ) loc ; D , b ( i ) ) i = 1 , . . . , n ( 8 ) µ ( i ) app , σ 2 ( i ) app = happ [ z ( i ) loc ] i = 1 , . . . , n ( 9 ) q ( z ( i ) app | z ( i ) loc , x ) = N ( z ( i ) app ; µ ( i ) app , diag ( σ 2 ( i ) app ) ) i = 1 , . . . , n , ( 10 ) where by v [ ek ] we denote the kth element of v ( with ek a standard basis vector ) , and the probability vector for location sampling at each step is computed iteratively : b ( i ) 0 = softmax ( hloc ) · ( 1D − ∑i−1 j=1 z ( j ) loc ) , b ( i ) = b ( i ) 0 · ||b ( i ) 0 || −1 1 . ( 11 ) The expectations in the variational bound are handled as follows : For n we use discrete categorical sampling . This gives a biased gradient estimator , but in practice this has not affected inference on n. We use the Gumbel-softmax relaxation ( Jang et al. , 2016 ; Maddison et al. , 2016 ) for zloc and the Gaussian reparameterization trick for zapp . 3 RESULTS . In all experiments we use the same architecture for LAVAE . See Appendix B for all implementation details . Briefly , the sprite decoder in the generative model p ( x | zloc , zapp ) consists of a fully connected layer and a convolutional layer , followed by 3 residual convolutional blocks . The appearance and location inference networks are fully convolutional , thus the feature maps happ and hloc have the same spatial size as x . As a baseline to compare the proposed model against we used a VAE implemented as a fully convolutional network : the absence of fully connected layers makes it easier to preserve spatial information , thereby allowing the model to achieve a higher likelihood and generalize more naturally . Moreover , this choice of baseline is closer in spirit to our model than a VAE with fully connected layers . We evaluate LAVAE on multi-MNIST and multi-dSprites data sets consisting of 200k images with 0 to 3 objects in each image . 190k images are used for training , whereas the remaining 10k are left for evaluation . We generated an additional test set of 10k images with 7 objects each . Examples with more than 3 objects are never used for training . 3.1 MULTI-MNIST . Each image in the multi-MNIST data set consists of a number of statically binarized MNIST digits scattered at random ( avoiding overlaps ) onto a black canvas of size 64 × 64 . The digits are first rescaled from their original size ( 28× 28 ) to 15× 15 by bilinear interpolation , and finally binarized by rounding . When generating images , digits are picked from a pool of either 1k or 10k MNIST digits . We call the two resulting data sets multi-MNIST-1k and multi-MNIST-10k , depending on the number of MNIST digits used for generation . We independently model each pixel as a Bernoulli random variable parameterized by the decoder : p ( x | zloc , zapp ) = D∏ d=1 p ( xd | zloc , zapp ) = D∏ d=1 Bernoulli ( xd ; λd ) , ( 12 ) where the parameter vector λ defined in Section 2 is the pixel-wise mean of the Bernoulli distribution . Figure 1 qualitatively shows the performance of LAVAE on multi-MNIST-10k test images . The inferred location of all objects in an image are summarized by a black canvas where white pixels indicate object presence . For each of those locations , the model infers a corresponding appearance latent variable . Each digit on the right is generated by fsprite from one of these appearance variables . Samples from the prior are shown in Figure 2 : note that the number of generated objects is consistent with the training set , since the prior p ( n ) is learned . Quantitative evaluation results are shown in Table 1 and compared with a VAE baseline . The inferred object count was correct on almost all images of all test sets—even with a held-out number of objects—and across all tested random seeds . A fundamental characteristic of disentangled representations is that a change in a single representational factor should correspond to a change in a single underlying factor of variation ( Bengio et al. , 2013 ; Locatello et al. , 2018 ) . By demonstrating that the appearance and location of single objects can be independently manipulated in the latent space , the qualitative disentanglement experiments in Figure 3 prove that objects are disentangled from one another , and so are the appearance and location of each object .

Summary: This paper introduces a compositional generative model of images, where the image is described by a variable number of latent variables. Moreover, the latent variables are disentangled, in the sense that they represent different parts of the scene, and where appearance and location are described separately. While the generative model is very similar to AIR [1], the central claim of the paper is that the proposed inference scheme can generalize to a much higher number of objects than seen during training, which is demonstrated empirically, and with which AIR struggles. Better generalization is achieved by removing the recurrent core of AIR and replacing it with a fully-convolutional inference network, which predicts appearance vectors at every location in the feature map. The appearance vectors are then stochastically sampled without replacement according to a location distribution. Unlike AIR, this approach does not model the object scale.