Unsupervised Causal Generative Understanding of Images
Abstract
We present a novel framework for unsupervised object-centric 3D scene understanding that generalizes robustly to out-of-distribution images.
To achieve this, we design a causal generative model reflecting the physical process by which an image is produced, when a camera captures a scene containing multiple objects.
This model is trained to reconstruct multi-view images via a latent representation describing the shapes, colours and positions of the 3D objects they show.
It explicitly represents object instances as separate neural radiance fields, placed into a 3D scene.
We then propose an inference algorithm that can infer this latent representation given a single out-of-distribution image as input – even when it shows an unseen combination of components, unseen spatial compositions or a radically new viewpoint.
We conduct extensive experiments applying our approach to test datasets that have zero probability under the training distribution.
These show that it accurately reconstructs a scene’s geometry, segments objects and infers their positions, despite not receiving any supervision.
Our approach significantly out-performs baselines that do not capture the true causal image generation process.
Introduction
Most machine learning approaches make the assumption that at test time, they are applied to data drawn from the same distribution as seen during training [14, 91, 73, 21, 33].
This means the generalization guarantees of statistical learning theory apply [127].
However, this does not apply to images drawn from a different distribution – recent works have shown this for images taken from unfamiliar viewpoints [4, 1, 8], shifted by few pixels [7] and showing scenes with an unseen composition of objects [11, 110, 115, 36, 37, 26, 70].
It has been suggested that this is because they learn spurious shortcuts [36] to achieve low training loss, but which do not capture the true causal relationship.
However, when deployed, machine learning methods often encounter observations drawn from a previously unseen distribution.
In this work, we consider the task of transforming a single observed image into a detailed representation of the scene it depicts, providing explicit information about its 3D structure such as object locations, shapes and appearances.
We focus on the challenging setting where at test time, we see images depicting scenes that have zero probability in the training distribution (Fig. 1).
We adopt an unsupervised approach to learning, avoiding the need for manual annotation of object masks, 3D positions, and similar – we require only a dataset of posed multi-view images.
Our approach is to build a generative model (Sec. 3) jointly over image pixels and the 3D world they depict, that can be robustly inverted to infer the latent factors that gave rise to an input image [114, 129, 76, 22].
If such a model is to support generalization to out-of-distribution (OOD) data, it should reflect the underlying causal model of the environment from which the data arose [102, 114].
In practice this means the conditional distributions in the model should correspond (as far as possible) to independent mechanisms [113, 56] in the world.
Thus, when mechanisms or physical processes in the environment change, we expect much of our model to remain applicable; this has been called sparse mechanism shift [114].
For example, even if objects are no longer placed according to the same patterns as at training time, we still expect the appearances of object types to remain consistent.
In general it is impossible to recover such a causal model purely from observational data [102, 114].
We therefore embed in our model the physical knowledge that the world is composed of 3D objects of different shapes, which may appear at different locations, imaged by a camera subject to the laws of 3D geometry and perspective projection.
Each object has separate appearance and position representations, ensuring disentanglement, so the same object can be represented invariantly in different locations.
The final image is synthesized by volumetric rendering of the latent 3D shapes placed in a 3D scene space according to their latent positions.
This structure means the model is compositional : having learnt about one type of object in one context, it can also represent and perform inference about a similar object in different contexts.
We perform inference on OOD data by intervening on our model [102, 105] – i.e. replacing certain conditional distributions (or mechanisms) that are no longer appropriate and then doing posterior inference.
To ensure that the inference method itself supports OOD generalization we cannot employ amortized variational inference [65, 109], as an encoder network is in general not robust to changes of distribution [89, 36].
Instead, we develop a novel Markov chain Monte-Carlo (MCMC) inference scheme, that finds posterior samples for a given test image without any encoder network.
Targeting OOD generalization requires important design choices in our model, that differ from other generative models of 3D scenes [97, 122, 72].
First, we use a non-learnt rendering mechanism, as this is guaranteed to generalise correctly to OOD data (e.g. novel viewpoints); it contrasts with prior works that learn the rendering process (e.g. with a CNN mapping rendered features to pixels [97, 94]).
Second, we use an explicit disentangled representation of the assumed underlying causal variables (e.g. object shapes and positions), which allows performing interventions and counterfactual inference; this contrasts with works that model the scene without disentangling the proper causal variables, such as spatial mixture models [33, 122].
Third, we separate the mechanisms (conditionals) for per-object appearances and for scene layouts, so the latter can be intervened on without affecting the former; this contrasts with methods that have a single global decoder [72, 94].
To evaluate our approach (Sec. 4) we create challenging test datasets that have zero probability under the corresponding training distribution, yet share aspects of its structure.
We show that our model can generalize to unseen numbers of objects, unseen compositions, and radically new camera viewpoints – all significantly better than existing works.
It successfully localizes objects in the 3D scene, reconstructs their 3D shapes, and hence predicts depth-maps and instance segmentation masks.
To summarise, our core contribution is the first unsupervised framework for inference of explicit object-centric 3D scene representations, that generalizes to out-of-distribution scenes .
Our secondary contributions are: (i) experimental evidence that current methods do not generalize to images with a radically new camera viewpoint, a different number of objects, or an unseen composition of objects; (ii) a novel generative model of 3D scenes based on multi-object radiance fields with explicit object positions and volumetric rendering; and (iii) a novel MCMC inference scheme exploiting the structure of our model, that allows inferring 3D scenes from a single OOD image.
Related Work
Recently numerous works have observed that modern learning-based computer vision methods fail to generalise on out-of-distribution data [4, 1, 46, 48, 51, 71, 7, 11, 110, 142, 115, 36, 37, 92, 87, 10, 123, 26, 141, 76] and constructed various benchmarks [68, 47, 8, 42, 62, 136, 131].
Others have tried to improve OOD generalisation, the most relevant by aiming to learn relations that are invariant across training domains [6, 119, 3, 2, 100, 74, 118, 104] or that maximize sparsity of interactions [39, 38].
These typically assume that multiple differently-distributed datasets are available during training, and mainly address the supervised setting (e.g. image classification).
Like the latter, our focus is on enabling OOD generalisation, but in contrast to them, we address unsupervised object-centric generative modelling.
Moreover, we do not rely on multiple training datasets to discover causal factors, but instead directly incorporate universal knowledge such as separation of the world into objects.
Our work is also connected to the vision as inverse graphics paradigm [41, 9, 63, 140, 27, 111, 90].
In this setting, it is assumed that we have access to (maybe parametric) 3D models of objects, and wish to find suitable pose and other parameters to explain an input image [83, 75, 53, 111, 112, 52].
Like our work, these typically use a test-time optimisation; unlike ours, they do not attempt to learn priors on object layout nor shapes – instead assuming these are known a priori .
Neural implicit scene representations [124] learn a continuous representation of a 3D scene from 2D images using neural rendering, either by explicit volumetric rendering [88, 85, 132, 135, 98] or with CNN post-processing of rendered features [121, 120].
These initial works fitted individual scenes without learning common characteristics, therefore requiring many images as input.
This was addressed by sharing models across different scenes [138, 125, 77, 35, 103, 96, 55], allowing inference of novel views from one or few images.
All these methods model a scene as a single entity without decomposing it into individual objects, meaning manipulating the scene (e.g. moving single objects) is not possible.
In contrast, [99] divides a scene into objects, but requires detailed manual annotations to do so, and does not support inference from few images; [28] relies on weaker supervision in the form of ground-truth object masks.
[139, 122] discover such a decomposition automatically (though with depth supervision for [122]), however they rely on spatial mixtures to assign each point in space to an object, without any explicit, controllable representation of object positions.
This entanglement of latent position and appearance means they are not guaranteed to generalise to OOD combinations of position and appearance [81].
Finally, [43] supports composition of multi-object scenes using neural scattering functions – but these must be learnt from multiple views of single objects (a form of weak supervision), with no probabilistic model over appearances nor layouts.
Other approaches extend neural rendering to the generative setting [116, 97, 94, 95, 72, 19, 25, 24], allowing sampling objects or scenes a priori .
However, these do not allow us to perform inference in the OOD setting, as they have components that do not reflect the causal, compositional structure of the world (e.g. a monolithic latent space lacking object-centric representations, or a learnt neural renderer that will not generalise to OOD viewpoints nor compositions).
In contrast, there are object-centric generative models that can sample plausible images and perform inference – but only in 2D, without reasoning over a latent 3D scene representation, and therefore without supporting 3D tasks such as depth prediction.
Some use a full-image spatial mixture model [33, 93, 67, 31, 32, 60] or alpha stacking [126, 128]; others model images as composed of smaller patches or sprites [34, 57, 5], learning both the distribution of appearances and compositions; others use compositional energybased models [29, 80].
A 3D extension of these latter is proposed by [45], but requires videos for inference at test time and does not support OOD generalisation.
Other methods take purely discriminative approaches to unsupervised segmentation [82, 61, 66].
In the supplementary, we provide a table comparing the capabilities of our method to closely related works.
Method
Our goal is to infer an explicit object-centric representation of a 3D scene from one or multiple images, even when they show scenes lying outside the distribution observed during training.
We achieve this by learning a compositional generative causal model jointly over multi-view images and the scenes they depict, as described in Sec. 3.1.
For training (Sec. 3.2) we use only posed multi-view images ( x 1 , v 1 ) , ..., ( x N , v N ) drawn from a training distribution P train ( x , v ) without any annotations such as depth-maps, bounding-boxes or segmentation masks.
At test time, input images are drawn from a distribution P OOD ( x , v ) with support disjoint from the training distribution.
Performing inference over the generative model using our proposed framework (Sec. 3.3) yields an explicit object-centric representation, including 3D object shapes, positions and appearances.
3.1 Compositional Generative Causal Model
We model a multi-view set of N images { x 1 . . .
x N } as caused by a single 3D scene S being rendered from viewpoints { v 1 . . .
v N } , by a function C ( S , v ) .
We now describe this scene representation, then the generative process by which it is sampled.
Scene representation.
The scene S = (cid:16) s bg , { ( s appi , s posi ) }
Oi =1 (cid:17) is composed of a 3D background component s bg
describing the background’s shape and color, and 3D objects indexed i = 1 . . .
O with shape and color described by s appi and explicit 3D positions s posi .
Each s appi explicitly represents the 3D appearance of an object as a neural radiance field (NeRF) [88] in a canonical space (e.g. with the object centered at the origin).
The positions s posi specify objects placement in the global 3D scene space .
While most prior work parametrizes 3D object positions as coordinates, we represent them by 1-hot vectors choosing from a set of plausible candidate locations to use as the center position of the object.
As we explain in Sec. 3.3, this makes gradient-based optimization easier.
Generative process for S .
We first sample a high-level latent scene embedding z g ∼ N ( 0 , I ) , that will model correlations between objects and learn the typical composition of a scene [57, 5].
The individual object appearances are specified by Gaussian latent variables z shapei and z coli that respectively encode the shape and color of the i th object; they are conditioned on z g , with mean and log-variance given by a fully-connected network ζ θ ( z g ) with weights θ .
The position z posi is specified by a categorical variable, with logits given by ξ θ ( z g ) .
We similarly introduce latents z shapebg and z colbg to encode the shape and color of the background.
For brevity, we will write z s = { z shapebg , z colbg , z shape 1 ...
O , z col 1 ...
O , z pos 1 ...O } .
The latents z s are mapped to the scene S by a function
S θ .
This sets s posi equal to z posi , and derives the object NeRF representations s appi from z shapei and z coli as described in the next paragraph.
The probability of an image x n given its camera viewpoint v n is then
[Equation]
where C ( S , v ) renders the scene described by S from viewpoint v , and f N represents a factored
Gaussian likelihood over the H × W × 3 pixels of the image, with fixed standard deviation σ .
The probability of a composition z s of objects and background in a scene is given by
[Equation]
We emphasise p θ ( · ) models different distributions for each object and variable; it does not assume the different scene variables constituting z s are I.I.D. Hence, p θ ( z s | z g ) can model any relationship among object locations, shapes and colors, which is necessary to sample scenes with plausible relationships among scene components.
The graphical model is illustrated in supplementary Fig. 9 in Sec. 7.
Rendering the scene S .
The rendering process C ( S , v ) outputs an image x for a camera viewpoint v , given our explicit compositional representation of a scene S .
Recall S contains a 3D background component s bg and a set of object components { s appi , s posi } Oi =1 ; for brevity we will identify the background as component i = 0 , with s posfixed to the origin.
We extend multi-component neural radiance fields (NeRF) [88, 85] to support explicit placement of objects in the 3D scene according to the position variables s posi .
Specifically, the latent codes z shapei and z coli for the i th object parametrize a learnt function f ∗ θ ( q ∗ ; z shapei , z coli ) , that maps points q ∗ in the canonical space of the object to a color c ∈ [0 , 1] 3 and density σ ∈ R + .
We place each object at its 3D position s posi by convolving its density and color functions with a one-hot location indicator:
[Equation]
where q is a position in scene space, and s posi ( q ) is an indicator function with a unit impulse if point q is chosen as the object center position by the 1-hot indicator s posi .
As in [139] we divide the scene space into foreground/background regions and only render the corresponding components in each.
Given the placed object densities σ i and colors
c i , we calculate the color of each pixel in the image x by casting a ray r ( t ) =
x 0 + t d ∈ R 3 from the pixel in direction d through a camera at position x 0 , summing the contributions from different objects [85, 86]:
[Equation]
where T ( t
[Equation]
= exp
[Equation]
Continuous relaxation of object placement.
To allow gradient-based training and inference, our generative process must be differentiable.
We therefore relax the categorical position variable to a
Gumbel-Softmax [54, 84].
This approach ensures we always receive non-zero gradients of the image with respect to every possible object position, easing optimisation.
This is in contrast to models based on spatial transformers [97, 134, 78], which can get stuck in local minima if the model has a poor initial prediction, as the gradient of pixels wrt position is zero if the predicted and true positions do not overlap.
Note that for a discretized representation (e.g. voxels [79] or triplanes [18]), the object placement operation Eq.
3 can be efficiently implemented as a convolution operation in the Fourier domain – exploiting the Fourier transform’s sifting property [16, 5].
3.2 Training
We train our generative model from a dataset of images containing K views for each of T scenes.
The model includes three learnable components, with parameters θ : (i) f ∗ θ ( q ; z shape , z col ) that represents a 3D object as a function from position to color and density conditioned on the object appearance embedding; (ii) f bgθ ( q ; z shapebg , z colorbg ) that similarly represents the 3D background; (iii) ζ θ and ξ θ , that map the global scene latent z g to parameters of the object and background latents z s .
We train the model using autoencoding variational Bayes [65, 109].
The posteriors over Gaussian latent variables are all diagonal Gaussians (parametrized by mean, and log-variance for z g ), whilst for positions the posterior is Gumbel-Softmax (parametrized by logits).
We use two encoder networks to parametrize these variational posteriors.
enc sϕ ( { x n , v n } Mn =1 ) parametrizes q ( z s |{ x n , v n } Mn =1 ) ; for efficiency, we pass it only a subset of M < K images.
It encodes each observed image and its viewpoint ( x n , v n ) independently then sums the results (as in [35]) before outputting the posterior parameters; this ensures the encoder is invariant to the ordering of images.
enc gϕ ( z s ) parametrizes q ( z g | z s ) , and takes the lower-level latent code z s as input.
For stable training, we adopt a two-stage approach.
We first train the model to reconstruct x 1 ...K , via the object-level latent space z s , ignoring the scene-level latent z g , i.e. maximizing the following loss:
[Equation]
After this has converged, we learn the scene-level latent space by maximizing
[Equation]
We use Adam for optimization [64], β -weighting of KL terms [50], and approximate each of the above expectations by a single sample.
We also further approximate L s by rendering only a random subset of pixels per minibatch.
More implementation details are in the supplementary material.
3.3 Inference for out-of-distribution (OOD) images
At test time, we assume images are sampled from a distribution disjoint from the training distribution.
This means that directly performing posterior inference under our model (which has learnt the training distribution, and ideally assigns zero probability to OOD test images) is not sound.
We therefore make appropriate interventions on our model [102], taking advantage of its causal nature.
For example, when the distribution of object arrangements is different at test time, we replace the learnt prior p θ ( z s ) on object arrangements with an uninformative uniform prior.
Moreover, the variational encoder networks enc sϕ and enc gϕ used during training are not suitable for use at test time, due to domain shift in their inputs.
Therefore, our framework instead directly samples the posterior distribution of latent variables given an observed image, using Markov chain Monte-Carlo (MCMC) inference [69].
Our novel MCMC scheme alternates Langevin dynamics (LD) [12, 130] and Metropolis-Hastings (MH) [44] steps, to infer the latent scene variables ( z s , z g ) from a single observed image x ∗ with viewpoint v ∗ .
The MH steps encourage the Markov chain to make large jumps between modes of the posterior, while the LD steps generate high-probability samples with less exploration.
Each LD step ascends the gradient of
[Equation]
Each MH step first picks an object slot i uniformly at random, then samples a new latents for that object from a proposal distribution ˜ p ( z shapei , z coli ) , accepting/rejecting it according to the usual MH criterion [44].
The proposal distribution ˜ p approximates 1 J (cid:80) Ji =1 p ( z shapei , z coli ) using a Gaussian mixture model fitted by expectation-maximisation [23]; it thus captures the distribution of object latent codes while disregarding the ordering of object indices.
Note that the compositionality of our model increases efficiency of the chain, as it allows MH proposals that affect only one object while keeping other variables fixed – in contrast to a monolithic latent embedding, which would require accepting or rejecting global modifications to the entire scene.
Thus, each MH step need not revert progress made on other variables: e.g. if the background is perfectly inferred but objects are not, then an MH proposal may change only an object, leaving the background intact.
Also, it allows caching computation and only re-rendering parts of the scene that need to be considered for a proposed change (e.g. just background).
In contrast, MCMC on non-structured models must render the entire scene from scratch.
Experiments
We conduct experiments on two synthetic datasets, using our model and three baselines.
We also include ablations of our model, without MCMC inference, and with an unstructured latent space.
We first evaluate performance in the standard setting where the test-set distribution matches the training distribution.
Then, we evaluate generalisation to OOD data, by using several OOD test splits for each dataset, which have zero probability under the training distribution.
We first describe these datasets and test splits (Sec. 4.1), the tasks, metrics and baselines used for evaluation (Sec. 4.2), and finally the performance of our proposed model, the baselines and ablations (Sec. 4.3).
Implementation details for all models (hyperparameters, hardware, etc.) are given in the supplementary material.
4.1 Datasets
GQN .
We render images of rooms containing several objects (cubes, cylinders, spheres), based on the ‘rooms ring camera’ dataset of [35]; similar datasets were used in [45, 33], but in all cases without
OOD test splits.
In the training split, the camera viewpoints are on a circular path around the center of the room, with the camera pointed at the center at fixed elevation angle.
Textures for the walls and colors for the objects are selected randomly from a finite set, with some combinations held out.
Three walls have the same texture as each other, with the fourth different.
There are 3–4 objects present; these are placed near the side of the room identified by the odd texture.
We define the following OOD test splits: (i) position : the 3–4 objects are now placed near a different wall; (ii) composition : unseen combinations of background and object textures; (iii) number of objects : 1/5/6 objects instead of
[Table] [Table]
ARROW .
We render images using a modified version of the CLEVR dataset [59], similar to those in [57].
These have four objects, of which one is always an arrow, two of which are the same as each other, and a fourth that is different.
The arrow always points at the odd (fourth) object.
Object colors are randomly sampled.
The camera has a random azimuth and shallow elevation.
We define the following OOD test splits: (i) position : the four objects are positioned in a line, and the arrow no longer points to the odd object; (ii) composition : all objects are the same shape and color, with no arrow present; (iii) number of objects : 1/5/6 objects instead of four; the arrow still points at one odd object; (iv) viewpoint : camera looks down steeply on the objects.
4.2 Evaluation
Tasks & metrics.
We consider the following discriminative tasks, all taking a single image as input: instance segmentation , measured by adjusted Rand index ( ARI ) and mean segmentation covering ( mSC ) [33], which in the PER - SCENE setting measures how well the model infers segmentation maps for novel viewpoints; depth prediction , measured by the mean relative error between predicted and true depths ( D.MRE ); and pixel reconstruction , measured by peak signal-noise ratio ( PSNR ), which in the PER - IMAGE setting measures how faithfully our latent representation can autoencode OOD images, and in the PER - SCENE setting how well the model performs novel view synthesis (NVS)
We report these metrics according to two different protocols: (i) PER - IMAGE , where we calculate the metrics only on the input image passed to the model; and (ii) PER - SCENE , where we calculate the metrics jointly over 10 images of the scene, taken from different viewpoints, but still having received only one of these as input.
Thus, the latter setting measures how well the model predicts appearance, depth and segmentation from novel viewpoints.
Finally, we evaluate a priori image generation, measured by the Fréchet Inception distance ( FID ) [49] and kernel Inception distance ( KID ) [15] between sampled and ground-truth image distributions.
Baselines.
We compare our approach to five existing works.
Slot Attention [82] is a recent unsupervised object segmentation model, with a spatial mixture representation.
It is purely 2D and not generative.
We additionally experiment with MCMC inference on the trained Slot-Attention decoder (instead of using their encoder), to see whether the decoder generalises to OOD data.
uORF [139] decomposes 2D images into 3D components represented as NeRFs.
Unlike ours, it is not generative, and does not explicitly represent position separate from appearance.
NeRF-VAE [72] is a generative method over 3D NeRFs, but which does not separate individual objects in its latent space.
IODINE [40] is a discriminative 2D method which performs iterative amortized inference over a spatial mixture model.
β -VAE [50] is an unstructured VAE aiming to learn disentangled representations, which has been hypothesized to help with OOD generalization.
Note that the baselines are either discriminative approaches that do not support accounting for distribution shifts, or model the scene with one latent variable which cannot be intervened on to model a different distribution of scenes.
Full details are in the supplementary Sec. 10.
4.3 Results
In-distribution data.
We first evaluate how each model performs on the distribution of images it was trained on.
Results on the discriminative tasks are given in Tab. 1; the top four rows show performance on the test split, which is drawn from the same distribution as the training data (all standard deviations are in the supplementary).
Qualitative results are displayed in Fig.
We see that all methods successfully reconstruct input images (high per-image PSNR).
The 3D-aware methods predict depths with a small relative error (low per-image D.MRE), with NeRF-VAE* slightly better on GQN , and ours on ARROW .
Our method performs similarly to uORF*, with ours slightly better on segmentation but slightly worse for novel view synthesis.
For GQN , Slot Attention performs best on segmentation under ARI scores, with ours better according to mSC and for both metrics on ARROW .
Results for image generation are given in Tab. 2 and Fig.
3 (see the supplementary for images from
NeRF-VAE*).
Note that Slot Attention, IODINE and uORF* are excluded, as they cannot generate images a priori .
Our model has clearly learnt the distribution of both datasets, including the space of likely object and background appearances.
Quantitatively, ours performs similarly to NeRF-VAE* on GQN (with FID slightly better and KID slightly worse), and out-performs it on both metrics for
ARROW .
In Fig.
4, we also show that our semantically-interpretable latent space allows users to edit the scene shown in an image, by counterfactual inference over the model.
[Table]
counterfactual inference.
We condition on the input scene, then modify (intervene on) certain latent variables and re-render the result
OOD data.
We now consider the more challenging setting where test data is drawn from a different distribution than the training set.
Quantitative results in this setting are given in the bottom half of
Tab. 1.
This shows the mean over all OOD splits; a full breakdown is given in the supplementary.
We see that in general, our method significantly out-performs the baselines on these splits – showing the benefit of our causal model and inference scheme.
In particular, our method performs best on all combinations of vision tasks and datasets, except for depth estimation on ARROW where uORF* is slightly better.
Slot-Attention with MCMC and IODINE with iterative amortized inference successfully reconstruct input images during their test-time optimization, but this comes at the cost of lower segmentation performance.
In contrast, our method performs well on both reconstruction and segmentation.
Thus, the best methods on the in-distribution test split are not the best on OOD data.
They leverage features that are predictive of the output in the training set [36] and easy to learn [133, 117, 106], but which are not reliable across different distributions [91].
Our qualitative results (Fig. 2) reinforce this interpretation – we see that our method still predicts accurate segmentations and depth-maps for OOD data, and novel viewpoints look plausible.
Notably, in all cases, it successfully reconstructs the input image via its latent space and synthesises new viewpoints, in spite of never having seen such an image during training.
In contrast, the baselines struggle with images that are different to their training distribution – e.g. mispredicting object colors (1st row, uORF* and
NeRF-VAE*), predicting an over-smoothed average scene (3rd row, NeRF-VAE*), or failing to separate walls and ceiling when they are in an unfamiliar pose (5th row, Slot Attention).
In the supplementary material (Sec. 6), we provide a more detailed analysis of how model performance generalizes on various axes of out-of-distribution images.
Particularly noteworthy is that the 2D baseline Slot Attention [82] has the most significant drop in reconstruction performance (measured by PSNR) on images taken from out-of-distribution camera viewpoints (Tab. 3)
In contrast, our 3D-aware model has a much smaller drop in performance.
Similarly important are the qualitative results in Fig.
5, showing that the non-compositional baseline model NeRF-VAE [72] fails on the straightforward task of reconstructing the input image – it predicts scenes similar to ones seen in the training dataset even when the OOD images are clearly different (e.g. still predicting the same object positions as in the training set despite objects being at the opposite side of the room).
In contrast, our method which uses test-time optimization accurately reconstructs out-of-distribution input images.
4.4 Ablation Study
Ablating MCMC inference.
We perform an ablation study on the effects of our novel MCMC scheme.
We compare it with the common approach of amortised inference (i.e. an encoder network predicts the posterior parameters).
This is denoted in Tab. 1 as Ours (w/o MCMC) .
We see that amortised inference is indeed a critical bottleneck for out-of-distribution generalization: though it performs well when the test distribution is identical to the training distribution, its performance drops significantly on out-of-distribution images, while MCMC holds up.
Ablating generative model structure.
We analyse the effects of using our compositional model compared to a non-compositional model, which has one global latent variable rather than one per object (similar to NeRF-VAE), but otherwise with the same architecture and inference as ours.
This is denoted in Tab. 1 as Ours (w/o structure) .
We see that our proposed compositional generative model significantly improves out-of-distribution generalization over the unstructured ablation: though both perform similarly on IID test data, the unstructured model performs significantly worse on out-of-distribution images.
Ablating p θ ( z s | z g ) .
We analyse the effects of our high-level prior over scene variables.
In Tab. 2, we evaluate samples generated by the model, comparing FID and KID with an ablated model that samples z s from a fixed prior.
The poorer scores from the ablated model demonstrate that our hierarchical approach with a high-level prior is necessary to correctly model the density of scenes: our model samples plausible scenes as it can model relationships between objects, while the ablated model performs much worse.
4.5 Limitations and Societal Impacts
Currently our method has several limitations.
• The scene-level prior distribution p θ ( z s | z g ) uses an unstructured latent space and decoder – thus, this particular component of the model is not interpretable.
It would be worthwhile to use a structured prior, that explicitly models sparse relations among objects.
• We have only tested it on synthetic data; it would be interesting to conduct experiments on natural images.
This would require modelling the image-formation process more faithfully, e.g. by modelling light transport and non-Lambertian materials.
• Our MCMC inference scheme is computationally expensive, as we render a complete image at each step.
The efficiency could be increased by using a stochastic gradient for the LD transitions.
• Our model lacks variables to explicitly model object presence – if an object slot is unused, it learns a latent representation of an empty, zero-density object.
• It would be preferable to use less dataset-specific prior information, e.g. candidate cells for objects.
• Optimizing the model required multiple stages of training, and optimising Eq.
5 was unstable.
It would be worthwhile to investigate joint training as in [57].
We do not anticipate any negative societal impacts from this work.
While all generative models have the potential to be used for creating fake content, our method requires further development to work on realistic images.
On discriminative tasks, support for OOD inference should reduce the susceptibility of models to dataset bias; thus, we see a significant long-term benefit to such methods.
Conclusion
We have presented a new object-centric generative model that captures the causal process by which images are produced, and incorporates universal physical knowledge such as 3D geometry.
We have shown that it is possible to intervene on our model then perform Monte-Carlo inference, in order to process OOD images.
On instance segmentation and novel view synthesis in this OOD setting, it significantly out-performs three state-of-the-art approaches.
Acknowledgments and Disclosure of Funding
We thank Professor Christopher K. I. Williams for detailed feedback on the draft and Professor
Iasonas Kokkinos and Hakan Bilen for fruitful discussions.
We also thank the NeurIPS reviewers for extensive feedback, discussion and experiment suggestions.
This work was mostly done during
T.A.’s internship at Snap Inc.
T.A. was partly supported by an EPSRC Doctoral Training Partnership.