Spatially Invariant Attend , Infer , Repeat

The physical world can be naturally sub-divided into discrete objects. Consequently, in the pursuit of constructing ever more intelligent agents, devising methods for reasoning and learning about objects should be regarded as an important goal. Indeed, recent machine learning literature is replete with examples of the benefits of object-like representations when modeling the physical world (e.g. Chang et al. (2016)). However, in order to reason in terms of objects, agents need a way of discovering and detecting objects from visual input a task which we call unsupervised object detection. This task has received significantly less attention in the literature than its supervised counterpart, especially in the case of large images containing many objects. In the current work, we develop a neural network architecture that effectively addresses this large-image, many-object setting. Through a series of experiments, we demonstrate a number of features of our architecture: that, unlike competing approaches, it is able to discover and detect objects in large, many-object scenes and that it has a significant ability to generalize to images that are larger and contain more objects than images encountered during training. We introduce Spatially Invariant Attend, Infer, Repeat (SPAIR), an architecture for unsupervised object detection that combines the unsupervised trainability of Attend, Infer, Repeat (AIR) (Eslami et al., 2016) with the spatial invariance properties of recent work in supervised object detection, particularly YOLO (Redmon et al., 2016). SPAIR is a Variational Autoencoder (VAE) (Kingma and Welling, 2013) with a highly structured, object-like latent representation z, a convolutional (and thus spatially invariant), object-detecting encoder network qφ(z|x), and a decoder network pθ(x|z) that “renders” detected objects into a reconstructed image. We now describe each of these components. Object-like Latent Representation. We first describe the representation scheme used by SPAIR for describing objects in its latent layer. Given an image with shape (Himg,Wimg, 3), it will be useful to spatially divide the image into an (H,W ) grid of cells, where H = dHimg/che, W = dWimg/cwe, and ch/cw are fixed integers giving the cell height/width in pixels. We employ a representation that allows a single object per cell (the extension to multiple objects per cell is straightforward). For a cell with indices (i, j), i ∈ {0, . ,H − 1}, j ∈ {0, . ,W − 1}, the corresponding object is described by the following variables: z what ∈ R A z depth ∈ R z ij pres ∈ {0, 1} z ij where ∈ R 4 z what is a vector with dimension A that stores appearance information for the object. z ij depth specifies the relative depth of the object; in the output image, objects with lower depth appear on top of objects with higher depth. z pres is a binary variable specifying whether the object exists and should be rendered to the output image. z where decomposes as z ij where = (z ij y , z ij x , z ij h , z ij w ), where z ij h and z ij w give the size of the object, and z y and z ij x give the position of the object relative to the current cell. Modeling the Physical World workshop at NeurIPS 2018, Montréal, Canada. Prior Distribution on Objects. A crucial component of a VAE is the prior distribution p(z) over the latent variables. For all real-valued variables, we assume independent Normal distributions; the moments of these distributions are model hyperparameters. For the binary random variables z pres, we design a prior that puts pressure on the network to reconstruct the image using as few objects as possible (i.e. as few z pres = 1 as possible), similar to the prior used by AIR. This pressure is necessary for the network to extract quality object-like representations; without it, the network is free to set all z pres = 1, and use multiple latent objects to explain each object in the image. Encoder Network. Our goal is to design an encoder network qφ(z|x) with spatially invariant properties. To this end, a convolutional neural network e φ (x) is first used to map from the input image x to a feature volume with spatial dimensions (H,W ). Next, this volume is processed sequentially cell-by-cell to produce objects, starting from the top left and proceeding row-by-row toward the bottom right. Processing a cell runs as follows. First, a multi-layer perceptron e φ produces parameters (μ where, σ ij where), (μ ij depth, σ ij depth) and β ij pres for distributions over z ij where, z ij depth and z ij pres, respectively. As input, e φ accepts the feature vector (output of e ) at the current cell as well as sampled objects at nearby cells that have already been processed. e can thus be thought of as encompassing a set of “lateral” connections which facilitate conditioning between nearby objects. Next, values are sampled from the distributions. Concretely: z where ∼ N(μ ij where, σ ij where) z ij depth ∼ N(μ ij depth, σ ij depth) z ij pres ∼ Bernoulli(β pres) The sampled value of z where is then used along with a spatial transformer T (Jaderberg et al., 2015) to extract a glimpse from the image. This glimpse is processed by an object encoder network e φ to yield parameters for a distribution over z what, which is subsequently sampled: μ what, σ ij what = e obj φ (T (x, z ij where)) z ij what ∼ N(μ ij what, σ ij what) Decoder Network. The decoder network is responsible for rendering the detected objects back into an image. First, an object decoder network d θ processes all z ij what to yield a reconstruction of the appearance of each object. Formally we have o , α = d θ (z ij what) where o ij is an RGB volume with shape (Hobj ,Wobj , 3) (for fixed integers Hobj , Wobj), and α is transparancy volume with shape (Hobj ,Wobj , 1). Appearances of all objects are then stitched together to yield a final image, with z where used along with a spatial transformer to give the objects the correct size and location. α is multiplied by z pres to ensure that objects with z pres = 0 are not drawn to the image. For objects that overlap spatially, z ij depth values are used to parameterize a convex combination between the objects, acting as a differentiable approximation of relative object depth. The output of rendering is an image xout, which parameterizes pθ(x|z) as a set of independent, pixel-wise Bernoulli variables. Training. In the VAE framework the Evidence Lower Bound (ELBO) is given by: L(θ, φ) := Ez∼qφ(z|x) [log pθ(x|z)]−DKL(qφ(z|x) ‖ p(z)) (1) It can be shown that log pθ(x) ≥ L(θ, φ); thus, to train the network, a sample-based estimate of the ELBO is maximized with respect to φ and θ using gradient ascent (Kingma and Welling, 2013). To backpropagate through the sampling process, we make use of the reparameterization trick (Kingma and Welling, 2013). For the Normally distributed random variables z where, z ij what and z ij depth this is straightforward. The discrete Bernoulli random variables z pres are replaced with Concrete variables, continuous relaxations of Bernoullis to which the reparameterization trick can be easily applied (Maddison et al., 2016). At validation and test time the Concretes are discretized via rounding.