Universal Mini-Batch Consistency for
Set Encoding Functions

Set encoding functions (Zaheer et al., 2017; Bruno et al., 2021; Lee et al., 2019; Kim, 2021) are becoming a broadly researched and cited topic in recent literature. This popularity can be partly attributed to natural set structures in data such as point clouds or even datasets themselves. Given a set of cardinality $N$, one may desire to group the elements (clustering), identify them (classification), or find likely elements to complete the set (completion/extension). Different from vanilla neural networks working on fixed input sizes, neural set functions must be able to handle dynamic set cardinalities for each input set. Additionally, sets are inherently unordered, so the function must make consistent predictions for any permutation of set elements.

Deep Sets (Zaheer et al., 2017) is a canonical work providing an in-depth investigation of the requirements and valid structures of neural set functions. Deep Sets utilizes traditional, permutation equivariant (creftype 3.2) linear and convolutional neural network layers in conjunction with featurewise permutation invariant (creftype 3.1) set-pooling functions (e.g. {min, max, sum, mean}) in order to satisfy the permutation consistency requirements and perform inference on sets. The Set Transformer (Lee et al., 2019) utilizes the power of Multi-Headed Attention (Vaswani et al., 2017) to construct multiple set-capable attention blocks, as well as an attentive pooling function. The previously mentioned works never explicitly consider the case where it may be required to process a set in multiple partitions, which can happen for a variety of reasons including resource constraints, prohibitively large or even infinite set sizes, and streaming set data.

Powerful models such as the Set Transformer cannot make consistency guarantees when updating pooled set representations, as self-attention blocks require all $N$ elements in one pass, and therefore do not satisfy MBC (i.e. batch processing of set partitions yields a different output than processing the whole set at once). Naively using such non-MBC set encoders in an MBC setting causes under-performance, as depicted in Figure 3 (\subrefst-single-point-\subrefst-chunk) where Set Transformer exhibits poor likelihood and inconsistent predictions. With an MBC guarantee Figure 3 (\subrefumbc-st-single-point-\subrefumbc-st-chunk), UMBC+Set Transformer gives consistent results, and vastly better likelihood (See Appendices B and 5 for details of the experiment) The quantitative effect of MBC vs non-MBC encoding on pooled set representations can be seen in Figure 2 which shows the variance between final representations of 100 random partitions of the same set. (See Appendix C for details).

The MBC property of set functions was identified by Bruno et al. (2021) who also proposed the Slot Set Encoder (SSE), a specific, constrained version of an attentive pooling mechanism, eliminating the need to store all $N$ set elements during the computation. The introduction of the MBC property naturally leads to the rise of a new dimension in the taxonomy of set functions, namely, those which satisfy MBC and those which do not. The main limitation of the SSE is the fact that it limits the number of valid MBC architectures, eliminating powerful models such as the Set Transformer, which can be the best choice for tasks which require leveraging pairwise set element relationships and self-attention. This is shown in Tables 2 and 4) where the Set Transformer outperforms SSE. In this work, we identify, prove, and verify that there is a universal way to convert arbitrary non-MBC set functions to MBC functions which can provably produce the same result for random partitioning schemes, allowing any set encoder to be used in an MBC setting. This result has large implications for all current and future set functions which do not natively satisfy MBC, as it unifies all set-functions into the MBC class of set-functions. This unification allows models to scale to sets larger than they could otherwise handle, and also allows them to be used in a wider range of settings (i.e.streaming data). Animations, code, and unit tests can be found in the supplementary file and also at: https://github.com/anonymous-subm1t/umbc

Our contributions in this work are as follows:

• In Theorem 4.1 we show that with a change in architecture, any arbitrary non-MBC set encoder can become MBC, guaranteeing that minibatch processing of sets gives the same result as processing the full set at once.

• We loosen the constraints of the SSE attention activations by showing many functions can be used. By factorizing the activation, we can maintain the MBC property and still normalize over $N$ (i.e. like traditional attention (Vaswani et al., 2017)).

• We uncover a connection between the pooling mechanism of the Set Transformer and SSE layers, which we show only differ in the attention matrix activation. We explore the effect of 5 different activation approaches.

• We explore an interesting dropout approach which arises as a consequence of UMBC’s structure, delivering improvements in calibration for both in-distribution and corrupted test sets.

Processing, pooling, and making a prediction for set structured data has been an active topic since the introduction of DeepSets Zaheer et al. (2017). Attention has been shown to be powerful in these tasks Lee et al. (2019), as simple independent row-wise operations may fail to capture pairwise interactions between the set elements. There have been subsequent works and variations of set attention which draw connections to optimal transport Mialon et al. (2020), and expectation maximization Kim (2021). Likewise, an efficient version of set-attention has been proposed which incorporates cross attention with lower dimensional self-attention in an iterative process Jaegle et al. (2021). Outside of attention, other approaches to set pooling functions include featurewise sorting Zhang et al. (2019), and canonical orderings to permutation sensitive functions Murphy et al. (2018).

Bruno et al. (2021), provide and especially important lens through which to view our work. Prior to the proposal of the MBC property, previous set function research never explicitly considered the mini-batched setting, which will likely become important with the ever increasing scales of models and data (Brown et al., 2020). Indeed most set functions do not satisfy creftype 3.3 (e.g. (Lee et al., 2019; Kim, 2021; Mialon et al., 2020; Jaegle et al., 2021; Zhang et al., 2019; Murphy et al., 2018)). Our work builds on the concepts established by Bruno et al. (2021), and ensures that any set functions proposed in the future, can be considered in terms of their MBC performance by incorporating UMBC.

Numerous prior works (Ovadia et al., 2019; Guo et al., 2017) focus on uncertainty quantification and improving probabilistic calibration, which can be crucial for tasks such as autonomous driving (Chen et al., 2017) and medical diagnosis (Zhou et al., 2021) where decisions can impact human well being. Guo et al. (2017), proposed quantifying uncertainty with the expected calibration error (ECE) metric measuring the mismatch between accuracy and confidence. Ovadia et al. (2019) used corrupted datasets made by Hendrycks and Dietterich (2019) (similar in form to ModelNet40-C (Ren et al., 2022) used in our experiments) to survey the landscape of neural network calibration. Guo et al. (2017); Ovadia et al. (2019) analyze variants of deep convolutional models, while Minderer et al. (2021) evaluate large Vision Transformers. To our knowledge, our work is the first to analyze set function calibration specifically, as most other works focus on general purpose classifiers.

For our setting, we define a neural set function $f$ which operates on a set $X=\{x_i{\}}_{i=1}^N$ with each set element $x_i\in R^d$. A dataset of set-structured data $D=\left\{\right(X_i,Y_i){\}}_{i=1}^M$ is itself a set composed of sets forming input sets $X_i$ and output sets $Y_i$ which can be learned by an appropriate function via mini-batch stochastic gradient descent. A set function $f:X\mapsto Y$ has a set-structured input space $X$ and output space $Y$, which may be discrete or continuous. An input to a set function, $X_i$ is a set. As the input is a set and the function must process any valid set, therefore any element of the powerset $P\left(X_i\right)$ also represents a valid input.

Deep Sets Zaheer et al. (2017) provided a crucial groundwork for neural set functions, formalizing the requirements of permutation equivariant architectures and invariant pooling mechanisms necessary for feature extraction and pooling of sets. Following these requirements, a function can assign a single output for each valid subset $X_i\in P\left(X_i\right)$ which is invariant to the permutations of the elements $x_{j\in N}\in X_{i\in M}$.

Permutation invariant layers are commonly referred to as set pooling functions, and have a stationary, fixed size output given any permutation, or cardinality of the input set, respectively. This stationary output can be seen as a Set to Vector function.

Additionally, Zaheer et al. (2017) prescribes that prior to any permutation invariant pooling, any composition of permutation equivariant layers may be used for feature extraction. Common linear and convolutional neural network layers are permutation equivariant layers when considering a batch of inputs as a set. For the remainder we assume $f$ contains both equivariant and invariant layers.

Lee et al. (2019) identified that Self-Attention (Vaswani et al., 2017) blocks satisfy creftype 3.2 and thus can be used as equivariant feature extractors for set functions, proposing the Set Transformer. SAB’s are defined as $\text{SAB}(X,X)=\text{Attention}(X,X)$. Additionally, the permutation invariant pooling layer of the Set Transformer (PMA), performs the attention operation between a learnable seed parameter $S\in R^{K\times d}$ and the input set, $\text{PMA}\left(X\right)=\text{SAB}(S,X)\in R^{K\times d}$.

Bruno et al. (2021) identified and formalized the MBC property, proposing the MBC Slot Set Encoder (SSE), adding a new dimension to the original view of creftype 3.1 from Zaheer et al. (2017). Instead of merely requiring that $f$ be permutation invariant for any permutation of the indices of a specific subset $X_i\in 2^{X_i}$, the MBC property also requires that sequential, mini-batched extraction/pooling, and subsequent aggregation of any partition of $X_i$ is also permutation invariant.

An SSE (Bruno et al., 2021) layer works in a similar fashion to the PMA layers of the Set Transformer, using parameterized slots $S\in R^{K\times d}$ as queries $Q$, and partition $X_{j\in P}$ as keys $K$ and values $V$, with an attention activation for a single $z_{i\in N}$ which does not depend on the other $N-1$ elements within the set. SSE uses a sigmoid activation with normalization $\sigma$ over the slot dimension $K$ (Bruno et al., 2021; Locatello et al., 2020) in the attention matrix $A=\sigma \left(Q{K}^{\top }\right)$. Then with $A_j\in R^{K\times |X_j|}$ and $X_j\in R^{|X_j|\times d}$

thereby allowing any partition scheme, and satisfying creftype 3.3. With the prior description in mind, an SSE can be phrased in terms of a PMA such that $\text{SSE}(S,X)=\widetilde{\text{SAB}}(S,X)$, with $\widetilde{\text{SAB}}$ signifying the use of the slot-normalized sigmoid activation used in order to satisfy creftype 3.3.

Originally, SSE acts as a S2V function, creating an encoded set representation for downstream handling by a task specific decoder. Decoders make different predictions given different representations, therefore creftype 3.3 need only be satisfied until the invariant S2V pooling function.

Put simply, Lemma 4.1 states that every module $f^{\ast }$ coming after a module which satisfies creftype 3.3 will continue to satisfy creftype 3.3, even though $f^{\ast }$ itself may not satisfy creftype 3.3. With this established, we can therefore use Lemma 4.1 in order to build a universally MBC set function.

In the simplest setting where the pooled representation is a vector $z$, $f^{\ast }$ receives a singleton set as input, which is valid, but may provide limited utility over a Deep Sets style encoder, as $f^{\ast }$ sees only a single element. SSE’s and PMA’s, however, output a set $Z\in R^{K\times d}$. Therefore, using a SSE/PMA layer with $K>1$ as the base module $f$, we can view SSE/PMA layer as a type of invariant feature extractor which takes a set of cardinality $N$ and maps it to a set of cardinality $K$. Our UMBC has features flowing through the model from the input space $X$ to the output space $Y$ as,

In this work, we have shown that composing a set function consisting of a mini-batch consistent base $f$, with an arbitrary set function head $f^{\ast }$, we can make the composition $F=f^{\ast }\circ f$ universally mini-batch consistent. We have provided proofs in Theorem 4.1, experiments Figure 2, and unit tests (included in the supplementary file) which prove our assertions. Likewise we have loosened the known constraints on the structure of the SSE Proposition 4.1, establishing an equivalency to the PMA layers of the Set Transformer. We have demonstrated that there are cases where a UMBC $F$ outperforms previous simpler MBC models, and explored an interesting dropout strategy which is made possible by our architecture and improved the calibration and NLL of UMBC. As the field of set-functions continues to widen, we look forward to seeing future research in the area of MBC set functions.