Scalable Set Encoding with Universal Mini-Batch Consistency and
Unbiased Full Set Gradient Approximation

Let $\mathfrak{X}$ be a $d_{x}$-dimensional vector space over $\mathbb{R}$ and let $2^{\mathfrak{X}}$ be the power set of $\mathfrak{X}$. We focus on a collection of finite sets $\mathcal{X}$, which is a subset of $2^{\mathfrak{X}}$ such that $\sup_{X\in\mathcal{X}}\lvert X\rvert\in\mathbb{N}$. We want to construct a parametric function $f_{\theta}:\mathcal{X}\to\mathcal{Z}$ satisfying permutation invariance. Specifically, given a set $X_{i}=\{{\mathbf{x}}_{i,j}\}_{j=1}^{N_{i}}\in\mathcal{X}$, the output of the function $Z_{i}=f_{\theta}(X_{i})$ is a fixed sized representation which is invariant to all permutations of the indices $\{1,\ldots,N_{i}\}$. For supervised learning, we define a task specific decoder $g_{\lambda}:\mathcal{Z}\to\mathbb{R}^{d_{y}}$ and optimize parameters $\theta$ and $\lambda$ to minimize the loss

on training data $((X_{i},y_{i}))_{i=1}^{n}$, where $y_{i}$ is a label for the input set $X_{i}$ and $\ell$ denotes a loss function.

We further assume that the cardinality of a set $X$ is sufficiently large, such that loading and processing the whole set at once is computationally prohibitive. For non-MBC models, a naïve approach to solve this problem would be to encode a small subset of the full set as an approximation, leading to a possibly suboptimal representation of the full set. Instead, Bruno et al. (2021) propose a mini-batch consistent (MBC) set encoder, the Slot Set Encoder (SSE), to piecewise process disjoint subsets of the full set and aggregate them to obtain a consistent full set representation.

Models which satisfy the MBC property can partition a set into subsets, encode, and then aggregate the subset representations to achieve the exact same output as encoding the full set. Due to constraints on the architecture of the SSE, however, on certain tasks the SSE shows weaker performance than non-MBC set encoders such as Set Transformer (Lee et al., 2019) which utilizes self attention. To tackle this limitation, we propose Universal MBC (UMBC) set encoders which are both MBC and also allow for the use of arbitrary non-MBC set functions while still satisfying MBC property.

In this section, we provide a formulation of our UMBC set encoder $f_{\theta}$. Given an input set $X\in\mathcal{X}$, we represent it as a matrix $X=[{\mathbf{x}}_{1}\cdots{\mathbf{x}}_{N}]^{\top}\in\mathbb{R}^{N\times d_{x}}$ whose rows are elements in the set, and independently process each element with $\phi:\mathbb{R}^{d_{x}}\to\mathbb{R}^{d_{h}}$ as $\Phi(X)=[\phi({\mathbf{x}}_{1})\cdots\phi({\mathbf{x}}_{N})]^{\top}$, where $\phi$ is a deep feature extractor. We then compute the un-normalized attention score between a set of learnable slots $\Sigma=[{\mathbf{s}}_{1}\cdots{\mathbf{s}}_{k}]^{\top}\in\mathbb{R}^{k\times d_{s}}$ and $\Phi(X)$ as:

where $\sigma$ is an element-wise activation function with $\sigma(x)\gneq 0$ for all $x\in\mathbb{R}$, LN denotes layer normalization (Ba et al., 2016), and $W^{Q}\in\mathbb{R}^{d_{s}\times d},W^{K}\in\mathbb{R}^{d_{h}\times d},W^{V}\in\mathbb{R}^{d_{h}\times d}$ are parameters which are part of $\theta$. For simplicity, we omit biases for $Q,K$, and $V$. With the un-normalized attention score $\hat{A}$, we can define a map

for $p=1,2$, where $\nu_{p}:\mathbb{R}^{k\times d}\to\mathbb{R}^{k\times d}$ is defined by either $\nu_{1}(\hat{A})_{i,j}=\hat{A}_{i,j}/\sum_{i=1}^{k}\hat{A}_{i,j}$ which normalizes the columns or the identity mapping $\nu_{2}(\hat{A})_{i,j}=\hat{A}_{i,j}$. The choice of $\nu_{p}$ depends on the desired activation function $\sigma$. Alternatively, similar to slot attention (Locatello et al., 2020), we can make the function stochastic by sampling ${\mathbf{s}}_{i}\sim\mathcal{N}(\mathbf{\mu}_{i},\text{diag}(\texttt{softplus}(\mathbf{v}_{i})))$ with reparameterization (Kingma & Welling, 2013) for $i=1,\ldots,k$, where $\mu_{i}\in\mathbb{R}^{d_{s}},\mathbf{v}_{i}\in\mathbb{R}^{d_{s}}$ are part of the parameters $\theta$. If we sample ${\mathbf{s}}_{1},\ldots,{\mathbf{s}}_{k}\stackrel{{\scriptstyle\text{iid}}}{{\sim}}\mathcal{N}(\mu_{1},\text{diag}(\texttt{softplus}(\mathbf{v}_{1})))$ with a sigmoid for $\sigma$ and $\nu_{1}$ for normalization, and then apply a pooling function (sum, mean, min, or max) to the columns of $[\hat{f}_{\theta}(X)^{\top}_{1}\cdots\hat{f}_{\theta}(X)^{\top}_{k}]^{\top}\in\mathbb{R}^{k\times d}$, we achieve a function equivalent to the SSE, where $\hat{f}_{\theta}(X)_{i}$ is $i$-th row of ${\hat{f}}_{\theta}(X)$.

However, SSE has some drawbacks. First, since the attention score of $\nu_{p}(\hat{A})_{i,j}$ is independent to the other $N-1$ attention scores $\nu_{p}(\hat{A})_{i,l}$ for $l\neq j$, it is impossible for the rows of $\nu_{p}(\hat{A})$ to be convex coefficients as the softmax outputs in conventional attention (Vaswani et al., 2017). Notably, in some of our experiments, the constrained attention activation originally used in the SSE, which we call slot-sigmoid, significantly degrades generalization performance. Furthermore, stacking hierarchical SSE layers has been shown to harm performance (Bruno et al., 2021), which limits the power of the overall model.

To overcome these limitations of the SSE, we propose a Universal Mini-Batch Consistent (UMBC) set encoder $f_{\theta}$ by allowing the set function $f_{\theta}$ to also use arbitrary non-MBC functions. Firstly, we propose normalizing the attention matrix $\nu_{p}(\hat{A})$ over rows to consider dependency among different elements of the set in the attention operation:

where $\bm{1}_{N}=(1,\ldots,1)\in\mathbb{R}^{N}$. We prove that a UMBC set encoder $f_{\theta}$ is permutation invariant, equivariant, and MBC.

Any strictly positive elementwise function is a valid $\sigma$. For an instance, if we use the identity mapping $\nu_{2}$ with $\sigma(\cdot)\coloneqq\exp(\cdot)$, the attention matrix $\text{diag}({\bar{f}}_{\theta}(X))^{-1}\nu_{p}(\hat{A})$ is equivalent to applying the softmax to each row of $\hat{A}$, which is hypothesized to break the MBC property by Bruno et al. (2021). However, we show that this does not break the MBC property in Appendix A.2. Intuitively, since

holds for any partition $\zeta(X)$ of the set $X$, we can iteratively process each subset $S\in\zeta(X)$ and aggregate them without losing any information of $f_{\theta}(X)$, i.e., $f_{\theta}$ is MBC even when normalizing over the $N$ elements of the set. Note that the operation outlined above is mathematically equivalent to the softmax, but uses a non-standard implementation. We discuss the implementation and list 5 such valid attention activation functions which satisfy the MBC property in Appendix I.

Lastly, we may consider the output of a UMBC set encoder $f_{\theta}(X)$ as either a fixed vector or a set of $k$ elements. Under the set interpretation, we may therefore apply subsequent functions on the set of cardinality $k$. To provide a valid input to subsequent set encoders, it is sufficient to view UMBC as a set to set function $\varphi(\Sigma;X,\theta):\Sigma\mapsto f_{\theta}(X)$ for each set $X\in\mathcal{X}$, which is permutation equivariant w.r.t. the slots $\Sigma$.

A key insight is that we can leverage non-MBC set encoders such as Set Transformer after a UMBC layer to improve expressive power of an MBC model while still satisfying MBC (Definition 3.2). As a result, with some assumptions, a UMBC set encoder used in combination with any continuously sum decomposable (Zaheer et al., 2017) permutation invariant deep neural network is a universal approximator of the class of continuously sum decomposable functions.

Although we use a non-MBC set encoder on top of UMBC, this does not violate the MBC property. Since we may obtain $f_{\theta}(X)$ by sequentially processing each subset of $X$ and the resulting set with cardinality $k$ is assumed small enough to load $f_{\theta}(X)$ in memory, we can directly provide the MBC output of UMBC to the non-MBC set encoder.

For notational convenience, we write $g_{\lambda}$ to indicate the composition $g_{\lambda}\circ\hat{g}_{\omega}$ of a set encoder $\hat{g}_{\omega}:\mathbb{R}^{k\times d}\to\mathcal{Z}$ and a decoder $g_{\lambda}:\mathcal{Z}\to\mathbb{R}^{d_{y}}$, throughout the paper. Similarly, the parameter $\lambda$ denotes $(\omega,\lambda)$.

Although we can leverage SSE or UMBC at test-time by sequentially processing subsets to obtain the full set representation $f_{\theta}(X)$, at train-time it is infeasible to utilize the gradient of the loss (equation 1) w.r.t. the full set. Computation of the full set gradient with automatic differentiation requires storing the entire computation graph for all forward passes of each subset $S$ from $\zeta(X)$ denoted as a partition of a set $X$, which incurs a prohibitive computational cost for large sets. As a simple approximation, Bruno et al. (2021) propose randomly sampling a single subset $S_{i,j}\in\zeta(X_{i})$ and computing the gradient of the loss $\ell((g_{\lambda}\circ f_{\theta})(S_{i,j}),y_{i})$ based on a single subset at each iteration.

In order to tackle this issue, we propose an unbiased estimation of the full set gradient which incurs a constant memory overhead. Firstly, we uniformly and independently sample a mini-batch $((\bar{X}_{i},{\bar{y}}_{i}))_{i=1}^{m}$ from the training dataset $((X_{i},y_{i}))_{i=1}^{n}$ for every iteration $t\in\mathbb{N}_{+}$. We denote this process by $((\bar{X}_{i},{\bar{y}}_{i}))^{m}_{i=1}\sim D[((X_{i},y_{i}))^{n}_{i=1}].$ Then, for each $\bar{X}_{i}$, we sample a mini-batch ${\bar{\zeta}}_{t}(\bar{X}_{i})=\{{\bar{S}}_{1},\ldots,{\bar{S}}_{\lvert{\bar{\zeta}}_{t}(\bar{X}_{i})\rvert}\}$ from the partition $\zeta_{t}(\bar{X}_{i})=\{S_{1},\ldots,S_{\lvert\zeta_{t}(\bar{X}_{i})\rvert}\}$ of $\bar{X}_{i}$, i.e., all ${\bar{S}}_{j}$ are drawn independently and uniformly from $\zeta_{t}(\bar{X}_{i})$. Denote this process by ${\bar{\zeta}}_{t}(\bar{X}_{i})\sim D[\zeta_{t}(\bar{X}_{i})]$. Instead of storing the computational graph of all forward passes of subsets in the partition $\zeta_{t}(\bar{X}_{i})$ of a set $\bar{X}_{i}$, we apply StopGrad to all subsets $S\notin{\bar{\zeta}}_{t}(\bar{X}_{i})$ as follows:

where, for any function $q:\theta\mapsto q(\theta)$, the symbol $\texttt{StopGrad}(q(\theta))$ denotes a constant with its value being $q(\theta)$, i.e., $\partial\texttt{StopGrad}(q(\theta))/\partial\theta=0$. For simplicity, we omit the superscript $\zeta_{t}$ if there is no ambiguity. Finally, we update both the parameter $\theta$ and $\lambda$ of the respective encoder and decoder using the gradient of the following functions, respectively at $t\in\mathbb{N}_{+}$:

We outline our proposed training method in Algorithm 1. Note that we can apply our algorithm to any set encoder for which a full set representation can be decomposed into a summation of subset representations as in equation 7 such as Deepsets with sum or mean pooling, or SSE which are in fact special cases of UMBC. Furthermore, we can apply the algorithm to any differentiable non-MBC set encoder if we simply place a UMBC layer before the non-MBC function. As a consequence of the $\texttt{StopGrad}()$ operation, if we set $\lvert{\bar{\zeta}}_{t}(\bar{X}_{i})\rvert=1$, our method incurs the same computation graph storage cost as randomly sampling a single subset. Moreover, $\partial L_{t,1}(\theta,\lambda)/\partial\theta$ and $\partial L_{t,2}(\theta,\lambda)/\partial\lambda$ are unbiased estimators of $\partial L(\theta,\lambda)/\partial\theta$ and $\partial L(\theta,\lambda)/\partial\lambda$, respectively.

Under mild conditions, Theorem 3.10 implies that the updating $\theta$ with our method makes progress towards minimizing $L(\theta,\lambda)$. We formalize this statement in Appendix B.

We consider amortized clustering on a dataset generated from Mixture of $K$ Gaussians (MoGs) (See Appendix C for dataset construction details). Given a set $X_{i}=\{{\mathbf{x}}_{i,j}\}_{j=1}^{N_{i}}$ sampled from a MoGs, the goal is to output the mixing coefficients, and Gaussian mean and variance, which minimizes the negative log-likelihood of the set as follows:

where $\Theta_{j}(X_{i})=(\mu_{j}(X_{i}),\Sigma_{j}(X_{i}))$ denotes a mean vector and a diagonal covariance matrix for $j$-th Gaussian, and $\pi_{j}(X_{i})$ is $j$-th mixing coefficient. Note that there is no label $y_{i}$ since it is an unsupervised clustering problem. We optimize the parameters of the set encoder $\theta$ to minimize the loss over a batch, $L(\theta)=1/n\sum_{i=1}^{n}\ell(f_{\theta}(X_{i}))$.

Setup. We evaluate training with the full set gradient vs. the unbiased estimation of the full set gradient. In this setting, for gradient computation, MBC models use a subset of 8 elements from a full set of 1024 elements. Non MBC models such as Set Transformer (ST), FSpool (Zhang et al., 2020), and Diff EM (Kim, 2022) are also trained with the set of 8 elements. We compare our UMBC model against Deepsets, SSE, Set Transformer, FSPool, and Diff EM. Note that at test-time all non-MBC models process every 8 element subset from the full set independently and aggregate the representations with mean pooling.

Results. In Figure 4, interestingly, the unbiased estimation of the full set gradient (red) is almost indistinguishable from the full set gradient (blue) for DeepSets and UMBC, while there is a significant gap for SSE. In all cases, the unbiased estimation of the full set gradient outperforms training with the biased gradient approximation with only the set of 8 elements per random sample (green), which is proposed by Bruno et al. (2021). Lastly, as shown in Table 2, we compare all models in terms of generalization performance (NLL), memory usage, and wall-clock time for processing a single subset. All non-MBC models show underperformance due to their violation of the MBC property. However, if we utilize ‘UMBC+’ compositions, the composition becomes MBC with significantly improved performance and little added overhead for memory and time complexity. In contrast, a composition of pure MBC functions, the Hierarchical SSE degrades the performance of SSE. Notably, UMBC with Set Transformer outperforms all other models whereas Set Transformer alone achieves the worst NLL. These results verify expressive power of UMBC in conjuction with non-MBC models.

In this task, we are given a set of $M$ RGB pixel values $y_{i}\in\mathbb{R}^{3}\times\cdots\times\mathbb{R}^{3}$ of an image as well as the corresponding 2-dimensional coordinates $({\mathbf{x}}_{i,1},\ldots,{\mathbf{x}}_{i,M})$ normalized to be in $[0,1]^{2}\times\cdots\times[0,1]^{2}$. The goal of the task is to predict RGB pixel values of all $M$ coordinates of the image. Specifically, given a context set $X_{i}=\{({\mathbf{x}}_{i,c_{j}},y_{i,c_{j}})\}_{j=1}^{N_{i}}$ processed by the set encoder, we obtain the set representation $f_{\theta}(X_{i})\in\mathbb{R}^{k\times d}$. Then, a decoder $g_{\lambda}:\mathbb{R}^{k\times d}\times[0,1]^{2}\times\cdots\times[0,1]^{2}\to\mathbb{R}^{6}\times\cdots\times\mathbb{R}^{6}$ which utilizes both the set representation and the target coordinates, learns to predict a mean and variance for each coordinate of the image as $((\hat{\mu}_{i,j,l})_{l=1}^{3},(\hat{\sigma}_{i,j,l})_{l=1}^{3})_{j=1}^{M}=g_{\lambda}({\mathbf{z}}_{i})$, where ${\mathbf{z}}_{i}=(f_{\theta}(X_{i}),{\mathbf{x}}_{i,1},\ldots,{\mathbf{x}}_{i,M})$. Then we compute the negative log-likelihood of the label set $y_{i}=((y_{i,j,l})_{l=1}^{3})_{j=1}^{M}$:

where $\mathcal{N}(\cdot;\mu,\sigma)$ is a univariate Gaussian probability density function. Finally, we optimize $\theta$ and $\lambda$ to minimize the loss $L(\theta,\lambda)=1/n\sum_{i=1}^{n}\ell(g_{\lambda}({\mathbf{z}}_{i}),y_{i})$.

Setup. In our experiments, we impose the restriction that a set encoder is only allowed to compute the gradient with 100 elements of a context set $X_{i}$ during training and the model can only process 100 elements of the context set at once at test time. We train the set encoders in a Conditional Neural Process (Garnelo et al., 2018) framework, using $32\times 32$ images from the CelebA dataset (Liu et al., 2015). We vary the cardinality of the context set size $N_{i}\in\{100,\ldots,500\}$ and compare the negative log-likelihood (NLL) of each model. For baselines, we compare our UMBC set encoder against: Deepsets, SSE, Hierarchical SSE, Set Transformer (ST), FSPool, and Diff EM. For our UMBC, we use the softmax for $\sigma$ in equation 3 and place the ST after the UMBC layer.

Results. First, as shown in Figure 5(a), our UMBC + ST model outperforms all baselines, empirically verifying the expressive power of UMBC. SSE underperforms in terms of NLL due to its constrained architecture. Moreover, stacking hierarchical SSE layers degrades the performance of SSE for larger sets. Note that all MBC set encoders (Deepsets, SSE, Hierarchical SSE and UMBC) in Figure 5(a) are trained with our proposed unbiased gradient approximation in Algorithm 1. On the other hand, we train non-MBC models such as Set Transformer (ST), FSPool, and Diff EM with a randomly sampled subset of 100 elements, and perform mean pooling over all subset representations at test-time to approximate an MBC model.

Additionally, Figure 5(b) shows GPU memory usage for each model while processing sets of varying cardinalities without a memory constraint. The marker size is proportional to the set cardinality. Notably, all four MBC models incur a constant memory overhead to process any set size, as we can apply StopGrad to most of the subset, and compute an unbiased gradient estimate with a fixed sized subset (100). However, memory overhead for all non-MBC models is a function of set size. Thus, Set Transformer uses more than twice the memory of UMBC to achieve a similar log-likelihood on a set of 500 elements.

Lastly, in Figure 5(c), we show how our proposed unbiased training algorithm (red) improves the generalization performance of UMBC models compared to training with a limited subset of 100 elements (green). Notably, the performance of our algorithm is indistinguishable from that of training models with the full set gradient (blue). We present similar plots for Deepsets and SSE in Figures 11(a) and 11(b). Across all models, our training algorithm significantly and consistently improves performance compared to training with random subsets of 100 elements, while requiring the same amount of memory. These empirical results verify both efficiency and effectiveness of our proposed method.

In this task, we are given a long document $X_{i}=({\mathbf{x}}_{i,1},\ldots,{\mathbf{x}}_{i,N_{i}})$ consisting of an average of $707.99$ words. The goal of this task is to predict a binary multi-label $y_{i}=(y_{i,1},\ldots,y_{i,c})\in\{0,1\}^{c}$ of the document, where $c$ is the number of classes. We ignore the order of words and consider the document as a multiset of words, i.e., a set allowing duplicate elements. Specifically, given a training dataset $((X_{i},y_{i}))_{i=1}^{n}$, we process each set $X_{i}$ with the set encoder to obtain the set representation $f_{\theta}(X_{i})\in\mathbb{R}^{k\times d}$. We then use a decoder $g_{\lambda}:\mathbb{R}^{k\times d}\to\mathbb{R}^{c}$ to output the probability of each class and compute the cross entropy loss:

where ${\mathbf{z}}_{i}=(z_{i,1},\ldots,z_{i,c})=(g_{\lambda}\circ f_{\theta})(X_{i})$ and $\tilde{\sigma}$ denotes the sigmoid function. Finally we optimize $\theta$ and $\lambda$ to minimize the loss $L(\theta,\lambda)=1/n\sum_{i=1}^{n}\ell({\mathbf{z}}_{i},y_{i})$.

Setup. All models are trained on the inverted EURLEX dataset (Chalkidis et al., 2019) consisting of long legal documents divided into sections. The order of sections are inverted following prior work (Park et al., 2022). To predict a label, we give the whole document to the model without any truncation. We compare the micro F1 of each model.

We compare UMBC to Deepsets, SSE, ToBERT (Pappagari et al., 2019), and Longformer (Beltagy et al., 2020). For Deepsets and SSE, we use the pre-trained word embedding from BERT (Devlin et al., 2019) without positional encoding and 2 layer fully connected (FC) networks for feature extractor $\phi$. We use another 3 layer FC network for the decoder. For UMBC, we use the same feature extractor as SSE but instead use the pre-trained BERT as a decoder, with a randomly initialized linear classifier. We remove the positional encoding of BERT for UMBC to ignore word order. For all the MBC models, we train them both with full set denoted as “w/ full” and with our gradient approximation method on a subset of 100 elements denoted as “w/ 100”.

Results. As shown in Table 3, our proposed UMBC outperforms all baselines including non-MBC models — Longformer and ToBERT which require excessive amounts of GPU memory for training models with long sequences. This result again verifies the expressive power of UMBC with BERT (a non-MBC model) for long document classification. Moreover, with significantly less GPU memory, all MBC models (Deepsets, SSE, and UMBC) trained with our unbiased gradient approximation using a subset of 100 elements, achieve similar performance to the models trained with full set. Lastly, in Figure 6, we plot the micro F1 score as a function of the cardinality of the subset used for gradient computation when training the UMBC model. Our proposed unbiased gradient approximation (red) shows consistent performance for all subset cardinalities. In contrast, training the model with a small random subset (green) is unstable, resulting in underperformance and higher F1 variance.

In MIL, we are given a ‘bag’ of instances with a corresponding bag label, but no labels for each instance within the bag. Labels should not depend on the order of the instances, i.e., MIL can be recast as a set classification problem. Specifically, given a set $X_{i}=\{{\mathbf{x}}_{i,j}\}_{j=1}^{N_{i}}$, the goal is to predict its binary label $y_{i}\in\{0,1\}$. For this task, we obtain two streams of set representations and compute the cross entropy loss from the decoder $g_{\lambda}:\mathbb{R}^{k\times d}\to\mathbb{R}$ output as:

where $w\in\mathbb{R}^{d_{h}}$ and $b\in\mathbb{R}$ are parameters and $\ell$ is the cross entropy loss described in equation 17 with $c=1$. We optimize all parameters $\theta,\lambda,w,\text{and }b$ to minimize the loss $1/n\sum_{i=1}^{n}\mathcal{L}_{i}$. At test time, we predict a label $y_{*}$ for a set $X_{*}$ as:

where ${\mathbf{z}}_{*,1}=\max\{w^{\top}\phi({\mathbf{x}}):{\mathbf{x}}\in X_{*}\}$, ${\mathbf{z}}_{*,2}=f_{\theta}(X_{*})$, $\tilde{\sigma}$ is the sigmoid function, $\mathbbm{1}$ is indicator function and $0<\tau<1$ is threshold tuned on the validation set.

Setup. We evaluate all models on the Camelyon16 Whole Slide Image cancer detection dataset (Bejnordi et al., 2017). Each instance consists of a high resolution image of tissue from a medical scan which is pre-processed into $256\times 256$ patches of RGB pixels. After pre-processing, the average number of patches in a single set is over 9,300 (7.3GB), making each input roughly equivalent to processing 1% of ImageNet1k (Deng et al., 2009). The largest input in the training set contains 32,382 patches (25.4 GB). We utilize a ResNet18 (He et al., 2016) which is pretrained on Camelyon16 (Li et al., 2021) via SimCLR (Chen et al., 2020) as a backbone feature extractor whose weights can be downloaded from this repository¹¹1https://github.com/binli123/dsmil-wsi. Our goal is to first pretrain MBC set encoders on the extracted features, and then use the unbiased estimation of the full set gradient to fine-tune the feature extractor on the full input sets. We evaluate the performance of UMBC against non-MBC MIL baselines: DS-MIL (Li et al., 2021) and AB-MIL (Ilse et al., 2018), as well as MBC baselines: DeepSets and SSE.

Results. As shown in Table 4, our UMBC model achieves the best accuracy and competitive AUROC score. Note that SSE shows the worst performance due to its constrained architecture, which even underperforms DeepSets in this task. These empirical results again verify the expressive power of our UMBC model. Moreover, we can further improve the performance of UMBC via fine-tuning the backbone network, ResNet18, which is only feasible as a consequence of our unbiased full set gradient approximation which incurs constant memory overhead. However, it is not possible for the non-MBC models to fine-tune with the ResNet18 since it is computationally prohibitive to compute the gradient of the ResNet18 with sets consisting of tens of thousands of patches with $256\times 256$ resolution.

To validate effectiveness of activation functions $\sigma$ for attention in equation 3, we train UMBC + Set Transformer with different activation functions listed in Table 15 for the amortized clustering and MIL pretraining tasks. As shown in Figures 7 and 7, softmax attention outperforms all the other activation functions whereas the slot-sigmoid used for attention in SSE underperforms. This experiment highlights the importance of choosing the proper activation function for attention, which is enabled by our UMBC framework.