Differentiable Random Partition Models
Authors: Thomas M. Sutter
Alain Ryser
Joram Liebeskind
Julia E Vogt
Venue: NeurIPS 2023 poster (2023)
Keywords: random partition model
continuous relaxation
reparameterization
generative models
vae
representation learning
weak supervision
variational clustering
deep learning
TL;DR: We propose a differentiable random partition model to be integrated into modern gradient-based optimization pipelines.

Abstract: Partitioning a set of elements into an unknown number of mutually exclusive subsets is essential in many machine learning problems.
However, assigning elements, such as samples in a dataset or neurons in a network layer, to an unknown and discrete number of subsets is inherently non-differentiable, prohibiting end-to-end gradient-based optimization of parameters.
We overcome this limitation by proposing a novel two-step method for inferring partitions, which allows its usage in variational inference tasks.
This new approach enables reparameterized gradients with respect to the parameters of the new random partition model.
Our method works by inferring the number of elements per subset and, second, by filling these subsets in a learned order.
We highlight the versatility of our general-purpose approach on three different challenging experiments: variational clustering, inference of shared and independent generative factors under weak supervision, and multitask learning.

Differentiable Random Partition Models
Thomas M. Sutter; Alain Ryser", Joram Liebeskind, Julia E. Vogt Department of Computer Science ETH Zurich
Abstract
Partitioning a set elements into an unknown number of mutually exclusive subsets is essential in many machine learning problems. However, assigning elements, such as samples in dataset or neurons in a network layer, to an unknown and discrete number of subsets is inherently non-differentiable, prohibiting end- to-end gradient-based optimization of parameters. We overcome this limitation by proposing a novel two-step method for inferring partitions, which allows its usage in variational inference tasks. This new approach enables reparameterized gradients with respect to the parameters of the new random partition model. Our method works by inferring the number of elements per subset and, second, by filling these subsets in a learned order. We highlight the versatility of our general- purpose approach on three different challenging experiments: variational clustering, inference of shared and independent generative factors under weak supervision, and multitask learning.
Introduction
Partitioning set of elements into subsets is classical mathematical problem that attracted much interest over the last few decades (Rota, 1964; Graham et al., 1989). A partition over given set collection of non-overlapping subsets such that their union results in the original set. In machine learning (ML., partitioning set of elements into different subsets essential for many applications, such clustering (Bishop and Svensen, 2004) classification (De la Cruz-Mesía et al., 2007).
Random partition models (RPM, Hartigan, 1990) define probability distribution over the space of partitions. RPMs can explicitly leverage the relationship between elements of a set, as they do not necessarily assume i.i.d. set elements. On the other hand, most existing RPMs are intractable for large datasets (MacQueen, 1967; Plackett, 1975; Pitman, 1996) and lack reparameterization scheme, prohibiting their direct use in gradient- based optimization frameworks.
In this work, we propose the differentiable random partition model (DRPM), a fullydlifferentiable relaxation for RPMs that allows reparametrizable sampling. The DRPM follows two-stage proce- dure: first, we model the number of elements per subset, and second, we learn an ordering of the elements with which we fill the elements into the subsets. The DRPM enables the integration of partition models into state-of-the-art ML frameworks and learning RPMs from data using stochastic optimization.
We evaluate our approach in three experiments, demonstrating the proposed DRPM's versatility and advantages. First, apply the DRPM to variational clustering task, highlighting how the reparametrizable sampling of partitions allows us to learn a novel kind of Variational Autoencoder (VAE, Kingma and Welling, 2014). By leveraging potential dependencies between samples in dataset, DRPM-based clustering overcomes the simplified i.i.d. assumption of previous works, which used categorical priors (Jiang al., 2016). In second experiment, demonstrate how retrieve sets of shared and independent generative factors of paired images using the proposed DRPM. In
"Equal Contribution. Correspondence to {thomas. sutter, alain. rraeen..sse)) ettzz.c
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

7
0
0
0
0
0
0
0
0
0
1
0
0
y22
0
1
0
0
0
0
0
0
0
1
0
0
m3
0
0
0
0
0
0
0/
Us
yo.
0
0
0
0
0
n5
0
0
0
0
Figure I: Illustration the proposed DRPM method. We first sample a permutation matrix æ and a set of subset sizes n separately in two stages. We then and ı to generate the assignment matrix Y, the matrix representation of a partition p.
contrast to previous works (Bouchacourt et al., 2018; Hosoya, 2018: Locatello et al., 2020), which rely on strong assumptions or heuristics, the DRPM enables end-to-end inference of generative factors. Finally, we perform multitask learning (MTL) by using the DRPM as building block in deterministic pipeline. We show how the DRPM learns to assign subsets of network neurons to specific tasks. The DRPM can infer the subset size per task based on its difficulty, overcoming the tedious work of finding optimal loss weights (Kurin et al., 2022; Xin et al., 2022).
To summarize, we introduce the DRPM, a novel differentiable and reparametrizable relaxation of RPMs. In extensive experiments, we demonstrate the versatility of the proposed method by applying the DRPM to clustering, inference of generative factors, and multitask learning.
2 Related Work
Random Partition Models Previous works on RPMs include product partition models (Hartigan, 1990), species sampling models (Pitman, 1996), and model-based clustering approaches (Bishop and Svensen, 2004). Further, Lee and Sang (2022) investigate the balancedness of subset sizes of RPMs. They all require tedious manual adjustment, are non-differentiable, and are, therefore, unsuitable for modern ML pipelines. A fundamental RPM application is clustering, where the goal is to partition given dataset into different subsets, the clusters. In contrast to many existing approaches (Yang et al., 2019; Sarfraz al., 2019; Cai et al., 2022), allowing us to treat clustering from a variational perspective. Previous works in variational clustering (Jiang al., 2016; Dilokthanakul et al., 2016; Manduchi et al., 2021) implicitly define RPMs to perform clustering. They compute partitions in a variational fashion by making i.i.d. assumptions about the samples in the dataset and imposing soft assignments of the clusters to data points during training. A problem related to set partitioning is the earth mover's distance problem (EMD, Monge, 1781; Rubner al., 2000). However, EMD aims assign a set's elements different subsets based on cost function and given subset sizes. Iterative solutions to the problem exist (Sinkhorn. 1964), and various methods have recently been proposed, e.g., for document ranking (Adams and Zemel, 2011) or permutation learning (Santa Cruz et al., 2017; Mena et al., 2018; Cuturi et al., 2019).
consider cluster assignments as random variables,
Differentiable and Reparameterizable Discrete Distributions Following the proposition of the Soumeel-Softax trick (GST, Jang et al., 2016; Maddison et al., 2017), interest in research around continuous relaxations for discrete distributions and non-differentiable algorithms rose. The GST enabled the reparameterization of categorical distributions and their integration into gradient- based optimization pipelines. Based on the same trick, Sutter et al. (2023) propose differentiable formulation for the multivariate hypergeometric distribution. Multiple works differentiable sorting procedures and permutation matrices have been proposed, e.g., Linderman et al. (2018); Prillo and Eisenschlos (2020); Petersen al. (2021). Further, Grover et al. (2019) described the distribution over permutation matrices p(T)) for a permutation matrix using the Plackett-Luue Luce distribution (PL, Luce, 1959; Plackett, 1975). Prillo and Eisenschlos (2020) proposed a computationally simpler variant of Grover et al. (2019). More examples of differentiable relaxations include the top-k elements selection procedure (Xie and Ermon, 2019), blackbox combinatorial solvers (Pogančié al., 2019), implicit likelihood estimations (Niepert al., 2021), and k-subset sampling (Ahmed al., 2022).

---

3 Preliminaries
{I..... n with clements is collection K subsets S₁ S [n] where K is a priori unknown (Mansour and Schork, 2016). For a partition p to be valid, it must hold that
Set Partitions A partiton p (St...., Sk) of a set [n]
(I) In other words, every element i E [m]] has to be assigned to precisely one subset Sk. We denote the size of the k-thh subset Sg ng | Alternatively, we can describe partition p through an assignment matrix Y = .y...... yK I E 10,1) Every row yıu E {0, 1] Ixn is multi-hot vector, where yk: 1 assigns element to subset Sg.
U.....U U....USE Sx = [n] and v1 S
Within the scope of our work, we view a partition of set of n elements as a special case of the urn model. Here, the urn contains marbles with different colors, where each color corresponds to subset in the partition. For each color, there are n marbles corresponding to the potential elements of their color/subset. To derive a partition, we sample marbles without replacement from the urn and register the order in which we draw the colors. The color of the i-th marble then determines the subset which element corresponds, Furthermore, we can constrain the partition to only K subsets by taking an urn with only K different colors.
Probability distribution over subset sizes The multivariate non-central hypergeometric distribu- tion (MVHG) describes sampling without replacement and allows to skew the importance of groups with additional importance parameter w (Fisher, 1935; Wallenius, 1963; Chesson, 1976). The MVHG an urn model and is described by the number of different groups K E N, the number of elements in the urn of every group m [mm,...., mR]] E NK, the total number of elements in the um Σk= my E N. the number of samples draw from the um n No, and the importance factor for every group w [wi....,.3p wk] E R (Johnson, 1987). Then, the probability of sampling (n....., nx) where ny describes the number of elements drawn from group K is
1 K
p(n: w, m) = Po
II A=1
m
W
where Po is a normalization constant. Hence, the MVHG p(n; w, m) allows us to model depen- dencies between different elements of set since drawing one element from the urn influences the probability of drawing one of the remaining elements, creating interdependence between them. For
the rest of the paper, we assume V m, E mg n. We thus use the shorthand p(n:; w) to denote the density of the MVHG. We refer to Appendix A.1 for more details.
Probability distribution over Permutation Matrices Let p(T) denote distribution over permuta- tion matrices π E {0, 1]=x", A permutation matrix æ is doubly stochastic (Marcus, 1960), meaning that its row and column vectors sum to 1. This property allows us to use π to describe an order over a set of n elements, where Ti = 1 means that element j ranked at position i the imposed order. In this work, we assume p(T) to be parameterized by scores s E R7, where each score s; corresponds an element i. The order given by sorting s in decreasing order corresponds the most likely permutation in p(T; s). Sampling from p(T; s) can be achieved by resampling the scores as &; = log 8; + 9: where 9: ~ Gumbel 3) for fixed scale 3, and sorting them in decreasing order. Hence, resampling scores s enables the resampling of permutation matrices и. The probability over orderings p(T; s) then given by (Thurstone, 1927; Luce, 1959; Plackett, 1975; Yellott, 1977)
(T)):
(T8)2
(TS)
p(:; s) p((T8) 2 (T:8m.)
Z Z (TS) 7 Σ [T))
where π is a permutation matrix and Z = E st The resulting distribution is Plackett-Luce (PL) distribution (Luce, 1959; Plackett, 1975) if and only if the scores are perturbed with noise drawn from Gumbel distributions with identical scales (Yellott, 1977). For more details, we refer to Appendix A.2).
4 A two-stage Approach to Random Partition Models
We propose the DRPM (Y; w, s), a differentiable and reparameterizable two-stage Random Partition Model (RPM). The proposed formulation separately infers the number of elements per subset
3

---

No. where Σk=1 7 = n, and the assignment of elements to subsets Sg by inducing an order on the n elements and filling Si...,Sp sequentially in this order. To model the order of the elements, we use permutation matrix π [ ....,........................... nnl" E (0, 1 1)*Xn, from which we infer Y by sequentially summing up rows according to n. Note that the doubly-stochastic property of all permutation matrices и ensures that the columns of Y remain one-hot vectors, assigning every element i to precisely one of the K subsets. At the same time, the k-th row of Y corresponds to an ngghoot vector ye and therefore serves as a subset selection vector, i.c.
Pk+AA y Σ Ts.
where U/ Σ n
(4)
such that Y = .y........ .K]". Additionally, Figure 1 provides an illustrative example. Note that K defines the maximum number of possible subsets, and not the effective number of non-empty subsets, because we allow Sg to be the empty set 0 (Mansour and Schork, 2016). We base the following Proposition 4.1 on the MVHG distribution p(12; w) for the subset sizes n and the PI. distribution p( 8) for assigning the elements to subsets. However, the proposed two-stage approach to RPMs is not restricted these two classes of probability distributions.
Proposition 4.1 (Two-stage Random Partition Model), Given probability distribution over subset a
sizes p(n; w) with n E No and distribution parameters w E E& and a PL probability distribution over random orderings p(T; 8) with π E [0,1)mxn and distribution parameters E I44 the probability mass function p(Y; oY; (w,8) of the two-stage RPM is given by
p(Y;w,8) p(y1........................ Vkiw, =p(n;w)
Σ p(T;8)
(5)
xElly
where IIy = f yk Ei TT 1...., K}. and yv and Vy in Equation (4).
In the following, we outline the proof of Proposition 4.1 and refer to Appendix B for a formal derivation. We calculate p(Y;w,s) a probability of subsets ply,,..., ; " 8), which compute sequentially over subsets, i.e.
pki = p(yliw, 8) ......-PP plyK 3)kiw,
(6)
where ycc ....... 9k-1] and
ply VeRiw, 8) = p(ne (eksw)
Σ
p(T ni, 8)):8),
#Ellys
where Hy in Equation (7) is the set of all subset permutations of elements E Sg. A subset permutation matrix represents an ordering over only nı out of the total n elements. The probability p(yx Y< w, s) describes the probability of a subset of given size probabilities of all subset permutations p( nk: V<k; Hence, the sum over all p( mk: y< 8) makes p(t) Yeck; s) invariant to the ordering of elements i E Se. (Xie and Ermon, 2019). Note that in a slight abuse of notation, we use p( nk, D<k: w, 8) as the probability of subset permutation म given that there are ni elements in Sy and thus स E {0, m
by marginalizing over the
The probability of subset permutation matrix p( Tok yyck 8) describes the probability of drawing the elements i E S₁ in the order defined by the subset permutation matrix म given that the elements in Se, are already determined. Hence, we condition on the subsets yzk- This property follows from Luce's 's choice axiom (LCA, Luce, 1959). Additionally, we condition on ng. the size of the subset Sp. The probability of subset permutation is given by
TTS) 2g Σ j= T8))
p( ne. Mckis)
II
(8)
In contrast to the distribution over permutations matrices p(T; s) in Equation (3), we compute the product over nig terms and have a different normalization constant Zu, which the sum over the scores 8; of all elements i E Sg. Although we induce an ordering over all elements by using a permutation matrix n, the probability p(y; Yekk w, 8) is invariant to intra-subset orderings of elements i E Sg. Finally, we arrive at Equation (5) by substituting Equation (7) into Equation (6).

---

and applying the definition of the conditional probability p(n; = II p(rz |nc w) and by reshuflling indices Evelly (T;;; 8) I ExEUL, p( nk. 8).
Note that in contrast to previous RPMs, which often need exponentially many distribution parameters (Plackett, 1975), the proposed two-stage approach to RPMs only needs (n + K) parameters to create an RPM for n elements: the score parameters E I and the group importance parameters w I
Finally, to sample from the two-stage RPM of Proposition 4.1 we apply the following procedure: First sample p(T;;& 8) and n p(n; w). From я and n, compute partition Y by summing the rows of π according to as described in Equation (4) and illustrated in Figure 1.
4.1 Approximating the Probability Mass Function
The number of permutations per subset [T scaaes factorially with the subset size ny. i.e. [g ny!. Consequently, the number of valid permutation matrices [Ty is given function of n, i.e.
K
ml
II
II ng!
Although Proposition 4. describes a well-defined distribution for p(Y; ;Y; w. 8). it in general com- putationally intractable due to Equation (9). In practice, we thus approximate p(Y; 8) using the following Lemma.
Lemma 4.2. pYY:,, (.ws) s) can be upper and lower bounded as follows
VT E Ily p((::)p((;s)) s) I p(Y:w,)) $ [1y max
(10)
We provide the proof in Appendix B. Note that from Equation (3) we see that max, p(t;; s) p( 8), where ny the permutation that results from sorting the unperturbed scores
=
4.2 The Differentiable Random Partition Model
To incorporate our two-stage RPM into gradient-based optimization frameworks, require that efficient computation gradients possible for every step of the method. The following Lemma guarantees differentiability, allowing us to train deep neural networks with our method in an end-to- end fashion:
Lemma 4.3 (DRPM). A two-stage RPM is differentiable and reparameterizable if the distribu- tion over subset sizes p(n; w) and the distribution over orderings p(T; s) are differentiable and reparameterizable.
We provide the proof in Appendix B. Note that Lemma 4.3 enables to learn variational posterior approximations and priors using Stochastic Gradient Variational Bayes (SGVB, Kingma and Welling, 2014). In our experiments, tions of the MVHG (Sutter et al., 2023) and the PL distribution (Grover et al., 2019), though other choices would also be valid.
apply Lemma 4.3 using the recently proposed differentiable formula-
5 Experiments
We demonstrate the versatility and effectiveness of the proposed DRPM in three different experiments. First, we propose novel generative clustering method based on the DRPM, which we compare against state-of-the-art variational clustering methods and demonstrate its conditional generation capabilities. Then, we demonstrate how the DRPM can infer shared and independent generative factors under weak supervision. Finally, we apply the DRPM to multitask learning (MTL.,, where the DRPM enables an adaptive neural network architecture that partitions layers based on task difficulty?.
We provide the code under https::/ //github//t...... com/thommssutterr/rrpp

---

Table 1: We compare the clustering performance of the DRPM-VC on test sets of MNIST and FMNIST between Gaussian Mixture Models (GMM), GMM in latent space (Latent GMM), and Variational Deep Embedding (VaDE). We measure performance terms of Normalized Mutual Information (NMI), Adjusted Rand Index (ARI), and cluster accuracy (ACC) over five seeds and put the best model in bold.
MNIST
FMNIST
NMI
ARI
ACC
NMI
ARI
ACC
GMM
0.32
0.22=0.02
0.41=0.01
0.49+0.01
0.33#0.00
0.44=0.01
LATENT GMM
0.86+0.02
0.03
0.88=+0.07
0.60=+0.00
0.47=0.01
0.62
VADE
0.84+0.01
0.76=0.05
0.82++004
0.56+0.02
0.40=0.01
0.56=0.03
DRPM-VC
0.8910.01
0.8810.03
0.94
0.64.10.00
0.51
0.65
Cluster
Generated Sample
Figure 2: A sample drawn from DRPM-VC model trained on FMNIST. On top is the sampled partition with the cluster assignments, and on the bottom are generated images corresponding to the sampled assignment matrix. The DRPM-VC learns consistent clusters for different pieces of clothing and can generate new samples each cluster with great variability.
5.1 Variational Clustering with Random Partition Models
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 æ dataset X generated by a latent vector R', where E N is the latent space size. Traditional VAEs would then assume that all latent vectors z are generated by a single Gaussian prior distribution N (O, 4r). Instead, we assume every z to be sampled from one of K different latent Gaussian distributions (Hu,, diag( (T)) = I..... K, with μkx d R Further, note that similar to an urn model (Section 3), if we draw batch from a given finite dataset with samples from different clusters, the cluster assignments within that batch are entirely independent. Since there is only a finite number of samples per cluster, drawing a sample from 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 variational clustering proposes to model the cluster assignment y E {0, 1 each sample æ through independent categorical distributions (Jiang 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 drawn from RPM. Hence, the generative process leading to X can summarized as follows: First, the cluster assignments are represented as a partition matrix Y and sampled from DRPM, i.c., Y p(Y; w.s). s). Given assignment y from Y, can sample the respective latent variable z, where z ~ N( diag (haag(ay))
TR. Note that we use the notational shorthand By = Marg max(y)) Like in vanilla VAEs, we infer æ by independently passing the corresponding through a decoder model. Assuming this generative process, we derive the following evidence lower bound (ELBO) for p(X):
LELBO
Σ Byi [[og((x|=)] #EX
- Σ E [KL]((] - KЛX(((X)))PX(
xEX
Note that computing K LL(g(Y|X) p()) directly is computationally intractable, and we need to upper bound according to Lemma 4.2. For an illustration of the generative assumptions and more details on the ELBO, 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 the shared and independent dimensions of the representation, predicting the respective generative factors. We report the results in adjusted balanced accuracy (Sutter al., 2023) across five seeds.
n2 =
11
R2 3
f2g =
I
S
LABEL
0.14
0.19 03
0.16#0.01
0.10+0.00
0.23
0.34.1..00
0.00=0..0
ADA
0.12 0.19
0.15
0.10+0.03
0.22#0.02
0.33+a.u3
0.00=#.00
HG
0.18
0.22=0.05
0.19
0.08
0.28=0.01
0.28
0.01
DRPM
0.26#0.02
0.39
0.2
0.15#
0.29
0.42#0.03
0.01
To assess the clustering performance, train our model 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 basclines 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 pretrained encoder as imput (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 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 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 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-supervisiee 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 cou- pled pairs while others are independent. The supervision is weak because we neither know the underlying gener- ative factors nor the number of shared and independent factors. In such a setting, we can use the DRPM to learn partition of the generative factors and assign them to be either shared or independent.
80
Label
60-
(Dum)
Mis(ee,
20
Number Sharod rol Factors iu,
The mean squared errors be- tween the estimated number of shared factors ft, and the true number of shared factors n, across five seeds for the Label- VAE, Ada-VAE, HG-VAE, and DRPM- VAE.
Figure 3:
In this experiment, we use paired frames X [x1, [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 ( X) and additionally prior distribution of the form p(Y). The model maximizes the following ELBO on the
7

---

marginal log- lieelikelihod of images through a VAE (Kingma and Welling, 2014):
2 By 2, xIX) [log E.g 2y)]
LELBO
(11)
j=1
B x) x) [KL [g( z1, 22 Y, X)Ip(( 21, x2)]] KL 1a X)II p(Y)
Similar to the ELBO for variational clustering in Section 5.1, computing KLL [g(Y X)I p(Y) directly is intractable, and we need to bound it according to Lemma 4.2.
We compare the proposed DREM-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) 18 assumes that the number of independent factors is known, the Ada- VAE (Locatello 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-VAI (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-gennrating process, we can control the number of shared n2 and independent n. factors. We compare the methods on four different datasets with n, E {0, 1, 5). In Figure 3, we demonstrate that the DRPM- 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 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-VAF 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 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. Lee. case of equal difficulty of the two tasks, classifier with equal weighting of the two classification losses serves an upper bound in terms of performance. However, finding suitable task weights 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 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 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 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 PI. distributions. We compare the DRPM-MTL to the unitary loss scaling method (ULS, Kurin et al.,
8

---

2022), which has a fixed urchitecture and scales task losses equally. Both DRPM-MTL and ULS use 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.,,,0177 2017), which overlaps two MNIST digits in one image, and we want to classify both numbers from 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 MuliiMMUtM 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
1.0
Accuracy
0.8
0.6
ULS Left
0.4
ULS Right
50
DRPM Left (Ours)
DRPM Right (Ours)
10
Partition Sizes
30
20
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.
0.0
0.2
0.1
0.6
0.8
Naise Ratio a
Figure Results for noisy MultiMNIST ex- periment. In the upper plot, 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 while it assigns the number of dimensions per task according to their difficulty.
We evaluate the DRPM-MTL concerning its classi- fication accuracy on the two tasks and compare the inferred subset sizes per task for different noise ratios dataset
a E {O.0..... 0.9} of the noisyMultiMNISs (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 o, the more the DRPM-MTL. tries to overcome the increased difficulty of the right task by assigning more dimensions it (lower part of Figure 4, noise ratio a 0.6-0.8). Note that for the maximum noise ratio of a = 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.
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 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(Y: w, 8) 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, will investigate better estimates for the PMF
Ultimately, we are interested 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.

---

Conclusion
In this work, we proposed the differentiable random partition model, novel approach to random partition models. Our two-stage method enables learning partitions end-to-end by separately control- ling 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-ban 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- sub-networks in applications where want to solve multiple tasks on single data point.
Acknowledgements
AR is supported by the Stimul grant #1-007811-009 and the Vontohel 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).
References
M. Abadi, A. Agarwal, Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean,
M. Devin, S. Ghemawat, J. Goodfellow, A. Harp. G. Irving, M. Isard, Y. Jia, R. Józefowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, 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/ 00 2016.
R. P. Adams and R. S. Zemel. Ranking via sinkhorn propagation. arXiv preprint arXiv: 196. 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 propagating gradients through stochastic neurons for conditional computation. arXiv preprint
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-Sccond AAAI Conference 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 Computer Vision and Pattern Recognition, pages 1-10, 2022.
R. Caruana. Multitask learning: A knowledge-based based source of inductive bias. In Machine learning: Proceedings of the tenth international conference, pages 41 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.
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 Fourleenth International Conference Artificial Intelligence and Statistics, pages 215-223, JMLR Workshop and Conference Proceedings, June 2011. ISSN: 1938-7228.
10

---

M. Cuturi, O. Teboul, and J.-P. Vert. Differentiable Ranks and Sorting using Optimal Transport, Nov. 2019. arXiv:1905.1 1885 [es, stat].
R. De la Cruz-Mesía, F. A. Quintana, and P. Müller. Semiparametric Bayesian Classification with longitudinal Markers. Journal of the Royal Statistical Society Series C: Applied Statistics, 56(2): 119-137, Mar. 2007. ISSN 0035-92544 doi: 10.1111/j.144 000777976220077 x00569.x.
J. Deng. W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
N. Dilokthanakul, P. A. M. Mediano, M. Gamelo, M. C. H. Lee, H. Salimbeni, K. Arulkumaran, and M. Shanahan. Deep unsupervised clustering with Gaussian mixture variational autoencoders. arXiv preprini arXiv: 1
2016.
R. A. Fisher. The logic of inductive inference. Journal of the royal statistical society, 98(1):39 82, 1935. Publisher: JSTOR.
A. Fog. Calculation methods for Wallenius' noncentral hypergeometric distribution. Communications in -imiuiats-iimaaiiion and ComputationO), 37(2):258-277 2008. Publisher: Taylor & Francis.
M. W. Gondal, M. Wuthrich, D. Miladinovic. F. Locatello, M. Breidt, V. Volchkov. J. Akpo, O. Bachem, B. Schölkopf, and S. Bauer. On the Transfer of Inductive Bias from Simulation to the Real World: New Disentanglement Dataset. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
R. I.. Graham, D. F. Knuth, O. Patashnik, S. L.-C. i. Physics, and undefined 1989. Concrete mathematics: foundation for computer science. aip.scitation.org. 3(5):15,, 1989. doi: 10.1063/1. 4822863. Publisher: AIP Publishing.
A. Grover, E. Wang, A. Zweig, and S. Ermon. Stochastic Optimization of Sorting Networks via Continuous Relaxations. In International Conference on Learning Representations, 2019.
A. Hartigan. Partition models. Communications in Statistics Theory and Methods, 19(8): 2745-2756, 1990. doi: 10. 00.188000660092290883333333 Publisher: Taylor & Francis.
K. He, X. Zhang, S. Ren, and J. Sun. Deep Residual Learning for Image Recognition, Dec. 2015.
I. Higgins, L. Matthey, A. Pal, C. Burgess, and X. Glorot. beta-vae: Learning basic visual concepts with constrained variational framework. 2016.
H. Hosoya. A simple probabilistic deep generative model for learning generalizable disentangled representations from grouped data. CoRR, abs/1809.0. 2018.
E. Jang, S. Gu, and B. Poole. Categorical reparameterizzation with gumbel-softmax. arXiv preprint 1rXiv:1611.01144, 2016.
Z. Jiang, Y. Zheng, H. Tan, B. Tang, and H. Zhou. Variational deep embedding: An unsupervised and generative approach clustering. arXiv preprint arXiv:1611.05148, 2016.
M. E. Johnson. Multivariate statistical simulation: A guide to selecting and generating continuous multivariate distributions, volume 192. John Wiley & Sons, 1987.
D. Kingma and M. Welling. Auto-Encoding Variational Bayes. In International Conference on Learning Representations, 2014.
V. Kurin, A. D. Palma, I. Kostrikov, S. Whiteson, and M. P. Kumar. In Defense of the Unitary Scalarization for Deep Multi-Task Learning. CoRR, abs/2201 0112200004122, 2022.
Y. LeCun, C. Cortes, and C. J. Burges. Gradient-based learning applied to document recognition. In Proceedings of the IEEE, volume 86, pages 2278-2324, 1998. Issue: 11.
C. J. Lee and Sang. Why the Rich Get Richer? the Balancedness Random Partition Models. arXiv preprint ar arX1v:2201.12697, 2022.

---

Y. Li, M. Yang, D. Peng, T. Li, J. Huang, and X. Peng. Twin Contrastive Learning for Online Clustering. International Journal of Computer Vision, 130(9):2205-222) -2211, Sept. 2022. ISSN 0920-5691, 1573-1405,
doi: 1010007/11260--022--166669
arXiv:2210.1 1680 [es].
S. W. Linderman, G. E. Mena, H. J. Cooper, L. Paninski, and Cunningham. Reparameterizing the Birkhoff Polytope for Variational Permutation Inference. In A. J. Storkey and F. Pérez-Cruz, editors, International Conference on Artificial Intelligence and Statistics, [AISTATS] 2018, 9-11 April 2018, Playa Blanca, Lanzarote, Canary Islands, Spain, volume 84 of Proceedings of Machine Learning Research, pages 1618-1627. PMLR, 2018.
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep Learning Face Attributes in the Wild, Sept. 2015. arXiv:1411.7766 [cs].
Locatello, B. Poole, G. Rätsch, Schölkopf, O. Bachem, and M. Tschannen. Weakly-supervised disentanglement without compromises. In International Conference Machine Learning, pages 6348-6359. PMLR, 2020.
I. Loshchilov and F. Hutter. Decoupled Weight Decay Regularization, Jan. 2019. arXiv:1711.05II [es, math].
R. D. Luce. Individual choice behavior: A theoretical analysis. Courier Corporation, 1959.
J. MacQueen. Classification and analysis of multivariate observations. In 5th Berkeley Symp. Math. Statist. Probability, pages 281-297, 1967.
C. Maddison, A. Mnih, and Y. Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations, 2017.
L. Manduchi, K. Chin-Cheong, H. Michel, S. Wellmann, and J. E. Vogt. Deep Conditional Gaussian Mixture Model for Constrained Clustering. ArXiv, abs/2106.0, 2021.
T. Mansour and M. Schork. Commutation relations, normal ordering, and Stirling numbers. CRC Press Boca Raton, 2016.
M. Mareus. Some Properties and Applications of Doubly Stochastic Matrices. The American Mathematical Monthly, 67(3):215-221, 1960. ISSN 00029890, 19300972. Publisher: Mathematical Association of America.
G. E. Mena, D. Belanger, S. W. Linderman, and J. Snoek. Learning Latent Permutations with Gumbel- Sinkhorn Networks, In 6th International Conference on Learning Representations, (ICLR) 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018.
G. Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math, Phys. Acad. Royale Sci., pages 666-704, 1781.
M. Nicpert, P. Minervini, and Franceschi. Implicit MLE; Backpropagating Through Discrete Expo- nential Family Distributions. In Advances in Neural Information Processing Systems, volume 34, pages 14567-14579. Curran Associates, Inc., 2021.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison,
Köpf, Z. Yang, Z. DeVito, M. Raison, A. Tejani, Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An Imperative Style, Hiig-Peeformance Deep Learning Library. CoRR, abs/1912.0, 2019.
F. Petersen, C. Borgelt, H. Kuchne, and O. Deussen. Monotonic Differentiable Sorting Networks, In International Conference on Learning Representations, 2021.
J. Pitman. Some Developments of the Mlackuell-Macqueen URN Scheme. Lecture Notos-Monograph Series, 30:245-267, 1996. ISSN 07492170. Publisher: Institute of Mathematical Statistics.
R. I.. Plackett. The analysis of permutations. Journal of the Royal Statistical Society: Series (Applied Statistics), 24(2):193-20) 1975. Publisher: Wiley Online Library.
12

---

M. V. Pogančić, A. Paulus, V. Musil, G. Martius, and M. Rolinek. Differentiation of Blackbox Combinatorial Solvers. In International Conference on Learning Representations, Sept. 2019.
Prillo and J. Eisenschlos. SoftSort: A Continuous Relaxation for the argsort Operator. In Proceedings of the 37th International Conference on Machine Learning, (ICML) 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 7793-7802. PMLR. 2020.
G.-C. Rota. The Number of Partitions of a Set. The American Mathematical Monthly, 71(5):498, May 1964. ISSN 00029890. doi: 10.2307/2312558 Publisher: JSTOR.
Y. Rubner, C. Tomasi, and L. J. Guibas. The earth mover's distance as metric for image retrieval. International journal of computer vision, 40(2):99, 2000. Publisher: Springer Nature BV.
Sabour, Frosst, and G. E. Hinton. Dynamic Routing Between Capsules. In Guyon, U. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
R. Santa Cruz, B. Fernando, A. Cherian, and S. Gould. Deeppermnet: Visual permutation learning. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 3949-3957, 2017.
S. Sarfraz, V. Sharma, and R. Stiefelhagen. Efficient parameter-free clustering using first neighbor relations. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9934-8943, 2019.
O. Sener and V. Koltun. Multi-tass task learning as multi-objective optimization. Advances in neural information processing systems, 31, 2018.
R. Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matrices. The annals of mathematical statistics, 35(2):876-879, 1964. Publisher: JSTOR.
T. M. Sutter, L. Manduchi. A. Ryser, and J. E. Vogt. Learning Group Importance using the Dif- ferentiable Hypergeometric Distribution. In The Eleventh International Conference on Learning Representations, 2023.
I. L. Thurstone. A law of comparative judgment. In Scaling, pages 81-92. Routledge, 1927.
K. Wallenius. Biased sampling: the noncentral hypergeometric probability distribution. Technical report, Stanford Univ Ca Applied Mathematics And Statistics Labs, 1963.
H. Xiao, K. Rasul, and R. Vollgraf. Fashion-MNIST: Novel Image Dataset for Benchmarking Machine Learning Algorithms, Sept. 2017, arXiv:1708 707.07747 [cs, stat].
S. M. Xie and S. Ermon. Reparameterizzbbl subset sampling via continuous relaxations. In Pro- ceedings of the 28th International Joint Conference Artificial Intelligence, IJCAI' 19, pages 3919-3922 3925, Macao, China, Aug. 2019. AAAI Press. ISBN 978-0-9992411-4-11
D. Xin, B. Ghorbani, A. Garg, O. Firat, and J. Gilmer. Do Current Multi-Task Optimization Methods in Deep Learning Even Help? arXiv preprint 2:Xiv:2209.11379, 2022.
X. Yang, C. Deng, F. Zheng. J. Yan. and W. Liu. Deep spectral clustering using dual autoencoder network. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4066 4075, 2019.
Yellott. The relationship between Luce's choice axiom, Thurstone's theory of comparative judgment, and the double exponential distribution. Journal of Mathematical Psychology. 15(2): 109, 1977. Publisher: Elsevier.
13

---

A Preliminaries
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 E N be the number of different classes or groups (e.g. marble colors the urn, m = [mm1,... .,..,...,... _x] k] E N describe the number of elements per class (e.g. marbles per color), NN = Σi --1 m be the total number of elements (e.g. all marbles in the urn) and n E ,O,...,,. J be the number of elements (e.g. marbles) to draw. Then, the multivariate hypergeometric distribution describes the probability of drawing n = [M1..... nx] E N marbles of class k.
where n2 is the number of drawn marbles by sampling without replacement such that Σ Sk=1 ng
In the literature, two different versions of the noncentral hypergeometric distribution exist, Fisher's s (Fisher, 1935) and Wallenius" (Wallenius, 1963; Chesson, 1976) distribution. Sutter 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 's noncentral hypergeometric distribution.
Definition A.1 (Multivariate Fisher's Noncentral Hypergeometric Distribution (Fisher, 1935)). A random vector X follows Fisher's 's noncentral maltivariate distribution, if its joint probability mass function is given by
K
1
mk
P(N- = =; = w)
wk
(12)
Po
k=
where
Po Σ
mk
with
(13)
7]
(m (M....R))E
k=1
The support S of the PMF is given by fn E NK Vk
n S
mk.
Ek=1
n
and
B(n-kj
The class importance w is 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 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 kt Every step includes merging operation, which leads to biased samples compared to groundtruth non-differentiable sampling with equal class weights w. Given that we intend to use the differentiable MVHG in settings where we want to learn the unknown class weights, we do not expect negative effect from this sampling procedure. For details on how merge the MVHG into sequence unimodal distributions, we refer to Sutter al. (2023).
The probability mass function calculation is based on unnormalized log- weights, which are interpreted as unnormalized weights of categorical distribution. The interpretation of the condittonai unimodal hypergeometric distributions Softmax trick (Jang et al., 2016; Maddison et al., 2017). Following the use of the Gumbel--oftmax trick, the class-conditional version of the hypergeometric distribution is differentiable and reparame- terizable. Hence, the MVHG has been made differentiable and reparameterizable as well. Again, for details refer to the original paper (Sutter et al., 2023).
categorical distributions allows applying the Gumbel-
14

---

A.2 Distribution over Random Orderings
Yellott (1977) show that the distribution over permutation matrices p( s) follows Plackett- Luce (PL) distribution (Plackett. 1975; Luce, 1959), if and only of the perturbed scores & are sampled independently from Gumbel distributions with identical scales. For each item i, sample Di Gumbel(O, (3) independently with zero mean and and fixed scale B. Let & be the vector of Gumbel perturbed log- weights such that 8; = B3 log 81 + Jt- Hence,
$1 82
g 2 ... kn) =
(14)
7 81
5;
We refer to Yellott (1977) or Grover al. (2019) for the proof. However, Grover et al. (2019) provide only an adapted proof sketch from Yellott (1977). The probability of sampling element first is given by its score s; divided by the sum of all weights in the set
8, ((s)) Z
(15)
For = log 8;s the right hand side of Equation (15) is equal to the softmax distribution softmax (x) exp(2i)/ / E, (2)( as already described in (Xie and Ermon, 2019). Hence, Equation (15) directly leads to the Gumbel-Softmax trick (Jang et al., 2016; Maddison al., 2017).
A.2.1 Differentiable Sorting
In the main text of the paper we rely on differentiable function f=()), which sorts the resampled version of the scores
π = (3) sort(
(16)
Here, we summarise the findings from Grover et al. (2019) on how to construct such differentiable sorting operator. As already mentioned in Section 2. there are multiple works on the topic (Prillo and Eisenschlos, 2020; Petersen al., 2021: Mena al.. 2018), but we restrict ourselves the work of Grover et al. (2019) as we see the differentiable generation of permutation matrices as tool in our pipeline.
Corollary A.2 (Permutation Matrix (Grover et al., 2019)). = )T be real-valued vector of length n. Let As denote the matrix of absolute pairwise differences of the elements of such that Ag [. = 18; 8yl. The permutation matrix corresponding to sort(s) is given by:
{
ifj arg max((n + 1 2i)s AA1]
π =
(17)
otherwise
where 1 denotes the column vector of all ones.
As know, the arg operator is non- non-differeetiable which prohibits the direct use of Corollary A.2 for gradient computation. Hence, Grover et al. (2019) propose to replace the arg max operator with softmax to obtain continuous relaxation (T) similar to the GS trick (Jang et al., 2016; Maddison et al., 2017). In particular, the ith row of x()) is given by:
[,]] = softmax((1 + 1 - 21) 1/1/T]
(18)
where + > 0 is a temperature parameter. We adapted this section from Grover et al. (2019) and we also refer their original work for more details on how to generate differentiable permutation matrices.
In this, work we remove the temperature parameter reduce clutter in the notation. Hence, only write т instead of (T), although it is still needed for the generation of the matrix и. For details on how select the temperature parameter in our experiments, we refer to Appendix C.
B Detailed Derivation of the Differentiable Two-Stage Random Partition Model
B.1 Two-Stage Partition Model
We want to partition n elements [m] = (1...... unknown.
........] n. into K subsets (Si,...,SKk where a priori
15

---

Definition B.1 (Partition). A partition of set of elements [n]] {I,...,n]] n) is a collection of subsets (Si...., .Kk such that
S1 U.U.UUU [m] and Vi + j S,nS
(19)
Put differently, every element i has to be assigned to precisely one subset S₁. We denote the size of the k-th subset Sk as n& |S&]. Alternatively, we describe partition P as an assignment matrix Y = ly...... MK] E {0, 1] 1j Every row ys E {0, 1|1x" is multi-hot vector, where yki = 1 assigns element to subset Sk.
In this work, we propose a new two-stage procedure to learn partitions. The proposed formulation
separately infers the number of elements per subset ni and the assignment of elements to subsets Sy by inducing an order on the n elements and filling Si,..., Sx sequentially in this order. See Figure 1 for example.
Definition B.2 (Two-stage partition model). Letn [mi,...., nxl E NK be the subset sizes in p. with No the set of natural numbers including and Σ ni = n, where n is the total number of elements. Let и E {0,1) 1]nxn be a permutation matrix that defines an order over the n elements. We define the two-stage partition model of n elements into K subsets as an assignment matrix Y yl y1..... yx] E (0, 1]Kx" with
Uk Mkt
k-1
yk
Σ Ti, 1
where Vk- =
n2
(20)
such that Y [[y) ni > 0]K
O= T.
Note that in contrast to previous work partition models (Mansour and Schork, 2016), allow Sk to be the empty set 0. Hence, K defines the maximum number of possible subsets, not the effective number of non-empty subsets.
To model the order of the elements, we use permutation matrix π [T].................... n., T., E (0, 1] which is square matrix where every row and column sums to 1. This doubly all permutation matrices x (Mareus, 1960) thus ensures that the columns of Y remain one-hot vectors. At the same time, its rows correspond to ny-hot vectors yı in Definition B.2 and therefore serve as subset assignment vectors.
property of
Corollary B.3. A two-stage partition model Y, which follows Definition B.2, is valid partition satisfying Definition B.1.
Proof. By definition, every row ग. and column т of is one-hot vector, hence every Σ nng results in different, non-overlapping ng-hot encodings, ensuring S; SS A i, and i # Further, since ng-hot encodings have exactly ny entries with 1, we have Σ De Σ/=1 Tij Tkk
2i=u
Hence, since
Σx=1 ni every element is assigned to yv: ensuring U... U....USp USk [n].
K
B.2 Two-Stage Random Partition Models
An RPM pYY) defines a probability distribution over partitions Y. In this section, we derive how to extend the two-stage procedure from Definition B.2 to the probabilistic setting to create two-stage RPM. To derive the two-stage RPM's probability distribution p(Y), we need to model distributions and x. We choose the MVHG distribution p(n:;) w) and the PL distribution p(T; s) (see Section 3).
We calculate the probability p(Y; w 8) sequentially over the probabilities of subsets Pys p(un
yek: w, 8). Pys itself depends on the probability over subset permutations Pa, := p( n y<ki 8). where subset permutation matrix म represents ordering over n, out of n elements.
Definition B.4 (Subset permutation matrix it.. A subset permutation matrix स E {0, 1 where nk S n, must fulfill
Vi < n4
Σ Tej
and Vj <:: Σ
3 1
16

---

We describe the probability distribution over subset permutation matrices PR. using Definition B.4 and Equation (3).
Lemma B.5 (Probability over subset permutations pra). The probability pnù of any subset permula- tion matrix मं = [ #nd E {0, 1,",." is given by
(T8);
psg = p( nu,, Mccki)))
(21)
Zx
2 21=1 (T8)
where yek = (91.................. yk 1.. 2 = Z Ejes., 8j and Sek= = U S
Proof. We provide the proof for pai. but it is equivalent for all other subsets. Without loss of generality, assume that there nı elements in Sj. Following Equation (3), the probability of permutation matrix p(T; 8) is given by
p( 8) (T8)
(T8)2
(T8))
(22)
Z (TS)1
Z - 3ーΣドー
1j=1 T8);
At the moment, we are only interested in the ordering of the first n1 elements. The probability of the first nı is given by marginalizing over the remaining n n1 elements:
p mi:w)
Σ p( $)
(23)
where IT, is the set of permutation matrices such that the top n1 rows select the elements in specific
ordering ग E {0, 11 xn i.e. IT, = ( [
n,,]² 1T
R). It follows
p( m1:w)
Σ p( 8)
(24)
mlll,
(TS))
Σ
(25)
Z
Σ j=1
(TS)
(is);
n-R1
(x8)ni+ +i
(26)
3=1 Σ (Ts))
Z R1 (T8); -E
(Ts),
xEIT;
im1
81
(ms);
(TS) niti
(27)
7. Σド (As))
7 Σ j= (mt)}
xelli
i=1
where 7 = ---
E}= ( follows
R2
(Ts)
p( miw
(28)
Z
(
Lemma B.5 describes the probability of drawing the elements E Sg in the order described by the subset permutation matrix म given that the elements in Sek are already determined. Note that in slight abuse of notation, we use p(# nig, Bicki w, s) as the probability of a subset permutation given that there are nı elements in Sş and thus म E {0, 1 X., Additionally, we condition on the subsets ye and nıg, the size of subset Sg. In contrast to the distribution over permutations matrices p(T;; s) in Equation (3), we take the product over n& terms and have a different normalization constant Zz. Although we induce an ordering over all elements in Definition B.2, the probability Py, is invariant to intra-subsel orderings of elements i E St.
Lemma B.6 (Probability distribution Py.). The probability distribution over subset assignments Pum is given by
Puc := p(y: yycki w. 8) p(n m)ckiw)
Σ
p( The Mccii) 8)
#Ell
where 11y. { 12 yk S ग.ं and p(
14, Yekk 8) in Lemma B.5.
17

---

Proof. We can proof the statement of Lemma B.6 as follows:
Pux
plyn yckiiw, 8)
Σ Plu s)k:w,
(29)
plnn yess s)p( mhs yers w. 8)
(30)
Σ p(n) meniw, s)p( m Dcus8)
(31)
"i
p(ax
nk. Mckis) s) nekkw, 8) s
(32)
p(ne
nexkww Σ
p( Mcki)
(33)
REllyR
Equation (29) holds by marginalization, where nig denotes the random variable that stands for the size of subset Sg. By Bayes' rule, from the fact that we can compute nick if y<v is given, as the assignments Yok hold information on the size of subsets Sek-- More explicitly, ny Σÿ=1 Mij- Further, yı is independent of if the size nig of subset S& is given, leading to Equation (31). We further observe that p(yk n/. M<ck s) is only non-zero, if n2 = EE=
then derive Equation (30). The next derivations stem
ng. Dropping all zero terms from the sum in Equation (31) thus results in Equation (32). Finally, by Definition B.2, know that yv: - Ti where k=1 1=1 n. and
P4 -'= need to marginalize over all permutations of the clements yı given that the elements in yek are already ordered, which corresponds exactly to marginalizing over all subset permutation matrices such that yı: Ee म, resulting in Equation (33).
{0, 1 permutation matrix. Hence, in order to get yı given Y-ck,
In Lemma B.6, we describe the set of all subset permutations of elements Sy by Hg Put differently, we make p(yk Y<k: w, s) invariant to the ordering of elements i S₁ by marginalizing over the probabilities of subset permutations Pa, (Xie and Ermon, 2019).
Using Lemmas B.5 and B.6, we propose the two-stage random partition o(Y:w,s) s). Since Y = .y,..... yk] |T calculate wY; s), the PMF of the two-stage RPM, sequentially ing Lemmas B.5 and B.6, where we leverage the PL distribution for permutation matrices p(T:s) to describe the probability distribution over subsets p(yx 8.).
(Two-Stage Random Partition Model). Given a probability distribution over subset sizes p(n; w) with n N and distribution parameters w E R.K and a PL probability distribution over random orderings p( s) with æ E (0, 1] ] and distribution parameters R the probability mass function p( p(Y;w, of the two-stage RPM is given by
Proposition 4.1
p(Y:,,s) p(v mk:w, 8) p(n;w) Σ p(T;s)
(34)
TETly
where 11y = 3k Disu+t1 Ii 1,..., K}. and ys and vy as in Definition B.2. ......,K
18

---

Proof. Using Lemmas B.5 and B.6, we write
p(Y) ply
yk:w, 8) = (() ..... (
p(m):w) Σ
p m1:s)
T
plnx (nsliekiw)
E
p(āk l (iglinsss
(35)
Axellyx
=p(n):w) ... p(nx (nkljekiw)
Σ p( 11:8
E p(ūk| (m))sskis)
(36)
Flellyi
ux Ex
=p(n;w)
Σ
Σ
p( 11;8) (.n;8)...ppp pp(px (ngbiskis)
(37)
Rielly;
xx Ex
w)(n;w) Σ
p(T n)
(38)
TETly
=p(n:w)
Σ p(T;:)) s)
(39)
relly
B.3
Approximating the Probability Mass Function
Lemma 4.2.
p(Y:w, s) can be upper and lower bounded as follows
Vī E Il ppp;:)p(T;s) I w(Y:w,s) s) s [y (p(::( maxpp(;;
(40)
Proof. Since p( s) is a probability know that Vz E 10, 1j p(T:s) 2 0. Thus, it follows
directly that:
Va E Il: p(Y:,,s) w(n:w)
Σ p(r';s) 2 p((;:))(TTss)
E'Elly
proving the lower bound of Lemma 4.2.
On the other hand, can prove the upper bound in Lemma 4.2 by:
8) = p(n; w)
Σ p(T':))
<'Elly
Sp(n:w)
Σ
max p( s)
TE
E'Elly
=p(n:w) nelly
p(T; 8)
Σ
E'elly
= p(n:w) max ((T:8) nllly
< [ ((); max (TT;
We can compute the maximum probability max, p(T; 8) with the probability of the permutation matrix fa (8), which sorts the unperturbed scores decreasing order.
B.4 The Differentiable Random Partition Model
We propose the DRPM p(Y: w, s), a differentiable and reparameterizable two-stage RPM.
19

---

Lemma 4.3 (DRPM). A two-stage RPM is differentiable and reparameterizable if the distribu- tion over subset sizes p(ne; w) and the distribution over orderings p( s) are differentiable and reparameterizable.
Proof. To prove that our two-stage RPM differentiable we need to prove that we can compute gradients for the bounds in Lemma 4.2 and provide reparameterization scheme for the two-stage approach in Definition B.2.
Gradients for the bounds: Since we assume that p(n:; w) and p(T;; s) are differentiable and repa- rametcrizable, we only need to show that we can compute [[ly] and max; p( s) in a differentiable manner to prove that the bounds in Lemma 4.2 are differentiable. By definition (see Section 4.1),
my
II mg!
Hence, | can be computed given a reparametrized version nyk which is provided by the reparametrization trick for the MVHG p(n; w). Further, from Equation (14), we immediately see that the most probable permutation is given by the order induced by sorting the original, unper- turbed scores from highest to lowest. This implies that max; p(#;; s) p p(=.:s) 8), which we compute due to p(T s) being differentiable according to our assumptions.
Reparametrization of the two-stage approach: Given reparametrized versions of n and x, we compute a partition as follows:
PE t
--1
Vk
Te.
where V/
n2
(41)
The challenge here is that we need to be able to backpropagate through wk. which appears an index in the sum. Let ar = {0, 1]", such that
{1
if u <is Mk+1 otherwise
(oox):
Given such oa., we can rewrite Equation (41) with
yk > (as)
(42)
While this solves the problem of propagating through sum indices, it is not clear how to compute Og in a differentiable manner. Similar to other works on continuous relaxations (Jang et al., 2016; Maddison al., 2017), we can compute relaxation of ag by introducing temperature т. Let us introduce auxiliary function f N + jo, 1)", that maps an integer to vector with entries
[((T;) ("-its)
such that fi(;;T Of s-Í and f((x;)) if -ーو 2 0. Note that (()) is the standard sigmoid function and <<< is a small positive constant to break the tie at a (0). We then compute an approximation of or, with
(T( = ((v=)) ((x-1:
@+(7) E [0, 1]". Then, for we have (x(TT) - ag. In practice, cannot set 0 since this would amount to division by O. Instead, we can apply the straight-through estimator (Bengio al., 2013) the auxiliary function (I;;T) in order get â E (0,1)" and use it to compute Equation (42).
Note that in our experiments, we use the MVHG relaxation of Sutter et al. (2023) and can thus leverage that they return one-hot encodings for ng. This allows a different path for computing os which circumvents introducing yet another temperature parameter altogether. We refer to our code in the supplement for more details.
20

---

Table 3: Total GPU hours per experiment. We report the cumulative training and testing hours to generate the results shown in the main part of this manuscript. We relied on our internal cluster infrastructure equipped with RTX2080Ti GPUs. Hence, we report the number of compute hours for this GPU-type.
Experiment
Computation Time (h)
Clustering (Section 5.1)
100
Partitioning of Generative Factors (Section 5.2)
480
MTL (Section 5.3)
100
C Experiments
In the following, we describe each of our experiments in more detail and provide additional ablations. All our experiments were run on RTX2080Ti GPUs. Each run took 6h-8h (Variational Clustering). 4h-6h (Generative Factor Partitioning). or ~ 1h (Multitask Learning) respectively. We report the training and test time per model. Please note that we can only report the numbers to generate the final results but not the development time.
com/thomassutter// drpm. Please note that the results reported in the main text slightly differ from the ones being generated from the official code. For the main paper, we based our own code for the experiments in Section 5.2 on the disentanglement ement_lib lib (Locatello et al., 2020). However, the library is based on Tensorflow v1 (Abadi et al., 2016), which makes it more and more difficult maintain and install. Therefore, we decided to re-implement everything in Py Torch (Paszke et al., 2019).
Code Release The official code can be found under https://gii
While the metrics of our method and the baselines slightly change, the relative performance between them remains the same.
The code and results for the remaining two experiments in Sections 5.1 and 5.3 are the same as in the main text.
C.1 Approximation quality of Lemma 4.2
To provide intuitive understanding of the bounds introduced in Lemma 4.2, we present an experiment in this subsection to demonstrate the behavior of the upper and lower bounds. It is important to note that RPMs are discrete distributions, as the number of possible samples finite for a given number of elements n and subsets K. Therefore, we can estimate the probability of a fixed partition Y under given w and ø by sampling M partitionn from p(Y: w, 8), counting the occurrences of in the samples, and dividing the count by M. As M approaches infinity, can obtain the true probability mass function (PMF) p(Y; 8) for every partition Y.
In our experiment, we set = 5 and aim to evaluate the quality of our bounds for all possible subset combinations of 5 elements, we thus set = n 5. To obtain a reliable estimate of the true PME, we set M = 10 From Lemma 4.2, we know that:
Vス E lly p((;:)p((z;)) 8) I p(Y;w,)) I [lly (p(p::(()))) max p(;; 8)
Let us define
Pu (Y: s) = Hy lp(n; w) max s)
Pr (Y max p(n; w )p(p; s)
HElly
In Figure 5, we present the estimated PMF along with the corresponding upper bounds (pv) and lower bounds (PL) for four different combinations of RPM parameters and s. We observe that when all scores $ are equal, as in the priors of the experiments in Sections 5.1 and 5.2, Pu (Y; w, s) approximates p(Y: w. s) well and serves as reliable estimate of the PMF However, when the scores vary, the upper bound becomes looser, particularly for lower probability partitions, as is dominated by the term max; p(ñ; s). Although Pr. (Y: w 8) appears looser than py (Y;w, 8) for certain configurations of and 8, provides more consistent results across all hyperparameter combinations.
21

---

-4
-5-
log(ptY; (.. 51)
log(ptY; W. s))
-12
-12
-14
-14
-15
-16
Upper Bound
Upper Bound
Estimated
Estimated
Lower Bound
Lower Bound
-20
Y- w.
Y p(Y ().))
(a) 1.
(b) ~ U(O,
-0
51
log(p(Y:
-12
109
-14
-14
-16
-15
Upper Bound
Upper Bound
Estimated Lower Bound
Lower Bound
-20
Y w.
Y p(Y; W. 5)
(c) w 1, s U[O,, 1])
(d) U[O, 1]), U 1)
Figure 5: Partitions with n = 5 and = 5 for different w and 8. Each point the plots corresponds to one of the n different partitions and their respective estimated probability mass and its upper/lower bounds according to Lemma 4.2.
C.2
Variational Clustering with Random Partition Models
C.2.1 Loss Function
As mentioned in Section 5.1, for given dataset X with N samples, let Z and Y contain the respective latent vectors and cluster assignments for each sample in X. The generative process then be summarized
follows: First, sample the cluster assignments Y from RPM, iee, Y ~ P(Y;w,,) 8). Given Y, we can sample the latent variables Z, where for each we have z (H of 1t). & R.. Finally, we sample X by passing each z through a decoder like in vanilla VAEs. Using Bayes rule and Jensen's 's inequality, can then derive the following evidence lower bound (ELBO):
log(p(X)) = log [|ExxxY Z)dZ
p p(XI2)2(2X3>000) X()
2 E
log
= Lx
We then assume that we can factorize the approximate posterior follows:
((2,\|x) x(|))) II <(2())
#EX
Note that while we do assume conditional independence between z given its corresponding x. model (X) with the DRPM and do not have to assume conditional independence between
22

---

X
X1
X2
IV
p(xx-1(ex-))
p(sx(es)
2
21
NN
15
pfeelao)
play-- 90-.1]
play mc)
40122)
Y
311
yN-1
UN
Figure 6: Generative model the DRPM clustering model. Generative paths are marked with thin arros,, whereas inference is in bold.
different cluster assignments. This allows leverage dependencies between samples from the dataset. Hence, we can rewrite the ELBO as follows:
CELBO(X) x [log(p(X|Z)]
- [K [[]][(][
Σ E [log p(+22)
=xX
X [KLIg(z|=))lp<</>> #
See Figure for an illustration of the generative process and the assumed inference model. Since computing P(Y) and (() is intractable, we further apply Lemma 4.2 to approximate the KL- Divergence term in LELBO, leading to the following lower bound:
LELBO >
Σ (=(=)))) [[og((=(())
(43)
x
Σ 1x( [KL|g<>])]]lx<]>] xex
(44)
Iy (())(())) ((n;w)((Tyis)
log
Ex(r(x))
(45)
log (maxq(";ssX))) (
(46)
where Ty is the permutation that lead to Y during the two-stage resampling process. Further, we want to control regularization strength of the KI. divergences similar to the 3-VAE (Higgins et al., 2016). Since the different terms have different regularizing effects, we rewrite Equations (45)
23

---

Prior Parameters
C
ee
(XX
3
Y
3
xXX
E
D
(4(2))
x
a(())
à
Figure 7: Autoencoder architecture of the DRPM-VC model.
and (46) and weight the individual terms as follows, leading to our final loss:
==
I (x(()) [log [(=2((
(47)
xEX
+ B
Σ BX( (1Zl((((=))\p(((1)))
(48)
|Ilr| w( w (X))
+ 7 (x(x)
log
(49)
p(n;w)
max; s())
+ 8 1x(r
log
(50)
p(Ty;:))
C.2.2 Architecture
The model for our clustering experiments is a relatively simple, fully-connected autoencoder with a structure seen in Figure 7. We have a fully connected encoder E with three layers mapping the input to 500, 500, and 2000 neurons, respectively. We then compute each parameter by passing the encoder output through linear layer and mapping the respective parameter dimension in the last layer. In our experiments, we use latent dimension size of / 10 for MNIST and = 20 for FMNIST, such that ((r)), (ax) E TR! To understand the architecture choice for the DRPM parameters, let us first take closer look at Equation (48), For each sample this term minimizes the expected KL divergence between its approximate posterior g x)) -N((xx))
(x))) and the prior index given by the partition Y sampled from the DRPM g(Y XX; 8, w). i.c., N. diag(oy)). Ideally, the most likely partition should assign the approximate posterior the prior that minimizes this KL divergence. We can compute such s(X) and w(X) given the parameters of the approximate posterior and priors follows:
Va, X s1(x;) u. (K arg
argmin min (K (lag (iag:)))))))(
1
Nx p4, diag(ox))
w(X)
X
Σ Nazlur (iiggoi
FEX
where u a scaling constant that controls the probability of sampling the most likely partition. Note that w and minimize Equation (48) if defined this way when given the distribution parameters of the approximate posterior and the priors. The only thing that left unclear is how much u should scale the scores s. Ultimately, we leave as learnable parameter but detach the rest of the computation of and w from the computational graph to improve stability during training. Finally, once we resample ~ ()((z), (zx)), we pass it through fully connected decoder D with four layers mapping to 2000, 500, and 500 neurons the first three layers and then finally back to the input dimension in the last layer to end up with the reconstructed sample c.
24

---

C.2.3 Training
As in vanilla VAES, we can estimate the reconstruction term in Equation (47) with MCMC by applying the reparametrization trick (Kingma and Welling, 2014) to ((2(x) to sample M samples x(1) {(2[2)) and compute their reconstruction error to estimate Equation (47). Similarly, we can sample from {{|)) I. times to estimate the terms in Equations (48) to (50), such that we minimize
M
Σ
log lggp(z)2(())
M
DEX
3
KT a2) (p<2)\\))
+
ØEX
1=1
L
[II ((n (X))
+
log
L
((n w)
1=1
maxx @( s(X))
+ T
log
((Ty(1;8)
In our experiments, we set M and L 100 since the MVHG and PL distributions are not concentrated around their mean very well, and more Monte Carlo samples thus lead to better approximations of the expectation terms. We further set 3 = for MNIST and B = 0.1 for FMNIST, and otherwise y and 5 = 0.01 for all experiments.
To resample n and ı we need to apply temperature annealing (Grover et al., 2019; Sutter al., 2023). To do this, we applied the exponential schedule that was originally proposed together with the Gumbel-Softman trick (Jang et al., 2016; Maddison et al., 2017), i.e., max(T]inal: exp(-rt)) where t is the current training step and r is the annealing rate. For our experiments, we choose log(T fimal) -log(Temee) in order to annealing over 100000 training step. Like Jang et al. (2016), 100000
set Timit = and Tfimol = 0.5.
Similar to Jiang et al. (2016), we quickly realized that proper initialization of the cluster parameters and network weights is crucial for variational clustering. In our experiments, we pretrained the autoencoder structure by adapting the contrastive loss of (Li et al., 2022), as they demonstrated that their representations manage to retain clusters in low-dimensional space. Further, we also added reconstruction loss to initialize the decoder properly. To initialize the prior parameters, fit GMM to the pretrained embeddings of the training set and took the resulting Gaussian parameters to initialize our priors. Note that used the same initialization across all baselines. See Appendix C.2.4 for an ablation where we pretrain with only a reconstruction loss similar to what was proposed with the VaDE baseline.
To optimize the DRPM-VC in our experiments, we used the Adam (Loshchilov and Hutter, 2019) optimizer with learning rate of 0.0001 with a batch size of 256 for 1024 epochs. During initial experiments with the DC, we realized that the pretrained weights the encoder would often lose the learned structure in the first couple of training epochs. We suspect this to be artifact of instabilities induced by temperature annealing. To deal with these problems, we decided to freeze the first three layers of the encoder when training the DRPM-VC, giving much better results. See Appendix C.2.5 for an ablation where we applied the same optimization procedure to VaDE.
Finally, when training the VaDE baseline and the DRPM-VC FMNIST, we often observe local optimum where the prior distributions collapse and become identical. We can solve this problem by refitting the GMM in the latent space every 10 epochs and by using the resulting parameters to reinitialize the prior distributions.
C.2.4 Reconstruction Pretraining
While the results of our variational clustering method depend a lot on the specific pretraining, we want to demonstrate that improvements over the baselines do not depend on the chosen pretraining method. To that end, we repeat our experiments but initialize the weights of our model with an autoencoder that has been trained to minimize the mean squared error between the input and the reconstruction. This initialization procedure was originally proposed in (Jiang et al., 2016). We
25

---

Table 4: We compare the clustering performance of the DRPM-VC on test sets of MNIST and FMNIST between GMM in latent space (Latent GMM) and Variational Deep Embedding (VaDE) initializing weights using an autoencoder trained on a reconstruction objective. We measure per- formance in terms of the Normalized Mutual Information (NMI), Adjusted Rand Index (ARI), and cluster accuracy (ACC) over five seeds and put the best model in bold.
MNIST
FMNIST
NMI
ARI
ACC
NMI
ARI
ACC
LATENT GMM
0.75=0.00
0.66=0.01
0.75+
0.56+0.002
0.41#0.03
0.57=0.02
VADE
0.77=0.02
0.62=0.01
0.69=0.01
0.53+
0.35+ 08
0.47+0.09
DRPM-VC
0.74=0.00
0.67=00001
0.75+0.00
0.59=0.01
0.47=0.02
0.62#0001
Table 5: We compare the clustering performance of the DRPM-VC on test sets of MNIST and FMNIST between GMM in latent space (Latent GMM), and Variational Deep Embedding (VaDE) when freezing the encoder. We measure performance in terms of the Normalized Mutual Information (NMI), Adjusted Rand Index (ART), and cluster accuracy (ACC) over five seeds and put the best model bold.
MNIST
FMNIST
NMI
ARI
ACC
NMI
ARI
ACC
LATENT GMM
0.86+0.02
0.83=0.06
0.88+0.07
0.60++0.00
0.47=0.01
0.62+0.01
VADE
0.90#0.02
0.88+0006
0.92+0.00
0.64+0.01
0.47=0.01
0.59=0.03
DRPM-VC
0.89
0.8810.03
0.94
0.64:10.00
0.51:0.01
0.65
present the results of this ablation in Table 4. Simply minimizing the reconstruction error does not necessarily retain cluster structures in the latent space. Thus, it does not come as a surprise that overall results get about 10% to 20% worse across most metrics, especially for MNIST, while results on FMNIST only slightly decrease. However, we still beat the baselines across most metrics, suggesting that modeling the implicit dependencies between cluster assignments helps to improve variational clustering performance.
C.2.5 Baselines with fixed Encoder
For the experiments in the main text, we wanted implement the VaDE baseline similar to the original method proposed in Jiang et al. (2016). This means, in contrast to our method, we used their optimization procedure, i.e., Adam with a learning rate of 0.002 with a decay of 0.95 every 10 steps, and did not freeze the encoder we do for the DRPM-VC. To ensure our results do not stem from this minor discrepancy, we perform an ablation experiment on VaDE using the same optimizer and learning rate with the DRPM-VC and freeze the encoder backbone. The results of this additional experiment can be found Appendix C.2.5. As can be seen, VaDE results do improve when adjusting the optimization procedure in this way. However, we still match or improve upon the results of VaDE in most metrics, especially in ARI and ACC, suggesting purer clusters compared to VaDE. We suspect this is because we assign samples to fixed clusters when sampling from the DRPM, whereas VaDE performs soft assignments by marginalizing over a categorical distribution.
C.2.6 Additional Partition Samples
In Section 5.1, we have seen a sample of a partition of the DRPM-VC trained on FMNIST. We provide additional samples for both MNIST and FMNIST at the end of the appendix in Figures 15 and 16. We can see that for both datasets, the DRPM-VC learns coherent representations of each cluster that casily allow us generate new samples from each class.
26

---

DRPM-VC
0.6
0.6
VaDE
Latent GMM
0.5
ARI
0.4
DRPM-VC
0.4
VaDE
0.3
Latent GMM
0.2
0.2
10
20
# Clusters
#Clusters
(a) NMI when varying K
(b) ARI when varying K
0.7
0.6
Accuracy
0.5
0.4
DRPM-VC
VaDE
0.3
Latent GMM
0.2
8
10
12
14
16
18
20
#Clusters
(c) Accuracy when varying K
Figure 8: Clustering performance of Latent GMM, VaDE, and DRPM-VC when varying the number of clusters K for seeds Fashion MNIST. DRPM-VC consistently outperforms the baselines except for = 20, where all methods perform similarly.
C.2.7 Samples per cluster
In addition to sampling partitions and then generating samples according the sampled cluster assigmments, we can also directly sample from each the learned priors. We show some examples of this for both MNIST and FMNIST the end of the appendix in Figures 17 and 18. We again see that the DRPM-VC learns accurate cluster representations since each the samples seems to correspond one of the classes in the datasets. Further, the clusters also seem to capture the diversity in each cluster, as we see a lot of variety across the generated samples.
C.2.8 Varying the number of Clusters
In previous experiments, we assumed that we had access to the true number of clusters of the dataset, which is, course, not true practice. We thus also want investigate the behavior of the DRPM-VC when varying the number of partitions K of the DRPM and compare it to our baselines. In Figure 8, we show the performance of Latent GMM, VaDE, and DRPM-VC for K E {6, 8, 10, 12, 14, 16, 18, 20) across 5 different seeds on FMNIST. DRPM-VC clearly outperforms the two baselines for all K, 20, where all models seem to perform similarly. Expectedly, DRPM-VC performs well when K is close to the true number of clusters, but performance decreases the farther we are from it. To investigate whether the model still learns meaningful patterns when are far from the true number of clusters, additionally generate samples from each prior of the DRPM-VCC VC Figure 9. Interestingly, we can see that DRPM- RCP still detects certain structures in the dataset but starts breaking specific FMNIST categories apart. For instance, it splits the clusters sandals/boots into clusters with (Priors 6/ 6/) and without (Priors 4 4/9) heels or the cluster T-shirt into clothing of lighter (Prior () and darker (Prior 3) color. Thus, ORPM-VC allows investigate clusters in datasets hierarchically, where low values of K detect more coarse and higher values of K more fine-grained grained patterns in the data.
except for the extreme case of
C.2.9 Clustering STL-10
include an additional variational clustering ablation on the STL-110 dataset (Coates et al., 2011) that follows the experimental setup of Jiang et al. (VaDE, 2016). As noted Jiang et al. (2016).
27

---

10
12
14
15
16
17
19
Sample
Figure 9: We generate 32 images from each prior when training DRPM-VC on Fashion MNIST with K 20. DRPM-VC starts detecting more fine-grained patterns in the data and splits some of the original clusters, such as sandals, boots, t-shirts, or handbags, into more specific sub-categories.
Table 6: We compare the clustering performance of the DRPM-VC and VaDE the test set of STL-10 10 across five seeds. We measure performance in terms of the Normalized Mutual Information (NMI), Adjusted Rand Index (ARI), and cluster accuracy (ACC) over five seeds and put the best model bold.
STL-10
NMI
ARI
ACC
VADE
0.80
0.7610.04
0.8810.02
DRPM-VC
0.83+0.01
0.80#0.02
0.91#0.00
variational clustering algorithms have difficulties clustering in raw pixel space for natural images, which is why we apply DRPM-VE VC to representations extracted using an Imagenet (Deng et al., 2009) pretrained ResNet-50 (He et al., 2015) as done Jiang al. (VaDE, 2016). Note that these results are hard to interpret, as STL-10 is a subset of Imagenet, meaning that representations are relatively easy to cluster as the pretraining task already encourages separation by label. For this reason, we list this experiment separately in the appendix instead of including it in the main text. We present the results of this ablation in Table 6, where we again confirm that modeling cluster assignments with our DRPM can improve upon previous work that modeled the assignments independently.
C.3 Variational Partitioning Generative Factors
We assume that we have access to multiple instances or views of the same event, where only subset of generative factors changes between views. The knowledge about the data collection process provides a form of weak supervision. For example, have two images of a robot arm as depicted here on the left side (see (Gondal et al., 2019)), which we would describe using high-level concepts such as color, position or rotation degree. From the data collection process, we know that subset of these generative factors is shared between the two views We do not know how many generative factors there are in total nor how many of them are shared. More precisely, looking at the robot arm, we do not know that the views share two latent factors, depicted in red, out of total of four factors. Please note that we chose four generative in Figure 10 only for illustrative reason as there are seven generative factors in the mpi3d toy dataset. Hence, the goal of learning under weak supervision is not
28

---

4
7
E
21
1
2
4
3 8
E
202
2
Figure 10: Motivation for the Partitioning of Generative Factors under weak supervision. The knowledge about the data collection process provides a weak supervision signal. We ohave access to dataset of pairs of images of the same robot with subset of shared generative factors (in red). We want to learn the shared and independent generative factors in addition to learning from the data. The images of the robot arms taken from Locatello et al. (2020) but originate from com/ com/rr-tu rr- srr-eearingg httgg/dseentnnng.. ement, ddttaset)....... The image is from Sutter al. (2023) and their ICLR 2023 presentation video (see https: / //iclr/ ///virtual/22233 ////2//2///ptte11107777777
the mpi3d toy dataset (see
only to infer good representations, but also inferring the number of shared and independent generative factors. Learning what is shared and what is independent lets reason about the group structure without requiring explicit knowledge in the form expensive labeling. Additionally, leveraging weak supervision and, hence, the underlying group structure holds promise for learning more generalizable and disentangled representations (see (e.g., Locatello al., 2020)).
C.3.1 Generative Model
We assume the following generative model for DRPM-VAE
x(X) f p(X,2)))) 2)dzz
(51)
[ p(X 2))(22dzz
(52)
where z = {2,, 21, }. The two frames share unknown number n, of generative latent factors 2,, and an unknown number, 71 and n2, of independent factors 21 and 22. The RPM infers ny and zz using Y. Hence, the generative model extends to
p(X) / (( 2) Σ p(z V))(Y)dz
/ p(01, x2 21 x2) Σ p(z Y)p(Y)dz
= p( Ze. 21)p(22
Zz. 22)
Σ 21, ())(Y)d2
(53)
Figure 11 shows the generative and inference models assumptions in graphical model.
29

---

Fi
32
X
p(a1 | Zv, z1)
p(az Rs. 22)
giY (sg mgl
U/x: (()))
olas J. mx)
pizs) #1, #2 (Y)
(s, IY, mi at)
(a) Generative Model
(b) Inference Model
Figure 11: Graphical Models for DRPM-VAE models in the weakly-supervised experiment.
C.3.2 DRPM ELBO
We derive the following ELBO using the posterior approximation g(2, Y X)
log Y2 log 9(2, Y |X)
x)]
LELBO (X) = [x(x
(54)
p(z, Y)
g(>
Y, x)e(
X)
logp(X (x) loog
(55)
p(z)p(Y)
g(= Y, Y,X)
loo ạ(v x)
log 111 022 |2) loog
(56)
p(z)
p(Y)
logp(am xx)
Es, 21)] - Eggx xIX)
[lgp(222 2s. 22)]
X| log g(2,, 21, #2
Y. X
1x( log q(Y X)
(57)
p(z,, 21, Z2)
p(Y)
Following Lemma 4.2, we are able to optimize DRPM- VAE using the following ELBO LELBO (X):
LELBO 2 {x)))
[logp(x;
2,, =1) Eg x) [logp(x22 Zss 22)]
(58)
log g(=,, 21, #2 Y, X)
Xx)
(59)
p(z,; 21, #2)
log |Il g(n X;w)
By(r1x)
(60)
p(n; wp)pp(pp( ( sp)
log (max ( x:
(61)
where Ty is the permutation that lead to Y during the two-stage resampling process. Further, want to control the regularization strength of the KL divergences similar to the 3-VAE (Higgins et al., 2016). The ELBO ((X) ) to be optimized can be written
C CELBO XI)) [log p(æ1
Zo. =1)l + Eg x) [log Z., 22]
(62)
log ((2e, 21, 22 X)X)
3 Ex(xY(
(63)
p(z,, 21. #2)
log | (nz; w( (XX))
7 [
(64)
p(n:: wp)
mmax q( s( X)
8 1g(Y(
log
(65)
p(Ty;8p)
where xXX and (XX) denote distribution parameters, which are inferred from (similar to the Gaussian parameters in the vanilla VAE).
30

---

View Aggrogation
22
(o(ss) Ra
(22)
Figure 12: Setup for the weakly-supervised experiment, The three methods differ only in the View Aggregation module.
As in vanilla VAEs, we can estimate the reconstruction term in Equation (58) with MCMC by applying the reparametrization trick (Kingma and Welling, 2014) to q(z Y, X) to sample samples x01) q(2 X) and compute their reconstruction error to estimate Equation (58). Similarly, we can sample from q( X) L times. We use L = 1 to estimate all expectations in LELBO.
C.3.3 Implementation and Hyperparameters
In this experiment, we use the disentanglements __i from Locatello et al. (2020). We use the same architectures proposed in the original paper for all methods we compare to. The bascline algorithms, Label (Bouchacourt et al., 2018; Hosoya, 2018) and AdaVAE (Locatello et al., 2020) are already implemented in di disentangl seeetanglemeet__11... ement_lis For details on the implementation of these methods we refer to the original paper from Locatello et al. (2020). HGVAE is implemented in Sutter et al. (2023). We did not change any hyperparameters or network details. All experiments were performed using 3 as this is the best performing / (according to Locatello et al. (2020). For DRPMVAE we chose = 0.25 for all runs. All models are trained on different random seeds and the reported results are averaged over the 5 seeds. We report mean performance with standard deviations.
We adapted Figure 12 from Sutter al. (2023). shows the baseline architecture, which is used for all methods. As already stated in the main part of the paper, the methods only differ in the View Aggregation module, which determines the shared and independent latent factors. Given a subset S of shared latent factors, we have
(0 æ3) = avg(g x1), g0(2) x2)) (2, x2)) (z 2,)
A eies else
(66)
(67)
where aug is the averaging function of choice (Locatello al., 2020; Sutter et al., 2023) and j E 21, The methods used (i. Label-VAE, Ada-VAE, HG-VAE, DRPM-VAE) differ in how to select the subset
For DRPM-VAE, we infer w from the pairwise KL-divergences K Lmn between the latent vectors of the views.
K KTpp (11, 222 = KL[((zi | m1))lq(Z2 xx2) + KT(((222 (g))lg(()))
(68)
where q(z; 2 are the encoder outputs of the respective images. We do not average or sum across dimensions in the computation of K LL such that the (. is d-dimensional, where d the latent space size. The encoder E in Figure 12 maps to μιæ;) and a (ιx of a Gaussian distribution. Hence, we can compute the KL divergences above in closed form. Afterwards, we feed the pairwise KI. divergence K KTpw to single fully-connected layer, which maps from d to K values
log w FC Lpu (x1,
(69)
where d = 10 and K 2 in this experiment. d is the total number of latent dimensions and K is the number of groups in the latent space. To infer the scores s(X) we again rely on the pairwise KL divergence K Lyw. Instead of using another fully-connected layer, we directly use the valuvs of the pairwise KL divergence
log = log K L (I1, 12)
(70)
Similar to the original works, we also anneal the temperature parameter for p(n; w) and p(T; 8) (Grover et al., 2019; Sutter et al., 2023). We use the same annealing function as in the clustering experiment (see Appendix C.2). We anneal the temperature т from 1.0 to 0.5 over the complete training time.
31

---

0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
O
P
7.
2
2
?
7
Figure 13: Samples from the noisyMultiiMNisS" dataset with increasing noise ratio in the right task.
C.4
Multitask Learning
C.4.1 MultiMNIST Dataset
The different tasks in multitask learning often vary in difficulty. To measure the effect of discrepancies in task difficulties on DRPM-MTL, we introduce the noisyMultMMNIST dataset.
The noisyMultiMNIST dataset modifies the MultiMNIST dataset (Sabour al., 2017) as follows. In the right image. we set each pixel value to zero with probability a E [O, 1]. This is done before merging the left and right image in order only affect the difficulty of the right task. Note that for = MoisyMultiMNNIST is equivalent to MultiMNIST and for a the right task can no longer be solved. This allows to control the difficulty of the right task, without changing the difficulty of the left. A few examples are shown in Figure 13.
C.4.2 Implementation & Architecture
The multitask loss function for the MultiMNIST dataset is
== w- + wnLnR
(71)
where wt and wz are the loss weights, and IL, and LR the individual loss for the respective tasks L and R. In our experiments, we set the task weights to be equal for all dataset versions, i.e. wt = w/ 0.5. We use these loss weights for the DRPM-MTL and ULS method. For the ULS method, it is by definition and to see the influence of mismatch in loss weights. The DRPM-MTL method on the other hand does not need additional weighting of loss terms. The task losses are defined as cross-entropy losses
C L = Σ gte log =-gt" log
(72)
where Ct CR = 10 for MultiNNIST, gt is a one-hot encoded label vector and is a categorical vector of estimated class assignments probabilities, i.e. >, Pe
The predictions for the individual tasks Pt are given as
P: he, (z), where
(73)
z enc (x)
(74)
for sample X (see also Figure 14). We use an adaptation of the LeNet-5 architecture LeCun et al. (1998) to the multitask learning problem (Sener and Koltun, 2018). Both DRPM-MTI. and ULS use the same network enco(.) with shared architecture up to some layer for both tasks, after which the network branches into two task-specific sub-networks that perform the classifications. Different to the ULS method, the task-specific networks in the DRPM-MTL pipeline predict the digit using only subset of z. DRPM-MTL uses the following prediction scheme
Pt = ho, (zz), where
(75)
2. z Yt
(76)
Ye = DRPM(W,, s): (z))
(77)
The DRPM-MTL encoder first predicts latent representation < enc,( where æ is the input image. Using the same encoder architecture but different parameters we predict partitioning encoding z' &- cz), With single linear layer per DRPM log-parameter log w and log are computed. Next we infer the partition masks yL. UR p(YL. yRi w. s). We then feed the masked latent representations 21. f- z @ yr. and žn +- 3 yn into the task specific classification networks ho, (21.) and hog (RR) respectively to obtain the task specific predictions. Since the two tasks in the MultiMNIST dataset are of similar nature, the sasc-f networks hag and hog share the same architecture, but have different parameters.
32

---

h0.
Task
ZL
ence
ZR
h0x
Task
enCo
W,
P(yL-YR:
Figure 14: Overview of the multitask learning pipeline of the DRPM-MTL method.
C.4.3 Training
For both the ULS and the DRPM-MTI. model, we use the Adam optimizer with learning rate 0.0005 and train them for 200 epochs with batch size of 256. We again choose an exponential schedule for the temperature т and anneal over the training time, as is explained Appendix C.2.3.
,0......... 0.9} and train each model with five different seeds. The reported accuracies and partition sizes are then means over the five seeds with the error bands indicating the variance and standard deviation respectively. We evaluate each model after the epoch with the best average test accuracy.
In our ablation we use E
C.4.4 CelebA for MTL
In addition to the experiment shown in Section 5.3, we show additional results for DRPM-MTL on the CelebA dataset (Liu et al., 2015). In MTL, each of the 40 attributes the CelebA dataset serves as individual task. Hence, using CelebA for MTL results is 40 task learning problem making the scaling of different task losses more difficult compared MultiMNIST (see Section 5.3) where we only need to scale two different tasks.
We again use the newly introduced DRPM-MTL method and compare the ULS model. We use the same pipeline as for MultiNNIST dataset but with different encoders and hyperparameters (see Appendices C.4.2 and C.4.3). We use the pipeline of Sener and Koltun (2018) with ResNet-based en- coder to map an image to representation of d = 64 dimensions. For architectural details, we refer to com/isl com/isl-org/Mulltae
Sener and Koltun (2018) and
Again, ULS inputs all d = 64 dimensions to the task-specific sub-networks whereas DRPM-MTL partitions the intermediate representations into nr different subsets, which are then fed to the respective task networks. Yrr the number of tasks.
Compared to the MultiMNIST experiment (see Appendix C.4.2), we introduce additional regular- ization for the DRPM-MTL method. The additional regularization based on the upper bound in Lemma 4.2 and is penalizing size of [ily for given Hence, the loss function changes to
L + X Log
(78)
nT
where
Lrog = log
II ni!
Σ log I'nn= +
(79)
For all versions of the experiment (i.e. By E 10, 20, 40), we set À = 0.015 ~ 64 which is the number of elements want partition. The task losses L, simple BCE losses similar to the MultiMNIST experiments but with two classes per task only.
We perform two different experiments based on the CelebA experiment. First, we use form MTL experiments using the first 10 attributes out the 40 attributes. Second, we increase the number of different tasks 20. Because sort the attributes alphabetically in both cases, the first 10 tasks
33

---

Table 7: Results for the MTL experiment on the CelebA dataset. We compare the DRPM-MTL again to the ULS method. We assess the performance of both methods on two cub-experiment of the CelebA experiment. In Table 7a, we form a MTI. experiment with 10 different tasks. Table 7b, we form experimet with 20 different tasks where the first 10 tasks are the same as in the 10 tasks experiment. And in Table 6c, MTL experiment with both methods for 50 epochs using learning rate of O. 0.0001 and batch size of 128. The temperature annealing schedule remains the same as in the MultiMNIST experiment. We report the per task classification accuracy in percentages (%) well the average task accuracy in the bottow row of both subtables.
tasks from the dataset. We train
(a) 10 Tasks
(b) 20 Tasks
ULS
DRPM
ULS
DRPM
TO
92.0+0.5
92.4i10.
TO
92.4
93.010.2
83.810.4
83.710.2
83.7
83.910.7
T2
80.210.5
80.210.4
T2
79.9
80.110.4
81.910.8
82.210.66
T3
82.4
83.010.7
T4
98.5+0022
98.5++0.1
T4 98.6=0.1
98.6+0.1
95.2+0.2
95.3+0.2
T5
95.2=0.1
95.5+0.0
T6
80.0+1.4
82.410.4
T6
12.0=1.3
84.410.4
T7
82.0+t.33
82.2+0.2
T7
82.520.1
82.8+0.2
T8
89.7+0.7
90.7+0.2
T8
90.1
91.010.4
T9
94.6:#0.5
95.0+0.2
T9
94.7
95.1+0.1
avg(Tasks)
87.840.3
88.310.1
T10
95.9=0.1
95.940.1
T11
84.910.1
84.620.3
T12
91.020.4
91.610.2
T13
94.710.1
94.910.1
T14
95.4=0.3
96.0+0.1
T15
99.2=0.0
99.2+0.1
T16
95.8=0.3
96.010.1
T17
97.3=0.3
97.5+0.2
TI8
91.2=0.3
91.240.1
T19
87.0
87.3+0.2
avg(Tasks)
90.7+0.2
91.110.1
shared between the two experiment versions. And third, we set ny = 40, where the first 20 tasks are the shared with the previous experiment.
Table 7 shows the results of all CelebA experiments for both methods, ULS and DRPM-MTL. We see that the DRPM-MTL scales better to a larger number of tasks compared to the ULS method, highlighting the importance of finding new ways of automatic scaling between tasks. Interestingly, the DRPM-MTL. ouperrformss the UI.S method on most tasks for the 20-tasks experiment even though has only access d/nr 64/ 64/20- = 3.2 dimensions on average. And even more extrem for nT 40, DRPM-MTL on average has only access to d/nt 64/40 1.6 dimensions. On the other hand, the ULS method can access the full set of 64 dimensions for every single task.
C.5 Supervised Learning
Given the true partition Y, we can also adapt the DRPM to learn partitions in a supervised fashion. One instance where we know the true partition elements in the case of classification. There, for given batch X containing B samples, we have Y == (yi,..., yB) where E _ is the one-hot encoding of the label of the i-th sample in the batch. In this ablation, we infer Y = DRPM (X), 80, (X)), where we compute w ((X) and s(X) as in the variational clustering experiment (Appendix C.2.2) and use the DRPM without resampling as in the multitask experiment (Section 5.3). We optimize
34

---

(c) 40 Tasks
ULS
DRPM
TO
92.9+00.6
92.6+0.5
T1
83.4
83.811.5
T2
80.640.8
80.640.3
T3
82.8+1.1
83.1+0.4
T4
98.6+0.1
98.7+0.1
T5
95.4+0.1
95.4+0.2
T6
81.6+2.2
84.4+1.2
T7
82.7
82.610.2
T8
90.110.8
90.7+0.5
T9
94.81102
95.0+0.1
T10
96.010.2
95.9+0.1
T11
85.0+0.5
85.010.3
T12
91.640.6
92.2410.1
T13
94.8+0.2
94.8+t0.2
T14
95.6+0.4
95.7
T15
99.3+0.1
99.3++0.1
T16
96.2+0.2
96.1+0.1
T17
97.410.2
97.5
T18
91.110.1
91.3+t0.3
T19
87.1+0.3
87.010.5
T20
98.610.0
98.5+0.1
T21
93.640.1
93.610.2
T22
96.010.1
95.940.2
T23
91.8+0.4
92.6+0.6
T24
95.4+0.2
95.5+0.1
T25
71.5+16
72.9+1.2
T26
96.0.1.4
96.410.2
T27
74.940.5
76.140.s
T28
93.5.10.5
94.0+0.3
T29
92.8+0.2
93.5+0.2
T30
96.310.3
96.1
T31
92.710.2
93.010.1
T32
81.0.0.4
82.1
T33
83.3+0.2
83.8+0.6
T34
89.2+0.2
89.3+0.1
T35
98.9+0.0
99.0+0.0
T36
91.710.11
91.840.2
T37
87.6:10.6
88.140.1
T38
95.5
95.640.2
T39
87.9
82.7+32
avg(Tasks)
90.640.1
90.840.1
35

---

parameters 01 and B2 by minimizing the following loss:
((X,Y)= - L)(X, + (x2((X,))
(x(,Y) 1 (
(2 = } |ln(X) K Σ yill"
where LCE denotes the standard cross-entropy loss, n denotes the output of the MVHG leading to Y, and L2 ensures that n, matches the number of appearances of label in the current batch. Using this simple training scheme, we achieve an fl-score of 96.43 + 0.02 and 82. = 0.03 on MNIST and FMNIST, respectively, further demonstrating the applicability and versatility of the DRPM to number of different problems.
36

---

PROP
3
O
4
9
4
4
3
S O
3
9
O
8
0
3
4
4
5
4
2
3
9
10
11
12
13
14
15
16
17
10
19
20
21
22
23
24
25
26
29
29
30
31
Sample
PHOOr
3
2
O
4
8
5
9
6
7
9
3
0
6 3
5
8
9
8
3
2
3
7
6
3
/
3
13
14
19
20
73
24
25
28
29
30
mor
2
6
8
8
2
O
8
2
3
0
7
4
2
6
9
4
8
3
2
6
4
9
Sample
Phor
7
O
6
7
9
4
O
9
O
8
5
2
3
O
19
25
30
Sample
Plor
4
6
/
3
4
3
O
8
7
9
8
4
O
5
7
4
3
8
9
10
13
14
15
18
19
20
22
23
24
25
26
28
29
30
31
Sample
Phlor
2
5
8
6
/
O
7
3
6
9
9
7
O
/
9
4
2
8
0
8
7
2
4
8
6
3
10
12
13
14
15
16
17
18
19
20
25
24
25
26
28
29
30
31
Sample
3
3
5
5
3
9
3
5
8
2
6
6
5
4
9
7
24
25
Figure 15: Additional partition samples from the DRPM- VC trained on MNIST. The different sets of each partition match each of the digits very well, even after repeatedly sampling from the model.
37

---

Pnior
Co
Sample
Proor
19
Sample
Fmorr
10
13
14
15
16
17
18
19
20
25
26
29
30
Sample
Piof
13
14
15
16
20
25
30
Phor
25
Sample
Phor
Sample
or
13
14
15
16
18
10
20
25
26
30
Sample
Figure 16: Additional partition samples from the DRPM-VC trained on FMNIST. Most clusters accurately represent one of the clothing categories and generate new samples very well. The only problem is with the handbag class, where the DRPM-VC learns different clusters for different kinds of handbags (cluster 5 and 6).
38

---

4
4
4
4
4
4
4
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
0
0
0
0
0
0
O
0
O
0
0
0
0
0
0
O
O
O
0
O
0
0
O
2
2
2
2
2
2
2
2-
2
2
2
2
7
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9 8
8
8
8
8
8
8
9
8
8
8
8
8
8
8
8
8
8 5
5
5
5
5
5
5
5
5
5
5
5
S 1
5
5
5
5
5
5
/
/
/
1
1
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
Samole
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4 6 0 8 9 8 5 1 7 3
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
O
0
0
0
O
O
0
0
0
0
0
0
0
0
0
0
0
0 2 9 8 S / 7 3
2 2
2
2
2
2
2
3
2
2
2
2
2
2
2
2
2
2
3
2
2 9
MO
9 9 9
9
9
9
9
9
9
9
9
9
9
9
9
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
5
5
5
5
5
5
5
5
5
5
5
5
S
5
5
5
5
S
5
5
5
5
5
5
5
5
5
/
/
/
(
I
/
/
1
I
/
/
I
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4 6 0
4 6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
0
0
O
0
0
O
O
0
0
0
0
0
0
0
0 2 9 8
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9 8
9 8 5
8 8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
5
5
5
5
S
5
5
5
5
5 / 7
S
5
5
5
5
5
5
5
5
5
5
/ 7 3
7
7
7
7
7
7
7
7
7
7
7
7
9 3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
Sample
4
4
4
4
4
4 6
4
4
4
4
4
4
4
4
4
4
4
4
4 6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0 2 9
0 2 9 8
2
2
3
2
2
2
8
2
2
2
2
2
2
2
2
2
2
2
2
2 9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8 5
5
5
6
5
5
5
5
5
5
5
S
5
5
5
S
5
5
5
/
/
/
I
1
7
7
7
7
7
7
7
3
3
3
3
3
3
3
3
3 3
3 3
3
3
3
3
3
3
3
3 3
Sample
Figure 17: Various samples from each of the generative priors. Each prior learns represent one of the digits. Further, we see a lot of variation between the different samples, suggesting that the clusters of the DRPM-VC manage capture of the diversity present in the dataset.
39

---

Sample
Sample
Sample
Sample
Figure 18: Various samples from each of the generative priors. Each prior learns to represent one of the digits. The DRPM-VC VC learns nice representations that provide coherent generations of most classes. For nieh- (cluster 4), generating new samples seems difficult due to the heterogeneity within that class.
40


Reviews:
[{"review_id": "vWhYnJrmnv", "reviewer": "Reviewer_WEvf", "summary": "This paper proposes a two-stage differentiable partitioning method where the first stage outputs number of elements in each cluster and  the second stage assigns samples by ordering them through a permutation matrix. The differentiability in their formulation is achieved via gumbel-softmax trick for distributions sampling. The paper shows the performance of their method for 1-a variational clustering task on MNIST nad FMNIST, 2-learn a partition of generative factors on mpi3d dataset, 3-neuron grouping task for multi-task learning on MultiMNIST. ", "questions": "- A quadratic relationship wrt n is mentioned in the limitations. How does your method's computational complexity compare with the baselines? Is this why you did not go beyond MNIST?\n\n- Why you experimented with fixed encoder and different initialization? (in the supplementary)\n\n", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed formulation and its usage in a differentiable setting is novel from my knowledge. Further applying this method for grouping neural networks for MTL and other tasks are novel from my knowledge. However, I am not an expert in this field, so I may not be aware of some of the previous work. \n\n- The paper is clearly written with a decent language and flow. Supplementary material is also very comprehensive and helpful for people who are not very familiar to the field (although I did not check all the details). For example, it introduces random partitioning well, carefully states the definitions and set up. Experiment set up and results are provided in detail. One suggestion is some of the material related to the method could be carried from supplementary to the main paper, such as final loss function, model plot, two stage algorithm details. \n\n- Although the applicability to a less trivial setting -- with larger datasets and harder tasks -- are not clear, the method shows promise to be improved and has sufficient contribution. Sampling from cluster numbers and permutation matrix distributions in a differentiable way to create partitioning is an interesting idea and application.", "weaknesses": "- The experiment results need more clarification for a couple of concerns I have to understand the impact and contribution better:\n\n    - Why some of the recent differentiable clustering papers are missing in terms of both citation and baselines(examples: [1,2,3])? Is it because only generative clustering is considered? If that is the reason, it needs to be mentioned somewhere in the paper, perhaps either in related work or introduction section.\n\n    - Why Table 1 especially ACC results are very different (low) than how the baselines' original papers reported (for example VADE paper)? Is the network used different, or set up? It needs to be clarified, and maybe another table with the same set up can be included.\n\n    - Why didn't you try any dataset other than MNIST for clustering tasks, for instance STL-10 or Reuters datasets at least?\n\n- If your main argument relies on contributing to the generative clustering methods, again introduction may be modified to reflect that. Otherwise, the paper needs to include other differentiable clustering/partitioning methods such as deep k-means and [1,2,3].\n\n\n \n\n\n\n\n\n1) Yang, Xu, et al. \"Deep spectral clustering using dual autoencoder network.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.\n2) Cai, Jinyu, et al. \"Efficient deep embedded subspace clustering.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n3) Sarfraz, Saquib, Vivek Sharma, and Rainer Stiefelhagen. \"Efficient parameter-free clustering using first neighbor relations.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for your detailed feedback. In the following, we would like to address your questions and concerns.\n\n>Why some of the recent differentiable clustering papers are missing in terms of both citation and baselines(examples: [1,2,3])? Is it because only generative clustering is considered? If that is the reason, it needs to be mentioned somewhere in the paper, perhaps either in related work or introduction section.\\\n>+\\\n>If your main argument relies on contributing to the generative clustering methods, [...]\n\nThank you for pointing this out. Our first experiment focuses on comparing generative clustering methods. Generative clustering approaches offer the unique advantage of enabling unsupervised conditional generation, with the tradeoff of having slightly worse clustering performance as opposed to more conventional methods. Consequently, a direct comparison with non-generative clustering methods might not be entirely fair. However, we would like to note that our DRPM offers the opportunity to develop other types of clustering algorithms, such as tree- or contrastive learning-based approaches if generative modeling is not desired. We will address this point in the camera-ready version and we will add the mentioned papers to our related work section.\n>Why Table 1 especially ACC results are very different (low) than how the baselines' original papers reported (for example VADE paper)? Is the network used different, or set up? It needs to be clarified, and maybe another table with the same set up can be included.\n\nFor the VADE experiments, we maintained consistency with the original VADE paper by using the same optimizer, architecture, and loss function. However, we would like to point out two discrepancies between our setup and what was reported in the original paper. First, the original VADE experiments did not have a test set but were trained and evaluated on the full dataset (train+test split) as can be seen from their code (https://github.com/slim1017/VaDE). Furthermore, in contrast to reporting average statistics, the original paper initialized the method ten times and then reported the performance of the model with the best objective value, as described in Section 4.3 of [6]. Both these aspects do not align with current best practices, and we hence used separate train/test splits in our experiments and decided to report average statistics, including standard deviations over several seeds. Finally, we added an ablation where we tuned the VADE baseline by employing certain optimizations that proved successful in the case of our DRPM-VC model. Using these optimizations, we were able to boost the average performance of VADE close to what was reported in the original paper, as shown in Table 5 of the appendix. \n>Why didn't you try any dataset other than MNIST for clustering tasks, for instance STL-10 or Reuters datasets at least?\n\nThe goal of our paper is to introduce the DRPM and demonstrate its versatility in a number of different scenarios. The focus of our experiments does not solely lie in raw performance on numerous datasets but also on demonstrating how the DRPM can be integrated into a variety of different scenarios as a handy tool that helps to incorporate new generative assumptions (Variational Clustering/Generative Factor Partitioning) or circumvent costly hyperparameter tuning (Multitask experiment).\nHowever, to demonstrate that we can easily extend our approach to other datasets, we conducted experiments on the STL-10 dataset following the setup of [6] for this rebuttal (see rebuttal pdf).\n>A quadratic relationship wrt n is mentioned in the limitations. How does your method's computational complexity compare with the baselines? Is this why you did not go beyond  MNIST?\n\nWe would like to note that for the experiments “Variational Partitioning of Generative Factors” and “Multitask Learning”, including DRPM does not increase memory complexity compared to baselines. Here, $n$ corresponds to the number of latent dimensions and number of neurons respectively. Since both model architectures employ fully connected layers to derive the respective intermediate representation, memory complexity is $\\Omega(n^2)$ for all baselines. In contrast, for the variational clustering experiment, additional memory complexity is added since $n$ denotes the batch size. However, $n$ is thus independent of the dataset and not a limiting factor since we can always choose it to be smaller if necessary.\n\nIn practice, we never encountered issues where we had to compromise on the value of $n$. For instance, we ran the variational clustering experiments with a batch size of $n=256$, partitioned the latent space $n=10$ on the mpi3d dataset, and added multitask experiments on CelebA with $n=64$ to the appendix (Section C.3). Note that all our experiments were conducted on an RTX2080Ti GPU. For this rebuttal, we also added an additional ablation for CelebA where we partitioned the neurons for all $40$ tasks without running into memory issues. These experiments highlight the versatility of the proposed method and that our approach is generally scalable to more difficult datasets.\n>Why you experimented with fixed encoder and different initialization? (in the supplementary)\n\nSimilar to the VADE paper, we realized that proper initialization of the cluster parameters and network weights is crucial for the success of variational clustering. We initially worked with the originally proposed autoencoder pretraining (Appendix C.2.4.) but found that incorporating more recent advances from the contrastive literature improves the clustering performance of both VADE and DRPM-VC. Further, we found that fixing early encoder layers leads to more stable training and better results (Appendix C.2.5). However, independent of the strategy used, we consistently outperform our baselines, suggesting that the DRPM may be a better prior model than i.i.d. Categorical for variational clustering.\n"}, {"review_id": "dYbRCIBrlt", "reviewer": "Reviewer_jtds", "summary": "The authors tackle the problem of partitioning a set of elements into a variable number of disjoint subsets in a differentiable manner. The method starts by first inferring the number of elements per subset. Consequently, they fill these subsets in a learned order.", "questions": "- In equation 2, what purpose does the superscript in $\\omega_k^{n_k}$ serve? It is my understanding that the weights are specific to each class, not each instance? Also, am I correct in my understanding that $P_0$ here is intractable?\n\n- I'm a bit confused as to how equation (2) induces a distribution over subset sizes (as per the paragraph title). I see that induces a distribution over the assignments to classes e.g. (2, 3, 0) would denote that we sampled 2 of class 0, 3 of class 1 and 0 of class 1. That should be enough to induce a distribution over subset sizes only if you assume somehow that each subset is comprised of items of a single class? What do we lose here, if anything, by making such an assumption?\n\n- If I'm not mistaken, Equation (3) is very similar to weighted reservoir sampling (see e.g. [1])?\n\n- It is not clear how you overcome the non-differentiability of sampling. Do you use Gumbel-Softmax when sampling both the permutation\nmatrix as well as an order?\n\n- Am I correct in my understanding that equation (5) is simply stating that the probability of subsets $\\mathbf{y}_1, \\ldots, \\mathbf{y}_n$ is simply the probability of the subsets sizes $\\mathbf{n}$ multiplied by the probability of all permutations of the elements of in each subset $\\mathbf{y}_i$?\n\n\nReferences:\n\n[1] Sang Michael Xie and Stefano Ermon. “Reparameterizable Subset Sampling via Continuous Relaxations.” International Joint Conference on Artificial Intelligence (2019).", "limitations": "Adequately addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The problem of partitioning a set of elements into variable number of disjoint sets while remaining amenable to back-prop is indeed an important and timely problem. The authors spend the time to attempt to develop an approach top the problem starting from first principles, which I appreciate.", "weaknesses": "- I found the writing to be quite a bit confusing and very notation heavy at times, with very little by way of intuition or meaning of the equations. A running example, e.g. an urn with marbles, would've greatly helped with the exposition.\n\n- After reading the paper, I'm still not really sure how the authors overcome the non-differentiability of sampling.\n\n- I'm not sure I get the point behind Lemma 4.3? It doesn't add much, and doesn't really mention how differentiability is achieved (I'm guessing through Gumbel-Softmax?)\n\n- The related work can be further improved as it's missing a lot of work regarding differentiable discrete distributions e.g. [1], [2], [3]\n\nReferences:\n\n[1] Marin Vlastelica, P., Anselm Paulus, Vít Musil, Georg Martius and Michal Rolinek. “Differentiation of Blackbox Combinatorial Solvers.” ArXiv abs/1912.02175 (2019)\n\n[2] Mathias Niepert, Pasquale Minervini and Luca Franceschi. “Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions.” Neural Information Processing Systems (2021).\n\n[3] Kareem Ahmed, Zhe Zeng, Mathias Niepert and Guy Van den Broeck. “SIMPLE: A Gradient Estimator for k-Subset Sampling.” ICLR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for your detailed feedback. In the following, we would like to address your questions and concerns.\n>I found the writing to be quite a bit confusing and very notation heavy at times, with very little by way of intuition or meaning of the equations. A running example, e.g. an urn with marbles, would've greatly helped with the exposition.\n\nWe like the proposition of including a marble example and propose integrating the following paragraph into the camera-ready version of our manuscript:\n\nWithin the scope of our work, a partition of a set of $n$ elements can be viewed as a special case of the urn model. Here, the urn contains marbles with $n$ different colors, where each color corresponds to a subset in the partition. For each color, there are $n$ marbles corresponding to the potential elements of their color/subset. To derive a partition, we sample $n$ marbles without replacement from the urn and register the order in which we draw the colors. The color of the $i$th marble then determines the subset to which element $i$ corresponds. Furthermore, we can constrain the partition to only $K$ subsets by taking an urn with only $K$ different colors.\n\nWe thank the reviewer for suggesting the idea of a running example. We would further like to note that we attempted to make the notation as concise and consistent as possible. However, it lies within the nature of the problem that a lot of notation is needed to adequately describe our work. We would be happy to take propositions regarding cleaner notation into consideration for the final manuscript. \n> After reading the paper, I'm still not really sure how the authors overcome the non-differentiability of sampling.\\\n> +\\\n> I'm not sure I get the point behind Lemma 4.3?[...] \\\n> +\\\n> It is not clear how you overcome the non-differentiability [...]\n\nWe outline the technical details to achieve differentiability in the proof of Lemma 4.3 in the appendix. In short, we apply variations of the Gumbel-Softmax trick and straight-through estimators [3] to get hard samples of the order/permutation $\\pi$ and elements per subset $\\mathbf{n}$ in the forward pass while retaining soft gradients in the backward pass. Given $\\\\nu_k:=\\\\sum\\_{i=1}^{k-1}\\\\mathbf{n}\\_i $\nand $\\\\nu\\_{k+1}$, we compute a vector $\\\\mathbf{\\\\tilde{\\\\alpha}}\\_k\\\\in[0,1]^n$ where $\\\\mathbf{\\\\tilde{\\\\alpha}}\\_k^{(i)}=1$ if $i\\in[\\nu_k,\\nu_{k+1}]$ and otherwise $\\\\mathbf{\\\\tilde{\\\\alpha}}\\_k^{(i)}=0$. We then use $\\\\mathbf{\\\\tilde{\\\\alpha}}$ to remove all entries from $\\pi$ that do not correspond to subset $k$ by multiplying it columnwise and then retrieve $\\mathbf{y}_k$ by adding up the resulting matrix row-wise, similar to the example in Figure 1.\n>The related work can be further improved [...]\n\nWe thank the reviewer for pointing out the additional references. We agree that they are related to our work and will add them to the camera-ready version of our manuscript.\n>In equation 2, what purpose does the superscript in $\\omega_k^{n_k}$ serve? It is my understanding that the weights are specific to each class, not each instance? Also, am I correct in my understanding that $P_0$ here is intractable?\n\nYes, the weights $\\omega_k^{n_k}$ are class/subset weights and not instance weights. The score $s_i$ for element $i$ could be interpreted as an instance weight. The superscript $n_k$ connects the number of drawn elements per class $n_k$ to the class importance $\\omega_k$ of every class $k$.\nRegarding the MVHG distribution, we follow the sequential procedure introduced in [4] instead of directly applying eq. (2). In eq. (2), $P_0$ is the distribution-wide normalization parameter. Using class-conditional distributions, the number of states per conditional distribution reduces to $n$, allowing tractable normalization. However, you are correct that $P_0$ is intractable if we directly apply eq. (2). For more details, we refer to [4].\n>I'm a bit confused as to how equation (2) induces a distribution over subset sizes [...]\n\nA partition model is defined as assigning every element $i$ to one of $K$ subsets. The content of each subset itself depends on the specific task we want to solve. In the proposed two-stage approach, we use the MVHG as a distribution over subset sizes and the PL distribution to model the content of each subset. Going back to the initial marble example, the MVHG models the number of marbles we draw per color, whereas the PL distribution models the order in which we draw the marbles. In this setting, a subset can be defined arbitrarily, and the MVHG still defines a distribution over subset sizes. The distribution over subset sizes does not make any assumptions about the content of each subset.\nIn our clustering experiment, the reviewer's observation is correct that we see the different classes of MNIST/FMNIST as subsets, whereas in the other experiments, subsets denote groups of shared/independent generative factors or neurons specialized to a single task.\nWe hope this addresses your question and are happy to answer any further questions regarding Equation (2).\n> If I'm not mistaken, Equation (3) is very similar to weighted reservoir sampling (see e.g. [1])?\n\nEquation (3) is indeed related to [1], as both originate from [2], where they show the relationship between Luce’s choice axiom and the Gumbel-max trick. We list [1] in our related work and reference [2] in the preliminaries. \n> Am I correct in my understanding that equation (5) is simply stating that the probability of subsets $\\mathbf{y}_1, \\ldots, \\mathbf{y}_n$ is [...]\n\nThe reviewers' observation is indeed correct and stems from our two-stage procedure. For each partition $Y$, there exists a single unique $\\mathbf{n}$ that leads to that partition. However, all permutations in $\\Pi_Y$ can also lead to $Y$, which is why we need to sum over their probabilities. For a more intuitive example and more details, we refer to Figure 1 and the proof of Proposition 4.1 in Appendix B.\n"}, {"review_id": "3Iy5XMGW8d", "reviewer": "Reviewer_6PXU", "summary": "The paper proposes an interesting approach for differentiable partitioning.", "questions": "-", "limitations": "Limitations are addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The writeup of the method is clean and clear and the method is sound.", "weaknesses": "While the experimental results demonstrate the strengths of the method, they seem to be pretty tailored for the presented method. \nIt would be great if there was a more \"conventional\" experiment where the method can be utilized; this would definitely strengthen the paper, make the approach more convincing, and I would like to offer increasing my score for an experiment that demonstrates the utility in more general settings.\n\nMaybe an experiment from https://arxiv.org/pdf/2305.16358.pdf could also work for your method. If applicable, demonstrating a more direct comparison like this would be very convincing.\n\nAs sinkhorn as well as differentiable relaxation of sorting and permutation matrices, it might be good to also include https://arxiv.org/abs/1905.11885 .", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for your detailed feedback. In the following, we would like to address your questions and concerns.\n>While the experimental results demonstrate the strengths of the method, they seem to be pretty tailored for the presented method. It would be great if there was a more \"conventional\" experiment where the method can be utilized; this would definitely strengthen the paper, make the approach more convincing, and I would like to offer increasing my score for an experiment that demonstrates the utility in more general settings. Maybe an experiment from https://arxiv.org/pdf/2305.16358.pdf could also work for your method. If applicable, demonstrating a more direct comparison like this would be very convincing.\n\nWe thank the reviewer for the suggestion. We decided to conduct the “Supervised Clustering” experiment from [5] and present some preliminary results here and in the rebuttal pdf. We train the supervised model by training a small CNN which maps each sample to a representation $\\\\mathbf{z}$. For each batch, we then use the vectors $\\mathbf{z}$ to learn a partition $\\\\tilde{M}\\_\\\\Omega$ that is trained to align with the actual cluster assignments $M_\\\\Omega$ similar to [5]. We train the network end-to-end by using the AdamW optimizer and applying the cross-entropy loss between the ground-truth $\\tilde{M}\\_\\\\Omega$ and $M\\_\\\\Omega$. In initial experiments, we achieved a precision of 96.47±0.25 on MNIST and 82.71±0.36 on Fashion MNIST without much tuning of network architecture or hyperparameters and running the experiment on 5 seeds each. Similar to [5], we also qualitatively inspected the learned representation by getting the UMAP embeddings in two dimensions. We include these plots in the rebuttal pdf and plan to include our findings in the appendix of the camera-ready version of our manuscript.\n\nWe further present additional ablations, such as for our PMF bounds, varying $K$ for clustering, FMNIST generations for $K=20$, clustering experiments on STL-10, and multitasking on CelebA for all $40$ labels in the rebuttal pdf.  We plan to include these in the appendix of the camera-ready version of our manuscript, together with the new supervised clustering ablation, which we compare to [5].\n>As sinkhorn as well as differentiable relaxation of sorting and permutation matrices, it might be good to also include https://arxiv.org/abs/1905.11885 .\n\nWe thank the reviewer for the additional reference. We will include it in the camera-ready version of our manuscript."}, {"review_id": "HRmUQXinZz", "reviewer": "Reviewer_Ynzj", "summary": "This paper proposes a probabilistic model for partitioning a non-i.i.d. set of points into a fixed number of subsets, such that the probabilistic model can be fit to the data by variational inference, using the reparameterization trick to enable computation of gradients.  The probabilistic model for partitioning is composed of two processes: one random process, using a multivariate hypergeometric distribution, for determining the sizes of the subsets, and a another random process, which assigns a random ordering to the points according to the Plackett-Luce distribution over permutations, which determines the assignment of points to subsets.  To implement the inference, an evidence lower bound for the marginal probability p(X) is derived.  Three additional methods utilizing the random partitioning model are presented.  One is a method for clustering the data, one is a method for inferring shared and unique latent factors for a pair of data points, and another is a method for partitioning final-layer neurons by task in a multi-task learning setup.", "questions": "* Rather than MVHG, wouldn't it be simpler to use a multinomial distribution instead?\n* Under what data generation conditions should the multivariate hypergeometric distribution be advantageous to the multinomial?\n* Is the success of MVHG distribution in cases such as MNIST due to the fact that the dataset is constructed to have a fixed number of exemplars per class?  Would we still see an RPM based on MVHG doing better than one based on multinomial distribution in an example where the data has been generated i.i.d.?", "limitations": "Yes. The authors mention the memory inefficiency and the possibility of improving the approximations used in the ELBO, which do seem be the most important limitations of the method.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper shows how random partitioning models have many applications in machine learning, and explains the advantages of this particular approach.  A large number of technical optimizations, such as lower bounds of the loss function, are developed, which seem essential for making the method usable.", "weaknesses": "The motivation for the multivariate hypergeometric distribution could benefit from more elaboration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for your detailed feedback. In the following, we would like to address your questions and concerns.\n>Rather than MVHG, wouldn't it be simpler to use a multinomial distribution instead? Under what data generation conditions should the multivariate hypergeometric distribution be advantageous to the multinomial?\\\n>+\\\n>Is the success of MVHG distribution in cases such as MNIST due to the fact that the dataset is constructed to have a fixed number of exemplars per class? Would we still see an RPM based on MVHG doing better than one based on multinomial distribution in an example where the data has been generated i.i.d.?\n\n\nWe chose the MVHG to model dependencies between elements, which would not be possible using the multinomial distribution. However, the multinomial distribution would also be valid, as the reviewer pointed out. One of the strengths of the proposed two-stage approach is its versatility. Any distribution over subset sizes that is differentiable could be integrated into our approach. Our paper relies on the MVHG distribution due to the mentioned advantages.\nThe advantage of using the MVHG distribution will probably vanish if the data is generated truly i.i.d. However, we can rarely guarantee truly i.i.d. data, and the MVHG distribution offers an elegant solution to this problem. In addition, the MVHG implicitly regularizes against collapsing classes, which could happen if using a sampling scheme with replacement.\n\n>The authors mention the memory inefficiency and the possibility of improving the approximations used in the ELBO, which do seem be the most important limitations of the method.\n\nFortunately, memory inefficiency has not been a problem for our experiments so far, since baseline methods often exhibit similar memory complexity. We provide more details regarding this in our answer to reviewer WEvf. However, for experimental setups where $n$ denotes, e.g., the input dimension, the memory complexity of the DRPM could become an issue and we plan to investigate how to scale our approach to such scenarios in future work.\nFurther, to investigate why the approximated ELBO often led to good results in our experiments, we conducted an ablation where we compare the PMF upper/lower bounds to the true PMF for small partitions where the PMF is tractable. We provide a figure showing this ablation in the rebuttal pdf. The results suggest that our approximations lead to a roughly constant offset to the true log PMF. By rescaling the corresponding loss terms with a hyperparameter in our ELBOs, we can thus compensate for the gap induced by our bounds.\n"}]

Author rebuttal:
We would like to thank all reviewers for providing comprehensive and valuable feedback. We particularly value the reviewers' recognition of the novelty (WEvf), importance (jtds), soundness (6PXU), and applicability (Ynzj) of our work. Your insightful comments and suggestions have been immensely beneficial in enhancing the overall quality of our paper.

Your thorough reviews have helped us identify potential areas for improvement and gain valuable perspectives that undoubtedly strengthen the contribution of our research. We have carefully addressed each of your comments and will incorporate your feedback into the camera-ready version of our paper. We believe that your inputs will contribute to the advancement of our work.

We welcome any additional suggestions, questions, or requests for information and encourage further discussions with all reviewers. We hope that our efforts in adequately addressing the reviewers' questions or concerns will lead them to consider adjusting their scores or confidences.

Once again, we thank the reviewers for their time and effort in evaluating our work.


**Rebuttal References:**

[1] Sang Michael Xie and Stefano Ermon. “Reparameterizable Subset Sampling via Continuous Relaxations.” International Joint Conference on Artificial Intelligence (2019).

[2] J. I. Yellott. The relationship between Luce’s choice axiom, Thurstone’s theory of comparative judgment, and the double exponential distribution. Journal of Mathematical Psychology, 15(2): 109–144, 1977. Publisher: Elsevier.

[3] Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.

[4] Sutter, Thomas M., et al. "Learning Group Importance using the Differentiable Hypergeometric Distribution." arXiv preprint arXiv:2203.01629 (2022).

[5] Stewart, Lawrence, et al. "Differentiable Clustering with Perturbed Spanning Forests." arXiv preprint arXiv:2305.16358 (2023).

[6] Jiang, Zhuxi, et al. "Variational deep embedding: An unsupervised and generative approach to clustering." arXiv preprint arXiv:1611.05148 (2016).

[7] Marin Vlastelica, P., Anselm Paulus, Vít Musil, Georg Martius and Michal Rolinek. “Differentiation of Blackbox Combinatorial Solvers.” ArXiv abs/1912.02175 (2019)

[8] Mathias Niepert, Pasquale Minervini and Luca Franceschi. “Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions.” Neural Information Processing Systems (2021).

[9] Kareem Ahmed, Zhe Zeng, Mathias Niepert and Guy Van den Broeck. “SIMPLE: A Gradient Estimator for k-Subset Sampling.” ICLR 2023.


Meta-review: This paper presents a novel and significant contribution, where a "differentiable" procedure of partitioning a given set of items into a variable number of blocks is presented. The presented algorithm would be useful for application in many areas of machine learning. The reviewers are generally happy with the novelty of the proposed method and the quality of the empirical validation.
Decision: Accept (poster)