Mini-Batch Consistent Slot Set Encoder for Scalable Set Encoding

1 Introduction . Recent interest in neural network architectures that operate on sets [ 14 , 6 ] has garnered momentum given that many problems in machine learning can be reformulated as learning functions on sets . Problems such as point cloud classification [ 13 ] , image reconstruction [ 3 , 4 , 7 ] , classification , set prediction [ 8 ] , and set extension can all be cast in this framework of learning functions over sets . Given that sets have no explicit structure on the set elements , such functions are required to conform to symmetric properties such as permutation invariance or equivariance for consistent processing . A defining property of many practical functions over sets involves an encoding of the input set to a vector representation , the set encoding . This set encoding is then used for downstream tasks such as reconstruction or classification . In DeepSets [ 14 ] , a sum-decomposable family of functions is derived for a class of neural network architectures that encode a given set to such a representation . However , the simplicity of the functions derived in DeepSets makes it ineffective for modeling pairwise interactions between the elements of the sets . Set Transformers [ 6 ] remedy this by using Transformers [ 12 ] to model higher order interactions resulting in rich and expressive set representations . In all these works , there is an implicit assumption in the experimental setup that the cardinality of the set is manageable , and that there are ample computational resources available for processing all the elements during the set encoding process . However in real-world applications such as large scale point cloud classification and data intensive applications in particle physics , the set size can be extremely large . In such cases , even if one has access to a set encoding function that is linear w.r.t . the number of elements in the set , it is still impossible to encode such sets in a single batch since we may not be able to load the full set into memory . Current encoding methods deal with this issue by sampling 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . subsets of the full set ( as a data preprocessing step ) and encoding them to a representative vector , as done in the ModelNet40 [ 13 ] experiments in [ 14 , 6 ] . However , this results in the loss of information about the full set , which is undesirable when huge amounts of monetary resources are invested in obtaining data such as in the Large Hadron Collider ( LHC ) experiments , and hence is imperative that the data is exploited for maximum performance on the given task . As such , it is necessary to have a set encoder that can take advantage of the full data , in a computationally efficient way . To this end , when the set is too large to fit in the memory , or given as a stream of data , we propose to iteratively encode it over multiple rounds , by randomly partitioning it , encoding each random subset ( mini-batch ) with the set encoder at each round , and aggregating all the encodings . However , naive mini-batch set encoding could compromise essential properties of the set encoder such as permutation invariance and equivariance . To use a concrete example , in Figure 2 we show the performance of Set Transformer ( in Negative Log-Likelihood ) on an image reconstruction task . By processing the full set ( in this case concatenations of pixel and coordinate values ) , the model performs as expected ( SetTransformer in Figure 2 ) . However when we instead partition the set elements into mini-batches , independently encode each batch and aggregate them to obtain a single set encoding , we observe performance degradation in the model as shown in Figure 2 ( SetTransformer ( MiniBatch ) in Figure 2 ) . This degradation stems from the attention module in Set Transformer where mini-batch encoding and aggregation results in mini-batch inconsistent set encoding . In this work , we identify and formalize a key property , Mini-Batch Consistency ( MBC ) , that is necessary for mini-batch encoding of sets . In what follows , we present the necessary constraints required for MBC and propose a slot based attentive set encoder , the Slot Set Encoder ( SSE ) which is MBC . SSE largely alleviates the performance degradation in mini-batch encoding of large sets . We depict MBC and non-MBC set encoders in Figure 1 . Our contributions in this work are as follows : • We identify and formalize a key property of set encoding functions called Mini-Batch Consistency ( MBC ) , which requires that mini-batch encoding of a set is provably equal to the encoding of the full set ( in Section 2.2 ) . • We present the first attention based set encoder that satisfies MBC . Our method can efficiently scale to arbitrary set sizes by removing the dependence of self-attention in the transformer through the use of slots which frees the transformer from requiring to process the full set in a single pass ( in Section 2.3 ) . • We perform extensive experiments on various tasks such as image reconstruction , point cloud classification and dataset encoding where we demonstrate that SSE significantly outperforms the relevant baselines ( Section 3 ) . 2 Mini-Batch Consistent Set Encoding with Slot Set Encoders . 2.1 Preliminaries . Suppose we are given a dataset of sets D = { Xi , . . . , XN } where |D| = N and each Xi ∈ Rni×d where ni is the number of elements in Xi and each element x ( j ) i ( that is , the jth element of Xi ) is represented by a d dimensional tensor . For each Xi , the ordering of the ni individual elements is considered arbitrary and could be any permutation of the indices . We further assume that both ni and d can be large enough such that processing an instance Xi is prohibitive both in terms of memory and computation . As an example , the processing task could be encoding each element Xi to a representative vector to be used for downstream tasks such as set reconstruction , set classification or set prediction . Current set encoding methods such as Zaheer et al . [ 14 ] and Lee et al . [ 6 ] deal with such large sets by sampling a smaller subset of the original set as a preprocessing step . In essence , they approximate the full set with a random subset which can be undesirable when we want to exploit all the data at our disposal for maximum performance . In this section , we present a method for encoding arbitrarily sized set data via a mini-batch encoding scheme . Our method can then iteratively encode subsets of Xi and aggregate the subset-encodings to obtain the full set representation . Our method is Mini-Batch Consistent , invariant to the order of the subsets , and invariant to permutations on the set elements . For clarity , when we say a set Xi is mini-batch processed , we mean that a random partition function can be applied to Xi to obtain proper subsets of Xi . The subsets , which we call mini-batches , are then processed independently and aggregated to obtain a representation for Xi . 2.2 Mini-Batch Consistent Set Encoding . We consider a single element Xi ∈ D with ni elements each with d dimensions . We truncate notation and write Xi , ni as X and n respectively since we present our algorithm for a single sample here . Our ultimate goal is to encode the set X to a representation Z ∈ Rd′ where d′ is the dimension of the encoding vector . We assume that both d and n are large such that it is not feasible to process X as a whole . To get around this computational problem , we propose to perform mini-batch encoding of X where the mini-batch samples are with respect to the number of elements n. Specifically , we partition X such that X can be written as a finite union X = X1 ∪X2 ∪ . . .∪Xp and each Xi in this partition has ni = |Xi| elements and p is the total number of partitions . We further assume that X is partitioned such that we can efficiently process a given mini-batch Xi in the partition . Each partition Xi is encoded with a set encoding function f ( Xi ) to obtain Zi . We then define the set encoding problem as follows : Z = g ( f ( X1 ) , f ( X2 ) , . . . , f ( Xp ) ) ( 1 ) where g ( · ) is an aggregation function that takes individual Zi ’ s and combines them to obtain Z . We further require that the functions f and g are Mini-Batch Consistent . Property 2.1 ( Mini-Batch Consistency ) . Let X ∈ Rn×d be partitioned such that X = X1 ∪X2 ∪ . . .∪Xp and f : Rni×d 7→ Rd ′ be a set encoding function such that f ( X ) = Z . Given an aggregation function g : { Zj ∈ Rd ′ } pj=1 7→ Rd ′ , g and f are Mini-Batch Consistent if and only if g ( f ( X1 ) , . . . , f ( Xp ) ) = f ( X ) Property 2.1 ensures that no information is lost by encoding the full set in independent mini-batches , and that the aggregation function g guarantees the same output as encoding the full set in a single batch . Additionally , Property 2.1 is general and places no restriction on the form of the functions f and g except that when used together , MBC must be satisfied . In what follows , we define an attention-based set encoding function f together with aggregation functions g that satisfy Property 2.1 . 2.3 Slot Set Encoder . In this section , we provide a formulation for an attention based set encoding function f in Equation 1 which utilizes slots [ 8 ] . Given an input set X ∈ Rn×d , we begin by sampling a random initialization for K slots S ∈ RK×h where h is the dimension of each slot . Specifically , S ∼ N ( µ , diag ( σ ) ) ∈ RK×h , ( 2 ) where µ ∈ R1×h and σ ∈ R1×h are learnable parameters . Optionally , instead of sampling random initialization for theK slots , we can instead designate deterministic learnable parameters Sχ ∈ RK×h . An advantage of random initialization over learnable parameters is that at test time , we can increase or decrease K for the randomly initialized model . In the ablation studies in section 3.1.1 , we explore the performance gains that result from the choice of slot initialization . The key design choice in the slot set encoder is that unlike attention based set encoding methods like Lee et al . [ 6 ] , we compute attention over S ( no self-attention ) instead of computing it over the n elements of X ( self-attention ) . This allows us to remove the dependence of the set encoding function on n , alleviating the requirement of processing the entire set at once , and allowing for MBC encoding of partitions of X since the same slots are used to encode all mini-batches . Specifically , we compute dot product attention between S and X according to the following formulation : attni , j : = σ ( Mi , j ) where M : = 1√ d̂ k ( X ) · q ( S ) T ∈ Rn×K ( 3 ) where k and q are linear projections of X and S to a common dimension d̂ and σ is the sigmoid activation function . There is a key difference in our attention formulation in that we use a sigmoid instead of a softmax to construct the attention matrix . Using a softmax , even when computed over S , breaks MBC as it requires a batch dependent normalization term which will vary across different partitions of X . By using the sigmoid function , we are able to free our model of this constraint and satisfy MBC . Finally , we weight the inputs based on the computed attention matrix to obtain the final slot set encoding Ŝ for X : Ŝ : =WT · v ( X ) ∈ RK×d̂ where Wi , j : = attni , j∑K l=1 attni , l ( 4 ) where v is also a linear projection applied to X and W ∈ Rn×K are the weights computed over slots instead of elements . The normalization constant in W occurs over the slot dimension K , and only contains a linear combination of a single ni ∈ X and slots S and therefore has no dependence on other elements within the partition Xi . The Slot Set Encoder is fully described in Algorithm 1 . The Slot Set Encoder in Algorithm 1 is functionally composable over partitions of the input set X . More concretely , for a given slot initialization and any partition of X , f ( X ) = g ( f ( X1 ) , f ( X2 ) , . . . , f ( Xp ) ) ( 5 ) Where the aggregation function g is chosen to satisfy Property 2.1 . This is convenient since it allows us to define g from the following set of operators : g ∈ { mean , sum , max , min } . We note that to satisfy MBC , we have chosen g to be an associative function but it is entirely possible to formulate a non-associative g and corresponding f that satisfies MBC and we leave that for future work . Proposition 2.1 . For a given input set X ∈ Rn×d and slot initialization S ∈ RK×d , the functions f and g as defined in Algorithm 1 are Mini-Batch Consistent for any partition of X and hence satisfy Property 2.1 . Proposition 2.2 . Let X ∈ Rn×d and S ∈ RK×d be an input set and slot initialization respectively . Additionally , let SSE ( X , S ) be the output of Algorithm 1 , and πX ∈ Rn×n and πS ∈ RK×K be arbitrary permutation matrices . Then , SSE ( πX ·X , πS · S ) = πS · SSE ( X , S ) Proofs of Proposition 2.1 & 2.2 involve showing that each component in Algorithm 1 are both MBC , and permutation invariant or equivariant . Details can be found in the Appendix with discussions on why the baselines fail or satisfy MBC .

Summary: The authors proposes the use of attention based set encoder which can scale to arbitrary set sizes. Attention in Transformers has the quadratic dependence on number of positions (i.e. cardinality of the set). The proposed method makes the problem easier by enabling the use of slots while preserving a property which method defines as mini-batch consistency. The proposed method is invariant to the order of the subsets, and invariant to permutations on the set elements.