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

Original: Abstract
Revised: Abstract

Original: 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.
Revised: Iterative reﬁnement – 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.

Original: 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.
Revised: 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.

Original: However, most prior works differentiate through the unrolled refinement process, which can make optimization challenging.
Revised: However, most prior works differentiate through the unrolled reﬁnement process, which can make optimization challenging.

Original: 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.
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 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.
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: 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 )) ... )) .
Revised: Conventionally, neural network models implement feedforward computation, transforming the input x into the output z through a ﬁxed series of operations corresponding to distinct layers as z = f N ( f N  1 ( ...f 2 ( f 1 ( x )) ... )) .

Original: 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 = f x ( z ) of an initial guess z 0 .
Revised: 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 ﬁxed point of an iterative reﬁnement z = f x ( z ) of an initial guess z 0 .

Original: 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].
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: 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.
Revised: Such iterative reﬁnement 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: 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].
Revised: In this work, we consider the case of iterative reﬁnement 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, 70].

Original: 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.
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.

Original: 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 
Revised: Iterative reﬁnement 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] c r o ss e n tr o p y [Figure]

Revised: a c

Original: [Figure]
Revised: [Figure]

Revised: i

Revised: [Figure]

Revised: s

Revised: [Figure]

Revised: l

Revised: [Figure]

Original: mapping from observation to representation.
Revised: mapping from observation to representation.

Original: The state-of-the-art of these iterative object-centric methods is the slot attention 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: [39], serving as the focus of our paper.
Revised: [39], serving as the focus of our paper.

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

Original: 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).
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).

Original: As a result, prior works train these methods by differentiating through the unrolled iterations of f .
Revised: As a result, prior works train these methods by differentiating through the unrolled iterations of f .

Original: Differentiating through this recurrence contributes to various training instabilities, as we can see by the growing spectral norm of f in Fig.
Revised: Differentiating through this recurrence contributes to various training instabilities, as we can see by the growing spectral norm of f in Fig.

Original: 1b.
Revised: 1b.

Original: 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.
Revised: Such instabilities result in sensitivity to hyperparameter choices (e g., number of reﬁnement 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 iterative refinement methods like slot attention have not taken full advantage of the fact that f can be viewed as a fixed point operation.
Revised: To approach this problem, we observe that previous iterative reﬁnement methods like slot attention have not taken full advantage of the fact that f can be viewed as a ﬁxed point operation.

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

Original: This paper investigates how advances in implicit differentiation in neural models can be applied to improve the training of iterative refinement methods.
Revised: This paper investigates how advances in implicit differentiation in neural models can be applied to improve the training of iterative reﬁnement methods.

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

Original: 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 .
Revised: First, we show that slot attention can be cast as a ﬁxed 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-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.
Revised: Second, we show on the latest state-of-the-art method of this kind, SLATE [57], that using the ﬁrst-order Neumann approximation of the implicit gradient for the slot attention module yields substantial improvement in optimization.

Original: 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.
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 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.
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, when integrated 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: [39], implicit differentiation substantially improves object property prediction and continues to predict intuitive segmentation 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

Original: 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: Much early work in artiﬁcial intelligence followed a paradigm of using an iterative learning procedure during execution time, whether it be a form of search, inference, or optimization [51].

Original: These include early work on variational inference  and energy-based models [12, 19, 37], Hopfield networks [31], Boltzmann machines [1], and associative memory [36].
Revised: These include early work on variational inference [12, 19] and energy-based models [14, 37, 65–68], Hopﬁeld networks [31], Boltzmann machines [1], and associative memory [36].

Original: 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: 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.

Original: 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].
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 work concerns the optimization of neural networks as update rules for these iterative refinement procedures.
Revised: Our work concerns the optimization of neural networks as update rules for these iterative reﬁnement procedures.

Original: 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.
Revised: Our methodological novelty is the adaptation of implicit differentiation techniques for training iterative reﬁnement procedures for representing of latent sets, where our key insight is that the iterative procedure used for symmetry-breaking reaches a ﬁxed point, thus allowing implicit differentiation to be used.

Original: 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.
Revised: We are not aware that this has been done before, as current methods that perform iterative reﬁnement for representing latent sets, often referred to as object-centric learning [24– 26, 35, 39, 57, 61, 62, 70], all differentiate through the unrolled dynamics of the ﬁxed point procedure, which as we show makes them difﬁcult to train.

Original: 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].
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 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].
Revised: Although we evaluate various techniques for implicit differentiation, our results highlight the beneﬁts of the simplest of these, which is that of using a truncated Neumann approximation [20, 21, 32, 55].

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

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

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

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

Original: 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.
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 refinement and implicit differentiation with deep networks.
Revised: Our work builds on prior works on iterative reﬁnement and implicit differentiation with deep networks.

Original: 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.
Revised: Iterative reﬁnement 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 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.
Revised: Because mixture models are also deﬁned 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 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 .
Revised: Under a clustering lens, the problem reduces 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 .

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

Original: 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 .
Revised: The slot attention module [39], e.g., computes ✓ nt + f ( ✓ nt , x n ) , where  n,m is updated as an intermediate step inside f .

Original: The θ n , called slots , serve as input to a downstream objective, e.g. image reconstruction, whose gradients are backpropagated through 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 [24] and videos [61].
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 a function defined 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 deﬁned in terms of satisfying a joint condition on the input and output.

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

Original: z
Revised: z

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

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

Original: 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 ∗ :
Revised: Letting f w be parameterized by weights w , with input x and ﬁxed 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 ( z ∗ ) is the Jacobian matrix of f w evaluated at z ∗ .
Revised: where J f w ( z ⇤ ) is the Jacobian matrix of f w evaluated at z ⇤ .

Original: Compared to backpropagating through the unrolled iteration of f , which is just one of many choices of the solver, implicit differentiation via 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 the 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)
Revised: Deep equilibrium models (DEQ)

Original: [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.
Revised: [7] represent the class of functions f parameterized by neural networks, which have successfully been trained with implicit differentiation and empirically produce stable ﬁxed 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 ( z ∗ )) − 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: [8], Pineda [48] and Almeida [3] propose to 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 fixed point of the linear system:
Revised: 1 as the ﬁxed 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 implicit differentiation to neural networks, this approach has in practice required some expensive regularization to maintain stability [8], 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: [20], Geng et al.
Revised: [20], Geng et al.

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

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

Original: [55] propose instead to approximate ( I − J f w ( z ∗ )) − 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 first-order approximation ( T = 1 ) amounts to applying f once to the fixed point z ∗ and differentiating through the resulting computation graph.
Revised: The ﬁrst-order approximation ( T = 1 ) amounts to applying f once to the ﬁxed point z ⇤ and differentiating through the resulting computation graph.

Original: 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.
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 sacriﬁcing performance.

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

Original: 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.
Revised: Our main contribution is to propose treating iterative reﬁnement algorithms used for inferring latent sets as ﬁxed-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 between iterative refinement for latent sets and the Expectation-Maximization algorithm [13], which also has been pointed out by Greff et al.
Revised: Our approach is inspired by the structural resemblance between iterative reﬁnement for latent sets and the Expectation-Maximization algorithm [13], which also has been pointed out by Greff et al.

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

Original: If f were to update the components θ n and assignments ϕ n,m in a way that 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 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 .
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 ﬁxed point operation whose ﬁxed point locally maximizes L n .

Original: 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 .
Revised: With this interpretation, such iterative reﬁnement 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 refinement methods are implemented differs from optimizing the per-datapoint
Revised: How these iterative reﬁnement 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: [24] and Locatello et al.
Revised: [24] and Locatello et al.

Original: [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.
Revised: [39] implemented ablations to their methods that do implement a variant of Expectation-Maximization on a well-deﬁned mixture model, but found that replacing the update rule with a recurrent neural network empirically performs better at extracting representations.

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

Original: 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.
Revised: While in general we still lack the theory to understand whether and how the slots of slot attention optimize a well-deﬁned objective like the ELBO, and thus have no theoretical guarantees that slot attention does converge to a ﬁxed point, we empirically observe from its stable forward relative residuals (Fig. 3a) that it does appear to approximately converge to a ﬁxed point.

Original: 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.
Revised: This suggests that slot attention can be understood as an instance of a DEQ that uses naive forward iteration to ﬁnd a ﬁxed point and backpropagates through the iteration to compute the gradient.

Original: 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.
Revised: The novel implication of this to inferring latent sets is that any root-ﬁnding solver and implicit gradient estimator can in principle be used to train slot attention, and we ﬁnd that many combinations do train slot attention effectively in Fig.

Original: 3b.
Revised: 3b.

Original: This insight is one of our main contributions, but in the next section we 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 [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).
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 ﬁrst-order Neumann approximation (Eq. 3).

Original: 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.
Revised: It can be implemented by simply differentiating the computation graph of applying the slot attention update once to the ﬁxed 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 slot attention 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 more complex 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 refinement methods for latent sets, specifically slot attention, can be trained as DEQs, and that consequently implicit differentiation would improve their optimization.
Revised: Our main hypothesis is that iterative reﬁnement methods for latent sets, speciﬁcally slot attention, can be trained as DEQs, and that consequently implicit differentiation would improve their optimization.

Original: We test this premise by replacing the backward pass of the slot attention module in the state-of-the-art
Revised: We test this premise by replacing the backward pass of the slot attention module in the state-of-the-art

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

Original: 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.
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: SLATE SLATE uses a discrete VAE [50] to compress an input image into 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, after applying 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 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.
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 flowing between the discrete VAE and the rest of the network (i.e. the slot attention module and the
Revised: Gradients are blocked from ﬂowing between the discrete VAE and the rest of the network (i.e. the slot attention module and the

Original: Image GPT decoder), but the entire system is trained simultaneously.
Revised: Image GPT decoder), but the entire system is trained simultaneously.

Original: 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.
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-Mirror directly 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: We pooled 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: [57] did and used the original train-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 the original train-validation-test splits.
Revised: The COCO-2017 dataset was downloaded from FiftyOne and used the original train-validation-test splits.

Original: [Figure]
Revised: [Figure]

Original: a
Revised: a

Original: [Figure]
Revised: [Figure]

Original: s i
Revised: s i

Original: [Figure]
Revised: [Figure]

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: l [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: [57, 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 96x96 due to computing constraints, rather than the 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 optimization appears 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
Revised: For a clean comparison

Original: [Figure]
Revised: [Figure]

Original: we simply integrated the official SLATE implementation 2 with the solvers and backward gradient hook from the official implementation of deep equilibrium models 3 .
Revised: we simply integrated the ofﬁcial SLATE implementation 2 with the solvers and backward gradient hook from the ofﬁcial implementation of deep equilibrium models 3 .

Original: 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.
Revised: Besides implementing 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.

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

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

Original: The first test to conduct is to see whether slot attention can be trained with implicit differentiation at all.
Revised: The ﬁrst 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 with implicit 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: There exists a (forward solver, implicit gradient estimator) that optimizes cross entropy no worse than backpropagating 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 enables significantly better optimization.
Revised: Yes, such a pair exists, and in fact it enables signiﬁcantly better optimization.

Original: [Table]
Revised: [Table]

Original: To conduct this test, we compare vanilla 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.
Revised: Broyden solver, used by Bai et al.

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

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

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

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

Original: We train these methods to model the CLEVR-Mirrors dataset with an 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 seven iterations 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 implicit differentiation, but three out of four of the implicit differentiation configurations from Table
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 conﬁgurations from Table

Original: [Figure]
Revised: [Figure]

Original: optimize significantly better than vanilla slot attention.
Revised: optimize signiﬁcantly better than vanilla slot attention.

Original: 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.
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 optimization improvement is largely due to whether we use implicit differentiation or not, rather than the choice of 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, and we hypothesize that this can be explained by the tendency for the fixed 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 ﬁxed point of Eq.

Original: 2 to become increasingly hard 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 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.
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 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.
Revised: We focus our attention on IN variant to conduct this test because it is the minimal modiﬁcation to the baseline slot attention for gaining the beneﬁts 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 correspondingly implicit 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 attention significantly improves optimization across all three datasets.
Revised: Hypothesis: Implicit slot attention signiﬁcantly improves optimization across all three datasets.

Original: Result:
Revised: Result:

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

Revised: metrics for image reconstruction through the slot bottleneck.

Original: Data
Revised: Data

Original: [Table]
Revised: [Table]

Original: Vanilla
Revised: COCO [Table] COCO

Original: [Table]
Revised: [Table]

Original: ShapeStacks
Revised: ShapeStacks

Revised: [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: ShapeStacks
Revised: ShapeStacks

Original: [Table]
Revised: [Table]

Original: MSE
Revised: MSE)

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table]

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

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

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

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

Original: 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.
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 significantly improves the quantitative image
Revised: Implicit differentiation signiﬁcantly improves the quantitative image

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

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

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: reconstruction metrics of SLATE across the test sets of CLEVR-Mirrors, Shapestacks, and COCO.
Revised: reconstruction metrics of SLATE across the test sets of CLEVR-Mirrors, Shapestacks [Figure] and COCO.

Original: In the case of MSE for CLEVR, this is almost a 7x improvement.
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 drops or 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, the generations 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 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.
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 ﬁnal evaluation that vanilla SLATE produced degenerated attention maps where one slot captures the entire foreground [Figure] and the background is divided among the other slots.

Original: 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.
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: 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.
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 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.
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.

Original: [Figure]
Revised: [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 benefits of implicit differentiation
Revised: To further understand the beneﬁts of implicit differentiation

Original: [Figure]
Revised: [Figure]

Original: 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.
Revised: 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 not require 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 of the 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 it also 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 we remove 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 for vanilla 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 also consistently makes vanilla SLATE’s training unstable, whereas it only affects the stability of our method 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 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.
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.

Original: 5.5 Does this mean that iterating slot attention for one iteration is enough?
Revised: 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 is sufficient to improve performance.
Revised: 6b to suggest that fewer iterations of vanilla slot attention is sufﬁcient to improve performance.

Original: 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.
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.

Original: 6c).
Revised: 6c).

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

Original: 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.
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 directly modified the released code from Locatello et al.
Revised: We directly modiﬁed the released code from Locatello et al.

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

Original: Implicit slot attention can instead scale to as many forward iterations as needed.
Revised: Implicit slot attention can instead scale to as many forward iterations as needed.

Original:  Does implicit slot attention still produce intuitive masks with a different architecture?
Revised: 5.6 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 segmentation masks 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: [39], which uses a spatial broadcast decoder [63] 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), suggesting that our findings are not specific to SLATE but apply to slot attention more broadly.
Revised: It indeed does (Fig. 8), suggesting that our ﬁndings are not speciﬁc 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 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.
Revised: The connection we made in this paper between slot attention and deep equilibrium models also highlights various other properties about iterative reﬁnement 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 [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.
Revised: First is the connection to the literature on fast weights [6, 33, 54]: interpreting slots as parameters that are modiﬁed 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 [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.
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: 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.
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.

Original: 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.
Revised: 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 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.
Revised: These different ﬁelds have their own conceptual and implementation tools that could potentially improve our understanding of how to build better iterative reﬁnement algorithms and inform how objects could potentially be represented in the mind.

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

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

Original: 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.
Revised: Because it is so simple to apply implicit differentiation to any ﬁxed 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 reﬁnement methods more broadly.

Original: Acknowledgements
Revised: Acknowledgements

Original: This work was supported ARL, W911NF2110097, and ONR grant number N00014-18-1-2873.
Revised: This work was supported ARL, W911NF and ONR grant number N00014-18-1-2873.

Original: Part of this work was completed while MC was an intern at Meta AI.
Revised: Part of this work was completed while MC was an intern at Meta AI.

Original: Computing support came from Google Cloud Platform and Meta.
Revised: Computing support came from Google Cloud Platform and Meta.

Original: 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 like to thank Shaojie Bai for help on implicit differentiation, Gautam Singh for help on SLATE, Alban Desmaison for help on PyTorch.

Original: We would also like to thank Yan Zhang, David Zhang, Thomas Kipf and Klaus Greff for insightful discussions.
Revised: We would also like to thank Yan Zhang, David Zhang, Thomas Kipf and Klaus Greff for insightful discussions and Jianwen Xie for pointing out previously missing references.

Original: Checklist
Revised: Checklist

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

Original: Please read the checklist guidelines carefully for information on how 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 justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
Revised: You are strongly encouraged to include a justiﬁcation 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 X. • 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 are proprietary.
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 instructions block 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 reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect 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 to them?
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 experimental 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 experiments multiple times)?
Revised: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [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)?
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] 4.
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
Revised: [No] (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: 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 identifiable information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

Original: [N/A] 5.
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?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
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: A Future Work

Original: Fig. 9 and Fig.

Original: 10 both highlight intriguing qualitative differences between vanilla and implicit

Original: SLATE and understanding what causes these differences would be valuable for future work.

Original: 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.

Original: This suggests that the optimization objectives for current object-centric models still underspecify the kinds of decompositions we seek to achieve in our models.

Original: The paradigm of decomposing static scenes, as opposed to interactive videos (e.g. [62, Fig. 7]), also contributes to this underspecification.

Original: [Figure]

Original: [Figure]

Original: B Further Experiments

Original: [Figure]

Original: [Figure]

Original: [Figure]