SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In our first experiment, we introduce a new version of a Variational Autoencoder (VAE, Kingma and Welling, 2014), the DRPM Variational Clustering (DRPM-VC) model. The DRPM-VC enables clustering and unsupervised conditional generation in a variational fashion. To that end, we assume that each sample  $\pmb{x}$  of a dataset  $X$  is generated by a latent vector  $\pmb{z} \in \mathbb{R}^l$ , where  $l \in \mathbb{N}$  is the latent space size. Traditional VAEs would then assume that all latent vectors  $\pmb{z}$  are generated by a single Gaussian prior distribution  $\mathcal{N}(\mathbf{0},\mathbb{I}_l)$ . Instead, we assume every  $\pmb{z}$  to be sampled from one of  $K$  different latent Gaussian distributions  $\mathcal{N}(\pmb{\mu}_k,\mathrm{diag}(\pmb{\sigma}_k))$ ,  $k = 1,\dots ,K$ , with  $\pmb{\mu}_k,\pmb{\sigma}_k \in \mathbb{R}^l$ . Further, note that similar to an urn model (Section 3), if we draw a batch from a given finite dataset with samples from different clusters, the cluster assignments within that batch are not entirely independent. Since there is only a finite number of samples per cluster, drawing a sample from a specific cluster decreases the chance of drawing a sample from that cluster again, and the distribution of the number of samples drawn per cluster will follow an MVHG distribution. Previous work on variational clustering proposes to model the cluster assignment  $y \in \{0,1\}^{K}$  of each sample  $x$  through independent categorical distributions (Jiang et al., 2016), which might thus be over-restrictive and not correctly reflect reality. Instead, we propose explicitly modeling the dependency between the  $y$  of different samples by assuming they are drawn from an RPM. Hence, the generative process leading to  $X$  can be summarized as follows: First, the cluster assignments are represented as a partition matrix  $Y$  and sampled from our DRPM, i.e.,  $Y \sim p(Y;\omega ,s)$ . Given an assignment  $y$  from  $Y$ , we can sample the respective latent variable  $z$ , where  $z \sim \mathcal{N}(\pmb{\mu_y},\mathrm{diag}(\pmb{\sigma_y}))$ ,  $z \in \mathbb{R}^l$ . Note that we use the notational shorthand  $\pmb{\mu_y} := \pmb{\mu}_{\arg \max(y)}$ . Like in vanilla VAEs, we infer  $x$  by independently passing the corresponding  $z$  through a decoder model. Assuming this generative process, we derive the following evidence lower bound (ELBO) for  $p(X)$ :
$$
\mathcal {L} _ {E L B O} = \sum_ {\boldsymbol {x} \in X} \mathbb {E} _ {q (\boldsymbol {z} | \boldsymbol {x})} [ \log p (\boldsymbol {x} | \boldsymbol {z}) ] - \sum_ {\boldsymbol {x} \in X} \mathbb {E} _ {q (Y | X)} [ K L [ q (\boldsymbol {z} | \boldsymbol {x}) | | p (\boldsymbol {z} | Y) ] ] - K L [ q (Y | X) | | p (Y) ]
$$
Note that computing  $KL[q(Y|X)||p(Y)]$  directly is computationally intractable, and we need to upper bound it according to Lemma 4.2. For an illustration of the generative assumptions and more details on the ELBO, we refer to Appendix C.2.
Table 2: Partitioning of Generative Factors. We evaluate the learned latent representations of the four methods (Label-VAE, Ada-VAE, HG-VAE, DRPM-VAE) with respect to the shared (S) and independent (I) generative factors. We do this by fitting linear classifiers on the shared and independent dimensions of the representation, predicting the respective generative factors. We report the results in adjusted balanced accuracy (Sutter et al., 2023) across five seeds.
ns=0	ns=1	ns=3	ns=5
I	S	I	S	I	S	I	I
LABEL	0.14±0.01	0.19±0.03	0.16±0.01	0.10±0.00	0.23±0.01	0.34±0.00	0.00±0.00
ADA	0.12±0.01	0.19±0.01	0.15±0.01	0.10±0.03	0.22±0.02	0.33±0.03	0.00±0.00
HG	0.18±0.01	0.22±0.05	0.19±0.01	0.08±0.02	0.28±0.01	0.28±0.01	0.01±0.00
DRPM	0.26±0.02	0.39±0.07	0.2±0.01	0.15±0.01	0.29±0.02	0.42±0.03	0.01±0.00
To assess the clustering performance, we train our model on two different datasets, namely MNIST (LeCun et al., 1998) and Fashion-MNIST (FMNIST, Xiao et al., 2017), and compare it to three baselines. Two of the baselines are based on a Gaussian Mixture Model, where one is directly trained on the original data space (GMM), whereas the other takes the embeddings from a pretrained encoder as input (Latent GMM). The third baseline is Variational Deep Embedding (VaDE, Jiang et al., 2016), which is similar to the DRPM-VC but assumes i.i.d. categorical cluster assignments. For all methods except GMM, we use the weights of a pretrained encoder to initialize the models and priors at the start of training. We present the results of these experiments in Table 1. As can be seen, we outperform all baselines, indicating that modeling the inherent dependencies implied by finite datasets benefits the performance of variational clustering. While achieving decent clustering performance, another benefit of variational clustering methods is that their reconstruction-based nature intrinsically allows unsupervised conditional generation. In Figure 2, we present the result of sampling a partition and the corresponding generations from the respective clusters after training the DRPM-VC on FMNIST. The model produces coherent generations despite not having access to labels, allowing us to investigate the structures learned by the model more closely. We refer to Appendix C.2 for more illustrations of the learned clusters, details on the training procedure, and ablation studies.
5.2 Variational Partitioning of Generative Factors
Data modalities not collected as i.i.d. samples, such as consecutive frames in a video, provide a weak-supervision signal for generative models and representation learning (Sutter et al., 2023). Here, on top of learning meaningful representations of the data samples, we are also interested in discovering the relationship between coupled samples. If we assume that the data is generated from underlying generative factors, weak supervision comes from the fact that we know that certain factors are shared between coupled pairs while others are independent. The supervision is weak because we neither know the underlying generative factors nor the number of shared and independent factors. In such a setting, we can use the DRPM to learn a partition of the generative factors and assign them to be either shared or independent.
Figure 3: The mean squared errors between the estimated number of shared factors  $\hat{n}_s$  and the true number of shared factors  $n_s$  across five seeds for the Label-VAE, Ada-VAE, HG-VAE, and DRPM-VAE.
In this experiment, we use paired frames  $\mathbf{X} = [\pmb{x}_1, \pmb{x}_2]$  from the mpi3d dataset (Gondal et al., 2019). Every pair of frames shares a subset of its seven generative factors. We introduce the DRPM-VAE, which models the division of the latent space into shared and independent latent factors as RPM. We add a posterior approximation  $q(Y \mid \mathbf{X})$  and additionally a prior distribution of the form  $p(Y)$ . The model maximizes the following ELBO on the
marginal log-likelihood of images through a VAE (Kingma and Welling, 2014):
$$
\begin{array}{l} \mathcal {L} _ {E L B O} = \sum_ {j = 1} ^ {2} \mathbb {E} _ {q \left(\boldsymbol {z} _ {s}, \boldsymbol {z} _ {j}, Y \mid \boldsymbol {X}\right)} \left[ \log p \left(\boldsymbol {x} _ {j} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {j}\right) \right] \tag {11} \\ - \mathbb {E} _ {q (Y | \boldsymbol {X})} [ K L [ q (\boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}, \boldsymbol {z} _ {2} \mid Y, \boldsymbol {X}) | | p (\boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}, \boldsymbol {z} _ {2}) ] ] - K L [ q (Y \mid \boldsymbol {X}) | | p (Y) ] \\ \end{array}
$$
Similar to the ELBO for variational clustering in Section 5.1, computing  $KL[q(Y\mid X)||p(Y)]$  directly is intractable, and we need to bound it according to Lemma 4.2.
We compare the proposed DRPM-VAE to three methods, which only differ in how they infer shared and latent dimensions. While the Label-VAE (Bouchacourt et al., 2018; Hosoya, 2018) assumes that the number of independent factors is known, the Ada-VAE (Locatello et al., 2020) relies on a heuristic-based approach to infer shared and independent latent factors. Like in Locatello et al. (2020) and Sutter et al. (2023), we assume a single known factor for Label-VAE in all experiments. HG-VAE (Sutter et al., 2023) also relies on the MVHG to model the number of shared and independent factors. Unlike the proposed DRPM-VAE approach, HG-VAE must rely on a heuristic to assign latent dimensions to shared factors, as the MVHG only allows to model the number of shared and independent factors but not their position in the latent vector. We use the code from Locatello et al. (2020) and follow the evaluation in Sutter et al. (2023). We refer to Appendix C.3 for details on the ELBO, the setup of the experiment, the implementation, and an illustration of the generative assumptions.
We evaluate all methods according to their ability to estimate the number of shared generative factors (Figure 3) and how well they partition the latent representations into shared and independent factors (Table 2). Because we have access to the data-generating process, we can control the number of shared  $n_s$  and independent  $n_i$  factors. We compare the methods on four different datasets with  $n_s \in \{0,1,3,5\}$ . In Figure 3, we demonstrate that the DRPM-VAE accurately estimates the true number of shared generative factors. It matches the performance of HG-VAE and outperforms the other two baselines, which consistently overestimate the true number of shared factors. In Table 2, we see a considerable performance improvement compared to previous work when assessing the learned latent representations. We attribute this to our ability to not only estimate the subset sizes of latent and shared factors like HG-VAE but also learn to assign specific latent dimensions to the corresponding shared or independent representations. Thus, the DRPM-VAE dynamically learns more meaningful representations and can better separate and infer the shared and independent subspaces for all dataset versions.
The DRPM-VAE provides empirical evidence of how RPMs can leverage weak supervision signals by learning to maximize the data likelihood while also inferring representations that capture the relationship between coupled data samples. Additionally, we can explicitly model the data-generating process in a theoretically grounded fashion instead of relying on heuristics.
5.3 Multitask Learning
Many ML applications aim to solve specific tasks, where we optimize for a single objective while ignoring potentially helpful information from related tasks. Multitask learning (MTL) aims to improve the generalization across all tasks, including the original one, by sharing representations between related tasks (Caruana, 1993; Caruana and de Sa, 1996). Recent works (Kurin et al., 2022; Xin et al., 2022) show that it is difficult to outperform a convex combination of task losses if the task losses are appropriately scaled. I.e., in case of equal difficulty of the two tasks, a classifier with equal weighting of the two classification losses serves as an upper bound in terms of performance. However, finding suitable task weights is a tedious and inefficient approach to MTL. A more automated way of weighting multiple tasks would thus be vastly appreciated.
In this experiment, we demonstrate how the DRPM can learn task difficulty by partitioning a network layer. Intuitively, a task that requires many neurons is more complex than a task that can be solved using a single neuron. Based on this observation, we propose the DRPM-MTL. The DRPM-MTL learns to partition the neurons of the last shared layer such that only a subset of the neurons are used for every task. In contrast to the other experiments (Sections 5.1 and 5.2), we use the DRPM without resampling and infer the partition  $Y$  as a deterministic function. This can be done by applying the two-step procedure of Proposition 4.1 but skipping the resampling step of the MVHG and PL distributions. We compare the DRPM-MTL to the unitary loss scaling method (ULS, Kurin et al.,
2022), which has a fixed architecture and scales task losses equally. Both DRPM-MTL and ULS use a network with shared architecture up to some layer, after which the network branches into two task-specific layers that perform the classifications. Note the difference between the methods. While the task-specific branches of the ULS method access all neurons of the last shared layer, the task-specific branches of the DRPM-MTL access only the subset of neurons reserved for the respective task.
We perform experiments on MultiMNIST (Sabour et al., 2017), which overlaps two MNIST digits in one image, and we want to classify both numbers from a single sample. Hence, the two tasks, classification of the left and the right digit (see Appendix C.4 for an example), are approximately equal in difficulty by default. To increase the difficulty of one of the two tasks, we introduce the noisyMultiMNIST dataset. There, we control task difficulty by adding salt and pepper noise to one of the two digits, subsequently increasing the difficulty of that task with increasing noise ratios. Varying the noise, we evaluate how our DRPM-MTL adapts to imbalanced difficulties, where one usually has to tediously search for optimal loss weights to reach good performance. We base our pipeline on (Sener and Koltun, 2018). For more details and additional CelebA MTL experiments we refer to Appendix C.4.
We evaluate the DRPM-MTL concerning its classification accuracy on the two tasks and compare the inferred subset sizes per task for different noise ratios  $\alpha \in \{0.0,\dots ,0.9\}$  of the noisyMultiMNIST dataset (see Figure 4). The DRPM-MTL achieves the same or better accuracy on both tasks for most noise levels (upper part of Figure 4). It is interesting to see that, the more we increase  $\alpha$  ,the more the DRPM-MTL
tries to overcome the increased difficulty of the right task by assigning more dimensions to it (lower part of Figure 4, noise ratio  $\alpha$  0.6-0.8). Note that for the maximum noise ratio of  $\alpha = 0.9$ , it seems that the DRPM-MTL basically surrenders and starts neglecting the right task, instead focusing on getting good performance on the left task, which impacts the average accuracy.
Figure 4: Results for noisyMultiMNIST experiment. In the upper plot, we compare the task accuracy of the two methods ULS and the DRPM-MTL. We see that the DRPM-MTL can reach higher accuracy for most of the different noise ratios  $\alpha$  while it assigns the number of dimensions per task according to their difficulty.
Limitations & Future Work
The proposed two-stage approach to RPMs requires distributions over subset sizes and permutation matrices. The memory usage of the permutation matrix used in the two-stage RPM increases quadratically in the number of elements  $n$ . Although we did not experience memory issues in our experiments, this may lead to problems when partitioning vast sets. Furthermore, learning subsets by first inferring an ordering of all elements can be a complex optimization problem. Approaches based on minimizing the earth mover's distance (Monge, 1781) to learn subset assignments could be an alternative to the ordering-based approach in our DRPM and pose an interesting direction for future work. Finally, note that we compute the probability mass function (PMF)  $p(\bar{Y};\omega ,s)$  by approximating it with the bounds in Lemma 4.2. While the upper bound is tight when all scores have similar magnitude, the bound loosens if scores differ a lot, leading Equation (10) to overestimate the value of the PMF. In practice, we thus reweight the respective terms in the loss function, but in the future, we will investigate better estimates for the PMF.
Ultimately, we are interested in exploring how to apply the DRPM to multimodal learning under weak supervision, for instance, in medical applications. Section 5.2 demonstrated the potential of learning from coupled samples, but further research is needed to ensure fairness concerning underlying, hidden attributes when working with sensitive data.
In this work, we proposed the differentiable random partition model, a novel approach to random partition models. Our two-stage method enables learning partitions end-to-end by separately controlling subset sizes and how elements are assigned to subsets. This new approach to partition learning enables the integration of random partition models into probabilistic and deterministic gradient-based optimization frameworks. We show the versatility of the proposed differentiable random partition model by applying it to three vastly different experiments. We demonstrate how learning partitions enables us to explore the modes of the data distribution, infer shared and independent generative factors from coupled samples, and learn task-specific sub-networks in applications where we want to solve multiple tasks on a single data point.
AR is supported by the StimuLoop grant #1-007811-002 and the Vontobel Foundation. TS is supported by the grant #2021-911 of the Strategic Focal Area "Personalized Health and Related Technologies (PHRT)" of the ETH Domain (Swiss Federal Institutes of Technology).
M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. J. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Jozefowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. G. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, I. Sutskever, K. Talwar, P. A. Tucker, V. Vanhoucke, V. Vasudevan, F. B. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and X. Zheng. TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems. CoRR, abs/1603.0, 2016.
R. P. Adams and R. S. Zemel. Ranking via sinkhorn propagation. arXiv preprint arXiv:1106.1925, 2011.
K. Ahmed, Z. Zeng, M. Niepert, and G. V. d. Broeck. SIMPLE: A Gradient Estimator for k-Subset Sampling. In The Eleventh International Conference on Learning Representations, Sept. 2022.
Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
C. M. Bishop and M. Svensen. Robust Bayesian Mixture Modelling. To appear in the Proceedings of ESANN, page 1, 2004. Publisher: Citeseer.
D. Bouchacourt, R. Tomioka, and S. Nowozin. Multi-level variational autoencoder: Learning disentangled representations from grouped observations. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
J. Cai, J. Fan, W. Guo, S. Wang, Y. Zhang, and Z. Zhang. Efficient deep embedded subspace clustering. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1-10, 2022.
R. Caruana. Multitask learning: A knowledge-based source of inductive bias. In Machine learning: Proceedings of the tenth international conference, pages 41-48, 1993.
R. Caruana and V. R. de Sa. Promoting poor features to supervisors: Some inputs work better as outputs. Advances in Neural Information Processing Systems, 9, 1996.
J. Chesson. A non-central multivariate hypergeometric distribution arising from biased sampling with application to selective predation. Journal of Applied Probability, 13(4):795-797, 1976. Publisher: Cambridge University Press.
A. Coates, A. Ng, and H. Lee. An Analysis of Single-Layer Networks in Unsupervised Feature Learning. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 215-223. JMLR Workshop and Conference Proceedings, June 2011. ISSN: 1938-7228.
A.1 Hypergeometric Distribution
This part is largely based on Sutter et al. (2023).
Suppose we have an urn with marbles in different colors. Let  $K \in \mathbb{N}$  be the number of different classes or groups (e.g. marble colors in the urn),  $\pmb{m} = [m_1, \dots, m_K] \in \mathbb{N}^K$  describe the number of elements per class (e.g. marbles per color),  $N = \sum_{k=1}^{K} m_K$  be the total number of elements (e.g. all marbles in the urn) and  $n \in \{0, \dots, N\}$  be the number of elements (e.g. marbles) to draw. Then, the multivariate hypergeometric distribution describes the probability of drawing  $\pmb{n} = [n_1, \dots, n_K] \in \mathbb{N}^K$  marbles by sampling without replacement such that  $\sum_{k=1}^{K} n_k = n$ , where  $n_k$  is the number of drawn marbles of class  $k$ .
In the literature, two different versions of the noncentral hypergeometric distribution exist, Fisher's (Fisher, 1935) and Wallenius' (Wallenius, 1963; Chesson, 1976) distribution. Sutter et al. (2023) restrict themselves to Fisher's noncentral hypergeometric distribution due to limitations of the latter (Fog, 2008). Hence, we will also talk solely about Fisher's noncentral hypergeometric distribution.
Definition A.1 (Multivariate Fisher's Noncentral Hypergeometric Distribution (Fisher, 1935)). A random vector  $\mathbf{X}$  follows Fisher's noncentral multivariate distribution, if its joint probability mass function is given by
$$
P (\boldsymbol {N} = \boldsymbol {n}; \omega) = p (\boldsymbol {n}; \omega) = \frac {1}{P _ {0}} \prod_ {k = 1} ^ {K} \binom {m _ {k}} {n _ {k}} \omega_ {k} ^ {n _ {k}} \tag {12}
$$
$$
\text {w h e r e} P _ {0} = \sum_ {\left(\eta_ {1}, \dots , \eta_ {K}\right) \in S} \prod_ {k = 1} ^ {K} \binom {m _ {k}} {\eta_ {k}} \omega_ {k} ^ {\eta_ {k}} \tag {13}
$$
The support  $S$  of the PMF is given by  $S = \{n\in \mathbb{N}^K:\forall k\quad n_k\leq m_k,\sum_{k = 1}^K n_k = n\}$  and  $\binom{n}{k}=\frac{n!}{k!(n-k)!}$ .
The class importance  $\omega$  is a crucial modeling parameter in applying the noncentral hypergeometric distribution (see (Chesson, 1976)).
A.1.1 Differentiable MVHG
Their reparameterizable sampling for the differentiable MVHG consists of three parts:
1. Reformulate the multivariate distribution as a sequence of interdependent and conditional univariate hypergeometric distributions.
2. Calculate the probability mass function of the respective univariate distributions.
3. Sample from the conditional distributions utilizing the Gumbel-Softmax trick.
Following the chain rule of probability, the MVHG distribution allows for sequential sampling over classes  $k$ . Every step includes a merging operation, which leads to biased samples compared to groundtruth non-differentiable sampling with equal class weights  $\omega$ . Given that we intend to use the differentiable MVHG in settings where we want to learn the unknown class weights, we do not expect a negative effect from this sampling procedure. For details on how to merge the MVHG into a sequence of unimodal distributions, we refer to Sutter et al. (2023).
The probability mass function calculation is based on unnormalized log-weights, which are interpreted as unnormalized log-weights of a categorical distribution. The interpretation of the class-conditional unimodal hypergeometric distributions as categorical distributions allows applying the Gumbel-Softmax trick (Jang et al., 2016; Maddison et al., 2017). Following the use of the Gumbel-Softmax trick, the class-conditional version of the hypergeometric distribution is differentiable and reparameterizable. Hence, the MVHG has been made differentiable and reparameterizable as well. Again, for details we refer to the original paper (Sutter et al., 2023).
A.2 Distribution over Random Orderings
Yellott (1977) show that the distribution over permutation matrices  $p(\pi; s)$  follows a Plackett-Luce (PL) distribution (Plackett, 1975; Luce, 1959), if and only if the perturbed scores  $\tilde{s}$  are sampled independently from Gumbel distributions with identical scales. For each item  $i$ , sample  $g_i \sim \mathrm{Gumbel}(0, \beta)$  independently with zero mean and fixed scale  $\beta$ . Let  $\tilde{s}$  be the vector of Gumbel perturbed log-weights such that  $\tilde{s}_i = \beta \log s_i + g_i$ . Hence,
$$
q \left(\tilde {s} _ {1} \geq \dots \geq \tilde {s} _ {n}\right) = \frac {s _ {1}}{Z} \cdot \frac {s _ {2}}{Z - s _ {1}} \cdot \dots \cdot \frac {s _ {n}}{Z - \sum_ {i = 1} ^ {n - 1} s _ {i}} \tag {14}
$$