Universal Mini-Batch Consistency for Set Encoding Functions
Authors: Jeffrey Ryan Willette
Bruno Andreis
Juho Lee
Sung Ju Hwang
Venue: Submitted to ICLR 2023 (2023)
Keywords: set

Abstract: Previous works have established solid foundations for neural set functions, complete with architectures which preserve the necessary properties for operating on sets, such as invariance to permutations of the set elements. Subsequent work has highlighted the utility of Mini-Batch Consistency (MBC), the ability to sequentially process any permutation of a set partition scheme (e.g. streaming chunks of data) while maintaining consistency guarantees on the output, although there are limited options for MBC architectures. We propose a framework which can convert an arbitrary non-MBC model to one which satisfies MBC. In doing so, we allow all set functions to universally be considered in an MBC setting (UMBC). Additionally, we explore a Monte Carlo dropout strategy made possible by our framework which allows performing Monte Carlo dropout on streaming sets while never seeing the entire set at once. We validate UMBC with theoretical proofs, unit tests, and also provide qualitative/quantitative experiments on Gaussian data, clean and corrupted point cloud classification, and amortized clustering on ImageNet. Additionally, we investigate the probabilistic calibration of set-functions under test-time distributional shifts. Our results demonstrate the utility of universal mini-batch consistency, and we further discover that our dropout strategy improves uncertainty calibration.

Under review as a conference paper at ICLR 2023
UNIVERSAL MINI-BATCH CONSISTENCY FOR SET ENCODING FUNCTIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Previous works have established solid foundations for neural set functions, com- plete with architectures which preserve the necessary properties for operating on sets, such invariance to permutations of the set elements. Subsequent work has highlighted the utility of Mini-Batch Consistency (MBC), the ability to sequentially process any permutation of set partition scheme (e.g. streaming chunks of data) while guaranteeing the same output as processing the whole set at once. Currently, there exists division between MBC and non-MBC architectures. We propose a framework which converts an arbitrary non-MBC model to one which satisfies MBC. In doing so, allow all set functions universally be considered in MBC setting (UMBC). Additionally, we explore set-based Monte Carlo dropout strategy which applies dropout to entire set elements. We validate UMBC with theoretical proofs, unit tests, and also provide qualitative/quaniittiiv experiments on Gaussian data, clean and corrupted point cloud classification, and amortized clustering on ImageNet. Additionally, we investigate the probabilistic calibra- tion of set-functions under test-time distributional shifts. Our results demonstrate the utility of UMBC, and we further discover that our dropout strategy improves uncertainty calibration.
INTRODUCTION
Set encoding functions (Zaheer al., 2017; Bruno et al., 2021; Lee et al., 2019; Kim, 2021) have become broad research topic in recent publications. 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 (completionn/xteesionn)) A key difference from vanilla neural networks, is that neural set functions must be able to handle dynamic set cardinalities for each input set. Additionally, sets are considered 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 investigation of the requirements and proposal of valid neural set function architectures. Deep Sets utilizes traditional, permutation equivariant (Property 3.2) linear and convolutional neural network layers in conjunction with per- mutation invariant (Property 3. set-pooling functions (e.g. {min, max, sum, mean)) in order to satisfy the necessary conditions and perform inference on sets. The Set Transformer (Lee et al., 2019) utilizes powerful multi-headed self-attention (Vaswani et al., 2017) to construct multiple set-capable transformer blocks, as well as an attentive pooling function. Though powerful, these works never explicitly considered the case where it may required process a set in multiple partitions at test time, which can happen for variety of reasons including device resource constraints, prohibitively large or even infinite test set sizes, and streaming data conditions.
The MBC property of set functions was identified by Bruno et al. (2021) who also proposed the Slot Set Encoder (SSE), a specific version of a cross-attentive pooling mechanism which satisfies MBC, guaranteeing will produce consistent output for all possible piecewise processing of set partitions. The introduction of the MBC property naturally leads to the rise of new dimension in the taxonomy of set functions, namely those which satisfy MBC and those which do not. The SSE is an example of one valid MBC architecture which comes at the cost of eliminating powerful self-attentive models such the Set Transformer. Self-attenti can be the best choice for tasks which require leveraging

---

Under review as a conference paper at ICLR 2023
NLL 1.67
1
Transformer single stream
Trunsformer
Transformat point
chunk Transformer
chunk
(a)
(b)
(c)
(d)
Figure 2: (o, +, x, D) correspond to classes in the input set. Ellipses are the model's clustering prediction. In streaming seftings, models must process the stream without storing streamed inputs. a-b: The Set Transformer delivers poor likelihood on different set streams. c-d: Set Transformer with UMBC module hecomes MBC function, yielding better likelihood, and consistent predictions regardless of the data stream. For a description of streaming settings, see Section 5; additional streams are shown Figure 8. 1. (Non-MBC) Set Encoder 2. (Non-MBC) Set Encoder
3. (MBC) Set Encoder
4. (MBC) Universal MBC Model
A
(MBC) Function
(Non-MBC) Set Function
Consistent Output
Not Consistant
Figure 3: Non-MBC set functions (1 & 2) are not MBC when sequentially processing random set partitions. MBC set functions (3) are MBC, but with limited valid architectures. Universal MBC (4) (UMBC) allows leveraging MBC+non-MBi set functions, widening the field of available MBC architectures.
pairwise relationships between set elements such as clustering (as we show later results in Figure and Table 2) where the Set Transformer outperforms SSE).
mini-batch pooling o²
Models such as the Set Transformer cannot make MBC guarantees when updating pooled set representations, require all N elements in a single pass, and therefore do not satisfy MBC (i.e. processing separate pieces of a set yields different output than processing the whole set at once). Naively using such non-MBC set functions in an MBC setting can cause severe degradation in performance, as depicted in Figures 2a and 2b where the Set Transformer exhibits poor likelihood and inconsis- tent clustering predictions. With the addition of a UMBC module UMBC+Set Transformer inherits an MBC guarantee, yielding con- sistent results, and much higher likelihoods (Figures 2c and 2d) (See Section and Appendix for details of the experiment). The quantitative effect of MBC vs non-MBC encoding on pooled set representation can be seen in Figure which shows the variance between pooled representations of 100 random partitions of the same set. (See Appendix for details).
Set Transformer
self-attention blocks
0.04
UMBC Set Transformer
mean
0.02
0.00
0
20
40
60
chunk size
Figure 1: 6² between encoded fea- tures of 100 random mini-batched set partitions. Set Transformer (not MBC) produces variance in out- put. UMBC+Set Transformer pro- duces consistent output for all 100 random mini-batched partitions.
In this work, we propose, and verify both theoretically and empirically that there exists uni- versal method for converting arbitrary non-MBC set functions into MBC functions, providing MBC guarantees for mini-batch processing of random set partitions. This allows for any set en- coder to be used in an MBC setting where it may have previously failed (e.g.streaming data). This result has large implications for all current and future set functions which not na- tively MBC, which can now be used in a wider variety of settings and under more restrictive conditions. Animations, code, and tests can be found in the supplementary file and also at: https: thtbbb...tttt/gg///g com/ anwaaanusnssssuuuuu suurmoouusuuumllee t umbc
Our contributions in this work are as follows:
In Theorem 4.1
show that any arbitrary non-MBC set encoder become MBC, guaranteeing that mini-batch processing of sets of any cardinality at test-time will give the same result as processing the full set at once.
2

---

Under review as a conference paper at ICLR 2023
We show that in tasks where non-MBC set functions outperform previous MBC functions, adding UMBC produces the strongest MBC model (Figure 4), highlighting the novelty and utiliiyy of UMBC.
With Proposition 4.1, we loosen the constraints on the MBC attention activation functions proposed by (Bruno et al., 2021) (which only uses sigmoid) by showing wide variety of functions can be used (e.g. traditional attention softmax (Vaswani et al., 2017)).
We explore an interesting MBC Monte Carlo dropout approach made possible by our UMBC module which delivers improvements in calibration for both in-distribution (Figure and corrupted test sets (Figure 6) while still maintaining MBC guarantees.
RELATED WORK
Table I: The MBC status of various set functions and the types of functions they contain.
Processing, pooling, and making a predictions for set structured data has been an active topic since the introduction of DeepSets (Zaheer et al., 2017). Attention has also been shown to be pow- erful in these tasks (Lee et al., 2019), as simple independent row-wise operations may fail to capture pairwise interactions between set ele- ments. Subsequent works have drawn connec- tions between set attention and optimal transport (Mialon et al., 2020), and subsequently expectation maximization (Kim, 2021). Likewise, efficient versions of set-attention have been proposed which incorporate cross attention with low dimensional self-attention in 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 (Murphy al., 2018) to tackle the problems posed by the required permutation invariance.
Model
MBC
Cioss-Att., Self-Attn.
Deep Sots (Zahoer al., 2017)
SSB (Bruno et al., 2021)
Set Transformer (Loc et 2019)
x
UMBC4Set Transfrmer
Bruno al. (2021), provide especially important lens through which view our work. Prior the proposal of the MBC property, previous works never explicitly considered the MBC 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 Property 3.3 (e.g. (Lee et al., 2019; Kim, 2021; Mialon et al., 2020; Jaegle 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 all set functions proposed in the future can be considered in MBC settings by incorporating UMBC.
Several prior works (Ovadia et al., 2019; Guo et al., 2017) highlight the problem of uncertainty quantification and probabilistic calibration, which can be crucial for tasks such as autonomous driving (Chen et al., 2017) and medical diagnoses (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 such as ImageNet-C (Hendrycks & Dietterich, 2019) to survey the landscape of neural network calibration. We take similar approach for set functions with ModelNet40-C (Ren et al., 2022) in our experiments. Guo et al. (2017); Ovadia et al. (2019) analyze calibration variants of deep convolutional models, while Minderer al. (2021) evaluate the calibration of 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.
3 PRELIMINARIES ON SET FUNCTIONS
For our setting, we define neural set function f X' ++ y with set-structured input space X' and output space y. I operates on sets = { with each set element x; E Idd A dataset of set-structured data D = {(X,, M consists of input sets X; and output sets Y,. the mapping of which can be learned by ff via stochastic gradient descent. Importantly, input X, is set and must process any set, therefore any element of the powerset (X;) also represents a possible input.
Deep Sets (Zaheer et al., 2017) provided crucial groundwork for neural set functions, formalizing the requirements of permutation equivariant (Property 3.2) architectures and invariant (Property 3.1) pooling mechanisms necessary for feature extraction, pooling, and predictions on sets. Following
3

---

Under review as a conference paper at ICLR 2023
these requirements. function assign (possibly different) output for each valid subset X, ((,), although the prediction should be invariant to permutations of the elements within X,.
Property 3.1 (Permutation Invariance). A function f ((t) + y acting on sets is permutation invariant to the order of objects the set iff for any permutation function ff f(x1,...,,,, = ((x,())....,x((())
Permutation invariant layers commonly referred to set pooling functions, and generally have fixed size output given any permutation or cardinality of the input set. Such a permutation invariant function can also be called Set to Vector function, it produces a vector of a fixed size given variable sized input.
{x} + Rd' is a set function satisfying Property 3.1, which projects a set (x; E Ji=1 to one or more vectors 2g E Rdd
Definition 3.1 (Set Vector Function). A Set to Vector Function (S2V) f
Additionally, Zaheer et al. (2017) prescribes that prior to any permutation invariant pooling, any composition of permutation equivariant layers may used for feature extraction. Common linear and convolutional neural network layers valid permutation equivariant functions when considering a batch of inputs as a set. For the remainder we assume contains both equivariant and invariant layers.
Property 3.2 (Permutation Equivariance). A function f P + y acting on sets is per- mutation equivariant to the order of objects in the set iff for any permutation function π ([xx<())-..,xx(())) Xx(N))) =
Lee et al. (2019) identified that self-attention layers (Vaswani et al., 2017) satisfy Property 3.2 and thus can be used equivariant feature extractors for set functions, creating the Set Transformer. Set Attention Blocks (SAB) are defined as SAB(X, X, X) = Attention X, X) = softmax (X,, (X, X Xx, (i.e. Self-Attention ( (Vaswani al., 2017)). Additionally, the permutation invariant pooling layer of the Set Transformer, Pooling by Multihead Attention (PMA), performs cross-attention between a learnable seed parameter S IRK and the input set, PMA(X) (S,
I Rxd
Bruno et al. (2021) identified the MBC property. proposing the MBC Slot Set Encoder (SSE), adding new dimension to Property 3.1 from Zaheer et al. (2017). Instead of merely requiring that fXX) be invariant to permutations of the indices i of x; E X, the MBC property also requires that sequential processing of partitions, with partition indices p for Xi.p X is also permutation invariant.
Property 3.3 (Mini-Batch Batch Consistency (MBC)). Let X X; Rdd Let X be partitioned such that = U Xp, and f {x} + Rd be a S2V function. Given an aggregation function g [ + Rd g and f are Mini-Batch Consistent iff the following holds for any permulation of any random partition scheme,
{x be set with each element
g(f(Xp)...... f(Xp.)) f(X)
The f function of an SSE (Bruno et al., 2021) layer cross-attention with slots S E k as queries, and partitions Xjer as keys and values. An SSE then utilizes an elementwise attention matrix activation which does not depend the other N elements within the set. SSE proposes sigmoid activation in the attention matrix A = (SX which is then normalized over the slot dimension K (Bruno et al., 2021; Locatello et al., 2020) such that A/., = Arj/ EM A With X, Rxx,|xd, and ē being the normalizzd sigmoid,
P Attention' (S, 6 Σ OSSX])),
1=1
where the sum is the partition cells, thereby consistently guarantecing the same output given any partition scheme and satisfying Property 3.3. The MBC aggregation function e (min, max, mean, sum) (Bruno et al., 2021) is represented by the
the RHS of Equation (1).
We omit the scaling of QK by va
for brevity throughout this text, but include the scaling our code.
4

---

Under review as a conference paper at ICLR 2023
4 BUILDING A UNIVERSALLY MBC SET FUNCTION
An important insight is that the SSE acts S2V function, creating encoded set representation for subsequent processing. However, S2V function which outputs K vectors 2, E IRd stated in Definition 3.1 can be re-interpreted as outputting set of cardinality K with each element z, R Therefore, by using a S2V function with K > as the base module f. we can view the S2V function function which takes set of cardinality N and maps downstream non-MBC module Property 3.3 need only be satisfied up until an invariant S2V pooling function.
set of cardinality K. Therefore,
Theorem 4.1. Let f be a set function satisfying either Property 3.2 or Property 3.1, and (g- f) be functions satisfying Property 3.3, which together form the functional composition F = f go f. For F to satisfy Property 3.3, it is sufficient to require the representation Z as input to I to satisfy Property 3.3.
Proof. Assume that g o f satisfies Property 3.3 and the composition F does not satisfy Property 3.3. updates Z as new partitions Xj arrive, yielding an MBC input to f", and therefore the same output of F for any permutation of random partition of X, contradicting the statement that F does not satisfy Property 3.3.
Algorithm Universal MBC Module with set func- tion f and slots S E R Kxd''
Put simply, Theorem 4.1 states that every mod- ule f coming after a module which satisfies Property 3.3 will continue to satisfy Property 3.3, even though f itself may not satisfy Prop- erty 3.3. With this established, we can therefore use Theorem 4. to build universally MBC set function.
1: Input: Partitioned Set X = (XI....,
2: Output: $ Rxdd
3: Set: Z 0
4: for
to do // to satisfy MBC, g=+ and f=Attention
5:
6: Z + Attention(S, X,) (Equation (3))
Maintaining attention normalization over N We now turn to the question of whether or not the constrained attention operation (i.e. avoiding normalization over N in the attention activation) described for SSE is in fact necessary in order to satisfy Property 3.3.
7: end for
8: return (2)
Proposition 4.1. By factorizing a normalization constant which depends on all N set elements from the attention matrix', normalization over N can be performed across mini- mathhbacheed partitions while still satisfying Property 3.3.
Proof. With o as elementwise activation function applied to the attention matrix, and å represent- ing the same function with normalization over N elements (i.e. like the traditional softmax in dot product attention),
Attention(S, X) ((SxT)x diag( (
Where diag(C) is a diagonal matrix containing the normalization constants CA - - where N is the set cardinality and S is single slot. Outside Φ(-). the final multiplication is associative, so we may simply evaluate ((SX") ).X, keeping a vector Ç with the incrementally updated normalization constants (e.g. the sum of the rows of exp(SX for softmax function). Factoring the attention in this way, we can update Ç and @(SXT) ] at the arrival of every partition
i=1 O(x sx)
Xy. normalize over N, and still satisfy Property 3.3.
Attention (S X) diag(()-1 diag(()))exp(SXX
diag
Σ
33
exp(SX)) )x
Interestingly, in our ablation study (Figure 5), find the softmax most effective, which requires the normalization over N as described above. For a note about about the numerical stability of the softmax calculated this way, see Appendix F.
We consider the dot product attention kemel, but any valid attention kemee function may be used.
5

---

Under review as a conference paper at ICLR 2023
Monte Carlo Slot Dropout As outlined in Section 4, our UMBC framework projects set of cardinality N to a fixed cardinality K. In doing so, there is a unique opportunity where we can treat each slot index { [=]=1 as Bernoulli random variable, dropping it with probability p (i.e. dropout (Hinton et al., 2012; Gal & Ghahramani, 2016) on the K set elements). This can be done while still satisfying Property 3.3, any dropout noise would be placed after the S2V function and thus be considered part of f*. Uniquely, a UMBC module can then perform Monte Carlo (MC) dropout (Gal & Ghahramani, 2016) on a streaming set, while never seeing the entire set at once. Dropping slot indices strictly decreases the cardinality of the set which is input to f". During training. this strategy could be useful for faster training, due to reduced set size for f (see Figure 18 for example). for regularization (Table 8 and Appendix J), or achieving a test-time ensemble of predictions by MC integration (Figure 7) as done by Gal Ghahramani (2016). Specifically, for test time MC integration, considering m as binary vector which selects the indices of set elements for processing by f", that is, f (X(), ) m) = f (iz; E f(X): m, = 1}), we perform the following approximation with dropout rate parameter p and sample size S.
p(ylX) - / ply X, m)p(m)dm
(f((), m), m, ~ Bernoulli(p)
i=1
Parallel UMBC Heads. We may also consider I. multiple parallel UMBC set projections layers, with each layer Si projecting the input of cardinality N to cardinality K to form an input to f" of cardinality LK. Specifically f" f(X)) Each f, would have independent multiheaded attention, allowing for independent representations of the same input set. Theorem 4.1 requires projecting a set of cardinality N to a set of cardinality K, which may be seen as bottleneck. Therefore parallel UMBC blocks could be used to add more expressiveness and reduce the overall effect of the UMBC bottleneck.
5 EXPERIMENTS
Metrics & Setup In the streaming settings of Figures 2 and 8 we placed arbitrarily hard MBC constraints (i.e. streaming settings described in Appendix B.1) show how failure occur in non-MBC model. However, it is unclear how to attempt to fairly compare against non-MBC models on MBC tasks. Additionally, the UMBC property has been proven both theoretically (Theorem 4.1) and empirically (Figures 1, and 8) Thus, in the following experiments, our aim is to evaluate the overall effect of the composition f (UMBCC.)), with the understanding that MBC models such as DeepSets and SSE extend to MBC settings while non-MBC models could be forced to perform in arbitrarily extreme MBC settings (like a single point stream).
We perform an ablation study on the components of UMBC, amortized clustering on Mixtures of Gaussians (MoG) and ImageNet- K, as well as ModelNet40 point cloud classification, Additionally, we analyze the calibration of popular set-functions on ModelNet40-C, which has not been examined in any prior work to our knowledge, We report accuracy, negative log likelihood (NLL), expected calibration error (ECE) (Guo et al., 2017), and Adjusted Rand Index (ARI) (Hubert & Arabie, 1985; Vinh et al., 2010). Standard settings of all UMBC models follow those shown in Table unless otherwise specified. All models trained over 5 different random initializations, with standard deviation error bars. We use open source code baseline code by Zaheer et al. (2017); Lee et al. (2019); Kim (2021) where applicable.
Amortized clustering We consider amortized clustering on similar Gaussian dataset as Lee et al. (2019) (See Appendix for dataset details). Figures and 8 contain a qualitative example of the task as well as a demonstration of how non-MBC models fail when used in an MBC setting. The task objective to maximize the likelihood (Equation (5)) of a set with K Gaussian components by predicting the component prior, mean, and covariance of each mixture component f(X) {T(X), (X)))
N
K
log (X;0)
Σ i=1 1=1
log n (a
(5)
We show the effect train/test set sizes in the full batch setting in Figure 4. Interestingly, the best performing MBC models (Figures 4e and 4f) are UMBC+Diff. EM. Set Transformer). Intuitively, this happens because the other MBC models (DeepSets and SSE) are not able to leverage pairwise

---

Under review as a conference paper at ICLR 2023
Dift.
Transformer
Deep Sets
2.0
1.90
2.72
train set size
1.9
1.85
2.70
1.80
2,68
1.0
1.75
120
2,66
256
1.7
1.70
1.65
2.64
512
1.6
1.60
2.62
200
400
600
800
1000
0
200
400
600
800
1000
200
400
600
800
1000
Lest set size
test set size
test set size
(a) Not-MBC
(b) Not-MBC
(c) MBC
Slot Set Encoder
UMBC Difr.
UMBC Set Transformer
2.20
2,00
1.90
2.15
1:95
1.85
2.10
1:90
1.80
2.05
2,00
1.85
1.75
1.95
1.80
1.70
200
400
600 800
1000
0
200
400
600
800
1000
200
400
600
800
1000
test Set size
test Set size
test set size
(d) MBC
(c) MBC
(f) MBC
Figure 4: Among models which satisfy Property 3.3 (c,d,e,f) UMBC+Set Transformer Shows the top performance. c-b show unmodified functions. and f' show UMBC augmented functions. UMBC+Set Transformer allows the power of a transformer to be used in an MBC setting. The legend in the top left figure applies to all plots. This experiment showcases the same task as seen in Figure 2
relationships between set elements. In UMBC, everything after the first function f can be arbitrarily complex, allowing I* to leverage self-attention. Crucially, this task requires a global view of relationships between points in order to predict cluster parameters. Therefore, it is intuitive to why f utilizing self-attention or expectation maximization perform better than simple f' (such as row-wise linear layers in Deepsets). Note that this is the same task as depicted in Figures 2 and 8, which shows how the better bottom line performance of the vanilla Set Transformer in Figure 4 evaporates in the MBC setting.
Table 2: Amortized Clustering on ImageNet fea- tures extracted with a pre-trained ResNet50.
We extended the amortized clustering to ImageNet-IK (Deng et al., 2009). using features extracted from pretrained ResNet50 (He et al., 2016) model (See details in Appendix E). Re- sults are shown in Table 2. The Empirical model in Table 2 is the NLL and ARI obtained using the actual prior, empirical mean, and diagonal covariance of each class cluster. In this task, UMBC+Set Transformer outperforms all other models. To account for UMBC's added parameters, we included UMBC on the baseline MBC models in Table 7, and UMBC+Set Transformer still shows the top performance. These results show that as UMBC may act as a bottleneck, there is no guarantee that the bottleneck will always negatively effect the performance of the underlying model.
Model
MBC
N
ARI
Empirical
1028.22:1.224
44.0910.11
Deep Sets (Zaheer et al., 2017)
531./410.15
6.18:1008
SSE (Bruno et al,, 2021)
520.2910063
22,91=::85
Dift. EM 2021)
524.7410.38
13.2210.16
Set Transfommr (Loce al., 2019)
512.5910.33
17.1313.67
UMBC-DiE, EM
518.56:10.99
13.011045
UMBC-SeI Set Transformer
503.899:10887
23.68: 1.85
Table 3: Valid UMBC attention activation functions. K norm. and N norm. refer to the the normalization constant over the slots K and instances 4 respectively.
SSE's Connection PMA'S With the introduction of Proposition 4.1,
it is easy to see that the only differ- ence between an SSE and PMA is
function (7)
K norm.
N' norm.
nome
reference
the choice of the attention activation function. Indeed any deterministic elementwise function which 1) maps the pre- activaaiiioon attention matrix to strictly positive values, and 2) has an
signnodd(:;))
slot sigmaid
(Bruna 2021)
exp(x,)
sum
sum
slot-softmax
(Locatello et al., 2020)
exp(ar.)
sum
solimax
(Lee al., 2019)
exp(x;)
maxx(()
sum
slot-exp
sigmaid(ar.)
sum
sigmoid
Diff. HM showed some instability on the ImageNet clustering task and failed to converge for one run.
Therefore variance is reported on 4/5 runs.
In cases where there is both K norm and an N norm, the K norm is performed first.
7

---

Under review as a conference paper at ICLR 2023
1.75
Actertion Actival ene
FT Layeroom
2.0
Siat Dropost
1.55
Siet Resicuel
1.50
coftrax
1.82
3,9
3.80
1.00
1.70
1.75
3.76
Test set dize
205t set size
Eest set size
rost set size
(a)
(b)
(c)
(d)
1.00
1.85
Type
cetemindddt
random
1.70
1.75
Kumber Faralle undersal Bodks
3.02
1.50
Numb Stats
1.000
1.5/3
Sut prop Samples
13
2.00
1.79
1:825
1.76
1.75
1.14
1/793
1.70
230
603
100
603
1000
200
400
600
800
1300
1.50
253
1003
(h)
Figure 5: Ablation study, analyzing the effects of different settings within the UMBC module. These experiments were performed on the MoG dataset with the UMBC+Set Transformer model.
Table 4: Point cloud classification (ModelNet40), Models are trained on a set size of 1000 randomly sampled points, and evaluated on 100, 1000, and 2048 (max) test set sizes. bold entries denote the best performance between models with and without UMBC. Models in the top row are for reference.
Aucusacy
NLL 1
BCE 1000
Model
MBC
100
1000
2048
100
1000
2048
100
2048
Deep Sets (Zabeer et al., 2017)
66.37.11.077
88.35_.0.22
58, 58.721021
1.510000
0.9010011
0.80.0011
17.58.0955
4.2110277
4.0210.16
SSE (Bruno et al. 2021)
31.09:0.51
87.85c11.39
$7.9289042
1.4260.30
0.5260005
0,5150006
16.9911111
5.9321006
5.89::1.17
Diff LM (Kim, 2021)
X
62.67.11.11
86.08 0.12
96.86.1..66
2.8010.11
0,711002
0.691000
22 161.99
5.15.0.111
4.96.0022
UMBC+DDef-BD EM
67.07:11.07
86.22
1:33
96.3711100
16110012
0,881006
0,5710086
13.91 1:11
4.32 1.37
4,38 1:27
Sei Transfomecs (Lax et sl., 2019)
x
74.21=1167
87.81-0.4
58.17-0.32
1.26=000
0.70-000
0.784000
17, 17.12200:4
7.45-0.22
7.57=0054
UMBC45e Transfurner
71.18=11.52
86
86:36-0.49
56. 56.77:00.29
1.2388115
0,53:1003
0,51=1083
10.37:2224
2.00:4119
2.35:0022
optional normalization constant over N which can be factored as in Proposition 4.1 is valid and will satisfy Property 3.3. With this in mind, we identify five functions, and explore their effects in Table and Figure 5.
Ablation Study Using the mixture of Gaussians dataset for amortized clustering, we analyze components of UMBC+Set Transformer in Figure 5. Of the five activation functions identified as valid in Section we found that the softmax used in traditional attention performs the best. In agreement with Bruno et al. (2021), we find that treating the slots as a Gaussian random variable, and learning them with reparameterization (outlined in Appendix H) leads to better overall results. We also find that layernorm on the post attention linear layer, residual connections on the slots before the FF layer (like the PMA of Lee et al. (2019)) to be beneficial. Increasing the number of slots (cardinality of input to f*), helped to point and then showed decrease in performance, likely due to overparameterization We used these settings to inform our base settings given in Table 6.
The effect of slot dropout at both train time and test time can be seen in Figure 5 (c. and h). Empirically, on the Gaussian clustering task, we found that using no slot dropout ultimately led to the best performance, which we think likely due the fact that the dataset can sample infinitely many instances and is therefore extremely resistant to overfitting, Using dropout on the ModelNer40 dataset (Figure 7), which prone to overfitting, led to increased performance on all metrics.
Point cloud classification We perform set classification experiments ModelNet40 (Wu et al., 2015) and analyze the robustness of different set encoders to dataset shifts and varying test-time set sizes using MedelNet4O-C (Ren et al., 2022) which contains 15 corruptions at 5 levels of intensity. Our experiments use the version of ModelNet40 and ModelNet40-C used by Ren et al. (2022) which contains 2048 points sampled from the original ModelNet40 (Wu et al., 2015) CAD models. Results are presented in Table 4 and Figure 6. For non-MBC models, augmenting the model with UMBC only shows a performance decrease for the larger set sizes, likely due to the bottleneck caused by the projection from cardinality N to K. Interestingly, on set sizes smaller than the model was trained on, UMBC augmented models show consistently higher accuracy and indicates that UMBC is resilient to smaller test set sizes. In terms of ECE and NLL., UMBC models outperform all non-MBC baselines.
8

---

Under review as a conference paper at ICLR 2023
Method
0.6
ECE
UMEC haadram
0.4
Text
Corruption Intensity (All Corruptions, Test Set Size: 100)
0.0
0.6
ECE
0.4
0.2
Corruption Intensity (All Corruptions, Test Set Size: 1000)
0.8
0.0
Corruption Intensity (All Comuptions, Test Size: 2048)
Figure 6: Expected Calibration Error on ModelNet4O-C which contains 15 corruptions at 5 different intensity levels. 'Test' corresponds to the uncorrupted test set. See Figures 13 and 14 for Accuracy and NLL and Figures 15 to 17 for results individual corruptions.
This increase ECE can be partly attributed to MC sampling slots test time (shown in Figures 11 and 12) and partly to Slot Dropout at train time (shown in Table 8).
0.5 9.4 0.3 0.2 S1 0.3
20.00
14:77
3.39
335
135
135
12.41
1244
12.5
3.35
Dropocut Rate
21.5
Drropouut Ratte
3.30
11.5
71
0.3
5.31
3.50
126
3.50
3.2
17:04
13:84
1843
12
13.55
11.0
3.25
14:00
71-37
71.00
14.00
70.5
3:4
125
121
1.5
128
3.20
14:22
14:38
14:33
425
18:00
10.0
100
100
Samples
Sam
Samples
Figure 7: Performing Monte Carlo Dropout on UMBC+Set Transformer slots leads to increases in accuracy, NLL, ECE. The top row corresponds to 0% dropout rate and is constant over dropout sample sizes. Experiment uses ModelNet40 with test set size of 100. Figures for set sizes 1000 and 2048 can be found in Figures 11 and 12
ModelNet40-C results can seen Figure 6. UMBC+ models give strong ECE performance in all test set sizes, improving over non-MBC baselines, especially for test set size 100 and UMBC+Set Transformer where the largest miscalibration baseline models is.
CONCLUSION
In this work, we have shown that composing a set function consisting of MBC base function f. with an arbitrary set function head f*, we can make the composition F= f o f universally mini-batch consistent (UMBC). We have provided proofs in Theorem 4.1, empirical experiments, 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 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 MBC dropout strategy made possible by UMBC which leads to improved calibration and NLL. As the field of set-functions continues to widen, look forward seeing future research in the of MBC set functions.

---

Under review as a conference paper at ICLR 2023
7 REPRODUCIBILITY STATEMENT
In addition to the details listed below, all code has been included as supplementary material and at the following URL where animations can be viewed in any broswer: https:/ / thub. com/ anonymous submlt t /umbc. Further descriptions of the construction of the Mixture of Gaussians dataset can be found in Appendix B, details ImageNet clustering can be found in Appendix E, further training details regarding optimization and model settings can be found in Appendices G and H respectively.
REFERENCES
Charles Blundell, Julien Corneiiee, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neurul information processing systems, 33:1877-1190 2020.
Andreis Bruno, Jeffree Willette, Juho Lee, and Sung Ju Hwang. Mini-batch consistent slot set encoder for scalable set encoding. Advances in Neural Information Processing Systems, 34, 2021.
Xiaozhi Chen, Huimin Ma, Ji Wan, Bo Li, and Tian Xia. Multi-view 3d object detection network for autonomous driving. In Proceedings of the IEEE conference Computer Vision and Pattern Recognition, pp. 1907-1915, 2017.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty deep learning. In international conference on machine learning, pp. 1050-1059. PMLR, 2016.
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. Advances in neural information processing systems, 30, 2017.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International Conference Machine Learning, pp. 1321-1330. PMLR, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness common corruptions and perturbations. arXiv preprint arXiv: /2261, 2019.
Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhutdinov: Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprini arXiv:1207.0580, 2012.
L Hubert and P Arabie. Comparing partitions journal of classification 2 193-218. Google Scholar, pp. 193- 128, 1985.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, pp. 4651-4664. PMLR, 2021.
Minyoung Kim. Differentiable expectation ionirattioommaiimzzation for set representation learning. In Interna- tional Conference on Learning Representattinss, 2021.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: framework for attention- based permutation invariian neural networks. In International Conference Machine Learning, pp. 7744-3753. PMLR, 2019.
10

---

Under review as a conference paper at ICLR 2023
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centrie learning with slot attention. Advances in Neural Information Processing Systems, 33:11525- 1151525588,,, 11538, 2020.
Grégoire Mialon, Dexiong Chen, Alexandre d' Aspremont, and Julien Mairal. A trainable optimal transport embedding for feature aggregation and its relationship to attention. arXiv preprint 12X1: 2020.
Matthias Minderer, Josip Djolonga, Rob Romijnders, Frances Hubis, Xiaohua Zhai, Neil Houlsby, Dustin Tran, and Mario Lucic. Revisiting the calibration of modern neural networks, Advances in Neural Information Processing Systems, 34, 2021.
Ryan L. Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Janossy pool- ing: Learning deep permutation-invarian functions for variable-size inputs. arXiv preprint arXiv: 1.11111111,000 2018.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. Advances in neural information processing systems, 32, 2019.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
Jiawei Ren, Liang Pan, and Ziwei Liu. Benchmarking and analyzing point cloud classification under corruptions. arXiv:2202.033: 03377, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017
Nguyen Xuan Vinh, Julien Epps, and James Bailey. Information theoretic measures for clusterings Variants, properties, normalization and correction for chance. The Journal of Machine Learning Research, 11 11:2837-2854, 2010.
comparison:
Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912-1920. 2015.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Yan Zhang, Jonathon Hare, and Adam Prügel-Benneth Fspool: Learning representations with featurewise sort pooling. arXiv preprint arXiv: /90.002995 2019.
S Kevin Zhou, Hayit Greenspan, Christos Davatzikos, James S Duncan, Bram Van Ginneken, Anant Madabhushi, Jerry L Prince, Daniel Rueckert, and Ronald M Summers. review of deep learning in medical imaging: Imaging traits, technology trends, case studies with progress highlights, and future promises. Proceedings of the IEFE, 2021.

---

Under review as a conference paper at ICLR 2023
A
APPENDIX
We will briefly describe the contents of each section of this appendix below:
Appendix B: Extra information and results related to MoG Amortized Clustering.
Appendix C: Details of the experiment depicted in Figure 1.
Appendix D: A note on MBC testing of the Set Transformer.
Appendix E: Details on ImageNet amortized clustering.
Appendix F: A note on the UMBC attention softmax stability.
Appendix G: Training parameters/setup.
Appendix H: Model hyperparamettes//sttup.
Appendix Additional ablation study results/discussion.
Appendix Additional results/discussion for ModelNet40 experiments.
Appendix K Extra results augmenting MBC models with UMBC.
Appendix Limitations and Future Work.
Appendix M Attention Activation Effects on Calibration.
B DETAILS ON THE MIXTURE OF GAUSSIANS AMORTIZED CLUSTERING EXPERIMENT
We used a modified version of the MoG amortized clustering dataset which was used by Lee al. (2019). We modified the experiment, adding random variance into the procedure in order to make more difficult dataset. Specifically, sample a single task for problem with classes,
Sample set size for the batch (trrin set size/2, / train set size).
2. Sample class priors π ~ (irichlet((11.... 1x)).
Sample class labels 2; ~ Categorical (T) for 1,.., N.
4. Generate cluster centers Pi,j ~ (--4(((44400 4) for i = 1,.., K and 1,2.
5. Generate cluster covariances ij (0.3, 0.6) for 1...., K and 1, Then make covariance matrix E; for each class with o as the diagonal.
6. Sample data X.j ~ N(H; E,)
In our MoG experiments, set K 4.
The Motivational Example in Figure 2 also used the MoG dataset, and performed MBC testing the set transformer corresponding to the procedure outlined in Appendix D
NEL 1.87
NLL1::1.7
class stream
each stream
class stream
mach stream
(a)
(b)
(c)
(d)
Figure 8: a-b: The Set Transformer struggles to perform well on different streams set. c-d: Our UMBC Module, UMBC+Set Transformer makes the Set Transformer an MBC function, yielding the same prediction regardless of the data stream. For a description of streaming settings, see Section 5, additional streams shown in Figure 2
12

---

Under review as a conference paper at ICLR 2023
B.1
STREAMING SETTINGS
The four total streaming settings in Figures and can be described as:
single point stream - streams each point in the set one by one. This causes the most severe underperformance by the Set Transformer.
class stream streams entire class once. The attention modules within Set Trans- former cannot compare the imput class with any other clusters, thereby degrading perfor- mance of Set Transformer.
chunk stream streams 8 random points at time from the dataset, Providing limited information to the Set Transformer's attention.
one each stream - streams a set consisting of a single instance from each class. Set Transformer
examples of each class, but with a limited sample size, the encoding fails to make accurate predictions.
C MEASURING THE VARIANCE OF POOLED FEATURES
In Figure 1, show the direct quantitative effect the pooled representation when using the original Set Transformer and with our UMBC module added, UMBC+Set Transformer. The UMBC model variance is always effectively O, while the Set Transformer gives different results for different set partition chunk sizes. The downward slope of the Set Transformer line can be explained by the fact that as the chunk size gets larger, the pooled representation will become closer to that of the full set. The procedure for MBC testing of the Set Transformer is outlined in Appendix D.
To perform this experiment, we used randomly initialized model with 128 hidden units, and sampled a random normal input with a set size of 1024, X E R 1024x 128 We then created 100 random permutations of the set elements of the input and split each permutation into partitions with various chunk sizes C; where the cardinality |C:| E {2') ] Jt=1- We then encode the whole set for each chunk size and report the observed variance between the 100 different random partitions the various chunk sizes in Figure 1. Note that the encoded set representation is vector and Figure 1 shows a scalar value. To achieve this, we take the feature-wise variance over the 100 encodings and report the mean over each feature. Specifically, with Z E IR 100x 128 representing all 100 encodings, = var (Z), with We then achieve the values in Figure by a simple mean over the feature dimension,
z E IR 128
1
Σ
24
(6)
128
D A NOTE ON MBC TESTING OF THE SET TRANSFORMER
In some illustrative experiments Figures and 2, we apply MBC testing to the Set Transformer to study the effects of using non-MBC model in an MBC setting. The Set Transformer does not have prescribed way to do this the original work, so we took the approach of processing each chunk up until the pooled representation that results from the PMA layer. We then performed mean pooling operation over the chunks in the following way, with Z representing the final mini-batch pooled features,
Z P
PMA(X,)
E DETAILS ON THE IMAGENET AMORTIZED CLUSTERING EXPERIMENT
For the ImageNet amortized clustering experiment outlined in Section 5, we first extracted the features up until the last hidden representation and before the final linear classifier layer of the pretrained
13

---

Under review as a conference paper at ICLR 2023
and frozen ResNet50. These features x, R2048 of large dimension which would create excessively large linear layers for this experiment. Therefore, we projected the features down to a lower dimension & E R2122 using random orthogonal Gaussian matrix. As this random Gaussian projection is suitable for random feature kernels (Rahimi & Recht, 2007). it should preserve the distances between points required for effective clustering with marginal effect on overall clustering performance. To validate this assumption, we ran the Empirical model (which computes the empirical cluster mean and diagonal covariance) both the original features and the projected features g(x) and present the results in the table above.
To construct the ImageNet dataset, we first initialized and saved the random Gaussian projection matrix, and proceeded to process the entire ImageNet training set with the saved matrix. From these extracted and projected features, we chose fixed 80/20 split for our train/test sets. Class indices for the train/text sets can be found in the supplementary file.
Version
ARI
R200
45.93:1.12
g(x,) 18812
44.09::0.11
F NUMERICAL STABILITY OF MBC SOFTMAX ATTENTION ACTIVATION
Numerical stability of the softmax requires that the values are not allowed to overflow. Generally this is done by subtracting the maximum value from all softmax logits which allows stable and equivalent computation.
max
max(x) (x)
(8)
Ex'Ex Ex ex (x)
(ax(x)
Extex
L
This poses problem when using the plain softmax attention activation, as the max(-) in Equation (8) requires max over the whole set of N items which is unknowable given the current mini-batch.
Originally, we had devised a special conditional update rule which would maintain the same form as in Equation (8), by tracking the overall max each row of the attention matrix and then conditionally updating either the current A and Ç or the previously stored values from the last processed partition. Those updates needed to be calculated in the exponential space which cause a propagation of numerical errors through the network, becoming large enough to interfere with inference. In experiments, we found it sufficient to calculate the softmax as a simple exponential activation with subsequent sum concern, one could also set a hyperparameter À for the model such that the softmax is calculated with an exponential function such e=- which should provide a reasonable solution.
N with consideration for numerical stability. If numerical stability is
G TRAINING SPECIFICATION
We use no L2 regularization, except for the ModelNet40 experiments, which use a small weight decay of le - 7. This was setting taken from previous experiments by Lee et al. (2019); Zaheer et al. (2017) which used dropout before and after the pooling layers and other regularization strategies such as gradient clipping to avoid overfitting.
The only experiment which utilized any kind of data augmentations was the ModelNet40 experiments which used random rotations of the point cloud as is common in the precedent experiments (Zaheer et al., 2017; Lee et al., 2019; Bruno et al., 2021)
All single runs of all of our experiments were able to fit on single GPU with 12GB of memory.
Table 5: The hyperparameter setup for all of our experiments involving UMBC modules.
Experiments
Setting
MoG
ImageNet
ModelNer40
Optimizer
Adam
Adam
Adam
Leamning Rate
le-3
le-3
le-3
Data Augmentation
Epochs
S0
50
1000
Iters/Epoch
1000
1000
9840
14

---

Bruno al. (2021), for each slot such that the set of slots S E x has a separate parameter for each K; E K. We do this because the original Slot Attention in (Locatello et al., 2020) used Under review as a conference paper at ICLR 2023
H UNIVERSAL MODEL SPECIHICATION
Unless otherwise specified, all universal modules were run with the following model hyperparameter settings in Table 6. The settings for the MoG dataset apply to those in Figure 4, and Figure 5 studies the effects of changing individual settings.
Table 6: The hyperparameter setup for all of our experiments involving UMBC modules, The hyperparameters were chosen as sensible default based previous architectures in Lee al. (2019); Zaheer et al. (2017); Kim (2021)
Experiments
Setting
MoG
ImageNet
ModelNet40
Embeddeer
128
Sloes Per Purallel UMBC
128 random
Stol-1ype
random
random
Stot LayerNomm
LayerNorm
Heads
Slot Dmpout Proh.
0%
50% softmax
50% softmax
Allention Activation
softmax
Slot Residual
/
UMBC Nom. Parallel
1
4
4
Test MC Samples
10
100
Slots Different from both Locatello et al. (2020) and unique initial slot parameters
Universal MBC Layer
(UMBC)
GRU in an inner loop adapt the single general slot into specific slots for given task, forcing them to 'compete' to capture different parts of the input. We cannot use a GRU, as it violates Property 3.3, so we instead let each slot K; E K learn to adapt to the overall data distribution. We always used the same dimension of inputs X and slots S.
Multi-Head Attention
A
Random Slots To initialize the random Gaussian slots, we use a similar initialization strategy as (Blundell al., 2015) and initialize E 4120.., 0.2] and log E 4[-5.00 -4.0). During training, we sample the distribu- tion with reparameterization S: u ak 6& with 6 ~N, ld).
FF
0
Dropout
Residual MBC Sum
Embedder We found useful to place a single layer embedding function at the base of UMBC modules which consists of a single linear layer and a ReLUU activation function. We used this embedder in all experiments except the ImageNet amortized clustering, as the ResNet feature extractor acted the embedding function in this case.
Matrix Mult.
Figure 9: The architecture of a UMBC layer. A" represents the unnormalized attention matrix a (SX,T )x; discussed in Proposition 4.1 and 'MBC Sum' rep- resents the summation in Equation (3)
1 ADDITIONAL ABLATION RESULTS
In addition to the results in Figure also did exper- iment looking at the effect of the number of attention heads in the UMBC layer in Figure 10. This result was uninformative, but we choose to use stock setting of 4 attention heads in our experiment as was common in the experiments performed by Lee et al. (2019).
J ADDITIONAL MODELLEET/MODELNET.. RESULTS
15

---

Under review as a conference paper at ICLR 2023
Table 8 shows extra results from the ModelNet point cloud classification task. In this table, we include results for 'UMBC+SSE' and *UMBC+Deep Sets' for complete- ness. While there is a slight decrease in accuracy for both "UMBC+SSE' and "UMBB++Dep Sets,' UMBC improves SSE in terms of NLL and ECE while lowering the perfor- mance of Deep Sets. This seems to generally agree with the results in Figure 4, indicating that it is likely unhelpful add a UMBC f to an already MBC f", and instead the model l should be chosen according to the given task first, and then UMBC considered if MBC treatment will be necessary.
Number of Attention Heads
1.76
1.75
1.74
10
15
Heads
Figure 10: Ablation study on the num- bers of attention heads in UMBC layers
Accuracy
NIL
FCE
0.0
Dropout Rate
BG:43
Dropout Rate
0.4 0.3 0.2 0.3
6:58
0.55
0.55
J:94
3.10
3.15
2.54
3.0
0.4 0.2
$1.42
86.6
6:46
0.54
0.55
1.14
0.55
0.54
Drogonu 0.4
7.44
3:72
2:71
2.44
233
2.8
86.5
0.53
2.6
86.99
86.4
<:54
2.79
85:00
86.55
64:92
66.65
06.79
86.3
0.5
6:53
0.53
0.52
9:52
0.54
0.52
2.58
1.28
2.30
2.39
2.43
10
25
50
100
5
10
25
50
100
10
25
50
100
Samples
Samples
Samples
Figure 11: Performing Monte Carlo Dropout UMBC+Set Transformer slots leads to increases in accuracy, NLL, ECE. The top row corresponds to a 0% dropout rate and is constant over dropout sample sizes. Experiment uses ModelNet40 with test set size of 1000.
Accuracy
NIL
FCE
0,0
0.54
86:85
2.8
Dropeut Rate
Drogout Rate
0.3 0.2 0.1
0.4 0.2
85.71
96.77
91.79
86.80
6:23
0.22
0.22
9:22
0.22
Dropout
8:73
1/72
3.74
2.44
2.5
2.7
2.6
86.77
86.29
05.22
6:33
0.33
0.22
9.24
1/13
1
2.46
234
0.5 0.4
0.51
65.24
$5.95
86,
70
6:52
0.51
0.5
2:43
2,37
2337
2.4
10
23
30
100
10
25
50
100
0.50
10
25
30
100
Samples
Samples
Samples
Figure 12: Performing Monte Carlo Dropout on UMBC+Set Transformer slots leads to increases in accuracy, NLL, ECE. The top row corresponds to 0% dropout rate and is constant over dropout sample sizes. Experiment uses ModelNet40 with test set size of 2048.
ModelNet40 is prone to overfitting, and previous experiments in Deep Sets (Zaheer et al., 2017) and Set Transformer (Lee et al., 2019) have used Dropout layers both before and after the pooling function in their encoders. To evaluate the regularization effect of our dropout strategy, the last block of Table 8 includes UMBC models trained without dropout. Training without dropout generally lowers test set performance in all metric categories.
For examples of the corrupted point clouds, we refer the reader to the original work which proposed ModeINet40-C (Ren et al., 2022). In Figures 13 and 14 we provide additional boxplots for accuracy and NLL metrics which correspond to the ECE metric reported in Figure 6. In Figures 15 to 17 we provide individual boxplots for each individual corruption
accuracy, ECE, and NLL respectively. The aggregate of all of these datapoints forms the boxplots seen in Figures 6, 13 and 14. Size is reduced to avoid excessive page length. Best viewed screen with high zoom.
16

---

Under review as a conference paper at ICLR 2023
Method
Accuracy
0.2
Text
Corruption intensity (All Corruptions, Test Set Size: 100)
Accursty
Corruption intensity (All Corruptions, Test Set Size: 1000)
Accurary
0.5
Corruption intensity (All Corruptions, Size:
Figure 13: Accuracy across all corruptions in the ModelNet40-C dataset. This figure corresponds to the ECE results presented Figures 6 and 14
Method
Comupttin Intensity Corruptions, Test 100)
Taxt
Corruption Intensity (All Corupttion,, 1000)
25
20
Tal
Corruption Intensity (All Corrpptions Test Set Size: 2048)
Figure 14: NLL across all corruptions in the ModelNet40-C dataset. This figure corresponds the ECE results presented in Figures 6 and 13
K ADDING THE UMBC MODULE To EXISTING MBC FUNCTIONS
L LIMITATIONS & FUTURE WORK
UMBC is a bottleneck UMBC projects the input set fixed size, and can therefore be bottleneck, causing possible loss information from the input set. An interesting line of research could be exploration of methods maximize mutual information between the input set of
cardinality N and the projected set cardinality K, or an exploration of other forms which UMBC may take, look forward to seeing future research in this area.
17

---

Under review as a conference paper at ICLR 2023
menelyy the Serrie
Figure 15: Accuracy boxplots for individual ModelNet-C test results. Size is minimized to avoid excessive page length. Best viewed on screen with high zoom
18

---

Under review as a conference paper at ICLR 2023
Figure 16: ECE boxplots for individual NeteI test results. Size minimized to avoid excessive page length. Best viewed on screen with high zoom
19

---

Under review as a conference paper at ICLR 2023
Figure 17: NLL boxplots for individual ModelNet-C test results. Size is minimized to avoid excessive page length. Best viewed on screen with high zoom
20

---

Under review as a conference paper at ICLR 2023
Inference Speed vs. Slot Dropout Rate
0.06
time
0.04
0.02
0.0
0.2
0.4
0.6
0.8
1.0
Because the model structure, higher slot dropout rates correspond to faster training times, and smaller set sizes input to the subsequent set encoder modules. A dropout rate of p= 0.5 in the function f will, in expectation, deliver a set size of K/2 to the subsequent function f" This figure was generated from UMBC+Set Transformer model with 128 hidden units and input x E R 200xx200dd__ The plotted line shows mean and standard deviation for 250 iterations at each p E [1, 99]. As safeguard against unstable training, we ensure that at least one slot remains after dropout is applied.
Figure 18:
Table 7: Amortized Clustering on ImageNet features extracted with pre-trained ResNet50.
Model
NLL J
ARI
Empirical
1028. 1088.22.1.22 1.24
44.09.10.11
Deep Sets
531.4410.11
6.08
SSE
$20.29-0.63
22.91=1185
Transformer
512.59-0.33
17.13+3.67
UMBC+Deep Sets
532.87=0.69
6,22=0.18
UMBC+SSF.
544.67+3.64
16.59+
UMBC+Set Transformer
503.89=0.87
23.68=+185
Train/Test Set Size Variability In Figure Deep Sets shows the tightest grouping between training set sizes, although giving the lowest overall performance, indicating that more complicated set functions which make pairwise comparisons may less robust varying training set sizes, which may provide interesting topic of future research.
Bayesian Slots our experiments, we used similar random slot parameter initialization as Blundell et al. (2015). Following Bruno al. (2021), we slots, so the increased performance of random slots likely due to randomness aiding in exploration of the parameter space rather than learning proper Bayesian posterior. Future work could explore the effects of incorporating prior distribution over slots or slot dropout rates (e.g. Concrete Dropout (Gal et al., 2017)). This could lead further increases in robustness to corruptions and varying set sizes.
Bayesian prior these random
Large Train-Tine Set Sizes The setting we have considered one where the train-time set size known and the test-time setting presents severely constrained computataional resources extremely large set sizes which require processing the set in chunks. An open problem, however, is how to handle tenee--tha--mmmooe each element the sets. We look forward to seeing future research which may tackle this problem, and allow for training with extremely large sizes.
sizes at train time, backpropagation requires memory allocation for
M ATTENTION ACTIVATIONS & CALIBRATION
To test the effect of training with different attention activation functions calibration, train and evaluate the UMBC+Set Transformer model on all corruptions of ModelNet4O-CC in Figures 19 21, and individual conruptions in Figures 22 24, Besides the change in attention activation, each
21

---

Under review as a conference paper at ICLR 2023
Table 8: Point cloud classification on ModelNet4O. All models are trained on size of 1000 randomly sampled points, and evaluated on 100, 1000, and 2048 (max) test set sizes. UMBC models in the second block are trained and tested with our slot dropout technique outlined in Section Models in the first block are trained without Slot Dropout, and use all available slots output from f at both train and test time
Accunaz
NLL.
BCE
Model
100
1000
2043
100
1000
2048
100
1000
2048
UMBCyDeep
09:96:10.64
$7.50:50221
87.5840.16
06641882
0664-0102
21.25:40.54
6.5969322
8:51-0.26
UMB wSSi (No Drigont haane)
02 (0)
64.81
84 $4.891110 1
1:100:00
0.5900044
054-002
11:091200
502-0.82
Tansfoeee (No Dropout wrinn)
71.5230.75
066320.14
13.3634.06
UMB MMCCDDop Sets
71.58411.00 183
67.5220.25
57 44.440445
14K-0.0
862-0085
16.3951 53
7.40-0.50
UMBCC+SSE
71.00:0072
19119.2
00. 16.10000 46
L11. 1.0-0.00
0:501001
9.67.42.01
2-021:0.77
2-17-1110
0.51-110
model was trained with the same settings as the UMBC+Set Transformer from the corresponding experiments in Figures 6, 13 and 14. Surprisingly, we find the slot-softmax, originally used by Locatello et al. (2020) delivers strong performance terms of NLL and ECE, although gives slightly lower accuracy the natural, uncorrupted test set.
Motthod
Accuracy
Telt
Corruption Intensity (All Corruptions, Test Set Size: 100)
Accuracy
04
Tux
Corruption Intensity (All Corruptions, Test Set Size: 1000)
Accuracy
Corruption Intensity (All Corruptions, Test Set Sibe: 2048)
Figure 19: Accuracy across all corruptions on the ModelNer40-C dataset for UMBC+Set Transformer with different attention activation functions. This figure corresponds to the results presented in Figures 20 and 21
22

---

Under review as a conference paper at ICLR 2023
Method
Corruption Intensity Corruptions, Test 100)
12.5
NL.
7.54
1elt
Corruption Intersity (All Comuptions, Test Set Size: 1000)
7.5
5.2
2.5
Corruption Intersity (All Corruptions, Test Size: 2048)
Figure 20: NLL across all corruptions on ModelNet40-C dataset for UMBC+Set Transformer with different attention activation functions. This figure corresponds to the results presented in Figures 19 and 21
Method
0.4
1/a
Corruption Intensity Corruptions,
Text
Corruption Intensity (All Corruptions, Test Set Size: 1000)
Texi
Corruption Intensity (All Corruptions, Test Set Size: 2048)
Figure 21: ECE across all corruptions on ModelNet40-C dataset for UMBC+Set Transformer with different attention activation functions. This figure corresponds to the results presented in Figures 19 and
23

---

Under review as a conference paper at ICLR 2023
Two bal
Figure 22: Accuracy boxplots for individual ModelNer-C tests with UMBC+Set Transformer and different attention activation functions. Size is minimized avoid excessive page length. Best viewed on screen with high zoom
24

---

Under review as a conference paper at ICLR 2023
Groulen veleee 100
Corvution
Giv lrmeeds
Figure 23: NLL boxplots for individual ModelNer-C tests with UMBC+Ser Transformer and different
attention activation functions. Size is minimized to avoid excessive page length. Best viewed screen with high zoom
25

---

Under review as a conference paper at ICLR 2023
Sele da bukkess Two biti S011
de
Figure 24: ECE boxplots for individual ModelNer-C tests with UMBC+Set Transformer and different
attention activation functions. Size is minimized to avoid excessive page length. Best viewed screen with high zoom
26


Reviews:
[{"review_id": "zVRfXYtWgY", "reviewer": "Reviewer_SNvt", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "This paper presents a new architecture to satisfy the mini-batch consistency (MBC) property, which is an important property required by set functions in the streaming fashion. Specifically, given a mini-batch consistent function $f$ (e.g., slot set encoder, SSE), the authors prove that for an arbitrary set function $f^*$, the composition $F = f^* \\circ f$ is also mini-batch consistency, resulting in the proposal architecture, universal MBC set function (UMBC). Additionally, a Monte Carlo dropout strategy is proposed to cast the UMBC module more robust on out-of-distribution data. Experimental results show promising results on amortized clustering and point cloud classification.", "summary_of_the_review": "Overall, although the problem is interesting and the method works well in some cases, I think the work needs to be improved. The major concern is the unclear contribution of the additional permutation invariant/equivariant set function $f^*$ on the top of SSE architecture. In this regard, I think the proposed method is somewhat incremental, and most claimed contributions are similar  as SSE.", "strength_and_weaknesses": "### Strengths\n- The problem raised in this paper is interesting. The attached code is helpful and reproducible.\n- The proposed MC dropout method is nice. Despite being straightforward, it is exactly the first time to consider uncertainty estimation and model calibration in the set function communities.\n### Weaknesses\n- The UMBC method proposed in this paper seems to be an extended version of SSE, which might limit the technical novelty.\n- It seems that the proposed UMBC architecture just composes the SSE ($g \\circ f$) and a permutation invariant/equivariant set function ($f^*$) on the top. It is unclear how the additional set function (i.e., $f^*$) contributes to the expressiveness of mini-batch consistency modeling. I am afraid that if we stack multiple layers of SSE [Bruno et al. 2021], we could achieve the same goal of using one layer of SSE and a set transformer on the top.\n\n- Proposition 4.1 seems to be an important contribution, as it releases the sigmoid constraint in SSE and shows that the softmax attention also satisfies mini-batch consistency. However, it is unclear about the contribution of softmax attention compared to the sigmoid attention. Actually, we can apply arbitrary kernel functions to calculate the attention weights between the inducing points $S$ and the data $X$. One could just replace the sigmoid function with an exponential kernel.\n\n\nBruno, Andreis, et al. \"Mini-Batch Consistent Slot Set Encoder for Scalable Set Encoding.\" *Advances in Neural Information Processing Systems* 34 (2021): 21365-21374.\n", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_AmWWHUQM0", "reviewer": "Reviewer_BZo6", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "**Background**. This work focuses on the problem of so-called *neural set* functions, where the input is a set, the task could be clustering, classification etc., and the map from a set to an output is a neural network.  \nSince the input is a set, there are certain invariances that are relevant for this case, e.g., the output should not change if the order changes, etc.\nAmong these, most relevant to this work is the so-called *Mini-Batch Consistency* (MBC) property, which is relevant for the case when we cannot pass the entire input to the model and we have to partition the input.\nThe MBC property requires that the final output of an auxiliary function that we use should be the same as if we passed the entire (non-partitioned) input.\nFor example, consider we want to take the max element of a large vector: we may partition the input, apply max, and then apply max on top of that, which gives overall an MBC map.\n\n**Summary**.\nLet us say that the overall map is a composition of two functions, say feature and prediction part, where the input of the latter is the output of the former. This paper points out that it suffices that the feature map is an MBC map so that the overall map is MBC (and the prediction map can be an arbitrary function). This insight simplifies the architectural constraints of the prior work of Bruno et al (2020), and allows for handling inputs of arbitrary and varying size/dimension. \nIn addition, dropout is considered on the output of the feature function, as well as several comparisons with some existing prior works.\n\n\n\n", "summary_of_the_review": "The paper points out a relevant abstract problem of neural set functions. It appears to me that the provided solution does not fully address the problem and the empirical results are insufficient (inconsistent benefit of the herein proposed UBMC, and lack of comparison with the simplest approach to address this).\nThis paper also considers dropout approaches on some parts of the model, but I do not think that these contribute to its novelty, since MCDropout is a general approach, especially since here it is applied to the remaining standard part of the neural net (after the aggregation). \n\n\nMy recommendation is only temporary, and I would be happy to raise my score if the authors address the major concerns or explain if I misunderstood.", "strength_and_weaknesses": "## Strengths\n- This work seems to identify an important problem that having an input of varying size may be hard to handle for neural set functions\n- the paper does a lot of empirical evaluations including comparison with baselines, dropout techniques, and ablation studies\n\n## Weaknesses\n- While this work seems to identify a relevant problem, I do not think it resolves it, please see the next part on motivation\n- although there are numerous experiments, I do not see these as conclusive, that is, it is still unclear to me if the proposed solution is outperforming the baselines, and moreover, how it performs against naive implementations when the input is of variable length, see part below on experiments\n- I find that the theoretical results are relatively simple arguments, see part below on theoretical results \n- [minor] writing: see the section below\n\n\n\n### Motivation\nIf I understand well, there are two main motivations to focus on this problem:\n1. memory issue: when having varying input sizes, if we can set a maximum size of the input we can always use padding with zeros for example, or with some negative number if we use max-pooling etc. However, this naive strategy may result in very large memory requirements. \n2. technical aspect: when we can not specify the maximum input size, the question that arises is how to technically deal with such inputs. \n\nRegarding the former, this paper does not actually resolve the problem. See for example the part on self-attention starting from Proposition 4.1. To perform the normalization one needs to first do a pass to compute the individual normalization constants and then run in mini-batches. However, when backpropagating during training, the resulting model will be the full graph (as if the entire input was passed). Please describe how your method works on self-attention models, and provide wall clock and memory footprint comparisons with the baselines.\n\nRegarding the latter, a baseline would be to simply pre-define a fixed input size and sample randomly that many samples. For example, if the input size is 1000, we can sample say pairs of elements, and we can fix to sample say 10000 such pairs (each time forming the pairs at random). This would increase significantly the dataset size, and the overall model (assuming the task is not simple to just do max-pooling) may perform well. Moreover, at inference time such a procedure would directly yield a confidence interval since multiple samples will be yielded from a full test input, thus the different predictions could be used for confidence intervals.\nPlease include such a baseline in your comparisons.\n\nLastly, I do not understand how, given an arbitrary aggregation function (not necessarily MBC), the herein proposed framework specifies that one proceeds  (as the abstract alludes to)? I understand that if that aggregation function is MBC then the entire map will be MBC as well. Nonetheless, often in practice, one does not know what is a suitable aggregation function (and standard ones may not be).  Hence I do not understand the claim of how *any* function can be transformed into an MBC map -- e.g. if there is prior knowledge that standard MBC functions are not suitable for the problem, but some other non-MBC one is, how could one proceed?\n\n\n### Experiments\n\nIt seems like UMBC does not always improve upon the baseline, e.g. Diff. EM + UMBC in Fig. 4. Similarly, in Tab. 3 sometimes it does not provide improvement.\nCould you provide more insights into this, e.g. additional datasets, or running UMBC with Deep Sets and Slot Set Enc?\n\nPlease also provide comparisons with a naive approach to solve this problem, see the part above on motivation.\n\nRegarding the MCDropout, a natural question that arises is if the baseline methods significantly improve when it is used. Could you provide such experiments for completeness?\n\n\n### Theoretical results\nI find that Lemma 4.1. is relatively obvious since it is easy to see that if we aggregate the results in a way that satisfies the MBC property, the overall output will not change with respect to if we were to pass the full input). \nI do not understand why Theorem 4.1. is not a simple corollary of Lemma 4.1, that is, since in Lemma 4.1  $f$ and $f^\\star$ are arbitrary maps, the statement in theorem 4.1 follows immediately as a special case of it (otherwise, a proof of it would be needed).\n\n\n", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ujbs4CKxES", "reviewer": "Reviewer_cnQu", "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": "", "summary_of_the_paper": "This paper considers the problem of training neural networks which operate on sets.  The main focus is on the mini-batch consistency (MBC) property, which roughly states the following: Let $X$ be a set, and $(X_1, \\dots, X_n)$.  Given a map $f$ from sets to vectors and an aggregation function $g$, the pair $(f,g)$ are MBC if $g(f(X_1), \\dots, f(X_n)) = f(X)$.  Informally, this means that the aggregated output from processing the partition of $X$ should give the same output as if $X$ was passed directly through $f$.\n\nThe authors show that non-MBC functions can be made MBC by pre-processing the data with an MBC function.  They then introduce a so-called \"universal\" MBC (UMBC) function.  They then present several experiments using their UMBC.", "summary_of_the_review": "To summarize, this paper is well-written and it has a broad array of experiments.  However, there are quite a few points of confusion in the notation and description of the setting.  Furthermore, I would argue that empirical results are inconclusive and that the theory does not constitute a significant contribution.  Therefore, I recommend that this paper not be accepted.\n\n**Post-rebuttal.**  The authors fixed a number of typos, clarified the notation, reworded some parts of the text, and re-captioned the plots of some of the experiments.  All of this has improved the paper, and therefore I will raise my score from 3 --> 5.  However, as discussed in my response, I think that there are still some fundamental issues which make me lean toward reject.", "strength_and_weaknesses": "## Strengths\n\n**Writing.**  The paper is well-written.  Of my batch of 7 papers, I would rank this as the second strongest in terms of writing.  The ideas are clearly articulated, and the flow of the paper makes sense.\n\n**Experiments in different settings.**  This paper conducts a range of experiments, from clustering to ablation studies to distribution shift robustness.  This is commendable, and a strong point of the paper.\n\n## Weaknesses\n\n**Defining notation and key terms.**  There were a few points in the paper where I had trouble parsing terms introduced by the authors.  For example, the authors talk about consistency: \n\n> \"Subsequent work has highlighted the utility of Mini-Batch Consistency (MBC), the ability to sequentially process any permutation of a set partition scheme (e.g. streaming chunks of data) while maintaining consistency guarantees on the output.\"\n\nAt this point in the paper (only a few lines in), the reader has no way of knowing what \"consistency\" means here.  It becomes somewhat inferable by the time the reader gets to Section 3, but I fear that some readers will be lost.  It would be worth explaining the main ideas of the paper in a way that does not require technical definitions in the abstract.\n\nOther feedback regarding similar weaknesses:\n\n* The authors seem to use MBC to refer to the term \"mini-batch consistency\" in the abstract and \"mini-batch consistent\" in the intro.  It would be clearer if the authors stuck to using MBC to refer to one or the other of these two phrases.\n\n* It's unclear what the authors mean by a \"valid\" set in Section 3.  \n\n* The presentation of set functions is confusing.  At the beginning of Section 3, set functions $f$ are said to map input sets $X$ to output sets $Y$.  However, later on, Property 3.1 tells us that set2vector functions map sets to one or several vectors.  And from this point onward, the authors seem to assume that all set functions are set2vector (e.g., in Thm. 4.1 -- the main result in this paper).  Given this, I think there needs to be clarification regarding what the input and outputs to the set functions are.  If we assume that $f$ maps sets to sets, then Thm. 4.1 seems to not apply.  If we assume that $f$ maps sets to vectors, then $g\\circ f$ is vector valued, and therefore the composition $f^\\star \\circ g\\circ f$ doesn't make sense if $f^\\star$ takes sets (not vectors) as input.\n\n* The notation in Property 3.2 is confusing.  The equation $f([x_{\\pi(i)}, \\dots x_{\\pi(N)}]) = [f_{\\pi(1)}(X_1), \\dots, f_{\\pi(n)}(x_n)]$ doesn't make sense to me.  Why does the permutation on the LHS go from $\\pi(i)$ to $\\pi(N)$ and why does the permutation on the RHS go from $\\pi(1)$ to $\\pi(n)$.  Should we assume that $N=n$?  What is $i$?\n\n* What do the authors mean when they say that MBC \"[added] a new dimension to the original view of Property 3.1?\"  In what way does this add a dimension?  Are we to interpret this mathematically or intuitively?  If the former, how does dimension come into play?  If the latter, it's not clear *how* this adds dimension, and it would be helpful if the authors could expand on this point.\n\n* Sets are often written as *elements* of $R^k$ for some $k$ or of $R^{n\\times k}$ for some $n$ and $k$.  I find this relatively confusing, because there are also various vectors in play which are truly elements of Euclidean space.  Is there a better way to denote which objects are sets and which are vectors?  For example, is it correct to say that $\\{1,2,3\\}\\in\\R^3$?  I would argue that the answer is no, because $\\{1,2,3\\}$ denotes the same set as $\\{3,2,1\\}$, but clearly as vectors $[3,2,1]$ is not the same as $[1,2,3]$.  So there is an identifiably issue.\n\n* The function Attention$(S,X)$ is used almost a full page before it is defined (on page 4).  This confused me as I was reading.\n\n* In the definition of Property 3.3, what are $n_i$, $d$, and $d'$?  Should they be inferable from context?\n\n* What is a slot-normalized activation function?  In general, after reading the paper, it was not clear to me what a \"slot\" is.  As this seems relatively important, it would be helpful if the authors could give intuition here.  One could certainly read the paper of Bruno et al., but for this paper to be self-contained and so as not to confuse readers, giving intuition for the most relevant related works seems important here.\n\n* What is $\\hat{\\sigma}$, i.e., what does the hat denote?\n\n* What do the authors mean when they say \"arbitrarily hard MBC constraints.\"  What makes one constraint harder than another?  In what sense can this get \"arbitrarily hard?\"\n\n* Why are MBC models like Deep Sets not able to leverage pairwise relationships? \n\n* What do the authors mean by \"simple\" $f^\\star$s?\n\n* In the intro, the authors mention infinite set sizes.  So it's reasonable to ask: How does all of this extend to settings where the cardinality of the input is infinite (countably? uncountably?)?  How does one even construct a partition?\n\n**Experiments.**  The experiments do not seem to support the conclusions made by the authors.  Indeed, I'm concerned that I missed something fundamental here, and if so, I hope that the authors will explain further.  However, it seems to me that baselines like Deep Sets ofter outperform the UMBC approach outlined by the authors.  For example, in Figure 4, the authors say that panels (e) and (f) show the best performance.  However, Diff. Em and Set Transformer (panels (c) and (d)) seem to reach lower values of the NLL for larger test set sizes.  Even for small test set sizes, (c) and (d) seem to do better.  Similarly, in Table 3, Deep Sets gets higher accuracy than all of the other methods on test sets of size 1000 and 2048.  Somewhat confusingly, these numbers are not bolded, despite the fact that (i) they perform better and (ii) (from my understanding) Deep Sets *is* MBC.  In the same categories, Deep Sets also outperforms UMBC on the NLL score.  And in Figure 6, it's not clear that UMBC does much better than the baselines.  So to summarize, the results seem relatively inconclusive as to why UMBC should be used when architectures like Deep Sets seem to offer strong performance.\n\n**Theory.**  The theory is also not a particular strength of this paper.  It's relatively strong to call the results lemmas and theorem in my opinion.  Thm. 4.1 says that if you preprocess the data in an MBC way, any function will be MBC which seems self-evident.  That is, if I have a set function $f^\\star$, if the input is in chunks, I can just use a separate MBC architecture $(g,f)$ to ensure that the output of the $g(f(X))$ is the same regardless of whether I chunked or not.  This holds directly by definition of MBC.  \n\nMore broadly, the question is: What is the impact of this theory?  Does it change our understanding or result in significantly better empirical results?  And based on my understanding of the experiments (see the discussion above), I'm not sure that it does.", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CRWXl61TGx5", "reviewer": "Reviewer_9w5z", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "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": "", "summary_of_the_paper": "This paper proposes a framework for converting any non-Mini-Batch-Consistent (non-MBC) models to an MBC model. Specifically, for any non-MBC set function, $f^\\star$, we can convert it to an MBC function by plugging in an MBC function before it. This framework also enables incorporating uncertainty estimation methods, such as MC-Dropout for neural set functions. The authors also conducted extensive experiments and ablation studies for demonstrating the effectiveness of their algorithms.\n", "summary_of_the_review": "Overall, I think the idea in this paper is interesting and novel, though it's not quite hard to derive. However, I have some doubts about the Lemmas and theorems in the paper as well as the empirical results, as detailed in the Strength&Weakness section. Also, I don't work in this area, so I am not able to assess the potential impact of this work. \nI will consider to increase the rating if the authors can address my concerns.", "strength_and_weaknesses": "Strengths:\n- The paper is mostly well-written, though there are some parts I feel are confusing (detailed later). \n\n- The proposed framework for converting any non-MBC functions to MBC functions is simple and general. \n\nWeakness:\n- I am a bit confused with Lemma 4.1 and Theorem 4.1. As in Proposition 3.3, the Mini-Batch-Consistency (MBC) is defined for a set function f and an aggregation function g, instead of just a single function. So, what do you mean by F satisfy property 3.3? What's the corresponding aggregation function for F?\n\n- Monte Carlo Slot Dropout: why does the composition of the functions $f^\\star$ and the dropout mask still a set function? Once the dropout mask is given, shouldn't the permutation of the slots affect the output? \n\n\n- Figure 4: It seems that UMBC + Diff. EM performs worse than Diff. EM alone, and this is also the case for UMBC + set transformer vs. Set transformer? I.e., UMBC + Diff. EM gives a larger negative log-likelihood, hence a lower log-likelihood. \n\n- Table 1: Why does Oracle perform the worst in terms of NLL? ", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject