Multiset-Equivariant Set Prediction with Approximate Implicit Differentiation
Authors: Yan Zhang
David W Zhang
Simon Lacoste-Julien
Gertjan J. Burghouts
Cees G. M. Snoek
Venue: ICLR 2022 Poster (2022)
Keywords: set prediction
permutation equivariance
implicit differentiation

Abstract: 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 beneficial 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 benefits made possible by implicit differentiation.

Published as a conference paper at ICLR 2022
MULTISET-EQUIVARIANT SET PREDICTION WITH APPROXIMATE IMPLICIT DIFFERENTIATION
Yan Zhang*!
David W. Zhang
Simon Lacoste-Julien' 13,4
Gertjan J. Burghouts
Cees G. M. Snoek?
Samsung SAIT AT Lab, Montreal
University of Amsterdam?
Mila, Université de Montreal
Canada CIFAR Al Chair
TNO
ABSTRACT
Most set prediction models in deep learning use set-equivariant operations, but they actually operate on multisets. We show that set-equivariant functions cannot repre- sent certain functions on multisets, so we introduce the more appropriate notion of mulliset-qquivvrianne We identify that the existing Deep Set Prediction Network (DSPN) can be multiset-cquivarian without being hindered by set- reteeiiivaeance and improve it with approximate implicit differentiation, allowing for better op- timization while being faster and saving memory. In a range of toy experiments, we show that the perspective of multistt-equivariance is beneficial and that 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 benefits made possible by implicit differentiation.
INTRODUCTION
An example of set prediction is detecting the objects in an image: we want 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 agents in an environment, instruments in a piece of music, or atoms in 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 arbitrary order. To still treat the list if it were multiset set, recent set prediction models avoid being affected by this arbitrary order by using list-to-list functions that are enforced to be permutation-equivrrae elements should change the order of the output elements in the same way.
(Zaheer al., 2017): changing the order of the input
Now, consider this function that takes multiset (represented as a list) of two scalars as input and "pushes them apart".
[a 1,b + 1] if a S b [a 1, otherwise
push.apart((o.b)
For example, push.apart ([1, = [0.3] 3] and push.apart(22.1] 1]) = [3. 0j. At first glance, appears that this function is permutation-equivarinn because swapping the inputs swaps the outputs. However, notice that push apart((0, 0)) = [11, 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 definition of permutation-qquiaarannce.. Is this problem in practice when there are rarely equal elements, but instead similar elements? We aim to address these question in our paper and present two key ideas:
Equull contribution

---

Published as a conference paper at ICLR 2022
Table Difference between equivariance properties in the context of multisets. Only when satisfying exclusive multiset-equivee equivariance (this paper) is it enforced that order does not matter while also allowing equal input elements to be mapped to different output elements like in apar_
is set- equivariant
is multiset- equivariant
f(la, b)) = [c. d] + *[([, al) = (d. cl
[((a, al) (c. d
guaranteed
not possible possible /
X
not guaranteed x
guaranteed
possible /
Multiset-qquivariince 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 the functions they can represent on multisets due to the limitation of being permutation-equivaiinn equivariant, which we call set-equivariance. Due to continuity, this is problem even when there are no exact duplicate elements. Thus, we develop a new notion of mulliset-equivarrianee, which we argue is a more suitable property for operations
multisets. To avoid the aforementioned limitation of set-equivariance, a crucial point is that it is possible to be multiset-equivarian without being sel-equivariant, which we call exclusive multiset-quuivarinnce We identify and show that deep set prediction networks (DSPN) (Zhang et al., 2019; Huang al., 2020) satisfy this property with specific choice of set encoder. This makes it the only set prediction model we are aware of that can satisfy exclusive multiset-equivarinnee
2. Implicit DSPN. Despite this beneficial property, DSPN is outperformed by the set-equivariant Slot Attention (Locatello al., 2020), which motivates us improve other aspects of DSPN. We propose implicit DSPN (iDSPN) in Section 3: 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 Hessian while still having the same benefits.
In our experiments in Section we start by evaluating the first 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 significantly less sample- efficient. Next, we evaluate the second point by testing the modeling capacity on autoencoding random sets, where iDSNN performs similarly to DSPN at the same iteration count and much better for the same computational cost. Lastly, iDSPN significantly raises the bar on CLEVR object property prediction, outperforming the state-of-the-as Slot Attention (Locatello et al., 2020) by a large margin of 69 percentage points the APo. metric while only training for 7% of the number of epochs.
LIMITATIONS OF SET-EQUIVARIANCE
In this section, we begin by identifying problems that limit the expressivity of a wide class of set prediction models satisfying type of cermutation-equivrrint set-equivariance is too strict and instead propose a relaxation of set-equivariance that more suitable We end by discussing how an existing method,
set- vet-equivariance. We argue that
for set prediction models: deep set prediction networks (Zhang et al., 2019; Huang al., 2020), is multiser-equivarian not set-equivariant and the associated benefits thereof (Table 1).
2.1 PRELIMINARIES
Sets in the context deep learning are typically unordered collections of real-valued vectors. They are represented in memory n x d matrices with the 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 short, when changing the order of the elements in the matrix, the output of function should either not change all (invariance) or change order similarly (equivariance).
-inittttiiii-i-rinnncccee and permutation-equuivriinn
(Zaheer al., 2017). In
2

---

Published as a conference paper at ICLR 2022
2.2 SET-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-equiivan set Yo. which is usually randomly generated or learnable matrix. Permutation Permutation-quuivarianne 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 casily composed build larger models that remain equivariant, which fits well into deep learning architectures as building blocks.
list-to-list) functions to refine an initial
However, without having any restrictions duplicate elements, it is accurate think of set prediction models as using multiset-to-multise 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-eqqivariieni for sets to multiset inputs and refer to it as ser-equivariance. This directly matches the usual definition of permutation permutation-equuvariaaee equivariance in the context of set imputs in the literature. Thus, many models already satisfy the following definition, including Slot Attention (Locatello et al., 2020), Set Transformers (Lee et al., 2019), TSPN (Kosiorek et al., 2020), and the equivariant DeepSets (Zaheer et al., 2017).
Definition (Set-equivariance for multiset inputs). A function f BRaxd -> Raxda with multiset input is set-equivariant iff for any permutation matrix P.
(PX = PS(X)
consequence is that elements in a multiset that are equal before set-equivariant function must remain equal after applying the function, as shown by the following:
Proposition 1. For all permutation matrices Pt, P2 such that P = P2X, set-equivariant function f satisfies:
= =
Proof.
-Pf(()
(4)
P-X can be satisfied by Pi + P2 when there are duplicate elements in X, so Pif(X) P2)(X) means that they must remain duplicates. This inability of separating equal elements demonstrates that the classical definition of permutation-equivarranc. is overly restrictive when directly applied to multisets as in Definition Therefore, we seek a more appropriate definition of permutation-qquiivrraanc for multiset inputs.
We now consider slightly different notion of cermutationeqquivrraaae that no longer forces equal elements to remain the same: mulliset-qquivariance. Recall the pushaapart function from Equation 1, which we suggest should be valid equivariant func- tion for multisets, but cannot be represented by a set-equivariant function. Therefore, we propose that there should be distinction between set-equivariance and corresponding to the appropriate equivariance property for the typical multiset-to to multiset mappings used in set prediction. The key difference of nuitiset-equivariance is that the output order for equal elements in the input multiset can be arbitrary. This leads us to the following definition:
Multiset-equivariance.
the latter
Definition 2 (Multiset-equivariance for multiset inputs). A function f Raxd - nxx with multiset input X is multiset-equiiaaiann iff for any permutation matrix Pt, there exists permutation matrix P2 such that P,X = P2X and:
f(, = Paf(XX
(5)
Any function that is set-equivariant is also multiset-equivariant (since we can set Pa = Pt), but the converse is not true.
The problem with set-equivariance. We argue that set-equivariance is too strong of property for set prediction. As we have shown in Proposition 1, it prevents equal elements from being processed differently (we demonstrate this experimentally in Section 4.1). More importantly, when the model is
3

---

Published as a conference paper at ICLR 2022
continuous', impossibility separate equal elements means difficulty in separating similar elements. For example, in order to learn the pushaapart function, a model has 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 the discontinuity (where elements are similar) the modeling error will likely be high. A function like push_apart can for example be relevant in turning a multiset of random 3d points that are close together (i.e., similar) into 3d shape with points that are farther away (Luo & Hu, 2021). Difficulty with separating similar elements can therefore lead to bad downstream performance. Similar elements could for example arise from a random initialization of Yo, nonlinearities throughout the model zeroing out (ReLU) or squashing elements (sigmoid and tanh), the model naturally tending towards equal elements (perhaps from oversmoothing as often seen graph neural networks (Li et al., 2018)), or simply from being part of the dataset used.
Summary. It is important to make the distinction between set-equivariance as has traditionally been considered in the literature and multiset-equivarianne Deep learning typically operates multisets rather than sets, so the less restrictive be the desirable property. Most existing set prediction models are continuous and set-equivariant, which limits their modeling capacity even if exact duplicates are never encountered. We refer to models that are multiset-qquvvarian but not set-equivariant as exclusively multiser-equivarian..
sellliistt-quivaiiannee without set- est-eevvvriiace should
2.3 DEEP SET PREDICTION NETWORKS
Here, we highlight the only set prediction model we aware of that can be exclusively multiset- equivariant: 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-invaiinl encoder (multiset-to- vector) to iteratively update a multiset Y. The objective is to find a Y that minimizes the difference between its vector encoding and the input vector (this vector could for example come from input image being encoded). This is expressed as the following optimization problem:
L(Y, 2, 0) = 11 0) =1|
(6)
DSPN 2)) = arg min L(Y.
(7)
Y
Because y is permutation-invarian the forward pass of the model, the arg min is approximated by running a fixed number of gradient descent steps. In the backward pass, the goal is differentiate Equation with respect to the input vector z and the parameters f 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.
any ordering for the elements in Y has the same value for L. In
Equivariance of DSPN. We now discuss the particular form of equivariance that DSPN takes, -iniaiariiiinniiaannt which has not been done before the context of multisets. The gradient of the permutation invariant encoder g with respect to the set input Y is always multiset-cquivariann but depending on the encoder, is not necessarily set-equivariant. Zhang et al. (2019) find that FSPool- encoders (Zhang et al., 2020) achieved by far the best results among the encoders they have tried. With FSPool, DSPN becomes exclusively santiiset-euvvaiiant its initialization Yo- This 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 -quuiiritteeiivaiiann,t,, but 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 specifically the use of the exclusively multiset-equivarinn gradient of sorting that makes DSPN exclusively multiset-equivariant
3 IMPLICIT DEEP SET PREDICTION NETWORKS
Despite being exclusively multiset-equivarian DSPN is outperformed by the set-equivariant Slot Attention (Locatello al., 2020). We begin by highlighting some issues in DSPN that might be the cause of this, then propose approximate implicit differentiation
solution.
'There is subtle difference between continuity in the usual sense and continuity multisets. While the discussion here the usual notion, elaborate the consequences of multiset-continutty in Appendix B.
4

---

Published as a conference paper at ICLR 2022
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 significantly more memory and compu- tation 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 negative effect on the modeling capacity due to insufficient minimization of Equation 7. We aim to address these problems in the following.
Implicit differentiation. Implicit differentiation is technique that allows efficient 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 differentiatiing 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' 0) is implicitly defined by the input vector z and the encoder parameters 0 through the fact that the gradient of L at an optimum is O:
Vy LYY 2,
The implicit function theorem allows us to differentiate Equation and obtain the Jacobians (z" 0) and (z, 00, which is all we need in order for iDSPN to fit into autodiff framework. Appendix D contains the full details on how this works.
In the implicit differentiation, @(Vy /) evaluated at Y (also explained in Appendix D). This can be problematic when H is ill-conditioned or singular, so we regularize it with the identity matrix through H' H H II with > 0 (Martens, 2010; Rajeswaran et al., 2019). In the limit of Y + 20, 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 the following Jacobians, which are equivalent to backpropagating through a single gradient descent step at the optimum:
need solve linear system involving the Hessian H
aY 8Y*(2.0 0) a v (Y", 2, 0))
(Y
a (V L 2, O))
dz
Oz
00
a0
Other motivations for this approximation are discussed in Fung et al. (2022), where they call this Jacobian-free backpropagation. Matching their results, we find that this is not only faster to run and easier to implement (no need solve a linear system involving H), but also leads better results than standard implicit differentiation in preliminary experiments. Hence, we apply this approximation to differentiate Equation with respect to = and 0 in all of our experiments.
Benefits. Implicit differentiation saves a significant 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 only take the derivative 27 reducing it by nearly a half. The combined improvements mean that we can minimize the objective in Equation better by increasing the number of optimization steps and using more sophisticated optimizers in the inner loop. The optimization procedure is treated black box now, so it does not even need be differentiable.
memory cost since updates to the set can be performed in- in-place. Because at the final set, we decrease the computational cost by a factor of around T-1
Further improvements. A simple improvement enabled by implicit differentiation is the use of Nesterov's 's Accelerated Gradient for estimating Y*. This has better convergence rate on convex problems than standard gradient descent ( (Nesterov, 1983) and has found wide applicability in deep learning (Sutskever al., 2013). Since the computation graph of the optimization is no longer kept in memory, it only has minor impact memory and computation: the momentum buffer is the same size as Y. updates can be performed in- 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, sample random initial set Yo (O I/10) to start the optimization with. Similar to DSPN, we can also use a learned initial set Y, to start closer to a solution. However, implicit differentiation treats the optimizer of Equation 7 as black box, so there is no gradient signal for Yo-
5

---

Published as a conference paper at ICLR 2022
We therefore need to include regularizer in Equation 6 to give gradient for Yo, for example by adding the regularizer from Rajeswaran et al. (2019): 2 Yo 12. We use this only in Section 4.2 to make the forward passes of iDSPN and DSPN the same that
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 enforce that every vector is in the probability simplex. As suggested by Blondel et al. (2021), 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 D):
class label, we would like to
proj(Y V 2, 0)) Y
(10)
proj is 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 satisfied. 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 V is set-equivariant. By using an exclusively multiset-equivariant Vy we avoid this problem because equal elements after projection can be separated again.
When we already know parts the desired Y (e.g. specific elements the set or the values along certain dimensions), we can help the model by keeping those parts fixed and not optimizing them in Equation 7. For example, in Section 4.1 we know that the first few dimensions in the output multiset should be the exactly same in the input multiset. If we fix 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) make iDSPN exclusively multiset-equivariant the input multiset, not just to the initialization Yo.
just
4 EXPERIMENTS
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). Appendix contains additional experiments on these datasets, including further evaluation of exclusive multiset-equivariane
We provide detailed descriptions of the experimental procedure in Appendix F, show example inputs and outputs Appendix G. and open-source the code to reproduce all experiments at https: // h////ittan github..b.... com/ dww.haaaauaaaaa momuuiiseett equivariiiiiieeee vani and in the supplementary material.
4.1 DIFFERENCE BETWERN 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-equivariine the aim 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 the input are mapped different elements in the output, so this is a task where being able to separate equal elements is directly necessary.
at all. Given a multiset like [a, a. b b, b],
Setup. Each element in the input set of size 64 generated by randomly picking a class out of choices 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 results. For the target output, a one-hot incremental to each element. A multiset matching loss is computed for all models (see details in Appendix F). We evaluate the accuracy of predicting every element in the multiset correctly.
that separate for each class is concatenated
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 applics even to transformers with random "position" embeddings, which similarly to TSPN (Kosiorek et al., 2020) and Slot Attention (Locatello et al., 2020) noise make elements

---

Published as a conference paper at ICLR 2022
Table 2: Sample efficiency of equivariance properties on task that requires separating equal elements. Accuracy in % (mean + standard deviation) over 6 seeds, higher is better. PE refers to positional encoding, random PE refers randomly sampled elements like in TSPN (Kosiorek et al., 2020).
Set- equivariant
Multiset- equivariant
Training samples
Model
10x
100x
DeepSets
0.0ga
0.0
0.0
Transformer without PE
V
0.0
0.0+a
0.0#o
Transformer with random PE Transformer with PE
X
X
28 28.5
22.6+s.2
22.7+7.7
X
X
0.0
46.4220.
94.421.8
BiLSTM
X
X
0.0go
25.1
97.9
DSPN iDSNN (ours)
X V
89.4+2.4
88.9+4.9
90.8+4.2
X
/ 96.8+o.a
98.1
97.9
not exactly equal. Meanwhile, BiLSTM and transformers with normal positional encoding require at least 100x more training samples to come close 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-efficient Appendix E.I shows the effect of data augmentation for the non-equivariant models.
4.2 AUTOMATIC VERSUS IMPLICIT DIFFERENTIATIO ON RANDOM MULTISETS
Next, we aim to evaluate our contributions in Section 3. To do this, we compare: 1) baseline DSPN, 2) iDSNN without momentum and 3) iDSPN with Nesterov's Accelerated Gradient. Comparing and the difference between backpropagation through the optimizer and the approximate implicit differentiation (Equation 9) when using the same forward pass, while comparing and allows
2 allows
the effects of using better optimizer for the inner optimization. In Appendix E.2 instead evaluate the difference between (exclusively multiset-equivarian)) and a modification to 1 where the FSPool is replaced with sum mean pooling (set-equivariant).
from.N(0, We propose the task of autoencoding random multisets: every element is sampled i.i.d. from (o I). The goal is reconstruct the input with a Dermutation-invariant latent vector as bottleneck. Varying the set size n and dimensionality of set elements d allows us to control the difficulty of the task. Note that while this appears like a toy task, is also likely harder than many real-world datasets with similar 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 solve this task exactly, which is especially difficult when want the latent space to be permutation-invarian This task also challenging because indirectly requires a version of ppsshapprrt similar elements in the initialization Yo might need to be moved far away from each other in the output.
memorize every element
Setup. All models use the same encoder, so DSPN and iDSPN without momentum have exactly the same forward pass. We train each model every combination of number of set elements n E {2, 4, 8, 16, 32). set element dimensionality d E (2, 4, 8, 16, 32), and optimization steps e (10, 20, 40). We evaluate the final validation reconstruction losses over 8 random seeds.
Results. In Figure I, we show selection of representative results. We show the full results in Appendix E.2 due 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
77 8, d
8, 8
T
16, d
loss
Baseline
0.02
0.02
0.2
iDSPN
iDSPN+mam
0.00
0.00
0.0
10
20
40
1x
2x
10
20
40
iterations
compute
ilerations
Figure 1: Implicit versus automatic differentiation for different dataset difficulties, reconstruction losses, lower is better. LDSPN+mon stands for iDSPN with Nesterov momentum. (left) and (right) compare models at equal iterations, (middle) compares models at roughly the same amount of computation. We select (right) as an atypical example where DSPN outperforms iDSPN.
7

---

Published as a conference paper at ICLR 2022
iDSPN momentum is 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 roughly the same amount of computation (Figure 1, middle), it is clear that iDSPN has significant practical benefits over DSPN regardless of momentum. These results apply to most dataset configurations (see Appendix F.2).
The only settings where iDSPN + momentum performed worse than DSPN is when n x d approaches the dimensionality of the latent vector, 512 (Figure 1, right). In these cases, the encoder runs into fundamental representation limits of set functions (Wagstaff et al., 2019). With performance being limited by the encoder, is natural that better optimization yields little no benefit. In Appendix G, it becomes clear that no model performs acceptubly on these dataset configurations. When comparing the exclusively multiset-equivariant iDSPN to the set-equivariant iDSPNs Appendix E.2, the former is consistently better, often with a lower loss by a factor of >> 10.
4.3 CLEVR OBJECT PROPERTY PREDICTION
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 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, 19th dimension is added as binary indicator for whether an element is present or not. Appendix E.3 compares an exclusively multist--equivarian iDSPN to set-equivariant iDSPNs and contains results for DSPN in like-to-like setting to iDSPN.
Setup. We mostly follow the experimental setup of DSPN where a ResNet (He et al., 2016) encodes the image into the input vector 2, which is decoded into set by iDSPN. We also include a few improvements that we elaborate on in Appendix F. Most notably, we use Nesterov's 's Accelerated Gradient as iDSPN optimizer and increase the number of optimization steps. The evaluation metric is a modified average precision (APrhrehooa)) where a predicted object is considered correct if every property correct and the 3d coordinates are within the given Euclidean distance threshold of the ground- truth element.
Results. Table 3 shows our results. iDSPN significantly improves over any configuration of Slot Attention (see caption for description on the configurations and f) on all AP metrics. It is better attribute classification when ignoring the 3d coordinates (AP oo. 96.4% + 98.8%) and improves especially for the metrics with stricter 3d coordinate thresholds (APo.125. 7.9% - 76.9%). This is despite Slot Attentionf using a three times higher weight on the loss for the coordinates than iDSPN. Note that stricter AP threshold is always upper bounded by looser AP threshold, so iDSPN is guaranteed to be better than Slot Attention on APo.0625. Again, we see in Appendix E.3 that the exclusively
iDSPN is better than the set-equivariant iDSPNs. Interestingly, iDSPN performs better than DSPN in a like-to-like setting (Appendix E.3) with the same number of iterations and same optimizer. This motivates thinking of iDSPN not just an approximation of DSPN, but as a separate model altogether.
Table 3: Performance on CLEVR object property set prediction task, average precision (AP) in % (mean + standard deviation) over 8 random seeds for our results, higher is better. Slot MLP and Slot Attention results are from Locatello et al. (2020). DSPN results are from Zhang et al. (2019). DSPN is trained at 10 iterations and evaluated at 30 iterations. For Slot Attention and t, Locatello et al. (2020) supervise the model at each iteration and scale up the loss on the coordinates by either or 6. We do not mark the iDSPN 256x image results in bold because of the difference in dataset setup.
Model
Image size
AP
AP
APas
APo25
APo.25
APo0625
Slot MI.P
128x
19.811.
1.4.aa
0.31oa
0.01oo
0.0ga.
DSPN
128x128
85.214
81.1252
47.4117a
10.81a9
0.62m7
Slot Attention
128x128
94.3111
86.721
56.01s.
10.8
0.92na
Slot *ttention*
128x128
96.4+o.a
93.6-o.s
80.4+2.1
26.5+2.7
2.6+0.s
Slot Attentiont
128x128
90.6+1a
89.1+2.1
84.4+2.s
50.3+s.
7.9+1.
iDSPN
128x128
98. 98.840.a
98.5
98.2+t..
95.820.r
76.9
32.3
iDSPN
256x256
19.4
99.220.4
99.020.5
97.8
86.921.
47.22a.
8

---

Published as a conference paper at ICLR 2022
Efficiency. 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.
5 RELATED 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). 1)). Since sets closely related to graphs, there are also natural extensions to molecule generation (Simonovsky & Komodakis, 2018) 8) and scene graph prediction (Krishna et al.. 2017: Xu et al., 2017). As we have covered in this paper, most existing set predictors
either not multiset- muluiset-equivariann (Vinyals
et al., 2016; Stewart & Andriluka, 2016; Vaswani et al., 2017: Achlioptas et al., 2018; I.i et al., 2018; Rezatofighi al., 2020) or set-equivariant (Lee al., 2019; Locatello et al., 2020; Kosiorek et al., 2020). The only set prediction model we are aware of that can be exclusively multiset-equivarian DSPN (Zhang et al., 2019; Huang et al., 2020). Note that Zhang al. (2021a) also make use of the exclusively multiset-equivarian
Jacobian of sorting, but instead focus on learning multimodal densities over sets, which is tangential to our work. We provide further background in Appendix C.
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 al., 2020) copies the same latent vector 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 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. the transformer with random position embeddings (similar application of noise as in TSPN) does not perform well.
6 CONCLUSION AND FUTURE WORK
We have established a more appropriate property for set prediction (exclusive nultiset--qqivaaiannce)) and that DSPN satisfies 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 significantly. 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 into a set. The benefit 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 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 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). We discuss the notion of multiset-continuity in Appendix B as a first step towards better understanding of these discontinuities. However, plenty of questions remain, for example: How can we obtain universal approximation of functions that are exclusively multiset-equivariant and multiset-continuuue How does multiset-continuity without continuity affect the training dynamics?

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
Proofs and derivations are contained in Appendix A and Appendix D. Detailed descriptions of experiments, including how to generate these datasets the case for Section 4.1 and Section 4.2, are given in Appendix F. We provide all the code, scripts to reproduce the experiments with the exact hyperparameters, and scripts to calculate the evaluation results at https: /qgttub........ com/ davzha / /mmlllleeeeeelmmmmmmmem t set equui vari ance. In the code repository, we make available every training run and the corresponding pre-trained weights used to produce the results in this paper through public Weights & Biases tables (Biewald, 2020).
ETHICS STATEMENT
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 3d space.
We propose improvements
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).
ACKNOWLEDGMENTS
We thank the reviewers for their reviews that let us improve several aspects of the paper. We especially thank reviewer 4sjy for pointing out issue with a definition that led us to think about it deeper and create Appendix B. The work of DWZ is part of the research programme Perspectief EDL with project number P16-25 project 3, which financed by the Dutch Research Council (NWO) domain Applied and Engincering Sciences (TTW). This research was partially supported by Calcul Québec (www.. ca 1 cul quebec. ca), Compute Canada (www. computecanate ca) and the Canada CIFAR AI Chair Program. Simon Lacoste-Julien is CIFAR Associate Fellow in the Learning Machines & Brains program.
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https: //t////////////////// htt///tttenstttee wwnsorrerrrroro ow org/. Software available from tensorflow.org.
Pierre Ablin, Gabriel Peyré, and Thomas Moreau. Super-efficiency of automatic differentiation for functions defined as minimum. In International Conference on Machine Learning (ICML), 2020.
Panos Achlioptas, Olga Diamanti, Ioannis Mitliagkas, and Leonidas J Guibas. Learning representa- tions and generative models for 3D point clouds. In International Conference on Machine Learning (ICML), 2018
10

---

Published as a conference paper at ICLR 2022
A. Agrawal, B. Amos, S. Barratt, S. Boyd. S. Diamond, and Z. Kolter. Differentiable convex optimization layers. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization as layer in neural networks. In International Conference on Machine Learning (ICML), 2017.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances Neural Information Processing Systems (NeurIPS), 2019.
Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Multiscale deep equilibrium models. In Advances in
Neural Information Processing Systems (NeurIPS), 2020.
David Belanger, Bishan Yang, and Andrew. McCallum. End-to-end learning for structured prediction energy networks. In International Conference on Machine Learning (ICML), 2017.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https: // /twww/wwww./ wandb. com/. Software available from wandb.com.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Phine Phillppi Vert. Efficient and modular implicit differentiation. arXiv preprint arXiv:2105. 2021.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision (ECCV), pp. 213-229, 2020.
Alexandre Carlier, Martin Danelljan, Alexandre Alahi, and Radu Timofte. Deepsvg: A hierarchical generative network for vector graphics animation. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Terrance DeVriss, Ishan Misra, Changhan Wang, and Laurens van der Maaten. Does object recogni- tion work for everyone. Computer Vision and Pottern Recognition (CVPR), 2019.
Laurent El Ghaoui, Fangda Gu, Bertrand Travacca, Armin Askari, and Alicia Tsai. Implicit deep learning. SIAM Journal Mathematics of Data Science, 3(3):930-958, 2021.
Yusuke Fujita, Naoyuki Kanda, Shota Horiguchi, Kenji Nagamatsu, and Shinji Watanabe. End-to-End Neural Speaker Diarization with Permutation-free Objectives. In Interspeech, 2019.
Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, and Wotao Yin. JFB: Jacobian-free backpropagation for implicit networks. In AAAI Conference on Artificial Intelligence, 2022.
Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint 1 2016.
Klaus Greff, Raphaël Lopez Kaufmann, Rishab Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In International Conference on Machine Learning (ICML), 2019.
Fangda Gu, Heng Chang, Wenwu Zhu, Somayeh Sojoudi, and Laurent El Ghaoui. Implicit graph neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Computer Vision and Pattern Recognition (CVPR), 2016.
Qian Huang, Horace He, Abhay Singh, Yan Zhang, Ser-Nam Lim, and Austin Benson. Better set representations for relational reasoning. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Peter J. Huber. Robust Estimation of a Location Parameter. The Annals of Mathematical Statistics, 35(1):33-10, 1964.
11

---

Published as a conference paper at ICLR 2022
Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Li Fei-Fei, C Lawrence Zitnick, and Ross Girshick. Clevr: A diagnostic dataset for compositional language and elementary visual reasoning. In Computer Vision and Pattern Recognition (CVPR), 2017.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Leurning Representations (ICLR), 2015.
Brendan F. Klare, Mark J. Burge, Joshua C. Klontz, Richard W. Vorder Bruegge, and Anil K. Jain. Face recognition performance: Role of demographic information. IEEE Transactions on Information Forensics and Security, 7(6), 2012.
Adam R Kosiorek, Hyunjik Kim, and Danilo J Rezende. Conditional set generation with transformers. Workshop Object-Oriented Learning ICML, 2020.
Steven G Krantz and Harold R Parks. The implicit function theorem: history, theory, and applications.
Springer Science & Business Media, 2012.
Ranjay Krishna, Yuke Zhu, Oliver Groth, Justin Johnson, Kenji Hata, Joshua Kravitz, Stephanie Chen, Yannis Kalantidis, Li-Jia Li, David A Shamma, Michael Bernstein, and Li Fei-Fei. Visual genome: Connecting language and vision using crowdsourced dense image annotations. International Journal of Computer Vision, 123(1):32-73, 2017.
Harold W Kuhn. The Hungarian method for the assignment problem. Naval Research Logistics Quarterly, 2(1-2):83-97, 1955.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: A framework for attention based permutation iitiiaatiioonnnniinn invariant neural networks. In International Conference Machine Learning (ICML), 2019.
Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper Insights into Graph Convolutional Networks for Semi-Supervised Learning. In AAAI Conference Artificial Intelligence, 2018.
Yujia Li, Oriol Vinyals, Chris Dyer, Razvan Pascanu, and Peter Battaglia. Learning deep generative models of graphs. In International Conference on Machine Learning (ICML), 2018.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Aaron Lou, Isay Katsman, Qingxuan Jiang, Serge Belongie, Ser-Nam Lim, and Christopher De Sa. Differentiating through the fréchet mean. In International Conference on Machine Learning (ICML), 2020.
Shitong Luo and Wei Hu. Diffusion probabilistic models for 3d point cloud generation. In Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
James Martens. Deep learning via hessian-free optimization. In International Conference on Machine Learning (ICMI.,, 2010.
Ninareh Mehrabi, Fred Morstatter, Nripsuta Saxena, Kristina Lerman, and Aram Galstyan. A survey on bias and fairness in machine learning. ACM Comput. Surv,, 54(6), 2021.
Yurii E. Nesterov. method for solving the convex programming problem with convergence rate of (1/kk)) Dold. Akad. Nauk SSSR, 269:543-5411 1983.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, highhpperfomance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Aravind Rajeswaran, Chelsea Finn. Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In Advances in Neural Information Processing Systems (NeurIPS). 2019.
12

---

Published as a conference paper at ICLR 2022
Hamid Rezatofighi, Roman Kaskman, Farbod T Motlagh, Qinfeng Shi, Anton Milan, Daniel Cremers, Laura Leal- Taixé, and lan Reid. Learn to predict sets using feed-forward neural networks. arXiv preprint ar arXiv:2001.11845, 2020.
Joseph P. Robinson, Gennady Livitz, Yann Henon, Can Qin, Yun Fu, and Samson Timoner. Face recognition: Too bias, or not too bias? In Computer Vision and Pattern Recognition (CVPR), 2020.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Tim Lillicrap. A simple neural network module for relational reasoning. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Martin Simonovsky and Nikos Komodakis. GraphVAE: Towards generation of small graphs using variational autoencoders. International Conference on Artificial Neural Networks (ICANN),
2018.
Russell Stewart and Mykhaylo Andriluka. End-to-end people detection in crowded scenes. In Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum deep learning. International Conference on Machine Learning (ICML., 2013.
Antonio Torralba and Alexei A. Efros. Unbiased look at dataset bias. In Computer Vision and Pattern Recognition (CVPR), 2011.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Oriol Vinyals, Samy Bengio, and Manjunath Kudlur. Order matters: Sequence to sequence for sets. In International Conference on Learning Representations (ICLR), 2016.
Edward Wagstaff, Fabian B. Fuchs, Martin Engelcke, Ingmar Posner, and Michael A. Osborne. On the limitations of representing functions on sets. In International Conference on Machine Learning (ICML), 2019.
Danfei Xu, Yuke Zhu, Christopher Choy, and Li Fei-Fei. Scene graph generation by iterative message passing. In Computer Vision and Pattern Recognition (CVPR), 2017.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations (ICLR), 2018
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
David W Zhang, Gertjan J Burghouts, and Cees GM Snoek. Set prediction without imposing structure as conditional density estimation. International Conference on Learning Representations (ICLR), 2021a.
Miao Zhang, Steven W. Su, Shirui Pan, Xiaojun Chang, Ehsan M Abbasnejad, and Reza Haffari. idarts: Differentiable architecture search with stochastic implicit gradients. International Conference Machine Learning (ICML), 2021b.
Yan Zhang, Jonathon Hare, and Adam Bennet. Deep Set Prediction Networks. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Yan Zhang, Jonathon Hare, and Adam Prügel-Bennee FSPool: Learning set representations with featurewise sort pooling. In International Conference on Learning Representations (ICLR), 2020.
13

---

Published as a conference paper at ICLR 2022
A THE JACOBIAN OF SORTING IS EXCLUSIVELY MULTISST-EQQIVARIINN
First, let us clarify what we mean by the Jacobian of sorting. We use sort: R - R to denote sorting of a list (representing a multiset) with scalar elements in ascending order. To be precise, we define sort(X) = Sx X where Sx is the permutation matrix that sorting X would be implemented by any sorting algorithm. When there multiple
apply. sort in this possible permutations that would lead to X being sorted the correct order (i.e. when there are duplicate elements X), Sx simply refers to the single permutation chosen by the algorithm. The details of the sorting algorithm are not relevant for the further discussion Sx can be considered an arbitrary choice within the possible permutations that would lead to a valid sorting.
With Jacobian of sorting, we mean:
Ösort 0x
0Sx 0X
SX
(11)
In the last step, we treat Sx as independent X because 0Sx/OX almost everywhere and undefined zero measure set. The definition 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 have different Sx for each dimension.
A.1 PROOF: NOT SET-EQUIVARIANT
To show that the Jacobian of numerical sorting is not set-equivariant, a simple counterexample is realize that the Jacobian does change between sort ((d 9 [0, O]) and sort (l2 [0,00] ))
The input does not change regardless of how we permute [0, 0] therefore the Jacobian of sorting does not change.
To make this point more concrete, we implement pwsh_apart from the main text using the Jacobian of sorting. We omit the transposes in the following for brevity. We begin by defining the function g((a, b)] = sort
(tt(((c, [a, b)) [-1, 1]. This multiplies the smaller value with -1 and the larger with I (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 ) so the -1 gets propagated to the smaller element and to the larger one. In other words, V Vaa,6]9 i-1, 1] if S b (alternatively
b), else [1, 1]. Now we can define push_apart ((a, b]) = [a, b] + V [a,b]9 to obtain exactly the same function as in the main text.
A.2 PROOF: MULTISET-EQQUVARIAN
We now show that the Jacobian of sorting is multiset-cquivariant Proposition 2. The Jacobian of sorting fXX = @ssor(X)/0XX = S is multiset-equivarian.
Proof. To show that for all Pt, there exists a P3 that satisfies the two conditions of multiset-
cquivariance, we set P2 = Sp, Sx. We use the fact that S is a permutation matrix so S s =
SS I. To show that the first condition is true:
Paf(X)= = PrS
(12)
Sx
(13)
=Sp,x
(14)
=(PiX)
(15)
We now use the fact that sorting is permuration-innariant so for any P,Sx = Spx PX. To show that the second condition true:
P&X = -x
(16)
Sp P,X
(17)
P.X
(18)
We showed the existence of P2 that fulfills both conditions for all P1, so is multiset-equiivaiann
14

---

Published as a conference paper at ICLR 2022
Note that in the definition of multiset-equivarianee it is important that the existential quantifier is on P2 acting on the output of f rather than P1 acting on the input. If the quantifiers were the other way around (A (PPP)),, then the multiset-equivarianne definition would not be useful: we have seen that permuting equal input elements does not affect the Jacobian of sorting, so there must be some P2 for which choice of P1 works.
B MULTISET-CONTINUEE
AND MULTIISET-RQUIVARINCC
In the main text. our discussion on the problems with separating similar elements concerns functions that are continuous and set-equivariant. In this section, we will make several aspects of this statement more precise. For example, one may ask whether there are functions that cannot separate equal elements, but can separate similar elements. This is an important point to address since exactly equal elements are rare in practice, while the behavior for similar elements is much more relevant. To talk about this properly, we first need to introduce the concept of multiset-continiity (Appendix B. where use the Hungarian loss metric multisets (Appendix B.2). Then, we discuss the various combinations of continuity with equivariance (A (Appendix B.3), and specifically that multiset- continuity combined with exclusive multiset-equivariaaee allows to guarantee that not just equal elements can be separated, but also similar elements.
B.1
MULTISET-COM CENTTISET-CNNINUITYY
While used the usual notion of continuity in the main text, it more appropriate to talk about con- tinuity on multisets
when working with multiset-to-multiset functions. Multiset- continuity is desirable for model when working with multisets for the same reasons that continuity is desirable when working with scalars or lists: it lets the model generalize from seen datapoints to nearby unseen datapoints predictably. The differing notion of continuity comes from the fact that order does not matter in multisets. This has the consequence that there are functions that are discontinuous when considered from an ordered list perspective, but they are continuous when con- sidered from an unordered multiset perspective. The list representation is an implementation detail; the property we care about when working with multisets (regardless of how they are represented) is continuity multisets.
Take the pusha example from Section 1. With e > O, push (0, 0]) = [-1, but push.apart ([c, 0)) = [1 + e, 1]. The list representation is discontinuous (a perturbation by changes [-1,1] to + c, -]] an L, distance of 4 + e) but when considering it as a multiset, appears continuous: [-1, 1] is not much different from [-1, 1 c). the L, distance is only since can change the "order" within the orderless multisets.
So. would like different notion of continuity multisets that takes the structure of multisets into account. We can do this by using the standard e-5 definition of continuity on a metric space, where our metric measures distances between multisets. The multiset metric will focus on the Hungarian loss.
Definition 3 (Permutation-invaritt metric), A metric D: Raxd x Raxd + R is called permutation- invariant iff for any X, Y E Ra and any permutation matrices Pı: Pa:
D(X, Y) D(PiX, PYY) Definition 4 (Muntiset-continntty... Let (Raxdd, D1) and (Raaxds Mixdd,DD) D2) be two metric spaces with permutation-inyarian metrics D1, D2. A function f: Raxd, - R is multiset-contnnuuus X I nxd iff for any 6 0, there exists (X), < The function f is multiset-continous iff it is multiset-continuuus for all X E gaxdi and we call ff multiset-disconinuuuss otherwise.
(19)
> such that D, Y) < implies
While multiple choices of cermutation-invarian metrics are possible, we will use the Hungarian loss in the following. With this definition in place, the Jacobian of sorting is another example of function that is not continuous but is caltiset-continuous. It is easy to see that it is not continuous: it can only take the value of discrete permutation matrices.
Proposition 3. The Jacobian of sorting f(X) = Ösort //x- = S is multiset-comtinuuuss
Proof. All permutation matrices S are equivalent up to permutation, hence for any permutation invariant metrie (((()) for all X, Y Rxd This is always less than > 0. n
15

---

Published as a conference paper at ICLR 2022
These definitions along with multiset- -quiiiistt-qvvviiannee give the necessary tools talk about the property being able to separate similar elements, the case most relevant in practice.
B.2 HUNGARIAN LOSS IS A VALID METRIC
For the sake of completeness, we show that the Hungarian loss combined with an elementwise metric Dy is a proper metric over multisets. The most relevant example for is the Euclidean distance over the vector space V = _R i.c., the elements of the multiset are d-dimensional vectors.
Proposition 4. The Dy-Hungarian loss for two multisets X, Y and a metric Dv: V x over the elements of the multiset, metric:
R
C( Y)
min xEll
Σ
(a (x()))
(20)
Proof. We show that L satisfies the three axioms: Identity of indiscernibles.
C(X, Y)
(21)
min Tlll
Σ
Dv(ww (2i, Vn(())
(22)
mm E 11: Vi: Dy (a, Vn(1))
(23)
3m 1I: Vi: T; yz(()
(24)
X =msel Y
(=mset denotes equality up to permutation)
(25)
2. Symmetry.
((X,Y) = min
Σ Dv Mr(1))
(26)
zell
min
Dy (y
(27)
xEIl
X)
(28)
Triangle inequality.
Let wxz arg min E. Dy (m,,, zx(1)) and ay arg arg min E, Dv These denote the optimal assignments from X to Z. and from Z to Y respectively.
C(X, Z) + L(
(29)
min TEll
Σ
Dy (wi, xx(i)
min Σ TEll
Dv Mr(i)
(30)
Σ Dv 2mze(t) (1)) + (zi,
Tazs Tay are the minimizers
(31)
Σ DV Znn (1) + Σ Dv Dv(zzaa Y (Teo
sum is commutative
(32)
Σ Dv 2mm)()) + Dv (Znzzlti) (re)(Ser))))
(33)
Σ (a Un (ray(
triangle inequality of Dy
(34)
min TEll
Dv yn())
(35)
C(X,Y)
(36)
n
16

---

Published as a conference paper at ICLR 2022
B.3 COMBINING MULTISET-CONTINUTTI WITH EQUIVARIANCE
We now have four combinations of the properties of interest:
1. multiset-continuous and set-equivariant,
2. multiset-continuuus and exclusively multiset-cquiivriiat,
3. muliiset-iiscontinuuus and set-equivariant,
4. multiset-discontinuuss and exclusively multiset-cquivvarint.
In the following, show why the second combination of properties (multiset-continuous and exclusively mulliset-qquivarian)) is most desirable.
B.3.1 MULTISET-CONTINNOO AND BET-EQUIVARIAN
Most existing set-equivariant functions have this combination of properties. This because they continuous in the usual sense by being composition of continuous functions, which gives us multiset-continitt.
Proposition 5. If function Ruxdd nxx continuous in metric space induced by some matrix norm, then it is also multiset-continuuus.
Proof. All matrix norms are equivalent in finite dimensions, thus any function that is continuous in a matrix-norm- nattix-norm-induced metric space is also continuous in all such metric spaces. In particular, f is also continuous for the metric Disa( [, X) EV ||2, yall- I.. We can show that Dia is an upper bound on the permutation invariant metric D mset.
iomutaatoo-invariiatt
Dus (X,Y)
Σ ||=; yill
(37)
2 min nell
Σ
[ y=(1)
(38)
L(X,Y)
(39)
Since Dús is an upper bound, Dis (S(X,, ((Y)) < everywhere (e-ō condition of continuity). then (( < everywhere well.
O
We know that set-equivariant models cannot separate equal elements, and by multiset-continuity. we know that a small perturbation equal elements will only result in a small change in the multiset output. Therefore, it will be difficult to separate similar elements.
B.3.2 MULTISET-CONTINCOUR AND EXCLUSIVELY MULTISET-EQUIVARIAN
We have already shown that the Jacobian of sorting satisfies both multiset-continuigg and exclusive multiset-cquivarianee This combination of properties is really the key for separating similar elements, as we will now show.
Proposition Let f: Ruxd Baxd denote a eguivariance is sufficient for there to exist a multiset containing equal elements that f can separate.
function. Exclusive multiset-
Proof. By definition, an exclusively multiset-equivarian function means that multiset-equivariant and not set-equivariant. Multiset-qquivarianee means that for all X and permutation matrices Pt. there exists permutation matrix P2 such that:
P ( =
(40)
(41)
No aet-equivariance means that there exists an X such that f(Pxtt) + Psf(Xt) for all permuta- tion matrices Py. Combining this with Equation 41, know that for Xt, have P1 * Pz. This implies that X must contain equals to satisfy Equation 40 and that these equals become separated since ((P x Paf(Xi) # (XX().). Therefore, f separates the equals in Xt.
17

---

Published as a conference paper at ICLR 2022
For all exclusively multiset-equivariann functions, there exists multiset with equal elements that can be separated. Combined with multiset-continuity. we know that similar elements around those equal elements can also be separated because small perturbation to the equal elements (obtaining similar elements) only results in a small change in the separated output elements, so they remain separated. This highlights the appeal of multisef-continuuus exclusive multiset-equivariant functions: all functions with this property can separate both equals and similar elements.
Note that while exclusive multiset-equivariance only requires there exist set containing equal elements that can be separated, functions like the Jacobian of sorting can actually separate all equals, and therefore separate all similar elements as well. This follows from the fact that each row in the permutation matrix S is always unique and a constant vector distance away from the other rows, regardless of how similar the input elements become.
B.3.3 MULTISET-DISCONTING
AND SET-EQUIVARIANT
The common theme among the combinations with multiset-discontinniyy is that we lose the guarantees we have on the behaviour for nearby points. For example, it is indeed possible to construct function that cannot separate equals, but can separate arbitrarily similar elements. Take a sort function and augment it so that for equal elements, instead of resolving the tie in the permutation arbitrarily, we equally distribute the equal elements. For example, sort2 ([1. 1]) 10.5 0.5 8:31 1] [L. 1]. The matrix we apply to the input is no longer permutation matrix, but doubly-stochastic matrix. Essentially, this behaves like a mean for equal elements and like a normal sorting otherwise. The Jacobian of this operation is set-equivariant since the mean prevents equals from being separated. For arbitrarily similar elements it behaves exactly like sorting, so it can separate similar elements but is
While there exist such functions that behave like a continuous exclusively multiset-equivariant function in practice (since exact equals are unlikely to be encountered), cannot obtain any guarantees on separating similar elements in general. The following function is a counterexample to the combination of multiset-discontinuie and set-equivariance being sufficient for separating similar elements. Consider the identity function s(X) = X that is augmented so that if all the entries are unique integers, it returns X + 1 instead. For example, f[I,,) 2]) [2, 3] but f([1,1) = [1. 1] because slightly perturbing [1. 2] as way that cannot be compensated by the matching, and
and f([1, 2 + <l) [1, 2 + e]. This is multiset-discontiuue
input makes large change the output set-equivariant because equal elements are mapped to equal outputs. Applying f on similar elements always just returns identity, so they can never be separated.
B.3.4 MULTISET-DISCONTINNGO AND EXCLUSIVELY MULTISET-EQUIVARRIA
Consider the identity function f(X) X that is augmented to separate duplicates in its input by concatenating a unique count to each duplicate element. For example, f1,1,1,2,2
2, 2 <)) [L1, (L, 1), (2, 0), (2 + c, 0)]. This function able to separate any equals in its input through the unique count, but it is multiset-discontinuur change in the output. In particular, it cannot separate arbitrarily similar elements that are not equal. This shows that exclusive multiset--quivariannce alone is not enough for separating similar elements. Multiset-continuity is needed to guarantee that separation of a multiset with equal elements implies separation of that multiset with similar elements as well.
because perturbing [2, 2] to [2, 2 E] results in large
C ADDITIONAL RELATED WORK
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-imarraace rather than equivariance: mean and max pooling lose information about duplicate elements while sum pooling does not. Equal elements easily arise in graphs (c.g. in molecules, the multiset of neighbors often contains multiple atoms of the same chemical element), making a clear distinction between sets and multisets for invariance and equivariance properties is especially important in such contexts.
Recent works in deep learning have successfully applied implicit differentiation to large-scale language modelling (Bai et al., 2019), semantic segmentation (Bai et al., 2020), meta-learning (Rajeswaran et al., 2019), neural architecture search (Zhang et al., 2021b), and more (Agrawal et al.,
18

---

Published as a conference paper at ICLR 2022
2019). The implicit function theorem allows for seamless integration of optimization problems into larger end-to-end trainable neural networks (Gould et al., 2016; Amos & Kolter, 2017; Blondel et al., 2021; El Ghaoui al., 2021). Beyond the sets and multisets considered in this work, implicit differentiation has been applied to other non-Fuclidean domains, such as hyperbolic space (Lou et al., 2020) and graphs (Gu al., 2020).
D DERIVATION OF IMPLICIT DSPN
The goal of differentiating Equation 7 is to obtain the gradient of the training objective with respect z and 0. If we can do that, then iDSPN fits into an autodiff framework like normal differentiable module. We start by deriving the Jacobians for z and O with implicit differentiation and present the approximation for the Hessian afterwards. Blondel et al. (2021) describe straightforward and efficient recipe for applying implicit differentiation to optimization problems. We follow their approach for the standard implicit differentiation derivation and efficiency concerns. We highly recommend their paper for further background and possible extensions that are made possible by implicit differentiation.
To simplify notation, combine the target vector with the encoder parameters into θ = (z, 0) and restate the optimality condition:
L(Y" 0) 0
(42)
This states that the gradient of the inner loss V T. when evaluated an optimum ()) should be a matrix of zeroes. This implicitly defines a minimizer function Y* (0)) The implicit function theorem (further background in Krantz & Parks (2012)) states that its Jacobian (x*(e)/00 exists if L is continuously differentiable and the Hessian 0(Vy L(Y", /)) is a square invertible matrix. Using Vy L rather than lets us apply the implicit function theorem: its Jacobian is a square matrix instead of vector and the optimality condition is fulfilled even when the root of V L only corresponds to 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 42 with respect to and reorder to obtain:
LVY θ) (Y) oY+
D
(43)
00
00
This is linear system of the form HX B where
use autodiff to compute H and B. Solving this system for X = (0) lets us compute the gradients for z and with VeL (Vy.L) X, where L is the loss of the prediction (as opposed to L).
Efficiency. In order to incorporate iDSPN into an autodiff framework, we only need the vector- Jacobian product (VJP) v' X for some vector v. Following Blondel al. (2021), We can efficiently compute this by solving the much smaller linear system Hu = Vy. C for u. This allows us to compute the desired gradient with VaL (Vr-L) X
HX B..
Instead of solving the linear system exactly, it is more efficient to use the conjugate gradient method (Blondel al., 2021). Rajeswaran et al. (2019) show that small number of conjugate gradient steps (5 in their experiment) can be enough to obtain sufficiently good solution. Another advantage of conjugate gradient 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 Tensorflow (Abadi et al., 2015) it also improves parallelization: to compute the full Hessian explicitly we would have to differentiate each entry in V L individually. The same advantage applies to B since we can also use a VJP to compute B.
Previous uses implicit differentiation in meta-lcarning (Rajeswaran et al., 2019) and neural architecture search (Zhang et al., 2021b) also involve solving 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 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 2. In preliminary experiments, Section 3 appears due regularizing I. by adding ||
19

---

Published as a conference paper at ICLR 2022
observed that even if L is not regularized, solving for the regularized Hessian (H/y+ + I) yielded better performance at the same number of conjugate gradient steps. Adding the identity matrix can be seen as adding
the cigenvalues to improve the conditioning of H.
Of course, in the main text the efficiency improvements are irrelevant because we take 7 - oo. Due to that, H I so Equation 43 simplifies to Equation 9. Fung al. (2022) provide different motivations for this approximation.
E EXTENDED EXPERIMENTAL RESULTS
E.1
NUMBERING
We perform additional experiments (Table 4) with data augmentation: by randomly permuting the input order, we can increase the effective number of training samples for the not-equivariant models. Since the other models are equivariant, data augmentation has no effect on them. We include additional runs for 8 classes (instead of 4 classes as in Section 4.1), in which we double the hidden dimensions of all models to account for the larger input. We use the same number of training steps in all runs (50,000).
We show these results Table 4. For 4 classes, adding data augmentation restores the performance for both transformer with PE and BiLSTM in the 1x and 10x setting. For 8 classes, adding data augmentation does not help as much. More classes make the task more challenging, because the expected number of different permutations per input sample increases with the number of classes. Note that the classes are uniformly sampled. Increasing the number of classes while keeping the number of training steps constant means that smaller fraction of all different permutations is during training.
Table 4: Extended experimental results for class-specific numbering with input set size 64 for 4 and classes. Data augmentation (DA) by randomly permuting the elements for the not-equivariant models closes the performance gap to the exclusively multiset-equivarian models on 4 classes, but not for the slightly more complicated case of classes. Accuracy in (mean standard deviation) over
seeds, higher is better.
Set- cquivariant
Multiset- cquivariant
Training samples
Model
1
10x
100x
4 classes
DeepSets
/
0.0+o
0.0+o
0.0+a
Transformer without PE Transformer with random PE Transformer with PE
/
0.0
0.02o
0.01e
X
28.5+s.1
22.6+s.2
22.7+7.7
X
X
0.0+o
46.4=20.2
94.4+1.s
BiLSTM
X
0.01o
25.1 12
97.9121
Transformer with random PE + DA Transformer with
X
26.83+o.s
24.42=4
25.52+0.7
X
94.93
96.35
96.5811.
BiLSTM + DA
X
X
98.12=a.a
99.04=a.a
98. 03.80
DSPN iDSPN (ours)
X
>
89.4+2.4
88.9*4.b
90.8+4.2
96.81o.s
18.1
97.91a.
8 classes
Transformer with random PE + DA Transformer with
X
6.95+4
12.03-a.4
11.23+2.
x
38.37112.
96.95
96.7911.7
BiLSTM + DA
X
61.56=16.
49.03=10.
61.67
DSPN
72.51 aT
73.87 ...
72.65
iDSPN (ours)
90.91
96.33-1.s
96.78+1.6
20

---

Published as a conference paper at ICLR 2022
E.2 RANDOM SETS AUTOENCODING
First, we show the full experimental results as mentioned in Section 4.2. We had to exclude 10 runs of DSPN due to significantly 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 rums of 2, 32 (out of 8 seeds),
run of n 2, d 16,
2 runs of n = 4, = 16,
run of n = 4, d = 32.
Figure and Figure 3 provide different views of the same underlying data in Table 5: each subplot in Figure 2 keeps n fixed and varies d, while each subplot in Figure 3 keeps d fixed and varies n. Since the scale of the losses varies across the different dataset configurations, we normalize each (n, d) combination by dividing the loss by the loss of the it=10 Basel leeelnne ne (DSPN) run. That is why it=10 Baself ne 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, Baseli tine 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 within the same n, d dataset configuration.
The impact of multiset-oquivarianee
We perform an additional experiment on this dataset in order to compare between an exclusively multiset- -ettiiittequvvaiiantt model (iDSPN with FSPool encoder and using momentum) and set-equivariant models (iDSPN with sum or mean pooling encoder and using momentum). The models and training settings are exactly the same aside from the difference in pooling method used and thus their equivariance property.
We show the results in Table 6. When comparing these results to the numbers in Table 5, it becomes clear that the exclusively multiset-equivariant iDSPN far better. For most dataset configurations, the loss of iDSNN with sum or mean pooling are worse by factor of 10 or more. For example, for (n=8, d=8) at 40 iterations, the loss improves from 6.4e-2 to 2.1e-3 when going from Sum to FSPool, lower loss by factor of around 30. Only when the task is extremely easy extremely hard is the difference smaller: for n=2, d=2 at 20 iterations, the loss FSPool is "only" better by a factor of 1.4 (Sum 2.7e-4, FSPool 1.9e-4), and for n=32, d=32 where every model essentially random guess, they are almost equal (Mean 3.4e-1, FSPool 3.2e-1).
Note that experienced significant instabilities and training divergence when using sum pooling for larger multiset sizes. We excluded these divergent runs manually (111 runs in total across all seeds), which sometimes leaves zero successful seeds for dataset configuration (e.g. n=32, d=16 for sum pooling). We did not observe any instabilities for FSPool and mean pooling.
21

---

Published as a conference paper at ICLR 2022
=
1.50
1.25
relative validation
1.00
0.75
0.50
0.25
0.00
70
YI 16
1.50
relative validation loss
1.25
1.00)
0.75
0.50
0.25
0.00
32
32
Y
1.50
i
relative validation loss
1.25
it=10 iDSPN
it=10 Baseline
1.00
it=20 IDSPN+
0.75
iDSPN
it=20 Baseline
0.50
it=40 iDSPN+mmo
0.25
it=40 IDSPN
it=40 Baseline
0.00
8
16
32
Figure 2: Relative reconstruction loss for autoencoding sets when fixing and varying d, loss relative to it=10 DSPN, lower is better.
22

---

Published as a conference paper at ICLR 2022
d d=2
d
1.50
1.25
relative validation
1.00
0.75
0.50
0.25
0.00
d
1.50
relative validation loss
1.25
1.001
0.75
0.50
0.25
0.00
8
32
8
32
1.50
i
relative validation loss
1.25
it=10 iDSPN
it=10 Baseline
1.00
it=20 IDSPN+
0.75
iDSPN
it=20 Baseline
0.50
it=40 iDSPN+mmom
0.25
it=40 IDSPN
it=40 Baseline
0.00
8
16
32
Figure 3: Relative reconstruction loss for autoencoding sets when fixing d and varying n, loss relative to it=10 DSPN, lower is better.
23

---

Published as a conference paper at ICLR 2022
Table 5: Full results for autoencoding sets every dataset configuration with set size n. and set element dimensionality d, every model, and varying number of iterations. Reconstruction Huber loss (mean + standard deviation) over rums for most results, lower is better. Some dataset configurations with DSPN use fewer runs, see Appendix E.2 for the explanation.
Model
Iterations
d=2
d=
=
d 32
iDSpN+mom
10
100-514e5
1.
2335
1.0e-3 1:
iDSPN
10
204
2.9c-42104
2-8
3
4.0
Baseline
10
1.8c-42204
4.9c-42404
1.94-42104
2.5e-3230.4
1.0e-2
iDSPN+mom
20
1.2
8.3e-527045
2.3c-42104
6.5c-421e4
7.1
iDSPN
20
1.7c-42764
1.2c-4260-5
4. 1114-42104
9.1e-4426-4
3.1
Baseline
20
1.2e-427055
2.7c-428o-
4.68-42204
1.0c-3322-4
1.60-3
iDSPN+mom
9.60-52805
1.1c-42104
1.9c-42965
3.3e-42865
2.0
iDSPN
40
6.50-52405
1.4c-42505
1.3
8.5c-42te4
2.70-32504
Baseline
40
1 1.0c-427645
2.1c-42104
1e
7.2c-42304
3.20-32mmn
71
4
iDSPN+mom
10
2.3
5.7c-429e5
2.2c-322e4
9.9c-32554
5.8c-2+2e3
iDSPN
10
1.7e-44104
1.6e-3
4.94-3
1.9e-2+364
8.8e-2+3e3
Baseline
10
4144
2.0e-3
5.8e-3
1.9e-2
6.7e-2+263
iDSPN+mom
20
1.9c-42le4
3.8c-42164
8.4c-417e5
3.60-3322-4
3.30-212e-
iDSPN
20
5.1e-4=1o4
9.6e-44204
2.1
8.3
5.2e-2+se-
Baseline
20
6.64-42164
1.1e-3
2.4e-3=865
9.0e-3
te5
iDSPN+mom
40
1.3e-42765
3.1e-42104
1.
2.7e-3
t.
iDSPN
40
4.00-4200-
1.0e-32204
1.94-323044
6.4
3.4
Baseline
40
428e-42104 1.2e-32204
2.20-3
1.0e-2
3.20-2+203
8
iDSPN+mom
10
4.2e-42305
2.5e-32104
1.36-2
5.0e-2
2.1e-1
iDSPN
10
1.4e-32704
7.0c-3
2.7c-2+964
8.2c-2
1
Baseline
10
1.7e-3#te4
8.5e-3
2.8
7.5e-2+5e-4
1.9e-1+20
iDSPN+mom
20
1.8c-42404
7.1c-4260-5
3.4c-32204
2.1e-224044
1.5e-1
iDSPN
20
442e-
2.3e-3
9.3e-3
4.7e-2t2263
1.98-1+203
Baseline
4
3.0e-3
1.le-2#5e-4
4.3e-2
1.5e-1+ee
iDSPN+mom
2.6e-44605
5.2e-4
2.1e-3422-4
9.4e-3
1.le-1+3e3
iDSPN
40
1. 1.4e33 442-4
1.6e-3#2e4
5.0e-3
2.5
1.6e-1+363
Baseline
40
1.le-3
1.9e-3
4.8
2.3e-2
1.
16
iDSPN+mom
10
1.0e-3#765
1.le-2#204
4.4e-2#2044
1.5e-1
1e0
iDSPN
10
2.8e-34364
2.0e-2
6.2e-2+te-4
1.7e-12163
1
Baseline
10
3.9e-3:1ed
2.20-2.15644
6.00-2:16-4
1.5e-1i2es
2.2c-1 114e3
iDSPN+mom
2.9e-42405
3.2e-3:2605
1.8e-2#444
1.0e-1 146-3
2.8e-1+te-
iDSPN
20
1.6e-3+144
9.0e-3
3.7e-2+8644
1203
2.8e-1+1e5
Baseline
20
1.6e-3ate4
1.0e-2
3.8e-2:55-4
1.2e-1
2.20-1
iDSpN+mom
40
3.2e-429665
1.4e-3
8.0e-32304
6.9e-2
2.6e-1+263
iDSPN
40
1.4e-3a9es
4.5e-3
2.0e-2a5e-4
1.le-1
2.7e-1i2e3
Baseline
5
5.3e-
2.2e-2
9.3e-2
2.2e-1ited
32
iDSpN+mom
10
6.
3.0e-2#3e4
9.3e-2#66-4
2.3e-1 +
3.3e-1
iDSPN
10
4.8e- 442e4
3.0e-2#dd44
1.0e- 1
2.3e-
+163
1
Baseline
10
5.6e-3 1e
2.9e-24364
9.3e-2+3364
2.2e-
+
1
iDSpN+mom
20
7.7e-44605
1.le-2
6.5e-2
2.3e-
+
3.2e-1+203
iDSPN
20
2.7e-347e5
2,0e-24:5-4
8.3e-244964
2.2e-1
1
3.2e-146e4
Baseline
20
2.10-2260-4
7.50-2226-4
2.0c- 1
3.0c-113e6
iDspN+mom
4.0e-443e5
5,4e-3
4.9e-242265
2. +:3-33
3.2e-1
iDSPN
40
1.20-2:22-4
6.9c-2+1e3
2. 19c-
1.26-1
Baseline
40
1.90-322e41
1.4e-2
1.20-2
2.0c- 1
2.96-1
24

---

Published as a conference paper at ICLR 2022
Table 6: Autoencoding random every dataset configuration with size n and element dimensionality d and varying number of iterations for set-equivariant iDSPN variants. Reconstruction Huber loss (mean standard deviation) over 8 runs for most results, lower is better. Some dataset configurations with iDSPN+mom+SSum use fewer runs due to divergent training behavior.
Model
Iterations
d 2
d = 4
d 8
16
d
32
n
2
SDSPN+mom+Sum
10
1.
1.7e-31se5
3.5e-3124 4
3
6.4e-2atte
iMsnN=mom+Mean
10
1.20-3tte4
2.1e-31204
3.4e-3ite4
8.3e-3
6.1e-243e3
iDSPN+mom+$Su
20
2.70-4430.5
5.7e-41905
1.5c-34104
4.8c-32204
1.4c-2+104
iDSPN-mom+Mean
20
3.70-411o4
7.0c-41805
1.60-31905
3.70-32104
1.4
IDSPN+mom+$Sum
40
2.90-41905
5.70-42204
7.5c-4
3.5c-32104
2.10-22604
IDSPN+mom+Mean
40
5:
5.9e-44304
8.9
3.5e-3+304
2.2
4
iDSPN+mom-Sum
10
1.3
9.9c-3122-4
3.20-213e4
1.1c-
2.4c-111t33
iDSnN-mom++Maan
10
5.4e-3t4e4 1.2e-2+4e4
3.6e-2+5e4
1.2e-1420.
2.3e-1#1ed
iDSPN-mom+Sum
2.7e-347e4
3.80-34204
1.6e-247e4
2
1.9e- +203
iDSPN+mom+Mean
20
3.76-3
5.4e-34304
1.6e-2+66-4
9.5e-2+3ea
2.0e-
1 +304
iDSPN+mom+Sum
40
2.30-311e4
1.2c-2:14
5.5e-
2:1-22te2
1.8c- 1.8e-1::2 12:3
iDSPN+mom+Mean
40
313e-4
3.1e-31264
1.2e-213e4
8.2e 2:2e3
2.0c 1i84
8
iDSPN+mom+SSum
10
2.
5.7e-2
9.5e-2
2.4e- 1 1:32
3.1e- ::ee
iDSPN-mom+Mean
10
1.5
4.4e-214e4
1.2e- 49944
2.4e-
1443
3.1e-1
48e4
iDSPN+mom+Sun
20
4.6e-3 mmn
1.3e-2+364
2.3e-
2.9e-
49e4
iDSPN+mom+Mcan
20
9.90-324e4
2.4c-21004
2
2.4c-
1764
2.9c-
2703
iDSPN+mom+Sum
40
2.90-318044
8.30-31304
6.40-21363
2.2c-
1
2.8c-
iDSPN+mom+Mean
40
6.70-3224-4
1.60-2160-
8.60-2+1e3
2.4c-
2.8c-
18c-1160-4 1
16
iDSPN+mom+Sum
10
7.8e-213e-2
1.50-113e2
1,50-
2.9c-
1222
3.5e-1 1
ims
10
4.20-2+20-3
1.2e-1
2.1e-1
2,9c-
+4e-3
3.5e- 1363
iDSPN+mom+SSum
20
8.40-3+304
4.0e-24504
2.3e- +
3.1e-I
#362
3.6e- 1.6e-1#362
iDSnN-mom=+Mean
20
2.1e-24764
7.1e-24304
1.9e- 22643
2.8e-
1404
3.5e-1 1
iDSPN+mom=4Sum
40
8.88-34203
1.3e-1 1443
2.4
Mes
40
1.89-219e4
2.5e-215e-4
1.8e- 1:233
2.7e-
1:3e
3.2e- 1:3e-
32
iDSNN=mom+Sum
10
4.4e-2+2e-
1.3e- 13e22
2.3e-1+tes
3,0e-
1 1
3.5e- 1663
Mesa
10
6.3e-2+4e.3
1.6e-1
2.6e- 4
2.9e- 4764
3,5e- 4564
iDSPN=mom+Sun
20
1.9e-2+mm
6.8e-24664
1.9e-
iDSPN-mmm+MMann
20
3.50-24263
1,40-111e2
2.50-1
2.9c-
1 1
3.4c- 1440
iDSPN+mom=Sum
40
2.8e-247-4
1.6e- 18c4
2.70-1
3.4c-
1
MDSNN+mom+Mean
40
3.0c- 244e-4
9,30-218e-4
2.1c-1
2.9c- 1
3.4c-1:50
25

---

Published as a conference paper at ICLR 2022
Table 7: Performance on CLEVR object property set prediction task, average precision (AP) in % (mean + standard deviation) over 8 random seeds, higher is better. The first block compares iDSPN models that use different pooling methods (resulting in different equivariance properties), the second block compares iDSPN to DSPN performance in an otherwise like-to-like training setup. iDSPN + FSPool + 0.9 momentum is the same setting as shown in the main paper.
Model
Pooling
Momentum
APx
API
APo.5
APo.5
APo.125
APo.0625
iDSPN
Mean
0.9
92.021a
88.5
29.
55.92so
22.2
5.221.0
iDSPN
Sum
0.9
97.4
96.1
94.221.6
85.82a9
49.4
12.7a1
iDSNN
FSPool
0.9
98.8
98.522.e
98. 98.2to.a
95.8
76.942.6
32.3=a.9
DSPN
FSPool
0.9
66.62a
51.920.7
29.9
7.8=s.o
1.421o
0.0ga.0
iDSPN
FSPool
0.5
48.4
98. 08...
97.8=o.
93.8*1.
61.3aa.
17.3
DSPN
FSPool
0.5
98.61a
98.01u.7
97.1
85.312.
34.814
6.111.
E.3
CLEVR OBJECT PROPERTY PREDICTION
On CLEVR, we perform two additional experiments with 128x128 image inputs. The results for both shown in Table 7.
The impact of muntiset-equivariince First, we evaluate the difference between the exclusively multiset-cquivarian iDSPN with FSPool and the set-equivariant iDSPNs with sum or mean pooling. In this setup, the models and training configurations are exactly the same except for the pooling method. The results show that the exclusively multiset-equivariant iDSPN still consistently outperforms the set-equivariant iDSPNs, though to lesser degree than in Appendix E.2. The largest differences can still be seen for the stricter AP thresholds. This can be explained by the one-hot-encoded attributes helping in keeping set elements farther apart so that separation of similar elements is less of concern.
DSPN in a like-to-like setting. While there are various experimental differences of our iDSPN to the original DSPN setup in Zhang et al. (2019) (more details in Appendix F), we can also run DSPN in a setting as close to iDSPN as possible. This allows us to evaluate the quality of the gradient estimation by the approximate implicit differentiation compared to differentiation through the unrolled optimization. This difference in how the gradient is calculated is the only difference between the models. All other details are exactly the same as in our normal iDSPN training setup.
First, we trained DSPN with the same momentum as iDSPN, which led to significantly degraded results. This suggests that iDSPN is more stable to the inner optimization hyperparameters than DSPN. After reducing the momentum of DSPN and iDSPN from 0.9 to 0.5, we obtain more reasonable results for DSPN. Still, the iDSPN results at 0.5 momentum are better than the DSPN results at 0.5 momentum. Thus, we believe that the approximate implicit differentiation should not be simply viewed as trying to estimate the automatic differentiation gradient and being at best equally as good as the autodiff gradient. iDSPN outperforming DSPN in like-to-like setting suggests that the approximate implicit gradient is actually better than the automatic differentiation gradient. This observation is inline with theoretical findings on strongly convex inner objectives (Blondel et al., 2021; Ablin et al., 2020), which indicate that implicit differentiation is better at approximating the Jacobian when the optimizer solution inexact. Thus, iDSPN should be considered different procedure to DSPN altogether.
Keep in mind that the better results of iDSPN in like-to-like setting compared DSPN also come with the other benefits of implicit differentiation such as faster training speed: in our training setup iDSPN takes around 2 hours 40 minutes to train on V100 GPU, while DSPN takes around hours 15 minutes.
26

---

Published as a conference paper at ICLR 2022
F DETAILED EXPERIMENTAL DESCRIPTION
Unless otherwise noted, 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 ground-truth set. Throughout the experiments, we make of the Hungarian loss (Kuhn, 1955):
C(Y (0)) min TEll
(1)
(2) yx(1) 1|
(44)
finds the optimal 1-to-1 mapping from each element in the prediction Y [y] (1)
each element in the ground-truth ( [v]*) ..... ...... yv J minimizing the total pairwise losses.
A pairwise loss is standard loss like mean squared eπror and cross-entropy. The Hungarian loss
permutation-invartee in both its arguments and able to map duplicate predicted elements to
different ground tround-truth elements and vice versa, which makes it well-suited for multiset prediction.
Set refiner networks. Set refiner networks (SRN) (Huang al., 2020) make few minor changes to DSPN. When we mention DSPN, we use the following modifications:
SRN found that learning rate of for the inner optimization (Equation 7) is sufficient. DSPN originally used 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.
F.1 CLASS-SPECIFIC NUMBERING
The following description corresponds to Section 4.1.
Dataset. This experiment can be thought of as an elementwise classification that requires global information: each input element needs to be classified (predicting the numerical IDD 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. To match the setup in the other experiments, we can implement this as a single call to the matching algorithm by increasing the cost of pairing up two elements of different classes. We use MSE pairwise loss for DSPNADSPN 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 64 with 4-dimensional elements corresponding to classes. This is represented i..d. from the multinomial distribution over the equally- 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 validation set of size 6,400 and a test set of size 64,000 for every run. In the extended experimental results in Appendix E.1 we also extend this setup to 8 classes, for which we use the same input multiset size and dataset sizes.
64 x 4 matrix where each row is one-hot vector that is sampled
Models. We have the prior information that every instance of class in the input will also appear in the desired output. To make the task casier for the models, modify them to take this into account. This makes sure that 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, they already satisfy the desired property. If the "first" element in the input is class c, the "first" predicted ID in the output will naturally correspond to class well.
27

---

Published as a conference paper at ICLR 2022
For the Transformer baselines, decoder is based on an affine transformation (Yo); = Wx; + b of the 1 I i S n input elements. Through this, each input element is paired up with output element without PE hecomes 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 significantly degraded the performance.
an additional modification: the input set to the transformer
well. The Transformer
We adapt the iDSPN and DSPN models to this setting by fixing the first four dimensions of Y to the values in the input multiset as we discussed in Section 3, which makes them multisel-equivariaat to the ipput. With the set encoder 9 architecture of Linear(4 + 64 + 256) ReLU-LLnearre 256)-FSPool, they become exclusively multiset-equivariant Since the initial set Yo 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 first four dimensions fixed. Additionally, we apply projected gradient descent with proj that projects each element back onto 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 significantly improved the performance over the unconstrained version.
-
Evaluation. We compute the accuracy at the sample-level, meaning 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.
F.2 RANDOM SETS AUTOENCODING
The following description corresponds to Section 4.2.
Dataset. Every set has fixed size with n elements of dimensionality d. Every element is sampled ii.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 final epoch without any early-stopping (we did not observe overfitting) 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 learned initial set Y, (Zhang al., 2019), we match this in iDSPN and add the regularizer 0. - Yoll 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 Yo at the start of training with elements sampled from A( I 1/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 > FSPool, where the linear layers are applied on each element in a multiset individually with shared weights across the elements.
- 512)-
Note that we do not compare against Slot Attention here due to the difference in setup. Slot Attention takes set directly as input, so giving it the input set directly would not give us a permutation-invarian bottleneck, which somewhat defcats the point of the autoencoder. It only has to learn to produce the identity function, which we verified it can easily do.
Training. All models are trained for 40 epochs at batch size 128 with Adam (Kingma & Ba, 2015) with learning rate of le-3 and default momentum parameters. This is enough to converge for almost every model and dataset combination. Only the hardest difficulty settings (32 x 16, 16 x 32, 32 x 32) would show any benefit from more training.
28

---

Published as a conference paper at ICLR 2022
F.3 CLEVR OBJECT PROPERTY PREDICTION
The following description corresponds to Section 4.3. Where possible, we use exactly the same setup as DSPN.
Every image is resized to 128x 28 (or 256x256) and scaled to be in the interval [0, 1]. For the ground-
truth, we scale down the target 3d coordinates from the interval [-3, 3] [-1, 1] and concatenate it to the attributes (cach 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) 12) feature map, on which we apply BatchNorm followed by 2x2 convolution with stride obtain a 2x2x128 (4x4x32) feature map. This is flattened to a latent vector with 512 dimensions, which is the input 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), 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 confidence), then computing the precision-recall curve by including more and more predicted elements, starting from the most confident 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 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 first - paired with FSPool, we skip the relation network entirely. This improves the complexity of the set encoder 9 used in iDSPN from O(n log to On n). Using the relation network approach would improve our results slightly (e-g. 3.5 percentage points improvement on APo. for 128 x 128), but is also bit slower. For simplicity, we therefore opted to not using relation networks.
The architecture of the set encoder g in iDSPN is L19 - 512) Ree---LLLLLLLr-e- 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².
+ 512)
Instead of ResNet34 to encode the input image, we use the smaller ResNet18. This did not appear
affect results.
Instead of using a learned initial set Yo as in DSPN, we find that it makes no difference to randomly sample the initial set for every example. We therefore use the latter for simplicity and to match the Slot Attention setup. In initial experiments we found that even initializing every element to 0 causes no problems.
While we do not use the initialization with Os in our experiments to better match the setup in Slot Attention, there are benefits to this initialization. In Slot Attention, the variance of the random initialization is learned, but the variance must strike a balance between being too low (elements are more similar, which poses difficulty in separating them) and too high (model has perform well over a wider value range, which is more difficult). Exclusively mullistt-quivaarian models can avoid this trade-off since separating similar elements poses problem. This is another benefit of exclusive multiset-equivariance
We increase the batch size from 32 128. There appeared 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.9 instead of standard gradient descent without momentum.
Instead of fixing the number of iterations 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 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 configurations.
We drop the learning rate after 90 epochs from 1e-3 to lc-4 for the last 10 epochs. This slightly improved training loss while also reducing variance in epoch-to-epoch validation loss.
29

---

Published as a conference paper at ICLR 2022
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 x 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 final feature map for CLEVR has 32x32 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 reflected by the fact that the time taken to process a sample is similar between R + iDSPN and CNN Slot Attention, even though use smaller batch size and could gain 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, IJ while we use an interval of [-1, 1]. In Table 3, Slot Attention* uses an interval of [-1, 1J (the same as iDSPN) while Slot Attentiont 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 striet thresholds but trades off classification performance at a threshold of infinity. We did not experiment with different coordinate scales and simply kept it the same as DSPN.
We observe some overfitting with 128x 128 images (1.6e-4 train loss, 5.4c-4 val 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 val loss). We thus believe that the overfitting is due to the ResNet18 image encoder rather than iDSPN.
G EXAMPLE 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 8 shows the worst failure case of cach model trained with 100x 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 misclassifications. Since make every model equivariant to the input, the classes shown in the prediction are guaranteed to be correct; the difficulty lies only in numbering them correctly. Note that the classes in the input 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 often somewhat subtle and best seen when zooming in on a digital display. It is most obvious on (n- 4, = 16) and (n 8, d = 16), but there are smaller differences visible with configurations like (n = 32, d = 2) as well, where iDSPN with momentum tends to be slightly more centered the ground-truth truth than DSPN. The examples where n x d 2 512 show that even if DSPN has a better loss than iDSPN on those configurations, the model quality is far from acceptable.
Table and Table 10 show predictions on CLEVR for different number of evaluation iterations (Section 4.3). We checked the first 128 images of the CLEVR validation set manually and selected examples that looked difficult 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 yellow cube that has its sides and bottom occluded for a cylinder (third image). The improvement seen when going from 128x 128 256x256 images is also quite natural, even for a human observer.
30

---

Published as a conference paper at ICLR 2022
Table 8: Worst failure cases for numbering elements within a class, selected by the highest amount of wrong numberings in prediction. Predictions are sorted for better readability. Wrong numberings are marked in red.
Transformer without PE
Transformer with random PE
Transformer with PE
DeepSets
BiL-TMM
DSPN
iDSPN
(a. 4)
(a, 0)
(a.
(a.
(a. 2)
(a,
(a, 4)
1)
(a, I)
(a. 1)
(a.
(a, 3)
(a, 2)
(a, 4)
(a, 1)
(a, 2)
(a. 2)
(a. 2)
(b. 0)
(a, 2)
(a, 4)
(a. 2)
(a, 3)
(a. 3)
(a. 3)
(b. 1)
(a, 2)
(a, 4)
(a, 3)
(a, 4)
(a, ()
(a, 4)
(b, 2)
(a. 2)
(a, 4)
(a. 4)
(a, 4)
(a.
(a. 5)
(b. 3)
(a,
(a, 4)
(a, 6)
(a, 5)
(a, 5)
(a, 6)
(b, 4)
(a,
(a, 4)
(a. 7)
(a, 7)
(a, 6)
(a. 7)
(b. 5)
(a.
(a,
(a.
(a, 8)
(a. 7)
(a.
(b. 6)
(a.
(a,
(a, 9)
(a,
(a,
(h, 7)
(a,
(a.
(b.
(a. 11)
11)
(a, 11)
(b, 9)
(a,
(a, 4)
(a, 12)
12)
(a, 12)
(a. 12)
(b, 10)
(b.
(b. 0)
(b. 0)
(a. 13)
(a, 13)
(a. 13)
(b. 11)
(b, 1)
(b, 0)
(b, 0)
(a, 14)
(b, 3)
(b, 4)
(b, 12)
(b, 1)
(b, 0)
(b, 1)
(a, 15)
(b,
(b, 6)
(b, 13)
(b, 1)
(b, 0)
(b.
(a. 16)
(c.0)
(c.0)
(c. 0)
(b, 1)
(b,
(b, 4)
(a, 17)
(c,
(c,
(c, I)
(b,
(b,
(b. 5)
(a. 18)
(c.2)
(c.2)
(c. 2)
(b,
(b,
(c. 0)
(a. 19)
(c,3)
(c.3)
(c. 3)
(b,
(b,
(c. 1)
21)
(c.4)
(c.4)
(c,
(b,
(c,
22)
(c,4)
(c,5)
(c,
0)
(c.3)
26)
(c.5)
(c.6)
(c.
(b,
(b, 0)
(c.4)
27)
(c.7)
(c
(c,
(b.
(b. 0)
(c.
27)
(c,8)
(c.7)
(c,
(b, 1)
(b, 0)
(c.
(a. 28)
(c,8)
(c.8)
(c.
(c, 5)
(c.
(c,
(a, 28)
(c.
(c,8)
(c,
(c,
(c,
(c.
(a, 28)
(c. 11)
(c.
(c. 10)
(c.5)
(c,
(c.
(a, 28)
(c. 12)
(c. 10)
(c. 11)
(c,5)
(c,4)
(c, 10)
(a, 29)
(c, 13)
(c, 10)
(c, (2)
(c.5)
(c,4)
(c, 11)
(a, 29)
(d. 0)
(c. (c)
(c. 13)
(c.5)
(c,4)
(c. 12)
29)
(d. 1)
(c. 14)
(c,5)
(c,4)
(c, 13)
(d, 2)
(c,
[c, 15)
(c.5)
(c,4)
(c. 14)
29)
(d.4)
[c. 16)
(c
(c,4)
(d.0)
(b.
(d.
17)
(c
(c,4)
(d, 2)
(b.
(d,
(c. 16)
(c, 18)
(c
(c,4)
(d.
(b. 2)
(d. 7)
(c. 17)
(c. 19)
(c.5)
(c)4)
(d.
(b, 3)
(d. 10)
(c, 17)
(c. 20)
(c.5)
(c,4)
(d.
(b. 4)
(d. 11)
(c, 18)
(c. 22)
(c.5)
(c,4)
(d, 5)
(b. 5)
(d. 12)
(c. 19)
(c. 23)
(c,
(c,4)
(d.5)
(b, 5)
(d, 13)
(c, 20)
(c. 24)
(c.
(c,
(d.6)
(b. 6)
(d. 15)
(c, 20)
(c. 25)
(d. 14)
(d.0)
(d.7)
(b. 7)
(d. 16)
(c,
(c. 26)
(d, 14)
()
(d,
(b,
17)
(c, 22)
(c,
(d,
(d.
(d.
(b.
(d. 19)
26)
(d, 14)
(d,
(d, 10)
(b, 10)
(d, 20)
24)
(c. 27)
(d, 14)
(d.
(d. 11)
(b, 11)
(d, 22)
(c. 26)
(c. 28)
(d. 14)
(d. 0)
(d, 11)
(b. 12)
(d. 23)
(7)
(c. 29)
(d. 14)
(d.
(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.
(d, 1)
(d, 14)
(d. 0)
(d, 18)
(c. 4)
(d. 27)
(d.
(d. 2)
(d, 14)
0)
(d, 19)
(c.
(d, 28)
().
(d, 3)
(d, 14)
(d.
(d,
(c.
4)
(d, 4)
(d,
()
(d, 20)
(d.
(d.
(d,
(d.
(d. 24)
(c.
().
(d. 6
(d, 14)
(d,
(d, 25)
(c.9)
(d, 29)
(d.
(d, 7)
(d, 14)
(d.
(d, 25)
10)
(d,
(d.
(d, 8)
(d. 14)
(d. 0)
(d, 25)
(c)
(d.
(d.9)
(d. 9)
(d, 14)
(d. 0)
(d, 27)
(c, 12)
(d,
(d, 10)
(d, 10)
(d, 14)
(d. 0)
(d, 27)
(c, 12)
(d,
(d. 11)
(d. 11)
(d. 14)
(d. 0)
(d, 27)
(d, 4)
(d,
(d. 12)
(d. 12)
(d, 14)
(d, 0)
(d,
(d, 5)
(9)
(d, 13)
(d. 13)
31

---

Published as a conference paper at ICLR 2022
Figure 4: Random set examples for each combination of n. and d. For d >> 2, only the first two dimensions are shown. All models use 20 iterations.
2,d
2,
2,
2,
16
d
32
2-
2-
2-
2
0-
0-
0-
2.5
0.0
2.5
2.5
0.0
2.5
2.5
0.0
2.5
2.5
0.0
2.5
2.5
0.0
255
4,d
4.d
4,
4, d
16
4,
32
22-
2-
2-
2-
2-
0-
g
2.5
0.0
2.5
2.5
0.0
2.5
2.5
0.0
2.5
2.5
0.0
2.5
2.5
0.0
25
8,
8,d
8, d
8, d
16
8, d
32
2-
2-
2-
2-
2
0-
0-
0A
-2.5
0.0
2.5
-2.5
0.0
2.5
-2.
0.0
2.5
-2.5
0.0
2.5
-2.
0.0
25
16. d 2
16.d
16, d
16.
16,
32
2-
2-
2-
2-
2-
83
0-
-2-
-2
-2
-2
-2
-2.5
0.0
2.5
-2.5
0.0
2.5
-2.5
0.0
2.5
-2.5
0.0
2.5
-25
0.0
2.5
32, d
32, d
32, d 8
32,
32, d
2-
2-
2-
2-
@
2-
2
24
2
-2.5
0.0
2.5
-2.5
0.0
2.5
-2.5
0.0
2.5
2.5
0.0
255
-25
0.0
255
Grvund-tuuth
DSPN
iDSPN moreenium
32

---

Published as a conference paper at ICLR 2022
Table 9: CLEVR examples with 128x 128 images, cherry-picked by difficulty (highly-occluded objects). Incorrect attributes and Euclidean distances (d) greater than 0.25 marked in red. Ground- truth objects are ordered going from top left to bottom right diagonally, predicted objects are matched to that ordering based on Equation 44, We choose an iDSPN run with median performance and evaluate at different iterations.
Image
Ground truth
IDSPN 10 inerations
ISPN 20 itenations
IDSPN 40 iocrations
(-2.95 -2. 0.70),
(-2.77 79 22-299 0.64) d-0.41
(-288 -2 58 0.79) d-0.13
(-2. -2-55 0.67) d- large perple metal
large purple metal cylinder
large purple
cylinder
large purple metal cylinder
(-2.26 0.66 0,35)
(-2.6 0.60 0. 26, den 24
(-2.3 32 0.41 0 y 26 small eyan metal sphere
(-2.31 0.45 0.35, d-n 21 small eyan metal sphere
cyan metal sphere
small eyan metal sphese
(-2.13 1622 0.35)
(-1 (-1462.00 2006 0.06) dunl 9p
(-204 L770 0,4 de0. 14
(2,03 1.34 0.37) 1
small brown mital cyliader
small beown metal eyfinder
small bownn metil cylinder
small beown metal cylinder
(- I 80 0.35)
(-0 69 33 35) den 55 small yellow rubber sphere
(-8) - $4 n 52) d=0 05 small red rubber sphere
(-0.8) -1 0 d=0 04 small red rubber sphere
small red ribber sphere
(-0.34 2 86 0.35)
(0.45 2.73 0.62) d 140.84
(0,08% 2 2.73 (.43) d 30
(-0.24 2.86 0.33) ded small eyan nilher cylinder
small eyan nilibrr cy cylinder
small yellow rilther ey linder
small yellow mibher sylinder
00.37 2.25 (1700)
(3.33 0.96 44) d=1.36
2.34 7.70) d-0 0o large eyan rubber cube
(0.42 2.28 0.00) d-0,06 large gyan rubber cube
lage syan rubber cube
small red rubber cube
25 25-2.004 0.70)
1,59 large brown metal cylaneer
-2.16
(1 36 d-1..3
d-n large yellow rubber cylinder
30
large yellow zubber cylinder
large yellow rubber cylindrr
0.35)
d:0.95 small purple metal cylinder
d=n.39 small pupple metil metilyylndds eylinder
0.05
0. small purple metal cytander small brown metal cylinder
small purple metal eylinder
2.29
2.23 des0,40
(1.07 0.46) den cylinder small brown
small brown cylinder
small metal
(2.93 2616 0,70) cuhe
(3,15 2.34 0.68) d:20.29
(2 9696 2.34 0.83) d=0 23 large
(2 90 2.20 0.20) d-0.05 large
large
large
(-1,73 2.38 0.35)
75 0.31) d=0.42
(-1 58 2.64 47) d=t..32
(-1 2.48 0.37) d-n 10 small green metal cylunder small rot motal linder
small green metal cylinder 122 small metal eylinneer
small metal cylunder d=0.25 small purple motal cyllnnder
small green metal cylinder
0.35)
small red metal cylinder
2009 0.35)
12
15
d=4.44
small brown rabber sphere
small brown nubber ey rylinder
small beown rubber cylimder
small brown rubber cylinder
(0
(0.48 71 (155) d5
(0.23 a
(o -283 0.66) d-0.09
largo purple
sphere
purple
sphere
large purple metal sphere
sphere purple
ro .002 (170)) large bise nibber
(.00s .1.79 24) d-0.90
(0.12 0z 0.60) d=i 28
10.29 aoos 0.71 d-0.09
large gray
sphere
large
(1.25 14 35)
(0.63 o 55 31) d= 7.70 small blue metal cylinder
(0.83 08 0.40) d=1 .4
(l 2020 15 o s0) d-p.06
small cyun metall cylinder
small metal cube
small cyan metal eylinder
(1 62 2.99 0.35) small yellow rublerr cuhe
(1.16 2 23 0.3 d=O 18
(1.57 2 90 o 10y d=0.0s
(1.58 2 0,33) d=0.07
small yellow subber cube
small yallow rubber cabe
small yellow
cube
(2.30 0.40 Q (0) lirge gray metal sphere
(203 0.63 0.54) d=0.37
(2.34 0.39 0,78) dw0.09
(2.20 0.39 0.69y d-0.01
lange gray metal sphere
large gny metal sphere
large gray metal sphere
(:2.3 0.66 0.70)
(-2.60 0. 0 75) d=1.13
(:2.99 0.05 o 70) d=0.66
(-2 66 0 o 21) d=0 30 hage oyan
lang eyan
sphore
large eyan
sphere
large eyan
sphere
sphere (:2.15 55 70) d-n
1-2.09 1.51 0.70)
(-1 07 -061) d=3
(-1.9 92 2.28 o 91) dei) 36
large blue
eyan metal
large
large
cube
(:2.00 2.71 0,70)
(227 2.48 0,30) den,54
(2.23 2 80 0.83) d 28 large cymn rilher sphere
(2.14 2 2,80 69) d=O 16 lange eyan nihhhr sphere
lango eyan nibher sphere
large eyan milher sphere
(-169 2344 0.35)
(-8.8 04 1,25) 1.1
(-1 33 2.81 0.55) d-0.41
(-1 2.75 0.30) d=o
small yellow metal cube
large yellow metal ey linder
small yellow metal cylinder
small yellow metal linder
(-0.18 2.91 0.70)
(-0.45 30 (164) de 29
(-0.34 34 3.02 o 73) d-0 20 lurge gray rubber cube
(-0.29 2.98 0.67) d-0.13 large gray rubber cube
large gray rubber cube
large gra rubber cube
(0 1212-2 2:40 0.35)
(0.33 -1 84 0.51) det 57
(-001 -2 30 0.38) den 17
(0.06 22 0,37) dan 24
small blee metal cylinder
small bise motal cylunder
small blue metall cylinder
small blue metal eylinder
del farge blue metall by cylinder
00 io 14
deni.35 tirge blue motal by cylinder
0.68) large yellow metal cylinder
(n
large yellow metal sylinder
(0.91 0.35)
0.29
d9 small blue metal cylinder
small bllce metal sylinder
small bise metal cyllnder
small blue metal cylinder
(1.97 0,70)
(2,00 2.79 d40
(1.88 2.63
(1.55
lange rol ribber cahe
largo ruhherr
nibber cube
lorge rod ruhber cuiee
(-2.82 2.97 0.70)
(-2.66 2 0 (.7%) d=O 18 lage eyun nutber cylinder
(-2.44 0.71) d-0
(-2.85 2.98 0,20) d=0,03
large eyan rubber eyfinder
larg eyan nubher cylinder
large eyun nulber eylinder
(-2.28 22 62 0,70)
(-2.5t 24 2.70 0.20) d-d jo
(-226 a 65 0.81) d-0 12
(-2 22.25 2,63
large purple metal sylinder
large purple metal ey lyndder
lange purple metal cylinder d-0.
large purple metal cylinder
(1177 -0.51
(-102 .n.54 0.63) d-11,74
(-1.74 40.3
-0.6 34) d-a small syan metal cylinder
small syan metal cylinder
small cyan metal cylinder
small eyan metal cylimber
36
small yellow rubber cylander
small yellow: rulber cylinder
small yellow nubiher cylinder
small yellow rubber cyllider
0.35)
(1 n d=
(0 a0 n d-i).42 brown rubber lyndeer
(0.22 -I zubber cylinder
small red rubber sylinder
small brown subber linder
(0.43 2.4.3..666 .666
(0.42 2.57 0.00 8.69) d=0 07 cylinder
(0.45 2.60 o 70) dan of cylunder (1.22 1.83 (2) dei 16 cy
46 62 0.67) d=0.05 large pupple
large people metal eylinder
large pepple
large purgle
eylinder
(1.32 -1.70 0.70) large (2.41 0.23 a 55)
(1.23 -1 89 0.67) dat 10 lange (2.56..35 a 10) d-0.27 metal sphere
(1.25 1.30 2.) da0 12 cylinder large brown
linder
large bnown
(2.44 0.22 30) d-0144 small red metal sphere
(2.44 2. a 0.36) d-0.04
small red metal sphere
sphere
6.2 0,35)
(-2 -2.86 den n9
(-280- -2 83 0.46) d-0. gray
(-2.5 88 de0 metal
gny metal
gray metal
(:2.7 2414 0.35)
(:2.61 2:22 d=n
(-2.0 70 2.33 o 29) d=0 21 purple
(-2.2 2.35 0.42) d-0.22 perple
purple
cube
perple metal
(-2.33 1 ju 0.70)
0:200 -162 0.57) d=0.54
(-2.36 4.8 0,26) d=0. 17 cy linder
(-2.42 -.95 75) dea 13 cy sylinder
large
cy linder
lange
cy gylindrr
large brown
lange brown
(-1 73 22 no 35)
(-144 A6 2.30 72) d=0,71
(-1.73 -2 71 (.4n) d-o to
(-1 1.64 2 so o 34) d-a 22 gray
small gray rubber sube
large gray
cube
small gray tubber cube
(-1.31 31 2.74 0,35)
(-0.48 65 0.4 d=2:26
(-1.15 1.80 o 25 den 96
(-1 2.24 0.34) d=0-52 small green motil sphere
small metal cabe
small green rubber cube
small gray motal sphere
small brown metal spteeee
(-1 1-1.8 1.67 0.35)
43 2 43 (1.4)) defUSt small bownn meti sphere
small brown mexal sphere
(-1 13 1 89 a 38 d-0 23
small brown metal spteee
(-1.) 17 1.60 36) d-0.06
(-0.0) 2.36 0,70)
(-0.06 3666 O 60) d=0 II
(0.01 2.47 0666) dan 12
10.01 2.36 a 69) d=0,02
large blue rubber cube
large biue nakber cale
large blue uhter cube
large bine rubber cube
10.45 -1.33 0.70)
10.66 2.07 0.51 d=0 55
00,50 1.76 0.s)) dei 13
10.49 1.36 0.68) d-0.05
lorge gray rakher osbe
small gny milher culs:
large gray nithee cuhe
large gray miber cole
33

---

Published as a conference paper at ICLR 2022
Table 10: CLEVR examples with 256x256 images, cherry-picked by difficulty (highly-occluderr objects), same examples as in Table 9. Incorrect attributes and Euclidean distances (d) greater than 0.25 marked in red. Ground-truth objects are ordered going from top left to bottom right diagonally, predicted objects are matched to that ordering based on Equation 44, We choose an iDSPN run with median performance and evaluate at different iterations.
Image
Ground truth
IDSPN 10 itteaaioons
IDSPN 20 iserations
IDSPN 40 incrations
(-2.95-2. -2, 0.70) large purple metal sylinder
(-2. 70 -1.53 0.53) d=1.07 purple
(-2.11 11 -2 49 0.69) d-0.84
(-3.01 so 60) d-00s
sphere
large rubber cylinder
lange purple metal cylinder
(:2.6 0.66 0. 0.35)
(-2.33 5.59 0. d-n 18
(-2.5 5i -0.4 n 4,46 de 1.18
(-2.22 0.56 0.35) d-n. II
small cyan metal sphere
small menil cylinder
small purple metal sphere
small evan metal sphere
(:2.33 1.52 0,35)
(-111 1.51 0.25) d=1.02
(-2.2 1.50 0.31) d-0 25
(2220 160 0.35, d=0.08
small hrown metal ey fyndddr
small brown metal cylinder
small brown metal cylinder
small brown metal cylinder
1.0.83- I 00.0.5)
(0.0) 22 7 (4 d=0.94
(-068 -2 2.00 0 38) des 26
(-0 78 76 36 de0 06
smail red nubber spiene
small cyan rubber cylinder
small cyan rubber sphere
small red rubber sphere
(-0.34 2,86 0.35)
(0.84 2,20 (.59) d=1.37 small yellow nibher cylinder
(0,02 0.97 0.36) d=1 93
(4.33 2.61 0 daa
small oyan nitieer cy gyniderr
small eyan nibher cabee
small gjan nuhher cube
(0.37 2:25 0.20)
d 74 large brown rubber cylinder
(0.26
(0.30 2.32 0.20) d-11.17
(0 2.30
d-11.6
lunge syan rubber cube
large syan rubber cube
large syan nubber cube
2.04 0.70)
2.51 0.75) d=2.41 large red rubber cylinder
98 d-0.38
d-n
large yellow rubber cylunder
large yellow rubber cylinder
large yellow rubber cylinder
0.35)
d-0.06
small parple metall eylinder
satall eyan metal eube
small purple metal eylander
small purple matal cytinder
(2002 0,35)
(2.44
d=1,43 brown metal linder
(1 0.36) d-f 06
de0.05
small metal eylinder
brown (2.66 2.41 0.66) d-0.26
small brown metal cylinder
22.93-2,160 0,70 large
(2 89 2.53 0.77) d=0038 large
(2.85 2221 0.68) d-0.10 10 large
cale
large
(-1.73 2:38 0.35)
(-1 51533 0,40) d=0.21
0. 59 0. d=1
(-1.79 2.43 0.37) d-n small green metal cylinder
small greas metal eylinder
small green metal cylinder
small metal cylinder
(-122] cylinder small metal mettllyylindee
2.78 0.33) d
small red motal cylinder
small motal sylinderr
small rod motal sylinder
2.09 0.35)
(0.19 0.54)
1.10000 2.05 d=0,33 rubber sphere
(-0.70
d-0 12 brown rubber sphere
small brown rubber aphere
brown metal sphere
small
88 large purple
(0-01 42 d-0.48 purple metal sphere
(3.13 -2.5 0.71) d-n 04
(0 2000 700 d-0.04
sphere
large purple
sphere
sphere purple
e0 -0.02 0,70) luge blue rulberr
to 29 0.16 cnn d-1
10 0.25 (u64) d=0 28
c0 39 0.09 ouoa d-1.11
large
large
large
(1.252 2 14 35)
(0.81 2 s2 o 51 d-11.50
(1.06 2.23 soy d-i) 22 small cyan metal cylinder
(1 23 12 o 35) d-0
small cyan metal cylinder
small eyan metal cylinder
small cyas metal eylinder
(1 6.22 2.93 0,35) small yellow nibiir cube
(1,76 3,24 0.37) d=0.3 34
(1 68 2.0 21 D, 3,36 d=0.07
(1.68 - 2.966 0.35) d=0.07
small yellow rubber cale
small yollow rubleer cule
small yollow rulber cuhe
2.30 0, large gray metal sphere
(2:23 0,19 0.80 d=0.24
(2.25 0.41 0.74) de0,05
(2.30 0.35 0,700 d=0.0s
large gny metal sphere
lirge gray metal sphere
large gny metal sphere
(:2.33 0.68 0.70)
(2771 -L83 0.69) d=1.15
(-2.75 - 59 70) d8
(-2.6 0.69 0 20) d20.03 lage oyan
large cyan sphere
large oyan
spliere
large eyan
sphere
aptire
(-2.99 1.5t 0.700
(2.55 1,25 ee d=0 54
(:2.06 56 0 60) d-n p4
(:2.04 54 0 21) d-n 05 cube 2200 2,69 69) d=0006 lagg cyan nihhher sphere
large
lange
large
cube
lange
(:2.00 2,71 0.50)
(0.6 1199 0.57) d=1.55
(-2.09 2.71 68) d=0.09
largo cyan rulberr sphere
large syan metal kyhhre
lorge cyan nither splire
(:1.6) 2,84 0.35)
(-1 35 2,36 0.30) d-0.59
(-1.5 16 2.78 0.57) d-i 15
(-) 1s so 0.30) d-0.1
small yellow metal cube
small yellow metal cylinder
small yellow metul cylinerr
small yellow metall ey linder
(-0.18 2.91 0. 70)
(-0m1 2.21 (00 dm0,72
(-0.17 2.87 a d-0 04.04
(-0.18 2.89 0.60) d-0.03
large gray zubber cubo
large pray rubber cube
large gray rubber cube
lurge gray rubber cube
(0, 12-2.40 35)
(0.62 -1 1,77 o (2) dea 84
(0.26 2 0.31) d=0
(0.14 1.34 0.35) d=0.07
small blue metal cylinder
small blue metal cylinder
smill blue metal cylinder
small blue metal cylinder
1.76 0,700
(-0 2
d-1 lange yellow metal cule
0.7%) 11 large yellow metal cy synderr
(0 31
large yellow metal cylinder
large yellow motall cylinder
0.35)
36
d-n
small biue metal cyllnder
small blue metal cylinder
smill beue metal cyllnder
small blue metal cylinder
(1.97 0.70)
(1.99 2252
red nubber cobe
large red nibber cube
lange rubber cabe
lingo red rubber cube
(-282 2.97 0.70)
(-2.2 3.07 (.555 d=0 21 large eyan nubber cylinder
(-2.78 3,30 02 0,26) d=0.08
(-2.82 2.92 0.69) d=0.05 large eyan nublerr cyllnner
lare epan nihher ey cylinder
large eyan ruliber cylinderr
(-2.28 2.62 (,20)
de1 202 large purple metal cylinder
(-1.6 86
(-1 22 2.70
d-il
(-2 24 2.56 0, deous lange purple motal by lylinerr
bege purple metal sylinder
large purple metal cyllnder
(-1.71 -0.51
(-0.65 -.98 n.
(-1.77 -0.4
d-o 20 small eyan metal cylinder
-047 17) d-0.0s small eyan metal cylinder
small eyun metal cylinder
small brown metal cylinder
d=)
small yellow suliher cylinder
small yellow nubher cylinder
small yellow nubber cyllinder
small yellow rubber cylinder
(0.001 small rubber cylinder
(1.11 -1.28
(8 in 0.31 11)
(0.07 -1.32
rubber cylinder
cyun
cylinder
zubber cylinder
(0.43 2464 70) large puple
(0 ..5k 0.60) d=n 17 eylunder
(0. 50 2.71 (.63) d=0, cylinder
(0.42 2.73 0,69) d-n
cylinder
large purple
large purple
large purple
cylinder
(1.32 -1,00 0, 70)
(- 1.00 .239 0.59) d=2 46 cylinder
(1.0d -1,48 0.51) d=o 41
(1.30 -1.66 0.68) d-0.00 cylinder
large (2.41 0.23 (.35) small red metal sphere
oy sylinder
lange purple
large
(2 24 0 36 0.4R) d-11.42 metal sphere
(2.57 0.17 (136) d-0.08 small nod metul sphere
(2.41 0.22 0.35) d-1.01
small
sphere
(-2 17.87-0 gray mital sphere
11 a d=90.9
(-2 86, -2 02 gny metal
d=
(-2.5 -203 0.35) d=002
small
(.2.8 2,14 0.35)
(-2.62 2.17 o (9) d-0,
(-2.76 2 06 0.35) d-0 09 parple
(-2.80 2.15 0.37) 4-0.03 purple
small puppic
euhe
purple
1:2.33 - 0,03 0.20)
(-1.72 -1 06 (156) d= 63
(-2. 26 -n ae 64) dwo 11
(-2 26 .0.5 U 68) dwo 10
large brown
by cylinder
large
cylineer
large brown
large
[-1 -2 79 o 35)
(-2. 30 98 52) d=0u62
(-1.78, -2 75 o. s0 d-
(-1 SS 78 n. 34) d-ons
small gray rubber cube
small gaay metal cube
small gray rubber cube
small gray rubber
(-1 31 0.35)
(-0 2.02 21 d=0.81
(-143 2.26 0 50) d small blue menal cabe
(-1 38 16 3.44 d=0.59 cube
small metal
small bloee metal cube
small bine
small beown metal sphere
(-118 1.77 0 15)
small brown motal sphere
(-1 67 1.90 o 48) d=0.56
(-1. 1.72 0.35) d-O 0.88 small brown metil sptiee
small brows motil sphere
(-1, 14.14 1.22 94) d-9:06
(-0.0) 2.36 0. 70)
40.62 2.39 6767) d=0.46
(-0.17 2 39 7.71) d=0
(-0.22 2.45 0.71) d=0
large blue ribber cube
large blue subher cube
large bise rubber cube
large blue nibber cube
(0.45 1.73 0.70)
(0.15 1.75 0.0 d=0.30
10.00 1.72 0.72) d-0.06
10.44 1.68 0.69) d-0,06
large goy witter cube
large gnay mubber cube
largo gray niber cube
lange gray mibber cuhe
34


Reviews:
[{"review_id": "ruSNPZYKmdW", "reviewer": "Reviewer_KAp5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n\n* The proposed methodology of using multiset equivariance makes sense and is a meaningful difference from standard set equivariance. This is a very good catch, as the difference is subtle but meaningful.\n* The empirical results are very strong. In particular, the results in the tables 2 and 3 show that the proposed model is able to actually learn (instead of being stuck at 0)\n* The contribution of using implicit differentiation for the DSPN is very useful for the field.\n\n## Weaknesses\n\n* The motivation for multiset equivariance is often confusing and/or very handwavy. Specifically, the push_apart function only serves to show what multiset equivariance is, but I'm not sure how this motivates any real world example. If possible, having a more real-world example would be much better. Similarly, the \"the problem with set-equivariance\" paragraph also suffers from a lack of specificity. A simple toy experiment here would be very helpful.\n\n## Nitpicks\n\n* The figures should be ordered correctly and stretch across too many pages. Currently, the text talks about, in order, the results in figure 1, table 2, and table 3. However, the actual figures appear in the order of table 2, figure 1, table 3. Furthermore, the figures do not necessarily all need to be at the top of a new page, especially as table 2 pushes too far away from its source text.\n* The paper should include a related work section discussing more implicit differentiation works. In addition to classical deep implicit layers works like [1, 2], I also recommend implicit layers work on other non-Euclidean domains beyond sets such as [3, 4].\n* The difference between the two \"Slot Attention\" rows in table 3 is not made clear.\n\n## References\n[1] https://arxiv.org/abs/1607.05447\n\n[2] https://arxiv.org/abs/1910.12430\n\n[3] https://arxiv.org/abs/2003.00335\n\n[4] https://arxiv.org/abs/2009.06211", "summary_of_the_paper": "The paper proposes the usage of multiset equivariance as a weaker constraint for deep sets. The authors then modify DSPN to be multiset equivariant and introduce the usage of Jacobian-free implicit differentiation to speed up computation. Finally, the authors compare on two test tasks and show general improvements.", "summary_of_the_review": "Overall, I found the paper to be a strong addition to the deep sets literature with many important contributions. However, I did have some overall complaints about the general structure given in the weaknesses and nitpicks sections above. As such, I am a bit reticent to accept. If these issues are resolved, I am more than happy to increase my rating.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "QlfUOJP_6qM", "reviewer": "Reviewer_MvMj", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- The paper is very well written and easy to understand. It is also correct as far as I can tell (apart from a tiny thing I noticed below).\n- Technically, it builds on top of the nicely carved out, yet somewhat subtle difference of multiset-equivariance and set-equivariance. The proposed improvement of DSPN with implicit differentiation is somewhat incremental: it combines existing ideas from implicit differentiation with an existing method, i.e. DSPN.\n- The experimental section illustrates large improvements in sample efficiency (e.g. vs. BiLSTMs or Transformers) on designed toy data. Also on already used data (CLEVR as in Locatello et al), significant improvements are demonstrated. As this is the only more realistic dataset where improvements are shown, it would be great to extend the analysis to another dataset.\n\nMinor Issues.\n- in eqs 8,9,10 : \\nabla_Y L(Y*,...) is trivially zero; what is meant is likely either \\nabla_{Y*} or L(Y,...) ?\n- pg. 8 Results. \"Slot Attention\\dagger\" is not clear what is meant from the main text. Please refer to Table 3.", "summary_of_the_paper": "The paper \"Multiset-equivariant set prediction with approximate implicit differentiation\" points out a difference between typically looked at set equivariance and multiset equivariance, the latter being able to produce different outputs for identical inputs. It is further shown that a method from the literature (DSPN) is multiset equivariant; to make this method perform better, this method is improved with an implicit differentiation approach. The superiority of the method is demonstrated on toy data, and also on CLEVR.", "summary_of_the_review": "The paper is well written, nicely connects a novel theoretical observation with a modeling improvement idea, which then culminates into a clear empirical improvement. More empirical investigations would be great, but I feel the paper is well rounded as is.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "BMD89a9RgDg", "reviewer": "Reviewer_4sjy", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths\n- The paper is well-written and understandable.\n- Implicit DSPN leads to very strong performance and is a good application of the recent progresses in deep learning with implicit differentiation.\n- The identification of the conceptual advantages of exclusive multiset-equivariance could serve as guidance towards which models future research will focus on. However, some issues have to be resolved, see weaknesses.\n\n# Weaknesses & Open Questions\n- Questions regarding importance of exclusive multiset-equivariance:\nIs exclusive multiset-equivariance really a deciding property in practice? The authors make the following two points: \n1) Set-equivariant networks cannot separate duplicate elements. This raises the question how often this happens in practice. The authors mention smoothing in graph problems (this is probably unlikely to produce exact duplicates), duplicates in the data (this is clearly the case in the toy experiment by design, but not obviously the case for the CLEVR benchmark), zeros from relu activations (to the best of my knowledge relu is not often used as the final nonlinearity) and duplicates from projections. It would be interesting to see how often duplicate elements actually show up, e.g. by reporting the percentage of instances for which this happens during the forward pass of a trained model on the CLEVR benchmark.\n2) ‘set-equivariance […] places a limitation on [the set prediction models] modeling capacity even if exact duplicates are never encountered’. I agree that very similar elements will be a problem for many set-equivariant networks that try to approximate functions such as push-apart, however, there are also set-equivariant networks for which this is not the case. An example would be to employ a sorting operation that instead of randomly breaking ties outputs the mean of the duplicate elements multiple times, which I believe makes the jacobian set-equivariant but is still able to separate arbitrarily similar (non-equal) elements. A version of DSPN/iDSPN with this modified sorting operation (possibly with noise added to inputs to break exact ties) would be an interesting comparison in the experiments as well, as it is the closest related architecture that breaks exclusive multiset-equivariance. If (even with added noise) this modified architecture fails to achieve similar performance this would really show the advantage of the exclusive multiset-equivariant architecture. \n\n- While producing impressive results, the technical contribution of iDSPN in itself is not very strong, as it is a straightforward application of implicit differentiation to DSPN. ", "summary_of_the_paper": "This paper the discusses the conceptual improvements of exclusive multiset-equivariance over set-equivariance in the context of set prediction networks. Deep Set Prediction Networks (DSPN) are identified as the only used architectures for set prediction that satisfy exclusive multiset-equivariance with a specific choice of encoder that employs sorting. As an answer to DSPN not being the leading architecture in terms of performance on set prediction benchmarks, iDSPN is proposed as an improved version of DSPN, which still satisfies exclusive multiset-equivariance and builds on implicit differentiation to employ better optimizers and reduce memory usage and computation time. The experiments highlight the benefits of exclusive multiset-equivariance in toy tasks and demonstrate that iDSPN significantly outperforms the state-of-the-art on the CLEVR object property prediction benchmark.", "summary_of_the_review": "The conceptual advantage of exclusive multiset-equivariance is interesting but lacks some additional comparisons in order to demonstrate that it is the deciding property in practice.\nImplicit DSPN works very well but on its own is just an application of existing methods to a new task, which is not necessarily bad but limits the technical novelty of this part of the contribution. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "8ywxITjZ3SX", "reviewer": "Reviewer_PV3V", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors consider an interesting problem in the construction of neural network models designed to be applied on set objects. The contribution of the authors is two-fold: 1) the authors formalize a discussion on multiset-equivariance and the limits of set-equivariance, and 2) the authors propose an improvement upon deep set prediction networks through the use of implicit gradients. In general, I found the treatment of (1) to be interesting, but somewhat inadequate as there appears to be remaining mathematical issues (see details below). Additionally, I was not convinced that multiset-equivariance is of practical importance in general problems of interest. On the other hand, I found the treatment of (2) to be a valuable contribution, with promising empirical results.\n\n**On the treatment of multiset-equivariance**\nI was confused by some claims on multiset-equivariant (but non-set-equivariant) functions made in the paper, as not all the constructions used in the paper appear well-defined. For example, the construction in Appendix A (Jacobian of sorting) appears particularly ill-defined, as the authors refer to “the permutation matrix”, which is not unique whenever there are equal elements. One possibility here would be to e.g. consider a stable sort, which breaks ties in the sequential order of the input, but this is a function on tuples and not sets, and I believe would be set-equivariant in any case.\n\nI was also uncertain about the treatment of element labels in the context of multiset-equivariant functions where equal elements may take different values. In particular, when we have labels attached to each input element, allowing non set-equivariant functions then would require a further permutation invariant matching step to reconcile labels of equal elements (as the outputs for equal elements are produced in arbitrary order). Given this, it would be great if the authors could address how this problem is tackled in their applications to counting equal elements in a multiset, and also provide a generic framework for tackling this important issue in using non set-equivariant functions.\n\n**On the empirical importance of multiset-equivariance**\nIt is not clear to me that typical set prediction problems are particularly constrained by the set-equivariance assumption used in the model. Indeed, as the authors note, slot attention (which is set-equivariant) performs better than DSPN on some of the tasks at hand. Perhaps a comparison of a set-equivariant and a strictly multiset-equivariant iDSPN architecture on a variety of tasks could help demonstrate the importance of being able to assign different values to identical inputs.\n\n**On the empirical results**\nI generally found the empirical results promising, although I believe it could be improved by including some of the following. 1) In table 2, a natural question would be if the apparent gap in required sample size could be mitigated or closed through data augmentation for non-equivariant models. 2) In table 3, it would be helpful to also include a like-for-like comparison of the iDSPN model with the same number of iterations as the DSPN model, to assess the potential degradation due to inaccurate gradient, and also compare the speed-up in a like-for-like scenario.\n", "summary_of_the_paper": "The authors examine the problem of building multiset-equivariant neural networks, and note that there exists multiset-equivariant functions which are not set-equivariant. Given this observation, using set-equivariant models for multiset prediction tasks may be limiting. The authors propose a strategy to construct multiset-equivariant functions (which are not set-equivariant) through the Deep set prediction network framework. The authors propose to improve upon that framework by making use of implicit differentiation to compute a backwards mode gradient without backpropagating through the optimization process which is memory- and compute- intensive. Some empirical results are provided for a simulated toy dataset, as well as a simulated object property prediction problem.\n", "summary_of_the_review": "The authors tackle an intriguing problem in modelling functions on multi-sets. The paper is overall promising, and I would be in favor of accepting if the authors could clarify some issues on the theory and add some more details in the empirical results.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper points out how set equivariant functions limit the types of functions that can be represented on multisets. They develop an new notion of multiset equivariance to address this limitation. The paper improves an existing multi-set equivariant Deep Set Prediction Network through implicit differentiation, which is an area of rising interest. The reviewers and I note that the paper is well written.
Decision: Accept (Poster)