O BJECT R EPRESENTATIONS AS E QUILIBRIA : T RAINING I TERATIVE I NFERENCE A LGORITHMS WITH I MPLICIT D IFFERENTIATION
Michael Chang, Sergey Levine & Thomas L. Griffiths ∗
A BSTRACT
Deep generative models, particularly those that aim to factorize the observations into discrete entities (such as objects), must often use iterative inference procedures that break symmetries among equally plausible explanations for the data.
Such inference procedures include variants of the expectation-maximization algorithm and structurally resemble clustering algorithms in a latent space.
However, combining such methods with deep neural networks necessitates differentiating through the inference process, which can make optimization exceptionally challenging.
In this work, we observe that such iterative inference methods can be made differentiable by means of the implicit function theorem, and develop an implicit differentiation approach that improves the stability and tractability of training such models by decoupling the forward and backward passes.
This connection enables us to apply recent advances in optimizing implicit layers to not only improve the stability and optimization of the slot attention module in SLATE, a state-of-the-art method for learning entity representations, but do so with constant space and time complexity in backpropagation and only one additional line of code.
I NTRODUCTION
[Figure]
Many generative models assume each observed datapoint x is generated from a set λ of latent variables and parameters, where the ordering of these variables does not matter.
In mixture models, for example, a single datapoint would have the same probability if the indices of the latent components were simply relabeled with a different permutation.
However, this symmetry results in many equally valid explanations of the same datapoint, making the inference problem ill-posed.
This phenomenon is especially relevant when learning representations of objects in visual scenes, where treating latent variables symmetrically enables the object representations to be composed in novel ways.
In these models, inference is typically done by breaking symmetry via a random initial guess λ 0 and then iteratively updating the value of λ , with variants of the classical expectation maximization (EM) algorithm
(Dempster et al., 1977) being prominent examples.
Recent works have replaced the analytical updates of EM with a neural network f that directly computes the update as λ =
f ( λ , x ) .
These iterative amortized inference algorithms (Marino et al., 2018a) have predominantly been applied to learning object-centric representations of visual scenes (Greff et al., 2017; Van Steenkiste et al., 2018; Greff et al., 2019; 2020;
Veerapaneni et al., 2020; Locatello et al., 2020; Kipf et al., 2021; Zoran et al., 2021; Singh et al., 2021).
All are trained by differentiating through the unrolled iterations of f .
Despite their conceptual elegance, it has been difficult to scale iterative amortized inference methods beyond modeling simple static scenes or short video sequences because differentiating through the unrolled forward iteration makes training unstable.
Fig.
5 shows that the spectral norm of the Jacobian of f gradually increases over the course of training, which has empirically been observed to cause training instabilities (Bai et al., 2021).
Such instabilities result in sensitivity to hyperparameter choices (e.g., number of inference iterations) and have motivated adding optimization tricks such as gradient clipping, learning rate warm-up, and learning rate decay, all of which make such models more complex and harder to use, restrict the model from optimizing its learning objective fully, and only temporarily delay instabilities that still emerge in later stages of training.
To approach this problem, we observe that previous methods have not taken advantage of the fact that f can be viewed as a fixed point operation.
Thus, f can be trained with implicit differentiation applied at the fixed point, without backpropagating gradients through the unrolled iterations.
We then leverage tools developed for implicit differentiation in neural models for improving the training of these methods.
Our primary contribution is to propose implicit differentiation for training the iterative amortized inference procedures of symmetric generative models, such as those used for learning object representations.
Namely:
We show that prior iterative amortized inference methods, including those on object-centric learning, can be cast as fixed point procedures that can be trained with implicit differentiation.
We call the resulting class of methods implicit iterative inference algorithms.
We show on the latest state-of-the-art of these methods, SLATE (Singh et al., 2021), that using the
first-order Neumann approximation of the implicit gradient for the slot attention module (Locatello et al., 2020) yields substantial improvement in optimization.
We show across three datasets that, compared to SLATE, our method for training achieves much lower validation loss in training, as well as lower Fr ´ echet inception distance (FID) (Heusel et al., 2017) and mean squared error (MSE) in image reconstruction.
Our method also removes the need for gradient clipping, learning rate decay, learning rate warmup, or tuning the number of iterations, while achieving lower space and time complexity in the backward pass, all with one additional line of code.
R ELATED W ORK
The equivalence of maximizing the evidence lower bound (ELBO) (Neal & Hinton, 1998; Bottou &
Bengio, 1995) and minimizing variational free energy (Dayan et al., 1995; Friston, 2010) signifies that any method for minimizing an energy function (LeCun et al., 2006) can be cast as a fixed point procedure.
Thus, iterative inference procedures can naturally be described as specifying the attractor dynamics of a dynamical system whose stable states correspond to the posterior or model parameters produced as a result of the inference, making these procedures especially useful for dynamically inferring representations of objects (Greff et al., 2020).
Thus far, however, the methods that instantiate these iterative inference updates with neural networks (see §3) are difficult to train because they all differentiate through the unrolled dynamics of the fixed point procedure.
Implicit differentiation can offer a solution to this difficulty.
In deep learning (Duvenaud et al., 2020), it has been applied to embedded optimization layers (Amos & Kolter, 2017; Agrawal et al., 2019), neural ordinary differential equations (Chen et al., 2018), meta-learning (Rajeswaran et al., 2019), implicit neural representations (Huang et al., 2021), declarative layers (Gould et al., 2019), and entire networks (Bai et al., 2019).
We focus on the novel application of implicit differentiation to iterative inference algorithms that approximate the optimization of the ELBO.
B ACKGROUND
Our work builds on prior works on iterative amortized inference and implicit differentiation with deep neural networks.
This section reviews recent advances in these two areas and introduces the formalism we use in the rest of the paper.
3. I TERATIVE AMORTIZED INFERENCE
Consider a generative model with observed variables X , local (per-observation) latent variables
Z , and global (across-observations) parameters θ , defining the joint distribution for a particular sample ( z, x ) as p ( z, x ; θ ) =
p ( z ; θ z ) p ( x | z ; θ x | z ) .
Given a datapoint x , the goals of statistical inference often involve estimating the parameters θ or inferring the posterior p ( z | x ) .
Both can be achieved via variational techniques (Neal & Hinton, 1998;
Dayan et al., 1995) that frame inference as a maximization of the evidence lower bound (ELBO) L with respect to θ and an approximate posterior q ( z | · [Equation]
[Equation]
Classical and amortized inference Classical approaches for maximizing L include variants of the expectation maximization (EM) algorithm
(Dempster et al., 1977), which alternates between optimizing max q L ( q, θ, x ) and max θ L ( q, θ, x ) , using incremental (e.g., gradient descent) or analytic approaches.
Given a dataset { x n } Nn =1 , and assuming that q can be parameterized by variational parameters ϕ , classical iterative methods employ a fixed learning rule (Hoffman et al., 2013), for improving ϕ or θ , e.g.
[Equation]
which is costly to scale to high-dimensional datasets.
Thus, techniques related to the variational autoencoder (Kingma & Welling, 2013; Rezende et al., 2014, VAE) amortize (Gershman & Goodman, 2014) the optimization of ϕ n for each x n via an encoder network that directly maps x to ϕ .
However, estimating ϕ without feedback from an iterative procedure results in decreased modeling performance (Krishnan et al., 2018; Cremer et al., 2018) and cannot break symmetry among exchangeable unobserved variables.
Iterative amortized inference Several works have proposed to combine the paradigms of iterative optimization and neural networks by replacing the fixed update rule with a update network f that is trained to optimize the unrolled iterative procedure (Kirsch & Schmidhuber, 2020; Andrychowicz et al., 2016) for improving ELBO.
All proposals so far have trained f by backpropagating gradients through the unrolled updates.
These can be categorized as performing posterior inference or parameter estimation via a meta-learning algorithm (Thrun & Pratt, 2012; Schmidhuber, 1987), with the former conducted across a single dataset (like Andrychowicz et al.
and the latter conducted across a dataset of mini-datasets (like Finn et al.
For methods that meta-learn posterior inference (Marino et al., 2018b;a; 2020; Greff et al., 2019;
Veerapaneni et al., 2020), instead of an encoder network that directly maps x n to ϕ n , an update network f improves a (initially random) previous estimate ϕ nt as ϕ nt +1 ← f ( ϕ nt , ∇ ϕ n t L t ) for each datapoint x n .
While ϕ n is updated per-datapoint, the model parameters θ and weights of f are updated across datapoints.
In contrast, methods that meta-learn parameter estimation (Greff et al., 2017; Van Steenkiste et al., 2018; Zoran et al., 2021; Locatello et al., 2020; Singh et al., 2021) treats each datapoint x n as itself a mini-dataset of M measurements x n,m (e.g. x n is an image and x n,m is a pixel or feature of x n ).
Each datapoint x n is generated from per-datapoint model parameters θ n with per-measurement latents z n,m , thus defining a per-datapoint ELBO L n .
The role of the update network f in this setting is to improve the (initially random) model parameters θ n as θ nt +
←
f ( θ nt , ∇ θ nt L nt ) , which generally also involves improving the per-measurement variational parameters ϕ n,m .
Object-centric learning A concrete application of iterative amortized inference has been in socalled object-centric learning , a research area that seeks to decompose observations x into a set of independent representations of entities without supervision on how to decompose.
Each datapoint x n (e.g. image or sensorimotor sequence) is a set of independent sensor measurements x n,m (e.g. pixels) which are generally posited as having been generated from a mixture model whose components represent the entities.
Under a clustering lens, the problem reduces to finding the K groups of cluster parameters θ n := { θ n,k } K k =1 and cluster assignments ϕ n,m := { ϕ n,m,k } K k =1 that were responsible for the measurements x n,m of the datapoint x n .
Modeling entities as cluster components encodes the assumption that entities are a priori independent and symmetric , thereby requiring a symmetry-breaking mechanism during inference.
To solve this problem, an update network f breaks symmetry among components by alternately updating θ n and ϕ n,m starting from independent randomly initialized θ n
,k ’s.
The state-of-the-art slot attention module (Locatello et al., 2020), e.g., computes θ nt +1 ← f ( θ nt , x n ) , where ϕ n,m is updated as an intermediate step inside f .
The θ n , called slots , serve as input to a downstream objective, e.g. image reconstruction, whose gradients are backpropagated through the unrolling of f .
Earlier works applied this meta-learned parameter estimation approach to binary images (Greff et al., 2017) and videos (Van Steenkiste et al., 2018).
Methods via meta-learned posterior inference have also been developed for images (Greff et al., 2017) and model-based planning (Veerapaneni et al., 2020) using mixture-density networks (Bishop, 1994) for the generative model.
3. I MPLICIT DIFFERENTIATION
Implicit differentiation is a technique for computing the gradients of a function defined in terms of satisfying a joint condition of the input and output.
For example, a fixed point operation f is defined to satisfy “find λ such that λ
= f ( x, λ ) ,” rather than through an explicit parameterization of f .
This fixed point λ ∗ can be computed by simply repeatedly applying f or by using a black-box root-finding solver.
Letting f w be parameterized by weights w , with input x and fixed point λ ∗ , the implicit function theorem (Cauchy, 1831) enables us to directly compute the gradient of the loss ℓ with respect to w , using only the output λ ∗ :
[Equation]
where J f w ( λ ∗ ) is the Jacobian matrix of f w evaluated at λ ∗ .
Compared to backpropagating through the unrolled iteration of f , which is just one of many choices of the solver, implicit differentiation via
Eq.
4 removes the memory cost of storing any intermediate results from the unrolled iteration.
Much effort has been put into approximating the inverse-Jacobian term ( I − J f w ( λ ∗ )) − 1 which has
O ( n 3 ) complexity to compute.
Geng et al.
; Fung et al.
; Huang et al.
(2021); Shaban et al.
(2019) propose instead to approximate ( I − J f w ( λ ∗ )) − 1 with its Neumann series expansion:
[Equation]
The first-order approximation ( T = 1 ) amounts to applying f once to the fixed point λ ∗ and differentiating through the resulting computation graph.
This is not only cheap to compute and easy to implement, but has also been shown empirically (Geng et al., 2021) to have a regularizing effect on the spectral norm of J f w without sacrificing performance.
I NFERENCE A S A F IXED P OINT I TERATION
To explain why it makes sense to train iterative amortized inference algorithms with implicit differentiation, we explicitly unify iterative amortized inference methods (§ 3.1) as solving a particular nested optimization problem whose inner optimization is that of maximizing the ELBO, thereby allowing them to be understood as fixed point procedures.
We describe the abstracted bi-level optimization problem (one-level of nesting), then show that meta-learned posterior inference and meta-learned parameter estimation instantiate this problem with one and two levels of nesting respectively.
4. T HE NESTED OPTIMIZATION PROBLEM
Consider the following bi-level optimization problem over a generic dataset { x n } Nn =1 with datapoints x n .
Define the parameters λ n as optimized per -datapoint, and the parameters w as optimized across datapoints.
With the ELBO L as the inner objective and a task objective J as the outer objective, we express the bi-level optimization problem as
[Equation]
When the inner optimization is conducted via a fixed update rule, the solution of the inner problem can be embedded as a differentiable optimization layer (Amos & Kolter, 2017) within a neural network with weights w .
Here, we partition w as w =
[ w e , w d ] , where w e are weights of an encoder that processes x n into ˜ x n , and w d are weights of a decoder that computes the outer objective J with the fixed point λ n ∗ as an input.
Special cases include the case where w e is the identity (i.e., no pre-processing of x n ) and the case where w d is the identity (i.e. no post-processing of λ n ∗ ).
Using a trainable network f w as the update rule instead, e.g. λ nt +1 ← f w ( λ nt , x n ) , implicitly parameterizes a constraint set C w ( x n ) .
The weights w from Eq.
6 now include the weights w u of the learnable update rule f w , yielding:
[Equation]
The constraint set C w ( x n ) implicitly depends on w e , which pre-processes x n , and w u , which updates λ n .
Any update rule that monotonically improves upon L is thus a fixed point operation whose fixed point locally maximizes L (Neal & Hinton, 1998; Wu, 1983).
It is in this sense that we can understand f w as trained to perform a fixed point operation.
4. P OSTERIOR INFERENCE AND PARAMETER ESTIMATION
Now we show that iterative amortized inference for posterior inference and parameter estimation implement fixed point procedures that solve the aforementioned nested optimization problem.
This is to our knowledge the first unification of both approaches under the same problem statement.
Meta-learned posterior inference Methods for meta-learned posterior inference (§3.1) train a
VAE decoder as the generative model with parameters θ and an update network f w that updates ϕ nt +1 ← f w ( ϕ nt , ∇ ϕ nt L t ) for each datapoint x n .
We recover the problem formulation in Eq.
7 by substituting the negative ELBO for the outer objective
J , the per-datapoint variational parameters ϕ n for λ n , and the model parameters θ for the subset of w d that compute the p ( x | z ) term of the negative ELBO.
Then the update network implements the fixed point operation ϕ
nt +1 ← f w ( ϕ nt , x n ) that computes ∇ ϕ nt L t from ϕ nt and x n as an initial pre-processing step.
Meta-learned parameter estimation Methods for meta-learned parameter estimation (§3.1) treat each datapoint x n as a mini-dataset of measurements x n,m .
We recover the problem formulation in Eq.
7 by substituting a per-datapoint ELBO L n for the inner objective and let λ n := (cid:0) θ n , { ϕ n,m } Mm =1 (cid:1) , meaning that the inner optimization jointly optimizes the per-datapoint model parameters θ n and all per-measurement variational parameters ϕ n,m .
Existing approaches implement this by using f w to compute an EM (Greff et al., 2017; Van Steenkiste et al., 2018) or modified soft K-means (Locatello et al., 2020; Zoran et al., 2021; Singh et al., 2021) step.
Since the variational inference problem (Eq. 1) is itself a bi-level optimization over θ and ϕ , meta-learned parameter estimation is actually a tri-level optimization, optimizing w across datapoints x n at the outer level, θ across measurements x n,m but per-datapoint at the middle level, and ϕ per-measurement at the inner level.
The inner two objectives are the ELBO L n defined for each datapoint and the outer objective J is a task objective specified for the dataset, such as image reconstruction or attribute classification (Locatello et al., 2020).
4. E NTITIES AS INDEPENDENTLY INITIALIZED FIXED POINTS
Having established the above formalism, the object-centric learning problem as studied so far represents a subset of instances of the nested optimization problem described in §4.1, where the inner optimization is of a set of independently initialized parameters λ n := { λ n,k } Kk =1 that are symmetrically updated by f w .
What this paper contributes is the explicit formulation of these entity representations λ n as fast weights that converge towards a set of fixed points during execution.
This gives us a unifying tangible problem statement for focusing object-centric learning research and enables us to improve the training of such methods with implicit differentiation, as we discuss in the next section.
I MPLICIT I TERATIVE I NFERENCE
[Figure]
Having shown how iterative amortized inference methods can be expressed as implementing learnable fixed point operations of the form λ n ← f w ( λ n , x n ) , it is simple to just substitute
J for ℓ in Eq.
4 to get the implicit gradient of J with respect to the weights w .
Because implicit differentiation decouples the forward and backward passes, any black box solver for computing the fixed point and black box gradient estimator for computing the implicit gradient can be used.
We refer to this family of algorithms for solving
Eq.
7 with implicit differentiation as implicit iterative inference algorithms.
Implicit object-centric learning To illustrate an example of such implicit inference algorithms, we propose implicit slot attention : a method for training the state-of-the-art slot attention module (Locatello et al., 2020), which performs meta-learned parameter estimation, with the simplest and most effective method that we have empirically found for approximating the implicit gradient, which is its first-order Neumann approximation (Eq. 5).
It can be implemented by simply differentiating the computation graph of applying the slot attention update once to the fixed point θ n ∗ , where θ n ∗ is computed by simply iterating the slot attention module forward as usual, but without the gradient tape.
The time and space complexity of backpropagation for our method compared to vanilla slot attention as a function of the number of slot attention iterations n , is shown below: vanilla
[Equation]
Our method is not only more efficient but also requires only one additional line of code (Fig. 2).
E XPERIMENTS
The main hypothesis behind this paper is that implicit differentiation can improve the training of iterative amortized inference methods for object-centric learning.
We test this hypothesis by replacing the backward pass of the slot attention module in SLATE (Singh et al., 2021) with the first-order
Neumann approximation of the implicit gradient, and measuring optimization performance.
For the task of image reconstruction, SLATE uses a discrete VAE (Ramesh et al., 2021) to compress an input image into a grid of discrete tokens.
These tokens index into a codebook of latent code-vectors, which, after applying a learned position encoding, serve as the input to the slot attention module.
An Image GPT decoder (Chen et al., 2020) is trained with a cross-entropy loss to autoregressively reconstruct the latent code-vectors, using the outputted slots from slot attention as queries and the latent code-vectors as keys/values.
Gradients are blocked from flowing in and out of the discrete VAE
[Figure]
to the rest of the network (i.e. the slot attention module and the Image GPT decoder)
, but the entire system is trained simultaneously.
We consider three datasets: CLEVR-Mirror (Singh et al., 2021), Shapestacks (Groth et al., 2018), and COCO-2017 (Lin et al., 2014).
We obtained CLEVR-Mirror directly from the SLATE authors and used a 70-15-15 split for training, validation, and testing.
We pooled all the data variants of
Shapestacks together as Singh et al.
(2021) did and used the original train-validation-test splits.
The
COCO-2017 dataset was downloaded from FiftyOne and used the original train-validation-test splits.
6. D OES IMPLICIT DIFFERENTIATION STABILIZE THE TRAINING OF SLOT ATTENTION ?
[Table]
Using the two primary metrics used in Singh et al.
(2021), images generated by SLATE trained with implicit differentiation achieve both lower pixel-wise mean-squared error and FID score (Heusel et al., 2017).
The FID score was computed with the PyTorch-Ignite (Fomin et al., 2020) library using the inception network from the PyTorch port of the FID official implementation.
All methods were trained for 250k gradient steps.
Table 1 compares the FID and MSE scores of the images that result from compressing the SLATE encoder’s set of discrete tokens through the slot attention bottleneck, using Image-GPT to autoregressively re-generate these image tokens one by one, and using the discrete VAE decoder to render the generated image tokens.
Implicit differentiation significantly improves the quantitative image reconstruction metrics of SLATE across the test sets of CLEVR-Mirrors, Shapestacks, and COCO.
In the case of MSE for CLEVR, this is almost a 7x improvement.
The higher quantitiatve metrics also translate into better quality reconstructions on the test set, as shown in Figure 3.
For CLEVR-Mirrors, vanilla SLATE sometimes drops or changes the appearance of objects, even simple scenes with three objects.
In contrast, the reconstructions produced from training with implicit differentiation match the ground truth very closely.
For Shapestacks, our method consistently segements the scene into constituent objects.
This is sometimes the case with vanilla SLATE on the training and validation set as well, but we observed for both of the seeds we ran that vanilla SLATE produced degenerated attention maps where one slot captures the entire foreground, and the background is divided among the other slots.
The visual complexity of the
COCO dataset is much higher than either CLEVR-Mirrors and Shapestacks, and the reconstructions on the COCO dataset are quite poor, for both SLATE’s discrete VAE and consequently for the
[Figure]
v a li d a t i o n c r o ss e n tr o p y
[Figure]
reconstruction through the slot bottleneck.
This may be expected because we did not attempt to tune
SLATE’s hyperparameters to COCO, but it does highlight the gap that still exists between using the state-of-the-art in object-centric learning out-of-the-box and what the community may want these methods to do.
The attention masks for both the vanilla SLATE and our method furthermore do not appear to correspond consistently to coherent objects in COCO but rather patches on the image that do not immediately seem to match with our human intuition of what constitutes a visual entity.
6. C AN WE SIMPLIFY THE NEED FOR OPTIMIZATION TRICKS ?
To further understand the benefits of implicit differentiation, we then ask whether it stabilizes the training of slot attention without the need for optimization tricks like learning rate decay, gradient clipping, and learning warmup.
Fig.
4 shows that these tricks generally help regularize spectral norm of the Jacobian of vanilla slot attention but are not required by our method.
Decaying the learning rate regularizes the Jacobian norm from exploding, but it also hurts optimization performance for both our method and vanilla SLATE, as expected.
When we remove gradient clipping, the Jacobian norm of vanilla SLATE explodes, whereas it stays stable for our method.
Lastly, removing learning rate warmup also consistently makes vanilla SLATE’s training unstable, whereas it only affects the stability of our method for one out of three seeds.
D ISCUSSION AND L IMITATIONS
Our results show clear signal that implicit differentiation can offer a significant optimization improvement over backpropagating through the unrolled iteration of slot attention, and potentially any iterative inference algorithm, with lower space and time complexity and only one additional line of code.
Despite our work pushing the optimization performance for a state-of-the-art model in object-centric learning, the discrepancy between the quantitative improvement in optimization and evaluation metrics on the one hand and the less intuitive qualitative attention masks on real world observations like COCO (Fig. 3) on the other hand still suggests a gap between what we optimize these methods to do and what we actually want them to do.
This paper proposes a novel conceptualization of object representations as fast weights that converge towards a set of fixed points during execution.
Because it is so simple to apply implicit differentiation to any fixed point algorithm, we hope this work inspires future work to leverage tools developed for implicit differentiation for improving object-centric learning and methods for learning latent structure more broadly. A CKNOWLEDGEMENTS This work was supported by ARL, W911NF2110097, with computing support from Google Cloud Platform.