Original: Object Representations as Fixed Points: Training Iterative Reﬁnement Algorithms with Implicit Differentiation
Revised: Object Representations as Fixed Points: Training Iterative Refinement Algorithms with Implicit Differentiation

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Object-centric learning has been motivated by developing algorithms that inferindependent and symmetric entities from perceptual input. This often requiresthe use of iterative reﬁnement procedures that break symmetries among equallyplausible explanations for the data, but most prior works differentiate through theunrolled reﬁnement process, which can make optimizationexceptionally challenging.
Revised: Iterative refinement – start with a random guess, then iteratively improve the guess – is a useful paradigm for representation learning because it offers a way to break symmetries among equally plausible explanations for the data. This property enables the application of such methods to infer representations of sets of entities, such as objects in physical scenes, structurally resembling clustering algorithms in latent space. However, most prior works differentiate through the unrolled refinement process, which can make optimization challenging.

Original: We observe that suchiterative reﬁnement methods can be made differentiableby means of the implicit function theorem, and develop an implicit differentiationapproach that improves the stability and tractability of training such models bydecoupling the forward and backward passes.
Revised: We observe that such 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 by decoupling the forward and backward passes.

Original: This connection enables us to applyrecent advances in optimizing implicit layers to not only improve the stability andoptimization of the slot attention module in SLATE, a state-of-the-art method forlearning entity representations, but do so with constant space and time complexityin backpropagation and only one additional line of code.
Revised: This connection enables us to apply advances in optimizing implicit layers to not only improve the 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.

Original:  Introduction
Revised:  Introduction

Original: One of the challenges for machine learning – particularly when working with complex real-worldproblems such as interpreting visual scenes – is deﬁning and performing inference in models thatare able to capture that complexity. Generative models attempt to do so by deﬁning probabilitydistributions that characterize the equivalently complex process by which data are generated, postulating latent parameters as intermediate steps in that process.
Revised: Conventionally, neural network models implement feedforward computation, transforming the input x into the output z through a fixed series of operations corresponding to distinct layers as z = f N ( f N − 1 ( ...f 2 ( f 1 ( x )) ... )) . However, a range of more sophisticated models implement iterative computation with the network, typically formulated and motivated as some sort of optimization procedure where the correct answer is the fixed point of an iterative refinement z

Original: We are interested in generative modelsthat assume each observed datapoint x is generated from a set  of latent parameters.
Revised: = f x ( z ) of an initial guess z 0 .

Original: For example,recent work has modeled objects in visual scenes by treating the representations of those objects asparameters of mixture model components [18, 28].
Revised: This includes diffusion models [30, 52, 58, 59], energy-based models [14, 37], deep equilibrium models [7], iterative amortized inference procedures [40–42], neural ordinary differential equations [11], meta-learning algorithms [5, 17, 23] and object-centric models [16, 35, 61, 62].

Original: In these latent variable models, inference is typically done by breaking symmetry via a random initialguess  0 and then iteratively updating the value of  , with variants of  the classical ExpectationMaximization (EM) algorithm [11] being prominent examples. Recent works have replaced theanalytical updates of EM with a neural network f that directly computes the update as  =
Revised: Such iterative refinement procedures may have a number of advantages over direct feedforward computation: they can serve to simplify the learned function (e.g., in the same way that a recursive program might be much simpler than an equivalent program implemented without recursion), introduce an inductive bias into the model that improves generalization, and break symmetries.

Original: f (  , x ) ,resulting in a class we call iterative reﬁnement algorithms . These algorithms have predominantly been applied to learningobject-centric representations of visual scenes [18–20, 26, 28, 37, 39, 40, 44].
Revised: In this work, we consider the case of iterative refinement applied to representation learning of latent sets, which has primarily been applied to learning representations of objects from pixels [24– 26, 35, 39, 57, 61, 62, 66].

Original: All are trained by differentiating through the unrolled iterations of f .
Revised: The particular challenge of this setting is that invariance of set elements to permutation means there are many latent sets that serve as equally plausible explanations for the data. Iterative refinement is especially useful in this context because it breaks the symmetry among these explanations with the randomness of the initial guess z 0 rather than encoding the symmetry-breaking mechanism in the weights of the network, as do conventional methods that learn a direct feedforward [Figure] mapping from observation to representation.

Original: The state-of-the-art  is the slotattention module from Locatello et al.
Revised: The state-of-the-art of these iterative object-centric methods is the slot attention module from Locatello et al.

Original: [28], which serves as the focus of our paper.
Revised: [39], serving as the focus of our paper.

Original: Despite their conceptual elegance, it has been difﬁcult to scale iterative reﬁnement methods beyond modeling simple static scenes or short video sequences because differentiating through the unrolledforward iteration makes training unstable.

Original: Unfortunately, these methods have been notoriouslydifﬁcult to train.
Revised: Unfortunately, these methods have been notoriously difficult to train.

Original: Fig.
Revised: Their nature as unsupervised representation learning methods means that we do not have ground-truth targets to supervise the outputs of each iteration of f , as in other settings (e.g. diffusion models). As a result, prior works train these methods by differentiating through the unrolled iterations of f .

Original: 1b shows that the  spectral norm ofthe Jacobian of f gradually increases over [Figure] the course of training, which has empirically been observed to cause training instabilities [7].
Revised: Differentiating through this recurrence contributes to various training instabilities, as we can see by the growing spectral norm of f in Fig. 1b.

Original: Suchinstabilities result in sensitivity to hyperparameter choices (e.g., number of inference iterations) andhave motivated adding optimization tricks such as gradient clipping, learning rate warm-up, andlearning rate decay, all of which make such models more complex and harder to use, restrict themodel from optimizing its learning objective fully, and only temporarily delay instabilities that stillemerge in later stages of training.
Revised: Such instabilities result in sensitivity to hyperparameter choices (e g., number of refinement steps) 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.

Original: To approach this problem, we observe that previous  methods  have not taken  advantage of the factthat f can be viewed as a ﬁxed point operation.
Revised: To approach this problem, we observe that previous iterative refinement methods like slot attention have not taken full advantage of the fact that f can be viewed as a fixed point operation.

Original: Thus, f can be trained with implicit differentiationapplied at the ﬁxed point, without backpropagating gradients through the unrolled iterations [8, 12].
Revised: Thus, f can be trained with implicit differentiation applied at the fixed point, without backpropagating gradients through the unrolled iterations [9, 15].

Original: We then leverage tools developed for implicit differentiation in neural models for improving thetraining ofthese iterative reﬁnement algorithms.
Revised: This paper investigates how advances in implicit differentiation in neural models can be applied to improve the training of iterative refinement methods.

Original: Our primary contribution is to propose implicit differentiation for trainingthe iterative reﬁnementprocedures of symmetric generative models, speciﬁcally slot attention.
Revised: Our primary contribution is to propose implicit differentiation for training iterative refinement procedures, specifically slot attention, for learning representations of latent sets.

Original: First, we show that slotattention can be cast as a ﬁxed point procedure that can be trained with implicit differentiation,resulting in what we call implicit slot attention .
Revised: First, we show that slot attention can be cast as a fixed point procedure that can be trained with implicit differentiation, resulting in what we call implicit slot attention .

Original: Second, we show on the latest state-of-the-artmethod of this kind, SLATE [37], that using the ﬁrst-order Neumann approximation of the implicitgradient for the slot attention module yields substantial improvement in optimization.
Revised: Second, we show on the latest state-of-the-art method of this kind, SLATE [57], that using the first-order Neumann approximation of the implicit gradient for the slot attention module yields substantial improvement in optimization.

Original: Third, weshow across three datasets that, compared to SLATE, our method for training achieves much lowervalidation loss in training, as well as lower Fréchet inception distance (FID) [23] and mean squarederror (MSE) in image reconstruction.
Revised: Third, 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échet inception distance (FID) [29] and mean squared error (MSE) in image reconstruction.

Original: Fourth, our method also removes the need for gradient clipping,learning rate decay, learning rate warmup, or tuning the number of iterations, while achieving lowerspace and time complexity in the backward pass, all with just one additional line of code.
Revised: Fourth, our method also removes the need for gradient clipping, learning rate warmup, or tuning the number of iterations, while achieving lower space and time complexity in the backward pass, all with just one additional line of code.

Original: Fifth, whenintegrated with the original slot attention encoders and decoders from Locatello et al.
Revised: Fifth, when integrated with the original slot attention encoders and decoders from Locatello et al.

Original: [28], implicitdifferentiation substantially improves object property prediction and continues to predict intuitivesegmentation masks as the vanilla slot attention.
Revised: [39], implicit differentiation substantially improves object property prediction and continues to predict intuitive segmentation masks as the vanilla slot attention.

Original:  Related Work
Revised:  Related Work

Revised: Much early work in artificial intelligence followed a paradigm of using an iterative learning procedure during execution time, whether it be a form of search, inference, or optimization [51].

Revised: These include early work on variational inference and energy-based models [12, 19, 37], Hopfield networks [31],

Revised: Boltzmann machines [1], and associative memory [36].

Revised: The rise of modern deep networks in the last decade shifted the method for computing solutions during execution time away from iterative procedures and towards producing the solution directly with a feedforward pass of a network.

Revised: Recent works have started to shift the paradigm of execution back to combining the best of function approximation and iterative search, with neural networks parameterizing initializations [17], update rules [5, 24, 40], search heurstics [34, 56], and evaluation functions [14].

Original: Our  methodological novelty is the adaptation of implicit differentiation techniques for trainingsymmetric factorized object-centric models, where our key insight is that the iterative procedure usedfor symmetry-breaking reaches a ﬁxed point, thus allowing implicit differentiation to be used.
Revised: Our work concerns the optimization of neural networks as update rules for these iterative refinement procedures. Our methodological novelty is the adaptation of implicit differentiation techniques for training iterative refinement procedures for representing of latent sets, where our key insight is that the iterative procedure used for symmetry-breaking reaches a fixed point, thus allowing implicit differentiation to be used.

Original: We arenot aware that this has been done before, as current methods that perform iterative reﬁnement forobject-centric learning [18–20, 26, 28, 37, 39, 40, 44] all differentiate through the unrolled dynamicsof the ﬁxed point procedure, which as we show makes them difﬁcult to train.
Revised: We are not aware that this has been done before, as current methods that perform iterative refinement for representing latent sets, often referred to as object-centric learning [24– 26, 35, 39, 57, 61, 62, 66], all differentiate through the unrolled dynamics of the fixed point procedure, which as we show makes them difficult to train.

Original: Our work draws upon innovations in implicit differentiation that have been applied in variousother applications besides object-centric learning, such as embedded optimization layers [1, 3], neural ordinary differential equations [10], meta-learning [32], implicit neural representations [24],declarative layers [17], and transformers [6].
Revised: Our work draws upon innovations in implicit differentiation that have been applied in various other applications besides object-centric learning, such as embedded optimization layers [2, 4], neural ordinary differential equations [11], meta-learning [49], implicit neural representations [32], declarative layers [22], and transformers [7].

Original: Although we evaluate various technique for implicitdifferentiation beyond any speciﬁc one, our results highlight the beneﬁts of the simplest of these,which is that of using a truncated Neumann approximation [15, 16, 24, 36].
Revised: Although we evaluate various techniques for implicit differentiation, our results highlight the benefits of the simplest of these, which is that of using a truncated Neumann approximation [20, 21, 32, 55].

Original: While this technique wasused in Zoran et al.
Revised: While this technique was used in Zoran et al.

Original: [44] without theoretical explanation, we propose implicit differentiation as anexplanation for why this technique is beneﬁcial.
Revised: [66] without theoretical explanation, we propose implicit differentiation as an explanation for why this technique is beneficial.

Original: The closest works to ours is the concurrent workof Zhang et al.
Revised: The closest works to ours is the concurrent work of Zhang et al.

Original: [43] which applies implicit differentiation to predicting properties of set elements andﬁnds similar beneﬁts to implicit differentiation as we do.
Revised: [65] which applies implicit differentiation to predicting properties of set elements and finds similar benefits.

Original: Our approach generalizes their results tothe unsupervised setting and scales to high dimensional outputs (e.g. images), taking a crucial stepforward toward improving self-supervised learning of symmetric generative models.
Revised: Our approach generalizes their results to the unsupervised setting and scales to high dimensional outputs (e.g., images), taking a crucial step toward improving representation learning of latent sets.

Original:  Background
Revised:  Background

Original: Our work builds on prior works on iterative reﬁnement and implicit differentiation with deep networks.
Revised: Our work builds on prior works on iterative refinement and implicit differentiation with deep networks.

Original: Iterative reﬁnement for object-centric learning Current work in object-centric learning has beenmotivated bydeveloping learning algorithms that infer independent and symmetric entities fromperceptual input.
Revised: Iterative refinement for inferring representations of latent sets Current work on inferring representations of latent sets is motivated by learning to represent objects – which we consider independent and symmetric entities – from perceptual input.

Original: Because mixture models are deﬁned to have a priori independent and symmetricmixture components,  these  methodsgenerally model each datapoint x n (e.g. image) as a set ofindependent sensor measurements x n,m (e.g. pixels) which are posited as having been generatedfrom a mixture model whose components represent the entities.
Revised: Because mixture models are also defined to have a priori independent and symmetric mixture components, they have been the model of choice for representing entities: thus these iterative methods model each datapoint x n (e.g. image) as a set of independent sensor measurements x n,m (e.g. pixels) which are posited as having been generated from a mixture model whose components represent the entities.

Original: Under a clustering lens, the problemreduces to ﬁnding the K groups of cluster parameters ✓ n := { ✓ n,k } Kk =1 and cluster assignments n,m := {  n,m,k } Kk =1 that were responsible for the measurements x n,m of the datapoint x n .
Revised: Under a clustering lens, the problem reduces to finding the K groups of cluster parameters θ n := { θ n,k } Kk =1 and cluster assignments ϕ n,m := { ϕ n,m,k } Kk =1 that were responsible for the measurements x n,m of the datapoint x n .

Original: A network f breaks symmetry among components by alternately updating ✓ n and  n,m startingfrom independent randomly initialized ✓ n,k
Revised: A network f breaks symmetry among components by alternately updating θ n and ϕ n,m starting from independent randomly initialized θ n

Original:  ’s.
Revised: ,k ’s.

Original: The slot attention module [28], e.g., computes ✓ nt +1f ( ✓ nt , x n ) , where  n,m is updated as an intermediate step inside f .
Revised: The slot attention module [39], e.g., computes θ nt +1 ← f ( θ nt , x n ) , where ϕ n,m is updated as an intermediate step inside f .

Original: The ✓ n , called slots , serve asinput to a downstream objective, e.g. image reconstruction, whose gradients are backpropagatedthrough the unrolling of f .
Revised: The θ n , called slots , serve as input to a downstream objective, e.g. image reconstruction, whose gradients are backpropagated through the unrolling of f .

Original: Earlier works applied this approach to binary images [18] and videos [39].
Revised: Earlier works applied this approach to binary images [24] and videos [61].

Original: Implicit differentiation Implicit differentiation is a technique for computing the gradients of afunction deﬁned in terms of satisfying a joint condition on the input and output.
Revised: Implicit differentiation Implicit differentiation is a technique for computing the gradients of a function defined in terms of satisfying a joint condition on the input and output.

Original: For example, a ﬁxedpoint operation f is deﬁned to satisfy “ﬁnd  such that = f
Revised: For example, a fixed point operation f is defined to satisfy “find z such that

Original:  ( x,  ),” rather than through an explicitparameterization of f .
Revised: z = f ( z, x ) ” (or written as z = f x ( z ) ) rather than through an explicit parameterization of f .

Original: This ﬁxed point  ⇤ can be computed by simply repeatedly applying f or byusing a black-box root-ﬁnding solver.
Revised: This fixed point z ∗ can be computed by simply repeatedly applying f or by using a black-box root-finding solver.

Original: Letting f w be parameterized by weights w , with input x andﬁxed point  ⇤ , the implicit function theorem [8] enables us to directly compute the gradient of theloss ` with respect to w , using only the output  ⇤ :
Revised: Letting f w be parameterized by weights w , with input x and fixed point z ∗ , the implicit function theorem [9] enables us to directly compute the gradient of the loss ℓ with respect to w , using only the output z ∗ :

Original: [Equation]
Revised: [Equation]

Original: where J f w (  ⇤ ) is the Jacobian matrix of f w evaluated at  ⇤ .
Revised: where J f w ( z ∗ ) is the Jacobian matrix of f w evaluated at z ∗ .

Original: Compared to backpropagating throughthe unrolled iteration of f , which is just one of many choices of the solver, implicit differentiationvia Eq.
Revised: Compared to backpropagating through the unrolled iteration of f , which is just one of many choices of the solver, implicit differentiation via Eq.

Original: 1 removes thememory cost of storing any intermediate results from the unrolled iteration.
Revised: 1 removes the cost of storing any intermediate results from the unrolled iteration.

Original: Deep equilibrium models (DEQ)Bai et al.
Revised: Deep equilibrium models (DEQ)

Original: [6] represent the class ofsuch functions f parameterizedby neural networks, which  empirically produce stable ﬁxed points even though the convergenceproperties of these models remains theoretically not well understood.
Revised: [7] represent the class of functions f parameterized by neural networks, which have successfully been trained with implicit differentiation and empirically produce stable fixed points even though their convergence properties remains theoretically not well understood.

Original: Much effort has been put into approximating the inverse-Jacobian term ( I  J f w (  ⇤ ))  1 which has
Revised: Much effort has been put into approximating the inverse-Jacobian term ( I − J f w ( z ∗ )) − 1 which has

Original: O ( n 3 ) complexity to compute.
Revised: O ( n 3 ) complexity to compute.

Original: Using notation from Bai et al.
Revised: Using notation from Bai et al.

Original: [7], Almeida [2], Pineda [31] proposeto approximate the u > term in Eq.
Revised: [8], Pineda [48] and Almeida [3] propose to approximate the u ⊤ term in Eq.

Original: 1 as the ﬁxed point of the linear system:
Revised: 1 as the fixed point of the linear system:

Original: [Equation]
Revised: [Equation]

Original: which can be solved with any black-box solver.
Revised: which can be solved with any black-box solver.

Original: However, in the context of applying implicitdifferentiation to neural networks, this approach has in practice required some expensive regularizationto maintain stability [7], so Fung et al.
Revised: However, in the context of applying implicit differentiation to neural networks, this approach has in practice required some expensive regularization to maintain stability [8], so Fung et al.

Original: [15], Geng et al.
Revised: [20], Geng et al.

Original: [16], Huang et al.
Revised: [21], Huang et al.

Original: [24], Shaban et al.
Revised: [32], Shaban et al.

Original: [36] propose instead to approximate ( I  J f w (  ⇤ ))  1 with its Neumann series expansion, yielding
Revised: [55] propose instead to approximate ( I − J f w ( z ∗ )) − 1 with its Neumann series expansion, yielding

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst-order approximation ( T = 1 ) amounts to applying f once to the ﬁxed point  ⇤ anddifferentiating through the resulting computation graph.
Revised: The first-order approximation ( T = 1 ) amounts to applying f once to the fixed point z ∗ and differentiating through the resulting computation graph.

Original: This is not only cheap to compute and easyto implement, but has also been shown empirically [16] to have a regularizing effect on the spectralnorm of J f w without sacriﬁcing performance.
Revised: This is not only cheap to compute and easy to implement, but has also been shown empirically [21] to have a regularizing effect on the spectral norm of J f w without sacrificing performance.

Original:  Implicit Iterative Reﬁnement
Revised:  Implicit Iterative Refinement

Original: Our main contribution is to propose treating iterative reﬁnement algorithms used for object-centriclearning as ﬁxed-point procedures, thereby motivating the use of implicit differentiation to train them.
Revised: Our main contribution is to propose treating iterative refinement algorithms used for inferring latent sets as fixed-point procedures, thereby motivating the use of implicit differentiation to train them.

Original: 4.1 Motivation
Revised: 4.1 Motivation

Original: Our approach is inspired by the structural resemblance betweenthese iterative reﬁnement methodsand the Expectation-Maximization algorithm [11], which also has been pointed out by Greff et al.[18] and Locatello et al.
Revised: Our approach is inspired by the structural resemblance between iterative refinement for latent sets and the Expectation-Maximization algorithm [13], which also has been pointed out by Greff et al. [24] and Locatello et al.

Original: If f were to update the components ✓ n and assignments  n,m in a waythat monotonically improves the evidence lower bound (ELBO)
Revised: If f were to update the components θ n and assignments ϕ n,m in a way that monotonically improves the evidence lower bound (ELBO)

Original: L n on the log-likehood of each imagex n with respect to ✓ n and  n,m , then we know from Neal and Hinton [29], Wu [42] that such anapproach is a ﬁxed point operation whose ﬁxed point locally maximizes L n .
Revised: L n on the log-likehood of each image x n with respect to θ n and ϕ n,m , then we know from Neal and Hinton [43] and Wu [64] that such an approach is a fixed point operation whose fixed point locally maximizes L n .

Original: With this interpretation,such iterative reﬁnement algorithms for object-centric learning can be viewed as performing a nestedoptimization with three levels, where the weights of the slot attention module are optimized acrossimages x n , the components ✓ n are optimized per-image x n but across measurements x n,m , and theassignments  n,m are optimized per-measurement x n,m .
Revised: With this interpretation, such iterative refinement algorithms for inferring latent sets can be viewed as performing a nested optimization with three levels, where the weights of the slot attention module are optimized across images x n , the components θ n are optimized per-image x n but across measurements x n,m , and the assignments ϕ n,m are optimized per-measurement x n,m .

Original: How these iterative reﬁnement methods are implemented differs from optimizing the per-datapoint
Revised: How these iterative refinement methods are implemented differs from optimizing the per-datapoint

Original: ELBO described above, however.
Revised: ELBO described above, however.

Original: Indeed, both Greff et al.
Revised: Indeed, both Greff et al.

Original: [18] and Locatello et al.
Revised: [24] and Locatello et al.

Original: [28] implementedablations to their methods that do implement a variant of Expectation-Maximization on a well-deﬁnedmixture model, but found that replacing the update rule with a recurrent neural network empiricallyperforms better at extracting representations.
Revised: [39] implemented ablations to their methods that do implement a variant of Expectation-Maximization on a well-defined mixture model, but found that replacing the update rule with a recurrent neural network empirically performs better at extracting representations.

Original: 4.2 Iterative reﬁnement as a ﬁxed point procedure
Revised: 4.2 Iterative refinement as a fixed point procedure

Original: While in general we do not yet have the theoretical tools to understand whether and how the slots ofslot attention optimize a well-deﬁned objective like the ELBO,  we empirically ﬁnd that the forwardrelative residual of the slots do appear stable (§5.2).
Revised: While in general we still lack the theory to understand whether and how the slots of slot attention optimize a well-defined objective like the ELBO, and thus have no theoretical guarantees that slot attention does converge to a fixed point, we empirically observe from its stable forward relative residuals (Fig. 3a) that it does appear to approximately converge to a fixed point.

Original: This suggests that slot attention can be understoodas an instance of a DEQ[6] that uses naive forward iteration to ﬁnd the ﬁxed point and backpropagatesthrough the iteration to compute the gradient.
Revised: This suggests that slot attention can be understood as an instance of a DEQ that uses naive forward iteration to find a fixed point and backpropagates through the iteration to compute the gradient.

Original: While the theoretical convergence properties of DEQsare an open question that is beyond the scope of this paper, they have been reported to empiricallyconverge to a ﬁxed point and have successfully been trained with implicit differentiation.

Original: The novelimplication of this to object-centric methods is that any root-ﬁnding solver and implicit gradientestimator can in principle be used to train slot attention, and we ﬁnd that many combinations do trainslot attention effectively in §5.2.
Revised: The novel implication of this to inferring latent sets is that any root-finding solver and implicit gradient estimator can in principle be used to train slot attention, and we find that many combinations do train slot attention effectively in Fig. 3b.

Original: This insight is one of our main contributions, but in the next sectionwe describe a particular combination that we empirically found effective.
Revised: This insight is one of our main contributions, but in the next section we describe a particular combination that we empirically found effective.

Original: 4.3 Implicit slot attention
Revised: 4.3 Implicit slot attention

Original: [Table]
Revised: [Table]

Original: We propose implicit slot attention : a method for training the state-of-the-art slot attention module [28] with thesimplest and most effective method that we have empirically found for approximating the implicit gradient, whichis its ﬁrst-order Neumann approximation (Eq. 3).
Revised: We propose implicit slot attention : a method for training the state-of-the-art slot attention module [39] 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. 3).

Original: It canbe implemented by simply differentiating the computationgraph of applying the slot attention update once to the ﬁxedpoint ✓ n ⇤ , where ✓ n ⇤ is computed by simply iterating the slotattention module forward as usual, but without the gradient tape, amounting to truncating the backpropagation.
Revised: 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, amounting to truncating the backpropagation.

Original: The time and space complexity of backpropagation for implicit slot attention compared to vanilla slot attention as a function of the number of slotattention iterations n , is shown in Table 1.
Revised: The time and space complexity of backpropagation for implicit slot attention compared to vanilla slot attention as a function of the number of slot attention iterations n , is shown in Table 1.

Original: While other methods for implicit differentiation are morecomplex to implement, this method requires only one additional line of code (Fig. 2).
Revised: While other methods for implicit differentiation are more complex to implement, this method requires only one additional line of code (Fig. 2).

Original:  Experiments
Revised:  Experiments

Original: Our main hypothesis is that iterative reﬁnement methods  can be trained as DEQs, and that consequently implicit differentiation would improve their optimization.
Revised: Our main hypothesis is that iterative refinement methods for latent sets, specifically slot attention, can be trained as DEQs, and that consequently implicit differentiation would improve their optimization.

Original: We test this premise by replacingthe backward pass of the slot attention module in the state-of-the-artSLATE [37] with the ﬁrst-order
Revised: We test this premise by replacing the backward pass of the slot attention module in the state-of-the-art

Original:  Neumann approximation of the implicit gradient.
Revised: SLATE [57] with the first-order Neumann approximation of the implicit gradient.

Original: This requires only one additional line of codeto the original slot attention implementation (Fig.2) and in one instance improves reconstructionmean-squared error by almost 7x.
Revised: This requires only one additional line of code to the original slot attention implementation (Fig.2) and in one instance improves reconstruction mean-squared error by almost 7x.

Original: 5.1 Experimental setup
Revised: 5.1 Experimental setup

Original: We summarize the SLATE architecture, datasets we considered, and relevant implementation details.
Revised: We summarize the SLATE architecture, datasets we considered, and relevant implementation details.

Original: SLATEFor image reconstruction, SLATE uses a discrete VAE [33] to compress an input imageinto a grid of discrete tokens.
Revised: SLATE SLATE uses a discrete VAE [50] to compress an input image into a grid of discrete tokens.

Original: These tokens index into a codebook of latent code-vectors, which, afterapplying a learned position encoding, serve as the input to the slot attention module.
Revised: 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.

Original: An Image GPTdecoder [9] 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.
Revised: An Image GPT decoder [10] 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.

Original: Gradients are blocked from ﬂowing between the discrete VAE and the rest of the network (i.e. theslot attention module and the Image GPT decoder), but the entire system is trained simultaneously.
Revised: Gradients are blocked from flowing between the discrete VAE and the rest of the network (i.e. the slot attention module and the Image GPT decoder), but the entire system is trained simultaneously.

Original: Data We consider three datasets: CLEVR-Mirror [37], Shapestacks [21], and COCO-2017 [27],the former two of which were used in the original SLATE paper.
Revised: Data We consider three datasets: CLEVR-Mirror [57], Shapestacks [27], and COCO-2017 [38], the former two of which were used in the original SLATE paper.

Original: We obtained CLEVR-Mirrordirectly from the SLATE authors and used a 70-15-15 split for training, validation, and testing.
Revised: We obtained CLEVR-Mirror directly from the SLATE authors and used a 70-15-15 split for training, validation, and testing.

Original: Wepooled all the data variants of Shapestacks together as Singh et al.
Revised: We pooled all the data variants of Shapestacks together as Singh et al.

Original: [37] did and used the originaltrain-validation-test splits.
Revised: [57] did and used the original train-validation-test splits.

Original: The COCO-2017 dataset was downloaded from FiftyOne and used theoriginal train-validation-test splits. 
Revised: The COCO-2017 dataset was downloaded from FiftyOne and used the original train-validation-test splits. [Figure] a

Original:  Implementation details In all of our experiments we used the same model and training hyperparameters as those used for the ShapeStacks experiment from Singh et al.
Revised: [Figure] s i [Figure] Implementation details In all of our experiments we used the same model and training hyperparameters as those used for the ShapeStacks experiment from Singh et al.

Original: [37, Table 6], run on an A100 GPU.
Revised: [57, Table 6], run on an A100 GPU.

Original: The only difference is in the image resolution (and consequently number of image tokens)because the largest image size we could train with was  128x128 used for their CLEVR experiment.
Revised: The only difference is in the image resolution (and consequently number of image tokens) because the largest image size we could train with was 96x96 due to computing constraints, rather than the 128x128 used for their CLEVR experiment.

Original: However, as we show for resolutions of both 96x96 and 64x64, the improvement in optimizationappears to be hold across image resolutions.
Revised: However, as we show for resolutions of both 96x96 and 64x64, the improvement in optimization appears to be hold across image resolutions.

Original: For a clean comparison, we simply integrated the ofﬁcial SLATE implementation 1 with the solversand backward gradient hook from the ofﬁcial implementation of deep equilibrium models 2 .
Revised: For a clean comparison [Figure] we simply integrated the official SLATE implementation 2 with the solvers and backward gradient hook from the official implementation of deep equilibrium models 3 .

Original: Besidesimplementing logging and evaluation code, as well as the ﬁrst-order Neumann approximation (Eq. 3),we did not change anything in the above two ofﬁcial implementations otherwise.
Revised: Besides implementing logging and evaluation code, as well as the first-order Neumann approximation (Eq. 3), we did not change anything in the above two official implementations otherwise.

Original: The SLATEimplementation serves as the baseline in all our experiments.
Revised: The SLATE implementation serves as the baseline in all our experiments.

Original: [Figure]

Original: 5.2 Training slot attention with implicit differentiation
Revised: 5.2 Training slot attention with implicit differentiation

Original: The ﬁrst test to conduct is to see whether slot attention can be trained with implicit differentiationat all.
Revised: The first test to conduct is to see whether slot attention can be trained with implicit differentiation at all.

Original: In summary: Test: Swap the forward pass with a different black-box solver, and train withimplicit gradients computed via various gradient estimation methods for Eq. 1.
Revised: In summary: Test: Swap the forward pass with a different black-box solver, and train with implicit gradients computed via various gradient estimation methods for Eq. 1.

Original: Hypothesis: Thereexists a (forward solver, implicit gradient estimator) that optimizes cross entropy no worse thanbackpropagating through unrolled inference.
Revised: Hypothesis: There exists a (forward solver, implicit gradient estimator) that optimizes cross entropy no worse than backpropagating through unrolled inference.

Original: Result:
Revised: Result:

Original: Yes, such a pair exists, and in fact it enablessigniﬁcantly better optimization.
Revised: Yes, such a pair exists, and in fact it enables significantly better optimization.

Original: [Table]
Revised: [Table]

Original: To conduct this test, we comparevanilla slot attention with four variants trained with implicit differentiation, shown in Table 2, labeled IB ,
Revised: To conduct this test, we compare vanilla slot attention with four variants trained with implicit differentiation, shown in Table 2, labeled IB ,

Original: BB , IN , and BN .
Revised: BB , IN , and BN .

Original: B stands for the
Revised: B stands for the

Original: Broyden solver, used by Bai et al.[6].
Revised: Broyden solver, used by Bai et al.

Original: We test conﬁgurations where the
Revised: We test configurations where the

Original: Broyden solver was both used to ﬁndthe ﬁxed point of slot attention and used to estimate the implicit gradient with Eq.2.
Revised: Broyden solver was both used to find the fixed point of slot attention and used to estimate the implicit gradient with Eq.

Original: N stands forthe ﬁrst-order Neumann approximation for computing the explicit gradient (Eq. 3).
Revised: N stands for the first-order Neumann approximation for computing the explicit gradient (Eq. 3).

Original: Our hypothesiswould be tentatively refuted if the cross-entropy learning curves of none of these variants convergesas efﬁciently as the baseline.
Revised: Our hypothesis would be tentatively refuted if the cross-entropy learning curves of none of these variants converges as efficiently as the baseline.

Original: We train these methods to model the CLEVR-Mirrors dataset withan image resolution of 64x64.
Revised: We train these methods to model the CLEVR-Mirrors dataset with an image resolution of 64x64.

Original: The slot attention baseline unrolls the slot attention cell for seveniterations and we also limit the maximum number of Broyden iterations to seven.
Revised: The slot attention baseline unrolls the slot attention cell for seven iterations and we also limit the maximum number of Broyden iterations to seven.

Original: Results and analysis Figure 3 shows that not only is it possible to train slot attention with implicitdifferentiation, but three out of four of the implicit differentiation conﬁgurations from Table 2optimize signiﬁcantly better than vanilla slot attention.
Revised: Results and analysis Figure 3 shows that not only is it possible to train slot attention with implicit differentiation, but three out of four of the implicit differentiation configurations from Table [Figure] optimize significantly better than vanilla slot attention.

Original: Whereas vanilla slot attention plateaus ata cross entropy of around 130, the BB , IN , and BN variants all achieve a 10x lower cross entropyloss.
Revised: Whereas vanilla slot attention plateaus at a cross entropy of around 130, the BB , IN , and BN variants all achieve a 10x lower cross entropy loss.

Original: Moreover, BB , IN , and BN all follow the same learning curve, suggesting that the optimizationimprovement is largely due to whether we use implicit differentiation or not, rather than the choiceof how we implement the implicit differentiation.
Revised: Moreover, BB , IN , and BN all follow the same learning curve, suggesting that the optimization improvement is largely due to whether we use implicit differentiation or not, rather than the choice of how we implement the implicit differentiation.

Original: The IB variant is less stable than the others, andwe hypothesize that this can be explained by the tendency for the ﬁxed point of Eq.
Revised: The IB variant is less stable than the others, and we hypothesize that this can be explained by the tendency for the fixed point of Eq.

Original: 2 to becomeincreasingly harder to estimate, an issue discussed in Bai et al.
Revised: 2 to become increasingly hard to estimate, an issue discussed in Bai et al.

Original: 5.3 Does this improvement generalize across datasets
Revised: 5.3 Does this improvement generalize across datasets?

Original: We now test whether this improved optimization holds across different datasets.
Revised: We now test whether this improved optimization holds across different datasets.

Original: If the improvementwas simply due to implicit differentiation being serendipitously useful for 64x64 images CLEVRMirror, then we should expect implicit differentiation to not help for other datasets.
Revised: If the improvement was simply due to implicit differentiation being serendipitously useful for 64x64 images CLEVRMirror, then we should expect implicit differentiation to not help for other datasets.

Original: We focus ourattention on IN variant to conduct this test because it is the minimal modiﬁcation to the baseline slotattention for gaining the beneﬁts of implicit differentiation, as it requires adding only one line of code[Figure] on top of the baseline slot attention implementation and does not require implementing any blackbox solvers.
Revised: We focus our attention on IN variant to conduct this test because it is the minimal modification to the baseline slot attention for gaining the benefits of implicit differentiation, as it requires adding only one line of code on top of the baseline slot attention implementation and does not require implementing any blackbox solvers.

Original: We henceforth refer to the IN variant as implicit slot attention (and correspondinglyimplicit SLATE), as described in §4.3.
Revised: We henceforth refer to the IN variant as implicit slot attention (and correspondingly implicit SLATE), as described in §4.3.

Original: In summary: Test: Apply implicit slot attention to CLEVRMirror, ShapeStacks, and COCO with 96x96 image resolution.
Revised: In summary: Test: Apply implicit slot attention to CLEVRMirror, ShapeStacks, and COCO with 96x96 image resolution.

Original: Hypothesis: Implicit slot attentionsigniﬁcantly improves optimization across all three datasets.
Revised: Hypothesis: Implicit slot attention significantly improves optimization across all three datasets.

Original: Result:
Revised: Result:

Original: Yes it does.
Revised: Yes it does.

Revised: Data

Revised: [Table]

Revised: Vanilla

Original: [Table]
Revised: [Table]

Revised: ShapeStacks

Revised: [Table]

Revised: ShapeStacks

Revised: [Table]

Revised: MSE

Original:  Results and quantitative analysis Using the two primary metrics usedin Singh et al.
Revised: [Table] [Table] [Table] Results and quantitative analysis Using the two primary metrics used in Singh et al.

Original: [37], images generatedby implicit SLATE achieve both lowerpixel-wise mean-squared error and FIDscore [23].
Revised: [57], images generated by implicit SLATE achieve both lower pixel-wise mean-squared error and FID score [29].

Original: The FID score was computedwith the PyTorch-Ignite [14] library using the inception network from the PyTorch port of the FID ofﬁcial implementation.
Revised: The FID score was computed with the PyTorchIgnite [18] library using the inception network from the PyTorch port of the FID official implementation.

Original: All methods were trained for250k gradient steps.
Revised: All methods were trained for 250k gradient steps.

Original: Table 3 comparesthe FID and MSE scores of the imagesthat result from compressing the SLATE encoder’s set of discrete tokens through the slot attentionbottleneck, using Image-GPT to autoregressively re-generate these image tokens one by one, andusing the discrete VAE decoder to render the generated image tokens.
Revised: Table 3 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.

Original: Implicit differentiation signiﬁcantly improves the quantitative image reconstruction metrics of SLATE across the test sets of
Revised: Implicit differentiation significantly improves the quantitative image [Figure] [Figure] [Figure]

Revised: [Figure]

Original:  CLEVR-Mirrors, Shapestacks, and COCO.
Revised: v a li d a t i o n c r o ss e n tr o p y [Figure] [Figure] reconstruction metrics of SLATE across the test sets of CLEVR-Mirrors, Shapestacks, and COCO.

Original: In the case of MSE for CLEVR, this is almost a 7ximprovement.
Revised: In the case of MSE for CLEVR, this is almost a 7x improvement.

Original: Qualitative analysis The higher quantitiatve metrics also translate into better quality reconstructions on the test set, as shown in Figure 4.
Revised: Qualitative analysis The higher quantitiatve metrics also translate into better quality reconstructions on the test set, as shown in Figure 4.

Original: For CLEVR-Mirrors, vanilla SLATE sometimes dropsor changes the appearance of objects, even in simple scenes with three objects.
Revised: For CLEVR-Mirrors, vanilla SLATE sometimes drops or changes the appearance of objects, even in simple scenes with three objects.

Original: In contrast, thegenerations produced from implicit SLATE match the ground truth very closely.
Revised: In contrast, the generations produced from implicit SLATE match the ground truth very closely.

Original: For Shapestacks,implicit SLATE consistently segments the scene into constituent objects.
Revised: For Shapestacks, implicit SLATE consistently segments the scene into constituent objects.

Original: This is sometimes thecase with vanilla SLATE on the training and validation set as well, but we observed for both of theseeds we ran for the ﬁnal evaluation that vanilla SLATE produced degenerated attention maps whereone slot captures the entire foreground, and the background is divided among the other slots.
Revised: 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 for the final evaluation that vanilla SLATE produced degenerated attention maps where one slot captures the entire foreground, and the background is divided among the other slots.

Original: Thevisual 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 [Figure] [Figure] [Figure]
Revised: 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 reconstruction through the slot bottleneck.

Original: v a li d a t i o n c r o ss e n tr o p y

Original: [Figure]

Original: [Figure]

Original: consequently for the reconstruction through the slot bottleneck.

Original: This highlights the gap that still existsbetween using the state-of-the-art in object-centric learning out-of-the-box and what the communitymay want these methods to do.
Revised: This highlights 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.

Original: The attention masks for both vanilla and implicit SLATE furthermoredo not appear to correspond consistently to coherent objects in COCO but rather patches on the imagethat do not immediately seem to match with our human intuition of what constitutes a visual entity. 
Revised: The attention masks for both vanilla and implicit SLATE 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. [Figure]

Original: 5.4 Can we simplify the need for optimization tricks
Revised: 5.4 Can we simplify the need for optimization tricks?

Original: To further understand the beneﬁts of implicit differentiation, we then ask whether it stabilizes thetraining of slot attention without the need for optimization tricks like learning rate decay, gradientclipping, and learning warmup.
Revised: To further understand the benefits of implicit differentiation [Figure] 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.

Original: In summary: Test: Remove learning rate decay, gradient clipping,and learning warmup each from vanilla SLATE and our method.
Revised: In summary: Test: Remove learning rate decay, gradient clipping, and learning warmup each from vanilla SLATE and our method.

Original: Hypothesis: Our method should notrequire these tricks to optimize stably.
Revised: Hypothesis: Our method should not require these tricks to optimize stably.

Original: Result:
Revised: Result:

Original: No, it generally does not require these tricks.
Revised: No, it generally does not require these tricks.

Original: Fig.
Revised: Fig.

Original: 5 shows that the ease and stability of training correlates with the spectral norm of the Jacobian ofthe slot attention cell.
Revised: 5 shows that the ease and stability of training correlates with the spectral norm of the Jacobian of the slot attention cell.

Original: Decaying the learning rate regularizes the Jacobian norm from exploding, but italso hurts optimization performance for both our method and vanilla SLATE, as expected.
Revised: 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.

Original: When weremove gradient clipping the Jacobian norm of vanilla SLATE explodes, as do its gradients (Fig. 6a),whereas both stay stable for our method.
Revised: When we remove gradient clipping the Jacobian norm of vanilla SLATE explodes, as do its gradients (Fig. 6a), whereas both stay stable for our method.

Original: This translates into a qualitative drop in performance forvanilla SLATE (Fig. 13) but not for implicit SLATE (Fig. 12).
Revised: This translates into a qualitative drop in performance for vanilla SLATE (Fig. 13) but not for implicit SLATE (Fig. 12).

Original: Lastly, removing learning rate warmup[Figure] also consistently makes vanilla SLATE’s training unstable, whereas it only affects the stability of ourmethod for one out of three seeds.
Revised: 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.

Original: Finally, Fig.
Revised: Finally, Fig.

Original: 6b shows that implicit slot attention is not sensitive tothe number of iterations with which to iterate the slot attention cell, whereas vanilla slot attention is,with more iterations being harder to train.5.5 Does this mean that iterating slot attention for one iteration is enough?
Revised: 6b shows that implicit slot attention is not sensitive to the number of iterations with which to iterate the slot attention cell, whereas vanilla slot attention is, with more iterations being harder to train. 5.5 Does this mean that iterating slot attention for one iteration is enough?

Original: One might be tempted to interpret Fig.
Revised: One might be tempted to interpret Fig.

Original: 6b to suggest that fewer iterations of vanilla slot attention issufﬁcient to improve performance.
Revised: 6b to suggest that fewer iterations of vanilla slot attention is sufficient to improve performance.

Original: This is not necessarily the case, as vanilla slot attention with oneiteration still achieves 2x worse reconstruction MSE than implicit slot attention with seven iterationson CLEVR (Fig. 6c).
Revised: This is not necessarily the case: although vanilla slot attention with one iteration trains more stably than its seven iteration counterpart, it still achieves 2x worse reconstruction MSE than implicit slot attention with seven iterations on CLEVR (Fig. 6c).

Original: Futhermore, Fig.
Revised: Futhermore, Fig.

Original: 7 shows that using only a single iteration is not enough to improve optimization ofvanilla slot attention for the object property prediction task used by Locatello et al.
Revised: 7 shows that using only a single iteration is not enough to improve optimization of vanilla slot attention for the object property prediction task used by Locatello et al.

Original: We directlymodiﬁed the released code from Locatello et al.
Revised: We directly modified the released code from Locatello et al.

Original: [28], which used three iterations, to use implicitdifferentiation.
Revised: [39], which used three iterations, to use implicit differentiation.

Original: Implicit slot attention can instead scale to as many forward iterations as needed.5.6 Does implicit slot attention still produce intuitive masks with a different architecture?
Revised: Implicit slot attention can instead scale to as many forward iterations as needed. Does implicit slot attention still produce intuitive masks with a different architecture?

Original: We sought to check whether implicit differentiation still preserves the quality of the segmentationmasks produced by the original slot attention architecture by Locatello et al.
Revised: We sought to check whether implicit differentiation still preserves the quality of the segmentation masks produced by the original slot attention architecture by Locatello et al.

Original: [28], which uses a spatialbroadcast decoder [41] rather than a transformer decoder as SLATE does.
Revised: [39], which uses a spatial broadcast decoder [63] rather than a transformer decoder as SLATE does.

Original: It indeed does (Fig. 8).
Revised: It indeed does (Fig. 8), suggesting that our findings are not specific to SLATE but apply to slot attention more broadly.

Original:  Discussion
Revised:  Discussion

Original: The connection we made in this paper between slot attention and deep equilibrium models alsohighlights various other properties about iterative reﬁnement procedures of symmetric generativemodels that suggest connections to other areas of machine learning that are worth theoreticallydeveloping in the future.
Revised: The connection we made in this paper between slot attention and deep equilibrium models also highlights various other properties about iterative refinement procedures for inferring latent sets that suggest connections to other areas of research that are worth theoretically developing in the future.

Original: First is the connection to the literature on fast weights [5, 25, 35], if weinterpret the slots as parameters that are modiﬁed during the inner optimization during executiontime.
Revised: First is the connection to the literature on fast weights [6, 33, 54]: interpreting slots as parameters that are modified during the inner optimization during execution time may give us novel formulation for how to represent and update fast weight memories.

Original: Second is the connection to the literature on meta-learning [4, 13, 34, 38], if we interpretthe slots as solutions to an inner optimization problem during execution time.
Revised: Second is the connection to the literature on meta-learning [5, 17, 53, 60]: interpreting slots as solutions to an inner optimization problem during execution time may give us a novel perspective on perception as itself a learning process.

Original: The third is theconnection to causality, if we interpret the independently generated and symmetrically processedslots as parameterizing independent causal mechanisms [22].
Revised: Third is the connection to causality [46, 47]: interpreting the independently generated and symmetrically processed slots as parameterizing independent causal mechanisms [28] may give us a novel approach for learning to represent causal models within a neural scaffolding. Fourth is the connection to dynamical systems [44]: interpreting slots as a set of attractor basins may offer a novel theory of how the error-correcting properties of discrete representations emerge from continuous ones.

Original: These different ﬁelds have their ownconceptual and implementation tools that could potentially improve our understanding of how tobuild better object-centric models and of how objects could potentially be represented in the mind.
Revised: These different fields have their own conceptual and implementation tools that could potentially improve our understanding of how to build better iterative refinement algorithms and inform how objects could potentially be represented in the mind.

Original: Conclusion We have proposed implicit differentiation for training iterative reﬁnement proceduresof symmetric generative models.
Revised: Conclusion We have proposed implicit differentiation for training iterative refinement procedures for inferring representations of latent sets.

Original: Our results show clear signal that implicit differentiation can offer a signiﬁcant optimization improvement over backpropagating through the unrolled iterationof slot attention, and potentially any other iterative reﬁnement algorithm, with lower space andtime complexity and only one additional line of code.
Revised: 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 other iterative refinement algorithm, with lower space and time complexity and only one additional line of code.

Original: Because it is so simple to apply implicitdifferentiation to any ﬁxed point algorithm, we hope our work inspires future work to leverage toolsdeveloped for implicit differentiation for improvingobject-centric learning  and  methodsfor learninglatent structure more broadly. 
Revised: Because it is so simple to apply implicit differentiation to any fixed point algorithm, we hope our work inspires future work to leverage tools developed for implicit differentiation for improving learning representations of latent sets and iterative refinement methods more broadly. Acknowledgements

Original: Checklist
Revised: This work was supported ARL, W911NF2110097, and ONR grant number N00014-18-1-2873. Part of this work was completed while MC was an intern at Meta AI. Computing support came from Google Cloud Platform and Meta. We would like to thank Shaojie Bai for help on implicit differentiation, Gautam Singh for help on SLATE, Alban Desmaison for help on PyTorch.

Revised: We would also like to thank Yan Zhang, David Zhang, Thomas Kipf and Klaus Greff for insightful discussions.

Original: The checklist follows the references.
Revised: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.
Revised: Please read the checklist guidelines carefully for information on how to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] or[N/A] .
Revised: For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .

Original: You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
Revised: You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.

Original: For example:
Revised: For example:

Original: • Did you include the license to the code and datasets?
Revised: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? . • Did you include the license to the code and datasets?
Revised: [Yes] See Section X. • Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.•
Revised: [No] The code and the data are proprietary.

Original:  Did you include the license to the code and datasets?
Revised: • Did you include the license to the code and datasets?

Original: [N/A]
Revised: [N/A]

Original: Please do not modify the questions and only use the provided macros for your answers.
Revised: Please do not modify the questions and only use the provided macros for your answers.

Original: Note that the
Revised: Note that the

Original: Checklist section does not count towards the page limit.
Revised: Checklist section does not count towards the page limit.

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
Revised: In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] (c) Did you discuss any potential negative societal impacts of your work?
Revised: [Yes] (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
Revised: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes]2.
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A] (b) Did you include complete proofs of all theoretical results?
Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

Original: [N/A]3.
Revised: [N/A] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] See Fig.
Revised: [Yes] See Fig.

Original: (b) Did you specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)?[Yes]
Revised: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original:  (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] 
Revised: [Yes] 4.

Original: If you are using existing assets (e g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: [No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A] 
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: Fig. 9 and Fig. 10 both highlight intriguing qualitative differences between vanilla and implicit SLATE and understanding what causes these differences would be valuable for future work.

Original: spent on participant compensation?
Revised: What these figures highlight is that multiple decompositions of a scene into components are possible, which may differ in how closely they reflect human intuition on what constitutes a visual entity. This suggests that the optimization objectives for current object-centric models still underspecify the kinds of decompositions we seek to achieve in our models. The paradigm of decomposing static scenes, as opposed to interactive videos (e.g. [62, Fig. 7]), also contributes to this underspecification. [Figure]

Original: [N/A]
Revised: [Figure]

Revised: B Further Experiments

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]