M ULTISET -E QUIVARIANT S ET P REDICTION WITH A PPROXIMATE I MPLICIT D IFFERENTIATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Most set prediction models in deep learning use set-equivariant operations, but they actually operate on multisets.
We show that set-equivariant functions cannot represent certain functions on multisets, so we introduce the more appropriate notion of multiset-equivariance .
We identify that the existing Deep Set Prediction Network (DSPN) can be multiset-equivariant without being hindered by set-equivariance and improve it with approximate implicit differentiation, allowing for better optimization while being faster and saving memory.
In a range of toy experiments, we show that the perspective of multiset-equivariance is beneﬁcial and that our changes to DSPN achieve better results in most cases.
On CLEVR object property prediction, we substantially improve over the state-of-the-art Slot Attention from 8% to 77% in one of the strictest evaluation metrics because of the beneﬁts made possible by implicit differentiation.
I NTRODUCTION
An example of set prediction is detecting the objects in an image: we want a collection of objects as output, but there is no inherent order to them.
Multisets and sets are natural representations for such data due to their orderless nature, with the difference between them being whether duplicate elements are allowed or not.
In general, they are useful for representing collections of discrete and composable elements , such as agents in an environment, instruments in a piece of music, or atoms in a molecule.
The usual approach to representing a multiset or set of vectors in memory is to store it as a list in which the elements are placed in an arbitrary order.
To still treat the list as if it were a multiset or set, recent set prediction models avoid being affected by this arbitrary order by using list-to-list functions that are enforced to be permutation-equivariant (Zaheer et al., 2017): changing the order of the input elements should change the order of the output elements in the same way.
Now, consider a function that takes a multiset (represented as a list) of two scalars as input and “pushes them apart”:
[Equation]
For example, push apart ([1 , 2]) =
[0 , 3] and push apart ([2 , 1]) =
At ﬁrst glance, it appears that this function is permutation-equivariant because swapping the inputs swaps the outputs.
However, notice that push apart ([0 , 0]) =
[ − 1 , 1] . Swapping the inputs has no effect in this case, so the outputs remain unchanged, thus push apart is not permutation-equivariant.
This is a bit strange – the function only violates the equivariance property for equal input elements, in which case the order of the equal input elements is irrelevant anyway.
How can we reconcile this observation with the standard deﬁnition of permutation-equivariance?
We aim to address this question in our paper and present two key ideas:
Multiset-equivariance.
We
begin Section by making the important observation that many existing models for set prediction thought to operate on sets actually operate on
multisets We then show that these models are too restricted in the functions they can represent on multisets due to the limitation of being permutation-equivariant, which we call set-equivariance .
This is a problem even when there are no exact duplicate elements.
Thus, we develop a new notion of multiset -equivariance, which we argue is a more suitable property for operations on multisets.
[Table]
To avoid the aforementioned limitation of set-equivariance, a crucial point is that it is possible to be multiset-equivariant without being set-equivariant, which we call exclusive multiset-equivariance .
We identify and show that deep set prediction networks (DSPN) (Zhang et al., 2019; Huang et al., 2020) satisfy this property with a speciﬁc choice of set encoder.
This makes it the only set prediction model we are aware of that can satisfy exclusive multiset-equivariance.
Implicit DSPN.
Despite this beneﬁcial property, DSPN is outperformed by the set-equivariant Slot Attention (Locatello et al., 2020), which motivates us to improve other aspects of DSPN.
We propose implicit DSPN (iDSPN) in Section 3: a version of DSPN that uses approximate implicit differentiation. Implicit differentiation enables us to use better optimizers and more iterations at a constant memory cost and less computation time.
The approximation makes this faster to run and easier to implement by avoiding the computation of a Hessian while still having the same beneﬁts.
In our experiments in Section 4, we start by evaluating the ﬁrst point by comparing exclusively multiset-equivariant models with both set-equivariant and non-equivariant models on synthetic data.
The former completely fail on this simple toy task while the latter are signiﬁcantly less sampleefﬁcient.
Next, we evaluate the second point by testing the modeling capacity on autoencoding random sets, where iDSPN performs similarly to DSPN at the same iteration count and much better for the same computational cost.
Lastly, iDSPN signiﬁcantly raises the bar on CLEVR object property prediction, outperforming the state-of-the-art Slot Attention (Locatello et al., 2020) by a large margin of 69 percentage points on the AP 0 .
125 metric while only training for 7% of the number of epochs.
L IMITATIONS OF SET - EQUIVARIANCE
In this section, we begin by identifying problems that limit the expressivity of a wide class of set prediction models satisfying a type of permutation-equivariance, set-equivariance .
We argue that set-equivariance is too strict and instead propose a relaxation of set-equivariance that is more suitable for set prediction models: multiset-equivariance .
We end by discussing how an existing method, deep set prediction networks (Zhang et al., 2019; Huang et al., 2020), is multiset -equivariant but not set -equivariant and the associated beneﬁts thereof (Table 1).
2. P RELIMINARIES
Sets in the context of deep learning are typically unordered collections of real-valued vectors.
They are represented in memory as n × d matrices with the n set elements in an arbitrary order, each element being a d -dimensional vector.
Since elements in the matrix representation are typically not restricted to be unique, the term multiset is more appropriate.
To still treat these matrices as multisets, operations must not rely on the arbitrary order of the elements.
The two main properties to accomplish this are permutation- invariance and permutation- equivariance (Zaheer et al., 2017).
In short, when changing the order of the elements in the matrix, the output of a function should either not change at all (invariance) or change order similarly (equivariance).
2. S ET - EQUIVARIANT MODELS
Recent set prediction models (Lee et al., 2019; Locatello et al., 2020; Carion et al., 2020; Kosiorek et al., 2020) make use of set-to-set (permutation-equivariant list-to-list) functions to reﬁne an initial set Y 0 , which is usually a randomly generated or learnable matrix.
Permutation-equivariance is desirable when processing sets because it prevents a function from relying on the arbitrary order of the set in its matrix representation.
Such functions can be easily composed to build larger models that remain equivariant, which ﬁts well into deep learning architectures as building blocks.
However, without having any restrictions on duplicate elements, it is more accurate to think of set prediction models as using multiset-to-multiset functions.
This difference subtly changes the appropriate notion of equivariance for sets versus multisets, which will become important for our discussion on their limitations.
Set-equivariance.
First, we generalize the standard notion of permutation-equivariance for sets to multiset inputs and refer to it as set-equivariance .
This directly matches the usual deﬁnition of permutation-equivariance in the context of set inputs in the literature.
Therefore, many existing models already satisfy the following deﬁnition, including Slot Attention (Locatello et al., 2020),
Set Transformers (Lee et al., 2019), TSPN (Kosiorek et al., 2020), and the equivariant version of
DeepSets (Zaheer et al., 2017).
Deﬁnition 1 (Set-equivariance for multiset inputs) .
A function f : R n × d 1 → R n × d 2 with multiset input X is set-equivariant iff for any permutation matrix P ,
[Equation]
A consequence is that elements in a multiset that are equal before a set-equivariant function must remain equal after applying the function, as shown by the following:
Proposition 1.
For all permutation matrices P 1 , P 2 such that P 1 X =
P 2 X , a set-equivariant function f satisﬁes:
[Equation]
Proof.
[Equation]
P 1 X =
P 2 X can be satisﬁed by P 1 (cid:54) =
P 2 when there are duplicate elements in X , so P 1 f ( X ) =
P 2 f ( X ) means that they must remain duplicates.
This inability of separating equal elements demonstrates that the classical deﬁnition of permutation-equivariance is overly restrictive when directly applied to multisets as in Deﬁnition 1.
Therefore, we seek a more appropriate deﬁnition of permutation-equivariance for multiset inputs.
Multiset-equivariance. We now consider a slightly different notion of permutation-equivariance that no longer forces equal elements to remain the same: multiset-equivariance .
Recall the push apart function from Equation 1, which we suggest should be a valid equivariant function for multisets, but cannot be represented by a set-equivariant function.
Therefore, we propose that there should be a distinction between set-equivariance and multiset -equivariance, the latter corresponding to the appropriate equivariance property for the typical multiset-to-multiset mappings used in set prediction.
The key difference of multiset-equivariance is that the output order for equal elements in the input multiset can be arbitrary .
This leads us to the following deﬁnition:
Deﬁnition 2 (Multiset-equivariance for multiset inputs) .
A function f : R n × d 1 → R n × d 2 with multiset input X is multiset-equivariant iff for any permutation matrix P 1 , there exists a permutation matrix P 2 such that P 1 X
= P 2 X
and:
[Equation]
Any function that is set-equivariant is also multiset-equivariant (since we can set P 2 = P 1 ), but the converse is not true.
The problem with set-equivariance. We argue that set-equivariance is too strong of a property for set prediction.
As we have shown in Proposition 1, it prevents equal elements from being processed differently.
These equal elements could for example arise from ReLU nonlinearities zeroing out elements, the model naturally tending towards equal
Even without exact duplicates in the input, very similar elements can still cause problems for set-equivariant models.
For example, in order to learn the push apart function, it is necessary for a model to decide which element is larger.
This decision is fundamentally discontinuous, but most models in deep learning cannot represent discontinuous jumps.
A continuous model must approximate the discontinuity, which means that close to the discontinuity (where elements are similar) the modeling error will likely be high. This is a general problem when trying to map similar elements in the input (e.g. due to random initialization of Y 0 ) to dissimilar elements in the output.
Summary. It is important to make the distinction between set-equivariance – as has traditionally been considered in the literature – and multiset-equivariance.
Deep learning typically operates on multisets rather than sets, so the less restrictive multiset-equivariance without set-equivariance should be the desirable property.
Most existing set prediction models are set-equivariant, which places a limitation on their modeling capacity even if exact duplicates are never encountered.
We refer to models that are multiset-equivariant but not set-equivariant as exclusively multiset-equivariant .
2. D EEP S ET P REDICTION N ETWORKS
Here, we highlight the only set prediction model we are aware of that can be exclusively multisetequivariant: deep set prediction networks (DSPN) (Zhang et al., 2019; Huang et al., 2020).
The core idea of DSPN is to use backpropagation through a permutation-invariant encoder g (multiset-tovector) to iteratively update a multiset Y .
The objective is to ﬁnd a Y that minimizes the difference between its vector encoding and
the input vector z (this vector could for example come from an input image being encoded).
This is expressed as the following optimization problem:
[Equation]
Because g is permutation-invariant, any ordering for the elements in Y has the same value for L .
In the forward pass of the model, the arg min is approximated by running a ﬁxed number of gradient descent steps.
In the backward pass, the goal is to differentiate Equation 7 with respect to the input vector z and the parameters θ of the encoder. To do this, Zhang et al. (2019) unroll the gradient descent applied in the forward pass and backpropagate through each gradient descent step.
Equivariance of DSPN. We now discuss the particular form of equivariance that DSPN takes, which has not been done before in the context of multisets.
The gradient of the permutation-invariant encoder g with respect to the set input Y is always multiset-equivariant, but depending on the encoder, it is not necessarily set-equivariant.
Zhang et al.
ﬁnd that FSPool-based encoders (Zhang et al., 2020) achieved by far the best results among the encoders they have tried.
With FSPool, DSPN becomes exclusively multiset-equivariant to its initialization Y 0 .
This is due to the use of numerical sorting in FSPool: the Jacobian of sorting is exclusively multiset-equivariant (Appendix A).
Note that DSPN is not always exclusively multiset-equivariant, but it depends on the choice of encoder.
A DeepSets encoder (Zaheer et al., 2017) – which is based on sum pooling – has the same gradients for equal elements, which would make DSPN set-equivariant.
It is speciﬁcally the use of the exclusively multiset-equivariant gradient of sorting that makes DSPN exclusively multiset-equivariant.
I MPLICIT D EEP S ET P REDICTION N ETWORKS
Despite being exclusively multiset-equivariant, DSPN is outperformed by the set-equivariant Slot
Attention (Locatello et al., 2020). We begin by highlighting some issues in DSPN that might be the cause of this, then propose approximate implicit differentiation as a solution.
Problems in DSPN. Increasing the number of optimization steps for solving Equation 7 generally results in a better solution (Zhang et al., 2019), but requires signiﬁcantly more memory and computation time and can lead to training problems (Belanger et al., 2017).
To be able to backpropagate through Equation 7, the activations of every intermediate gradient descent step have to be kept in memory.
Each additional optimization step in the forward pass also requires backpropagating that step in the backward pass.
These issues limit the number of iterations that are computationally feasible (DSPN uses only 10 steps), which can have a negative effect on the modeling capacity due to insufﬁcient minimization of Equation 7.
We aim to address these problems in the following.
Implicit differentiation. Implicit differentiation is a technique that allows efﬁcient differentiation of certain types of optimization problems (Krantz & Parks, 2012; Blondel et al., 2021).
We replace the normal backward pass of DSPN with implicit differentiation and call our method implicit deep set prediction networks ( iDSPN ).
Instead of differentiating through the entire forward optimization, implicit differentiation allows us to only differentiate the (local) optimum obtained at the end of the forward optimization in Equation 7.
The minimizer Y ∗ ( z , θ ) is implicitly deﬁned by the input vector z and the encoder parameters θ through the fact that the gradient of L at an optimum is 0:
[Equation]
The implicit function theorem allows us to differentiate Equation 8 and obtain the Jacobians ∂ Y ∗ ( z , θ ) /∂
z and ∂ Y ∗ ( z , θ ) /∂
θ , which is all we need in order for iDSPN to ﬁt into an autodiff framework.
Appendix B contains the full details on how this works.
In the implicit differentiation, we need to solve a linear system involving the Hessian H = ∂ ( ∇ Y L ) /∂
Y evaluated at Y ∗ (also explained in Appendix B).
This can be problematic when
H is ill-conditioned or singular, so we regularize it with the identity matrix through H (cid:48)
=
H /γ
+ I with γ >
(Martens, 2010; Rajeswaran et al., 2019).
In the limit of γ → ∞ , the Hessian is simply replaced by the identity matrix.
The limit can be interpreted as pretending that the inexact optimization of L found an exact (local) minimum and that L curves up equally in all directions around that point.
Setting the Hessian equal to the identity matrix results in the following Jacobians, which are equivalent to backpropagating through a single gradient descent step at the optimum:
[Equation]
Other motivations for this approximation are discussed in Fung et al. (2022), where they call this Jacobian-free backpropagation.
Matching their results, we ﬁnd that this is not only faster to run and easier to implement (no need to solve a linear system involving H ), but also leads to better results than standard implicit differentiation in preliminary experiments.
Hence, we apply this approximation to differentiate Equation 7 with respect to z and θ in all of our experiments.
Beneﬁts.
Implicit differentiation saves a signiﬁcant amount of memory in training because memory no longer scales with the number of iterations; only the size of the set and its optimizer determine the memory cost since updates to the set can be performed in-place.
Because we only take the derivative at the ﬁnal set, we decrease the computational cost by a factor of around T − 1 2 T , reducing it by nearly a half.
The combined improvements mean that we can minimize the objective in Equation 7 better by increasing the number of optimization steps and using more sophisticated optimizers in the inner loop.
The optimization procedure is treated as a black box now, so it does not even need to be differentiable.
Further improvements. A simple improvement enabled by implicit differentiation is the use of
Nesterov’s Accelerated Gradient for estimating Y ∗ .
This has a better convergence rate on convex problems than standard gradient descent (Nesterov, 1983) and has found wide applicability in deep learning (Sutskever et al., 2013).
Since the computation graph of the optimization is no longer kept in memory, it only has a minor impact on memory and computation: the momentum buffer is the same size as Y , updates can be performed in-place, and the additional momentum calculation does not need to be differentiated.
We have also tried gradient descent with line search, but found that the slowdown due to additional evaluations of L was not worth the small improvement in performance.
By default, we sample a random initial set Y 0 ∼ N ( 0 , I / 10) to start the optimization with.
Similar to DSPN, we can also use a learned initial set Y 0 to start closer to a solution.
However, implicit differentiation treats the optimizer of Equation 7 as a black box, so there is no gradient signal for Y 0 .
We therefore need to include a regularizer in Equation 6 to give us a gradient for Y 0 , for example by adding the regularizer from Rajeswaran et al.
λ 2 || Y − Y
|| 2 .
We use this only in Section 4.2 to make the forward passes of iDSPN and DSPN the same so that we can compare them fairly.
Incorporating constraints. We would like to have the ability to impose constraints on the outputs of iDSPN.
For example, when each set element corresponds to a class label, we would like to [Table] [Table] enforce that every vector is in the probability simplex.
As suggested by Blondel et al.
[Table]
we can incorporate such constraints through constrained optimization in implicit differentiation by changing
Equation 8 to the following and then following the standard derivation in Appendix B:
[Equation]
(10) proj is a function that projects its input onto the feasible region.
We can then use projected gradient descent with this projection to ensure that the constraints are always satisﬁed.
For example, proj could map each element onto the closest point in the probability simplex, which is what we use in Section 4.1.
Note that projections can map different elements to the same value, which can cause problems when ∇ Y L is set-equivariant.
By using an exclusively multiset-equivariant ∇ Y L , we avoid this problem because equal elements after a projection can be separated again.
When we already know parts of the desired Y (e.g. speciﬁc elements in the set or the values along certain dimensions), we can help the model by keeping those parts ﬁxed and not optimizing them in
Equation 7.
For example, in Section 4.1 we know that the ﬁrst few dimensions in the output multiset should be the exactly same as in the input multiset. If we ﬁx these dimensions during optimization, the model only needs to learn the remaining dimensions that we do not already know.
We use this in
Section 4.1 (but not the other experiments) to make iDSPN exclusively multiset-equivariant to the input multiset, not just to the initialization Y 0 .
E XPERIMENTS
In this section, we evaluate three different aspects of our contributions: the usefulness of exclusive multiset-equivariance (Section 4.1), the differences between automatic and our approximate implicit differentiation (Section 4.2), and the applicability of iDSPN to a larger-scale dataset (Section 4.3).
We provide detailed descriptions of the experimental procedure in Appendix D, show example inputs and outputs in Appendix E, and open-source the code to reproduce all experiments at https: //github.com/<redacted>/<redacted> and in the supplementary material.
4. D IFFERENCE BETWEEN EQUIVARIANCES ON CLASS - SPECIFIC NUMBERING
First, we propose a task that highlights the difference between models that are set-equivariant, exclusively multiset-equivariant, and not multiset-equivariant at all.
Given a multiset like [ a, a, b, b, b ] , the aim is to number the elements for each class separately:
[( a, 0) , ( a, 1) , ( b, 0) , ( b, 1) , ( b, 2)] .
Since these are multisets, the numbered elements do not need to be in any particular order.
Equal elements in the input are mapped to different elements in the output, so this is a task where being able to separate equal elements is directly necessary.
Setup. Each element in the input set of size 64 is generated by randomly picking a class out ofchoices and one-hot encoding it.
We have also considered different set sizes and number of classes, but varying these only affected the number of training steps and model parameters required for similar
[Figure]
v a l l o ss n = 8 , d =1 × 2 × compute 0 .0 .n
= 8 , d =  40 iterations 0 .0 .n
= 16 , d =
Baseline iDSPN iDSPN+mom results.
For the target output, a one-hot incremental ID that is separate for each class is concatenated to each element.
We evaluate the accuracy of predicting every element in the multiset correctly.
Results. Table 2 shows our results.
The two exclusively multiset-equivariant models DSPN and iDSPN perform well at any training set size.
As expected, the set-equivariant models are unable to solve this task because they cannot map the equal elements in the input to different elements in the output.
This applies even to transformers with random “position” embeddings, which similarly to TSPN (Kosiorek et al., 2020) and Slot Attention (Locatello et al., 2020) use noise to make elements not exactly equal.
Meanwhile, BiLSTM and transformers with normal positional encoding require at least 100 × more training samples to come close to the performance of iDSPN.
This is because they lack the correct structural bias of not relying on the order of the elements, which makes them less sample-efﬁcient.
Note that the non-equivariant models are unlikely to beneﬁt signiﬁcantly from even more training data because they overﬁt on 1 × and 10 × , but no longer on 100 × .
4. A UTOMATIC VERSUS IMPLICIT DIFFERENTIATION ON RANDOM MULTISETS
Next, we aim to evaluate our contributions in Section 3.
To do this, we compare: 1) baseline DSPN, 2) iDSPN without momentum and 3) iDSPN with Nesterov’s Accelerated Gradient.
Comparing 1 and 2 allows us to see the difference between backpropagation through the optimizer and the approximate implicit differentiation (Equation 9) when using the same forward pass, while comparing 2 andallows us to see the effects of using a better optimizer for the inner optimization.
We propose the task of autoencoding random multisets: every element is sampled i.i.d.
from N ( 0 , I ) .
The goal is to reconstruct the input with a permutation-invariant latent vector as bottleneck.
Varying the set size n and dimensionality of set elements d allows us to control the difﬁculty of the task.
Note that while this appears like a toy task, it is also likely harder than many real-world datasets with similar n and d .
This is because the independent nature of the elements means that there is no global structure for the model to rely on; the only way to solve this task is to memorize every element exactly, which is especially difﬁcult when we want the latent space to be permutation-invariant. This task is also challenging because it indirectly requires a version of push apart : similar elements in the initialization Y 0 might need to be moved far away from each other in the output.
Setup. All models use the same encoder, so DSPN and iDSPN without momentum have exactly the same forward pass.
We train each model on every combination of number of set elements n ∈ { 2 , 4 , 8 , 16 , 32 } , set element dimensionality d ∈ { 2 , 4 , 8 , 16 , 32 } , and optimization steps ∈ { 10 , 20 , 40 } .
We evaluate the ﬁnal validation reconstruction losses over 8 random seeds.
Results. In Figure 1, we show a selection of representative results.
We show the full results in
Appendix C due to the large number of dataset and model combinations.
When comparing the models at the same number of iterations (Figure 1, left), iDSPN performs similarly to DSPN, but iDSPN + momentum is a major improvement.
This shows that the iDSPN gradients are of similar quality despite only using the approximate implicit differentiation, and that the improved optimizer leads to consistently better results.
If we instead compare the models at roughly the same amount of computation (Figure 1, middle), it is clear that iDSPN has signiﬁcant practical beneﬁts over DSPN regardless of momentum.
These results apply to most dataset conﬁgurations (see Appendix C).
The only settings where iDSPN + momentum performed worse than DSPN is when n × d approaches the dimensionality of the latent vector, [-
-]512 (Figure 1, right).
In these cases, the set encoder runs into [Table] fundamental representation limits of set functions (Wagstaff et al., 2019).
With performance being limited by the encoder, it is natural that better optimization yields little to no beneﬁt.
In Appendix E, it becomes clear that no model performs acceptably on these dataset conﬁgurations.
4. CLEVR O BJECT P ROPERTY P REDICTION
Finally, we evaluate iDSPN on the CLEVR (Johnson et al., 2017) object property prediction task that was used in DSPN (Zhang et al., 2019) and Slot Attention (Locatello et al., 2020).
Given an image of a synthetic 3d scene containing up to ten objects, the goal is to predict the set of their properties: 3d coordinate, size, color, and material (18 dimensions in total).
To account for different set sizes in the dataset, a 19th dimension is added as a binary indicator for whether an element is present or not.
Setup.
We mostly follow the experimental setup of DSPN where a ResNet (He et al., 2016) encodes the image into the input vector z , which is decoded into a set by iDSPN.
We also include a few modiﬁcations, which we elaborate on in Appendix D.
Most notably, we use Nesterov’s Accelerated Gradient as iDSPN optimizer and increase the number of optimization steps, both changes that would make DSPN even slower and use more memory.
Aside from these, most changes we make are simpliﬁcations that speed up iDSPN slightly without improving results, so they are also not likely to improve the DSPN baseline results.
The evaluation metric is a modiﬁed average precision (AP threshold ) where a predicted object is considered correct if every property is correct and the 3d coordinates are within the given Euclidean distance threshold of the ground-truth element.
Results. Table 3 shows our results.
iDSPN signiﬁcantly improves over any conﬁguration of Slot
Attention on all AP metrics.
It is better at attribute classiﬁcation when ignoring the 3d coordinates (AP ∞ , 96.4%
→ 98.8%) and improves especially for the metrics with stricter 3d coordinate thresholds (AP 0 .
125 , 7.9% → 76.9%).
This is despite Slot Attention † using a three times higher weight on the loss for the coordinates than iDSPN. Note that a stricter AP threshold is always upper bounded by a looser AP threshold, so iDSPN is guaranteed to be better than Slot Attention † on AP 0 .
We observe some overﬁtting with 128x128 image inputs (1.6e-4 train loss, 5.4e-4 validation loss), which did not appear in preliminary experiments when training an autoencoder with the ground-truth set as input.
We reduce this generalization gap by increasing the image size to 256x256 while keeping the latent vector size the same, which results in further performance improvements (1.1e-4 train loss, 2.5e-4 validation loss).
We thus believe that the overﬁtting is due to the ResNet18 image encoder rather than iDSPN.
Efﬁciency. iDSPN also compares favorably against Slot Attention when considering training times.
Locatello et al.
(2020) train Slot Attention for 200k steps with batch size 512 while we only train for 100 epochs with batch size 128, which is equivalent to 55k steps.
Their model takes around 51 hours to train on a V100 GPU, while our model only takes 2 hours and 40 minutes on a V100 GPU (though with different infrastructure around it).
Ignoring possible speed-ups from increasing the batch size for our model, the per-sample iteration speed is similar.
R ELATED WORK
Set prediction is ubiquitous due to the versatility of representing data as sets.
Example applications include speaker diarization (Fujita et al., 2019), vector graphics generation (Carlier et al., 2020), object detection (Carion et al., 2020) and point cloud generation (Achlioptas et al., 2018).
Since sets are closely related to graphs, there are also natural extensions to molecule generation (Simonovsky & Komodakis, 2018) and scene graph prediction (Krishna et al., 2017; Xu et al., 2017).
As we have covered in this paper, most existing set predictors are either not multiset-equivariant (Vinyals et al., 2016; Stewart & Andriluka, 2016; Vaswani et al., 2017; Achlioptas et al., 2018;
Li et al., 2018; Rezatoﬁghi et al., 2020) or set-equivariant (Lee et al., 2019; Locatello et al., 2020; Kosiorek et al., 2020).
The only set prediction model we are aware of that can be exclusively multiset-equivariant is DSPN (Zhang et al., 2019; Huang et al., 2020). Note that Zhang et al.
(2021a) also make use of the exclusively multiset-equivariant Jacobian of sorting, but instead focus on learning multimodal densities over sets, which is tangential to our work.
One topic we have not covered in detail so far is the use of noise to separate similar or identical elements.
For example, TSPN (Kosiorek et al., 2020) copies the same latent vector z over all elements and adds Gaussian noise to still be able to separate them with set-equivariant self-attention.
While this makes the elements different, the other problems we point out in Section 2.2 still apply, such as the fact that elements can become similar in later parts of the network and that there can be modeling problems with similar (not necessarily equal) elements.
Indeed, in Section 4.1 the transformer with random position embeddings (similar application of noise as in TSPN) does not perform well.
The difference between operating on sets versus multisets has also been studied in the context of graph neighborhood aggregation (Xu et al., 2018).
There, the concern is about set and multiset- invariance rather than equivariance: mean and max pooling lose information about duplicate elements while sum pooling does not.
Equal elements easily arise in graphs (e.g. in molecules, the multiset of neighbors often contains multiple atoms of the same chemical element), so making a clear distinction between sets and multisets for invariance and equivariance properties is especially important in such contexts.
C ONCLUSION AND F UTURE W ORK
We have established a more appropriate property for set prediction (exclusive multiset-equivariance) and that DSPN satisﬁes this property.
We showed that after improving DSPN with approximate implicit differentiation, iDSPN can outperform a state-of-the-art set predictor like Slot Attention signiﬁcantly.
However, there are still some open problems.
Lack of inductive bias to spatial inputs. Most structured data can be encoded into a vector, which makes iDSPN with its vector input versatile.
However, it also limits the inductive bias towards spatial domains.
When working on CLEVR with image inputs and set outputs, we had to compress all spatial information into a single vector before decoding it into a set.
The beneﬁt of methods like Slot
Attention is that they take another set as input, possibly including positional encoding.
This allows them to work with feature map inputs rather than only feature vectors, which we believe is especially helpful with tasks like object discovery where having a spatial bias is crucial.
In preliminary experiments on the Tetrominoes object discovery task from Greff et al.
(2019), we only inconsistently reached above 98% on the ARI evaluation metric, unable to match the 5-run average of 99.5% ARI for Slot Attention (Locatello et al., 2020) and 99.2% ARI for IODINE (Greff et al., 2019).
We believe that in order for iDPSN to scale to larger and more complex tasks, an architecture that makes better use of spatial information is crucial.
We also propose that better set prediction is a requirement for better object discovery; if a model cannot predict a set with direct supervision well, why should we expect it to be able to infer a similar one without direct supervision?
Multiset prediction is fundamentally discontinuous.
Lastly, we made the point in Section 2 that multiset-equivariant functions like push apart can require discontinuous changes when elements are equal.
This is a concrete example of the responsibility problem (Zhang et al., 2020; Huang et al., 2020), which shows that certain set prediction tasks require discontinuous changes.
It is an open problem on how to best incorporate such discontinuities into a model while still being trainable.
Optimization as is used in iDSPN is a promising step towards dealing with this problem, but there is certainly plenty of room for improvement in both theory and practice.
R EPRODUCIBILITY S TATEMENT
Proofs and derivations are contained in Appendix A and Appendix B. Detailed descriptions of experiments, including how to generate these datasets in the case for Section 4.1 and Section 4.2, are given in Appendix D.
We provide all the code, scripts to reproduce the experiments with the exact hyperparameters, and scripts to calculate the evaluation results.
In the de-anonymized version, we intend to make every training run and the corresponding pre-trained weights used to produce the results in this paper available through public Weights & Biases tables (Biewald, 2020).
E THICS STATEMENT
We propose improvements to a general vector-to-set model and evaluate it largely on synthetic datasets. We have not yet shown the applicability of iDSPN on realistic large-scale datasets. The most immediate potential application is in object detection, which can be used for improving safety (better understanding of the world in self-driving cars) but also for surveillance (face detection and tracking). This is shared with all the general set prediction models.
We showed that the main improvement on CLEVR property prediction for iDSPN is the better localization of the objects in the 3d world coordinates. This should be more useful for self-driving cars to locate objects and predict their trajectories more accurately, rather than surveillance where presence is likely more important than the precise position in 3d space. In general, it is important to be mindful of biases when using such machine learning systems (Torralba & Efros, 2011; Mehrabi et al., 2021). For example, models trained on general object recognition datasets can be discriminatory towards lower-income households (DeVries et al., 2019).
It is of course especially important to be careful when the system concerns humans, for example with biases relating to human ethnicity and gender in facial recognition systems (Klare et al., 2012; Robinson et al., 2020). A T HE J ACOBIAN OF SORTING IS EXCLUSIVELY MULTISET - EQUIVARIANT First, let us clarify what we mean by the Jacobian of sorting. We use sort : R n × 1 → R n ×to denote sorting of a multiset with scalar elements in ascending order.
To be precise, we deﬁne sort ( X ) =
S X X where S X is the permutation matrix that sorting X would apply.
With Jacobian of sorting, we mean: [Equation] In the last step, we treat S X as independent of X because ∂ S X /∂ X is 0 almost everywhere and undeﬁned on a zero measure set.
The deﬁnition of sorting can be extended to vector elements (the setting in FSPool (Zhang et al., 2020)) by independently sorting each dimension across the multiset so that we have a different S X for each dimension.
No set-equivariance.
To show that the Jacobian of numerical sorting is not set-equivariant, a simple counterexample is to realize that the Jacobian does not change between sort (cid:0) [ 1 0 0 1 ] [0 , 0] (cid:62) (cid:1) and sort (cid:0) [ 0 1 1 0 ] [0 , 0] (cid:62) (cid:1) .
The input does not change regardless of how we permute [0 , 0] (cid:62) , therefore the Jacobian of sorting does not change.
To make this point more concrete, we implement push apart from the main text using the Jacobian of sorting.
We omit the transposes in the following for brevity.
We begin by deﬁning the function g ([ a, b ]) =
sort ([ a, b ]) · [ − 1 , 1] .
This multiplies the smaller value with − 1 and the larger with(ties broken arbitrarily), then adds them together.
Taking the gradient of g with respect to the input set means following the permutation that the sort applied in reverse ( S (cid:62) X ), so the − 1 gets propagated to the smaller element and 1 to the larger one.
In other words, ∇ [ a,b ] g = [ − 1 , 1] if a ≤ b (alternatively a < b ), else [1 , − 1] .
Now we can deﬁne push apart ([ a, b ]) =
[ a, b ] + ∇ [ a,b ] g to obtain exactly the same function as in the main text.
Multiset-equivariance. We now show that the Jacobian of sorting is multiset-equivariant.
Proposition 2.
The Jacobian of sorting f ( X ) =
∂ sort ( X ) /∂
X = S (cid:62) X is multiset-equivariant.
Proof.
To show that for all P 1 , there exists a P 2 that satisﬁes the two conditions of multisetequivariance, we set P 2 = S (cid:62) P 1 X S X .
We use the fact that S is a permutation matrix so S (cid:62) S =
SS (cid:62)
= I .
To show that the ﬁrst condition is true:
[Equation]
We now use the fact that sorting is permutation-invariant, so for any P , S X X = S PX P X .
To show that the second condition is true:
[Equation]
[Equation] We showed the existence of a P 2 that fulﬁlls both conditions for all P 1 , so f is multiset-equivariant.
Note that in the deﬁnition of multiset-equivariance, it is important that the existential quantiﬁer is on
P 2 acting on the output of f rather than P 1 acting on the input.
If the quantiﬁers were the other way around ( ∀ P 2 ∃ P 1 ), then the multiset-equivariance deﬁnition would not be useful: we have seen that permuting equal input elements does not affect the Jacobian of sorting, so there must be some P 2 for which no choice of P 1 works.
B D ERIVATION OF IMPLICIT DSPN
The goal of differentiating Equation 7 is to obtain the gradient of the training objective with respect to z and θ .
If we can do that, then iDSPN ﬁts into an autodiff framework like a normal differentiable module.
We start by deriving the Jacobians for z and θ with implicit differentiation and present the approximation for the Hessian afterwards.
Blondel et al. (2021) describe a straightforward and efﬁcient recipe for applying implicit differentiation to optimization problems. We follow their approach for the standard implicit differentiation derivation and efﬁciency concerns. We highly recommend their paper for further background and possible extensions that are made possible by implicit differentiation.
To simplify notation, we combine the target vector with the encoder parameters into Θ =
( z , θ ) and restate the optimality condition:
[Equation]
This implicitly deﬁnes a minimizer function Y ∗ (Θ) .
The implicit function theorem (further background in Krantz & Parks (2012)) states that its Jacobian ∂ Y ∗ (Θ) /∂
Θ exists if ∇ Y L is continuously differentiable and the Hessian ∂ ( ∇ Y L ( Y ∗ , Θ)) /∂
Y ∗ is a square invertible matrix.
Using ∇ Y L rather than L lets us apply the implicit function theorem: its Jacobian is a square matrix instead of a vector and the optimality condition is fulﬁlled even when the root of ∇ Y L only corresponds to a local minimum of L .
In practice, we will estimate Y ∗ with an optimizer (e.g. gradient descent) that approximately minimizes L .
Using the chain rule, we differentiate both sides of Equation 19 with respect to Θ and reorder to obtain:
[Equation]
This is a linear system of the form HX = B where we can use autodiff to compute H and B .
Solving this system for X = ∂ Y ∗ (Θ) /∂
Θ lets us compute the gradients for z and θ with ∇ Θ L = ( ∇ Y ∗ L ) (cid:62) X , where L is the loss of the prediction (as opposed to L ).
Efﬁciency. In order to incorporate iDSPN into an autodiff framework, we only need the vectorJacobian product (VJP) v T X for some vector v . Following Blondel et al.
(2021), We can efﬁciently compute this by solving the much smaller linear system Hu = ∇ Y ∗ L for u .
This allows us to compute the desired gradient with ∇ Θ L = ( ∇ Y ∗ L ) (cid:62)
X
= u (cid:62) HX
= u (cid:62) B . Instead of solving the linear system exactly, it is more efﬁcient to use the conjugate gradient method (Blondel et al., 2021). Rajeswaran et al.
(2019) show that a small number of conjugate gradient steps (5 in their experiment) can be enough to obtain a sufﬁciently good solution.
Another advantage of conjugate gradient is that we only need to compute the Hessian-vector product Hu instead of the full Hessian H .
Not only does this avoid storing H explicitly (decreasing memory requirements), with reverse-mode autodiff like in PyTorch (Paszke et al., 2019) and Tensorﬂow (Abadi et al., 2015) it also improves parallelization: to compute the full Hessian explicitly we would have to differentiate each entry in ∇ Y L individually.
The same advantage applies to B since we can also use a VJP to compute u (cid:62) B .
Previous uses of implicit differentiation in meta-learning (Rajeswaran et al., 2019) and neural architecture search (Zhang et al., 2021b) also involve solving a linear system, but their Y corresponds to the neural network parameters and thus usually has millions of entries.
In contrast, in our setting we work with a much smaller Y , which for example only contains 190 entries in Section 4.3. Regularization. In Rajeswaran et al.
(2019), the regularization of the Hessian like we described in
{+Section 3 appears due to regularizing L by adding λ 2 || Y − Y 
+}
|| 2 .
In preliminary experiments, we observed that even if L is not regularized, solving for the regularized Hessian ( H /γ + I ) yielded better performance at the same number of conjugate gradient steps. Adding the identity matrix can be seen as adding 1 to the eigenvalues to improve the conditioning of H . Of course, in the main text the efﬁciency improvements are irrelevant because we take γ → ∞ . Due to that, H =
I
so Equation 20 simpliﬁes to Equation 9.
Fung et al.
(2022) provide different motivations for this approximation.
C R ANDOM SETS FULL RESULTS
We had to exclude 10 runs of DSPN due to signiﬁcantly worse results making their average uninformative.
We did not observe any stability issues with iDSPN, so we did not have to exclude any iDSPN runs.
The following excluded DSPN runs were all using 40 iterations:
• 6 runs of n = 2 , d = 32 (out of 8 seeds), • 1 run of n = 2 , d = 16 , • 2 runs of n = 4 , d = 16 , • 1 run of n = 4 , d = 32 .
Figure 2 and Figure 3 provide different views of the same underlying data in Table 4: each subplot in
Figure 2 keeps n ﬁxed and varies d , while each subplot in Figure 3 keeps d ﬁxed and varies n . Since the scale of the losses varies across the different dataset conﬁgurations, we normalize each ( n, d ) combination by dividing the loss by the loss of the it=10 Baseline (DSPN) run. That is why it=10 Baseline always has a relative loss of 1. The runs are re-ordered from the main text to be in the order iDSPN with momentum, iDSPN,
Baseline to make visual comparison for equal computation easier. The correspondence between color and model is still the same. To make a comparison for equal iterations, compare the baseline
DSPN (orange bar) to the two bars to the left of it (in the table, the two rows above).
To make a comparison for similar computation, compare the baseline DSPN (orange bar) to the two bars to the right of it (in the table, the two rows below), but only if within the same n, d dataset conﬁguration.
[Figure]
r e l a t i v e v a li d a t i o n l o ss
[Figure]
n
[Figure]
l a t i v e v a li d a t i o n l
[Figure]
l a t i v e v a li d a t i o n l
[Figure]
ss
[Figure]
[Figure]
r e l a t i v e v a li d a t i o n l o ss
[Figure]
l a t i v e v a li d a t i o n l
[Figure]
l a t i v e v a li d a t i o n[-l o ss d = it=
-]
[-it=
-]
[-iDSPN it=10 Baseline it=20 iDSPN+mom it=20 iDSPN it=20 Baseline it=
 iDSPN+mom it=40 iDSPN it=40 Baseline [Table]-][Figure]
[Table]
D D ETAILED EXPERIMENTAL DESCRIPTION
In each of the experiments, we make use of FSPool (Zhang et al., 2020) in the set encoder g for iDSPN and DSPN, which allows them to be exclusively multiset-equivariant.
Multiset loss. When the task is to predict sets, we need to compute a loss between a predicted and a ground-truth set.
Throughout this paper, we will make use of the Hungarian loss (Kuhn, 1955):
[Equation]
It ﬁnds the optimal 1-to-1 mapping from each element in the prediction Y (1) =
[ y (1)1 , . . .
, y (1) n ] (cid:62) to each element in the ground-truth Y (2) =
[ y (2)1 , . . .
, y (2) n ] (cid:62) , minimizing the total pairwise losses.
A pairwise loss is a standard loss like mean squared error and cross-entropy.
The Hungarian loss is permutation-invariant in both its arguments and is able to map duplicate predicted elements to different ground-truth elements and vice versa, which makes it well-suited for multiset prediction.
Set reﬁner networks.
Set reﬁner networks (SRN) (Huang et al., 2020) make a few minor changes to DSPN.
When we mention DSPN, we use the following modiﬁcations:
• SRN found that a learning rate of 1 for the inner optimization (Equation 7) is sufﬁcient.
DSPN originally used a learning rate of 800 due to normalizing by the set size in the set encoder and a bug that made the effective learning rate depend on the batch size.
• SRN found that with the better learning rate, the additional supervised loss term that DSPN proposes is not necessary for successful training.
D. C LASS - SPECIFIC NUMBERING
The following description corresponds to Section 4.1.
Dataset.
This experiment can be thought of as an elementwise classiﬁcation that requires global information: each input element needs to be classiﬁed (predicting the numerical ID of the element), but the prediction needs to take the other predictions into account to avoid duplicate IDs.
We use a Hungarian loss to match the predictions and ground-truth within each class, because we do not care which instance of a class gets which ID as long as all the correct IDs are predicted.
We use
MSE as pairwise loss for DSPN/iDSPN and cross-entropy as pairwise loss for the other models. The baselines perform worse with MSE as pairwise loss. For each example, the input multiset has a size of 64 with 4-dimensional elements corresponding toclasses. This is represented as a 64 × 4 matrix where each row is a one-hot vector that is sampled i.i.d.
from the multinomial distribution over the equally-weighted 4 classes.
We generate the target multiset of the same set size with 64-dimensional elements ( 64 × 64 matrix), each element being a one-hot vector that represents a number. For each class, we number the corresponding elements sequentially.
The setting with 1 × samples has a training dataset of size 640.
We additionally use a validation set of size 6,400 and a test set of size 64,000 for every run.
Models. We have the prior information that every instance of a class in the input will also appear in the desired output. To make the task easier for the models, we modify them to take this into account. This makes sure that a model does not need to learn to preserve the input class in the output; it only needs to predict the ID associated to the class.
Both the equivariant DeepSets version (as opposed to the more commonly-used invariant version) and BiLSTM naturally pair up each input with each output element, so they already satisfy the desired property.
If the “ﬁrst” element in the input is class c , the “ﬁrst” predicted ID in the output will naturally correspond to class c as well.
For the Transformer baselines, we use an additional modiﬁcation: the input set to the transformer decoder is based on an afﬁne transformation ( Y 0 ) i
=
W x i + b of the 1 ≤ i ≤ n input elements.
Through this, each input element is paired up with an output element as well.
The Transformer without PE becomes set-equivariant to the input this way. For the Transformers with PE and random PE, we provide the position encoding only to the transformer decoder because giving it to the transformer encoder signiﬁcantly degraded the performance.
We adapt the iDSPN and DSPN models to this setting by ﬁxing the ﬁrst four dimensions of Y to the values in the input multiset as we discussed in Section 3, which makes them multiset-equivariant to the input.
With the set encoder g architecture of Linear( 4 + 64 → 256 )–ReLU–Linear( 256 → 256 )–FSPool, they become exclusively multiset-equivariant.
Since the initial set Y 0 is sampled i.i.d. , we do not need to pay special attention to which input element gets associated with which random initialization.
During the inner optimization of iDSPN, we only update the 64 dimensions corresponding to the output set and keep the ﬁrst four dimensions ﬁxed.
Additionally, we apply projected gradient descent with a proj that projects each element back onto a point in the probability simplex – the convex hull relaxation of the one-hot encoded IDs – after each gradient descent update. This ensures that the vectors in Y remain valid probability distributions after each update.
We found that this signiﬁcantly improved the performance over the unconstrained version.
Evaluation. We compute the accuracy at the sample-level, meaning a predicted set is considered correct only if every element is correct.
The predicted ID for each predicted element is obtained by taking the argmax over the elements’ dimensions in the output.
The baselines are very volatile during training, which results in very large variances at the end of training.
To reduce this variance, we pick the best model according to the validation accuracy, evaluated every 500 training steps.
We report the accuracy on the test set for the best model.
Each run samples new datasets based on the random seed, so we evaluate all models using the same set of random seeds.
D. R ANDOM S ETS A UTOENCODING
The following description corresponds to Section 4.2.
Dataset. Every set has a ﬁxed size with n elements of dimensionality d .
Every element is sampled i.i.d.
from N ( 0 , I ) .
We assign 64,000 samples to the training set and 6,400 samples to the test set.
Changing the random seed also changes the dataset.
As loss and evaluation metric, we use Hungarian matching with the Huber loss (Huber, 1964) (quadratic for distances below 1, linear for distances above 1) as pairwise loss.
We always use the results for a model after the ﬁnal epoch without any early-stopping (we did not observe overﬁtting) or selection based on best loss.
Models. We try to make the DSPN baseline and the iDSPN models as similar as possible.
Since
DSPN was intended to be used with a learned initial set Y
(Zhang et al., 2019), we match this in iDSPN and add the regularizer 0 .
1 || Y − Y
|| 2 to Equation 8 as we discussed in Section 3.
We apply the same regularization to the objective in DSPN to make the objectives match up.
We initialize Y 0 at the start of training with elements sampled from N ( 0 , I / 10) for all models.
For the same seed, the weight initializations are the same for all models.
We use the same set encoder to encode the autoencoder input between the models.
We use the same set encoder g between the models (but we do not share the weights with the encoder for the autoencoder input).
The architecture of these set encoders is Linear( d → 512 )–ReLU–Linear( 512 → 512 )–
FSPool, where the linear layers are applied on each element in a multiset individually with shared weights across the elements. Note that we do not compare against Slot Attention here due to the difference in setup.
Slot Attention takes a set directly as input, so giving it the input set directly would not give us a permutation-invariant bottleneck, which somewhat defeats the point of the autoencoder.
It only has to learn to produce the identity function, which we veriﬁed it can easily do.
Training. All models are trained for 40 epochs at batch size 128 with Adam (Kingma & Ba, 2015) with a learning rate of 1e-3 and default momentum parameters.
This is enough to converge for almost every model and dataset combination.
Only the hardest difﬁculty settings ( 32 × 16 , 16 × 32 , 32 × 32 ) would show any beneﬁt from more training.
D. CLEVR O BJECT P ROPERTY P REDICTION
The following description corresponds to Section 4.3.
Where possible, we use exactly the same setup as DSPN.
Every image is resized to 128x128 (or 256x256) and scaled to be in the interval [0 , 1] .
For the groundtruth, we scale down the target 3d coordinates from the interval [ − 3 , 3] to [ − 1 , 1] and concatenate it to the attributes (each of which is one-hot encoded) and the binary mask feature (1 if element is present, 0 if not).
The input image is processed by a ResNet encoder, which scales the input down to a 4x4x512 (8x8x512) feature map, on which we apply BatchNorm followed by a 2x2 convolution with stride 2 to obtain a 2x2x128 (4x4x32) feature map.
This is ﬂattened to a latent vector withdimensions, which is the input z to iDSPN. We use the Hungarian loss with Huber loss (Huber, 1964) as pairwise loss. We train the model for 100 epochs with the Adam optimizer (Kingma & Ba, 2015), a learning rate of 1e-3, and default momentum hyperparameters. For evaluation, we adapt the average precision evaluation code to PyTorch.
It is calculated by sorting the predicted elements in descending order by the dimension corresponding to the binary mask feature (acting as conﬁdence), then computing the precision-recall curve by including more and more predicted elements, starting from the most conﬁdent elements.
A prediction is considered correct only if all attributes are correct (after taking the argmax for each attribute individually) and the 3d coordinate is within a given threshold.
We make some changes to the training procedure of DSPN:
• Instead of using a relation network (Santoro et al., 2017) – expanding the set into the set of all pairs ﬁrst – paired with FSPool, we skip the relation network entirely.
This improves the complexity of the set encoder g used in iDSPN from O ( n 2 log n ) to O ( n log n ) .
Using the relation network approach would improve our results slightly (e.g. 3.5 percentage points improvement on AP 0 .for 128 × 128 ), but is also a bit slower.
For simplicity, we therefore opted to not using relation networks.
The architecture of the set encoder g in iDSPN is Linear( 19 → 512 )–ReLU–Linear( 512 → 512 )– FSPool.
The main difference to DSPN is that since there is no concatenation of pairs, so the input dimensionality is 19 instead of 38 and everything is applied on sets of size n rather than sets of size n 2 .
• Instead of ResNet34 to encode the input image, we use the smaller ResNet18.
This did not appear to affect results.
• Instead of using a learned initial set Y 0 as in DSPN, we ﬁnd that it makes no difference to randomly sample the initial set for every example.
We therefore use the latter for simplicity.
In initial experiments we found that even initializing every element to 0 causes no problems.
• We increase the batch size from 32 to 128.
There appeared to be no difference in results between the two, with 128 being faster by making better use of parallelization.
• We use Nesterov’s Accelerated Gradient (Nesterov, 1983) with a momentum parameter of 0.
instead of standard gradient descent without momentum.
• Instead of ﬁxing the number of iterations at 10 like DSPN, we set the number of iterations to 20 at the start of training and change it to 40 after 50 epochs.
This had slightly better training loss than starting training with 40 iterations.
We have tried a few other ways of increasing the number of iterations throughout training (going from 10 to 20 to 30 to 40 iterations, smooth increase from 1 to 40 over the epochs, randomly sampling an iteration between 20 and 40 every batch), which had little impact on results. iDSPN training was stable in all of these conﬁgurations.
• We drop the learning rate after 90 epochs from 1e-3 to 1e-4 for the last 10 epochs.
This slightly improved training loss while also reducing variance in epoch-to-epoch validation loss.
• In preliminary experiments, we rarely observed spikes in the training loss.
Clipping the gradients in the inner optimization to a maximum L2 norm of 10 seemed to help.
Note that we use a ResNet18 image encoder, while Slot Attention uses a simpler 4-layer convolutional neural network with 5 × 5 kernels.
This difference is not so easily compared: we compress the image into a 512d vector while they operate on a feature map.
Their ﬁnal feature map for CLEVR has 32x spatial positions with 64d each, so it can be argued that their latent space is 65536-dimensional.
It is natural that a tighter bottleneck requires more processing.
ResNet18 also applies several strided convolutions early, so the amount of processing between it and the Slot Attention image encoder are not too dissimilar.
This is reﬂected by the fact that the time taken to process a sample is similar between ResNet18 + iDSPN and CNN + Slot Attention, even though we use a smaller batch size and could gain a speed-up from better parallelization for higher batch sizes. A small difference between the default Slot Attention setup and our setup is that they normalize the 3d coordinates to be within [0, 1] while we use an interval of [-1, 1]. In Table 3, Slot Attention* uses an interval of [-1, 1] (the same as iDSPN) while Slot Attention † uses the interval of [-3, 3] (the default coordinate range for this dataset). This increases the weight on the coordinates versus the other attributes, so it improves AP for strict thresholds but trades off classiﬁcation performance at a threshold of inﬁnity.
We did not experiment with different coordinate scales and simply kept it the same as DSPN.
E E XAMPLE INPUTS AND OUTPUTS
In the following, we show some example inputs of each dataset and corresponding model outputs.
The focus is not on showing representative examples, but highlighting some qualitative differences. Table 5 shows the worst failure case of each model trained with 100 × samples when numbering elements within a class (Section 4.1). For each model, the sample is selected from the dataset by picking the prediction with the highest number of misclassiﬁcations. Since we make every model equivariant to the input, the classes shown in the prediction are guaranteed to be correct; the difﬁculty lies only in numbering them correctly.
Note that the classes in the input are ordered arbitrarily and the order of the numberings within each class is arbitrary as well; we sort the outputs in the table so that it is easier to tell what elements are missing or duplicated.
Interestingly, the worst input for transformer with PE, BiLSTM, DSPN, and iDSPN all contain a class that only occurs twice, which is atypically few.
Figure 4 shows the autoencoding of randomly-generated sets (not in training set nor validation set) for each n and d combination (Section 4.2).
We leave out iDSPN without momentum because its performance is so similar to DSPN.
Both models apply 20 iterations in the forward pass.
The visual difference between model qualities is often somewhat subtle and best seen when zooming in on a digital display.
It is most obvious on ( n = 4 , d = 16) and ( n = 8 , d = 16) , but there are smaller differences visible with conﬁgurations like ( n = 32 , d = 2) as well, where iDSPN with momentum tends to be slightly more centered on the ground-truth than DSPN.
The examples where n × d ≥show that even if DSPN has a better loss than iDSPN on those conﬁgurations, the model quality is far from acceptable.
Table 6 and Table 7 show predictions on CLEVR for different number of evaluation iterations (Section 4.3).
We checked the ﬁrst 128 images of the CLEVR validation set manually and selected examples that looked difﬁcult due to occluded objects, then evaluated the model on these examples.
For easier examples, iDSPN is usually correct on all objects at 40 iterations.
The mistakes of the model are reasonable, such as confusing a yellow cube that has its sides and bottom occluded for a cylinder (third image).
The improvement seen when going from 128x128 to 256x256 images is also quite natural, even for a human observer.
[Table]
) (c, 4) (d, 2) (b, 1) (d, 6) (c, 16) (c, 18) (c, 5) (c, 4) (d, 3) (b, 2) (d, 7) (c, 17) (c, 19) (c, 5) (c, 4) (d, 4) (b, 3) (d, 10) (c, 17) (c, 20) (c, 5) (c, 4) (d, 5) (b, 4) (d, 11) (c, 18) (c, 22) (c, 5) (c, 4) (d, 5) (b, 5) (d, 12) (c, 19) (c, 23) (c, 5) (c, 4) (d, 5) (b, 5) (d, 13) (c, 20) (c, 24) (c, 5) (c, 4) (d, 6) (b, 6) (d, 15) (c, 20) (c, 25) (d, 14) (d, 0) (d, 7) (b, 7) (d, 16) (c, 21) (c, 26) (d, 14) (d, 0) (d, 8) (b, 8) (d, 17) (c, 22) (c, 26) (d, 14) (d, 0) (d, 9) (b, 9) (d, 19) (c, 23) (c, 26) (d, 14) (d, 0) (d, 10) (b, 10) (d, 20) (c, 24) (c, 27) (d, 14) (d, 0) (d, 11) (b, 11) (d, 22) (c, 26) (c, 28) (d, 14) (d, 0) (d, 11) (b, 12) (d, 23) (c, 27) (c, 29) (d, 14) (d, 0) (d, 14) (c, 0) (d, 24) (c, 27) (c, 29) (d, 14) (d, 0) (d, 15) (c, 1) (d, 25) (c, 28) (c, 29) (d, 14) (d, 0) (d, 17) (c, 2) (d, 26) (d, 0) (d, 0) (d, 14) (d, 0) (d, 18) (c, 3) (d, 26) (d, 1) (d, 1) (d, 14) (d, 0) (d, 18) (c, 4) (d, 27) (d, 2) (d, 2) (d, 14) (d, 0) (d, 19) (c, 5) (d, 28) (d, 3) (d, 3) (d, 14) (d, 0) (d, 20) (c, 6) (d, 28) (d, 4) (d, 4) (d, 14) (d, 0) (d, 20) (c, 7) (d, 28) (d, 5) (d, 5) (d, 14) (d, 0) (d, 24) (c, 8) (d, 28) (d, 6) (d, 6) (d, 14) (d, 0) (d, 25) (c, 9) (d, 29) (d, 7) (d, 7) (d, 14) (d, 0) (d, 25) (c, 10) (d, 29) (d, 8) (d, 8) (d, 14) (d, 0) (d, 25) (c, 11) (d, 29) (d, 9) (d,
[Table]
[Figure]
[Table]
[Table]
[Table]