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