Multiset-Equivariant Set Prediction with Approximate Implicit Differentiation

1 INTRODUCTION . An example of set prediction is detecting the objects in an image : we want a collection of objects as output , but there is no inherent order to them . Multisets and sets are natural representations for such data due to their orderless nature , with the difference between them being whether duplicate elements are allowed or not . In general , they are useful for representing collections of discrete and composable elements , such as agents in an environment , instruments in a piece of music , or atoms in a molecule . The usual approach to representing a multiset or set of vectors in memory is to store it as a list in which the elements are placed in an arbitrary order . To still treat the list as if it were a multiset or set , recent set prediction models avoid being affected by this arbitrary order by using list-to-list functions that are enforced to be permutation-equivariant ( Zaheer et al. , 2017 ) : changing the order of the input elements should change the order of the output elements in the same way . Now , consider this function that takes a multiset ( represented as a list ) of two scalars as input and “ pushes them apart ” . push apart ( [ a , b ] ) = { [ a− 1 , b+ 1 ] if a ≤ b [ a+ 1 , b− 1 ] otherwise ( 1 ) For example , push apart ( [ 1 , 2 ] ) = [ 0 , 3 ] and push apart ( [ 2 , 1 ] ) = [ 3 , 0 ] . At first glance , it appears that this function is permutation-equivariant because swapping the inputs swaps the outputs . However , notice that push apart ( [ 0 , 0 ] ) = [ −1 , 1 ] . Swapping the inputs has no effect in this case , so the outputs remain unchanged , thus push apart is not permutation-equivariant . This is a bit strange – the function only violates the equivariance property for equal input elements , in which case the order of the equal input elements is irrelevant anyway . How can we reconcile this observation with the standard definition of permutation-equivariance ? Is this a 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 : ∗Equal contribution Table 1 : Difference between equivariance properties in the context of multisets . Only when satisfying exclusive multiset-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 push apart . f is setequivariant f is multisetequivariant f ( [ a , b ] ) = [ c , d ] ⇔f ( [ b , a ] ) = [ d , c ] f ( [ a , a ] ) = [ c , d ] X X guaranteed 3 not possible 7 7 7 not guaranteed 7 possible 3 7 X guaranteed 3 possible 3 1 . Multiset-equivariance . We begin Section 2 by making the important observation that many existing models for set prediction thought to operate on sets actually operate on multisets . We then show that these models are too restricted in the functions they can represent on multisets due to the limitation of being permutation-equivariant , which we call set-equivariance . Due to continuity , this is a problem even when there are no exact duplicate elements . Thus , we develop a new notion of multiset-equivariance , which we argue is a more suitable property for operations on multisets . To avoid the aforementioned limitation of set-equivariance , a crucial point is that it is possible to be multiset-equivariant without being set-equivariant , which we call exclusive multiset-equivariance . We identify and show that deep set prediction networks ( DSPN ) ( Zhang et al. , 2019 ; Huang et al. , 2020 ) satisfy this property with a specific choice of set encoder . This makes it the only set prediction model we are aware of that can satisfy exclusive multiset-equivariance . 2 . Implicit DSPN . Despite this beneficial property , DSPN is outperformed by the set-equivariant Slot Attention ( Locatello et al. , 2020 ) , which motivates us to improve other aspects of DSPN . We propose implicit DSPN ( iDSPN ) in Section 3 : a version of DSPN that uses approximate implicit differentiation . Implicit differentiation enables us to use better optimizers and more iterations at a constant memory cost and less computation time . The approximation makes this faster to run and easier to implement by avoiding the computation of a Hessian while still having the same benefits . In our experiments in Section 4 , 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 sampleefficient . Next , we evaluate the second point by testing the modeling capacity on autoencoding random sets , where iDSPN performs similarly to DSPN at the same iteration count and much better for the same computational cost . Lastly , iDSPN significantly raises the bar on CLEVR object property prediction , outperforming the state-of-the-art Slot Attention ( Locatello et al. , 2020 ) by a large margin of 69 percentage points on the AP0.125 metric while only training for 7 % of the number of epochs . 2 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 a type of permutation-equivariance , set-equivariance . We argue that set-equivariance is too strict and instead propose a relaxation of set-equivariance that is more suitable for set prediction models : multiset-equivariance . We end by discussing how an existing method , deep set prediction networks ( Zhang et al. , 2019 ; Huang et al. , 2020 ) , is multiset-equivariant but not set-equivariant and the associated benefits thereof ( Table 1 ) . 2.1 PRELIMINARIES . Sets in the context of deep learning are typically unordered collections of real-valued vectors . They are represented in memory as n × d matrices with the n set elements in an arbitrary order , each element being a d-dimensional vector . Since elements in the matrix representation are typically not restricted to be unique , the term multiset is more appropriate . To still treat these matrices as multisets , operations must not rely on the arbitrary order of the elements . The two main properties to accomplish this are permutation-invariance and permutation-equivariance ( Zaheer et al. , 2017 ) . In short , when changing the order of the elements in the matrix , the output of a function should either not change at all ( invariance ) or change order similarly ( equivariance ) . 2.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-equivariant list-to-list ) functions to refine an initial set Y0 , which is usually a randomly generated or learnable matrix . Permutation-equivariance is desirable when processing sets because it prevents a function from relying on the arbitrary order of the set in its matrix representation . Such functions can be easily composed to build larger models that remain equivariant , which fits well into deep learning architectures as building blocks . However , without having any restrictions on duplicate elements , it is more accurate to think of set prediction models as using multiset-to-multiset functions . This difference subtly changes the appropriate notion of equivariance for sets versus multisets , which will become important for our discussion on their limitations . Set-equivariance . First , we generalize the standard notion of permutation-equivariance for sets to multiset inputs and refer to it as set-equivariance . This directly matches the usual definition of permutation-equivariance in the context of set inputs 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 1 ( Set-equivariance for multiset inputs ) . A function f : Rn×d1 → Rn×d2 with multiset inputX is set-equivariant iff for any permutation matrix P , f ( PX ) = P f ( X ) ( 2 ) A consequence is that elements in a multiset that are equal before a set-equivariant function must remain equal after applying the function , as shown by the following : Proposition 1 . For all permutation matrices P1 , P2 such that P1X = P2X , a set-equivariant function f satisfies : P1f ( X ) = P2f ( X ) ( 3 ) Proof . P1f ( X ) = f ( P1X ) = f ( P2X ) = P2f ( X ) ( 4 ) P1X = P2X can be satisfied by P1 6= P2 when there are duplicate elements inX , so P1f ( X ) = P2f ( X ) means that they must remain duplicates . This inability of separating equal elements demonstrates that the classical definition of permutation-equivariance is overly restrictive when directly applied to multisets as in Definition 1 . Therefore , we seek a more appropriate definition of permutation-equivariance for multiset inputs . Multiset-equivariance . We now consider a slightly different notion of permutation-equivariance that no longer forces equal elements to remain the same : multiset-equivariance . Recall the push apart function from Equation 1 , which we suggest should be a valid equivariant function for multisets , but can not be represented by a set-equivariant function . Therefore , we propose that there should be a distinction between set-equivariance and multiset-equivariance , the latter corresponding to the appropriate equivariance property for the typical multiset-to-multiset mappings used in set prediction . The key difference of multiset-equivariance is that the output order for equal elements in the input multiset can be arbitrary . This leads us to the following definition : Definition 2 ( Multiset-equivariance for multiset inputs ) . A function f : Rn×d1 → Rn×d2 with multiset input X is multiset-equivariant iff for any permutation matrix P1 , there exists a permutation matrix P2 such that P1X = P2X and : f ( P1X ) = P2f ( X ) ( 5 ) Any function that is set-equivariant is also multiset-equivariant ( since we can set P2 = P1 ) , but the converse is not true . The problem with set-equivariance . We argue that set-equivariance is too strong of a property for set prediction . As we have shown in Proposition 1 , it prevents equal elements from being processed differently ( we demonstrate this experimentally in Section 4.1 ) . More importantly , when the model is continuous1 , impossibility to separate equal elements means difficulty in separating similar elements . For example , in order to learn the push apart function , a model has to decide which element is larger . This decision is fundamentally discontinuous , but most models in deep learning can not represent discontinuous jumps . A continuous model must approximate the discontinuity , which means that close to the discontinuity ( where elements are similar ) the modeling error will likely be high . 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 a 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 Y0 , nonlinearities throughout the model zeroing out ( ReLU ) or squashing elements ( sigmoid and tanh ) , the model naturally tending towards equal elements ( perhaps from oversmoothing as is often seen in 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-equivariance . Deep learning typically operates on multisets rather than sets , so the less restrictive multiset-equivariance without set-equivariance should be the desirable property . Most existing set prediction models are continuous and set-equivariant , which limits their modeling capacity even if exact duplicates are never encountered . We refer to models that are multiset-equivariant but not set-equivariant as exclusively multiset-equivariant .

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