Title: Differentiable Random Partition Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: random partition model

continuous relaxation

reparameterization

generative models

vae

representation learning

weak supervision

variational clustering

deep learning
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

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 endto-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 generalpurpose approach on three different challenging experiments: variational clustering, inference of shared and independent generative factors under weak supervision, and multitask learning.

1 Introduction

Partitioning a set of elements into subsets is a classical mathematical problem that attracted much interest over the last few decades (Rota, 1964; Graham et al., 1989). A partition over a given set is a collection of non-overlapping subsets such that their union results in the original set. In machine learning (ML), partitioning a set of elements into different subsets is essential for many applications, such as clustering (Bishop and Svensen, 2004) or classification (De la Cruz-Mesía et al., 2007).

Random partition models (RPM, Hartigan, 1990) define a 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 a reparameterization scheme, prohibiting their direct use in gradient-based optimization frameworks.

In this work, we propose the differentiable random partition model (DRPM), a fully-differentiable relaxation for RPMs that allows reparametrizable sampling. The DRPM follows a two-stage procedure: 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, we apply the DRPM to a 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 a dataset, DRPM-based clustering overcomes the simplified i.i.d. assumption of previous works, which used categorical priors (Jiang et al., 2016). In our second experiment, we demonstrate how to retrieve sets of shared and independent generative factors of paired images using the proposed DRPM. In

Figure 1: Illustration of the proposed DRPM method. We first sample a permutation matrix π and a set of subset sizes n separately in two stages. We then use n and π to generate the assignment matrix Y , the matrix representation of a partition ρ .

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 a building block in a 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 a given dataset into different subsets, the clusters. In contrast to many existing approaches (Yang et al., 2019; Sarfraz et al., 2019; Cai et al., 2022), we consider cluster assignments as random variables, allowing us to treat clustering from a variational perspective. Previous works in variational clustering (Jiang et 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 et al., 2000). However, EMD aims to assign a set's elements to different subsets based on a 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).

Differentiable and Reparameterizable Discrete Distributions Following the proposition of the Gumbel-Softmax 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 gradientbased optimization pipelines. Based on the same trick, Sutter et al. (2023) propose a differentiable formulation for the multivariate hypergeometric distribution. Multiple works on differentiable sorting procedures and permutation matrices have been proposed, e.g., Linderman et al. (2018); Prillo and Eisenschlos (2020); Petersen et al. (2021). Further, Grover et al. (2019) described the distribution over permutation matrices p ( π ) for a permutation matrix π using the Plackett-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 topk elements selection procedure (Xie and Ermon, 2019), blackbox combinatorial solvers (Poganˇ ci´ c et al., 2019), implicit likelihood estimations (Niepert et al., 2021), and k -subset sampling (Ahmed et al., 2022).

3 Preliminaries

Set Partitions A partiton ρ = ( S 1 , . . . , S K ) of a set [ n ] = { 1 , . . . , n } with n elements is a collection of K subsets S k ⊆ [ n ] where K is a priori unknown (Mansour and Schork, 2016). For a partition ρ to be valid, it must hold that

̸

In other words, every element i ∈ [ n ] has to be assigned to precisely one subset S k . We denote the size of the k -th subset S k as n k = |S k | . Alternatively, we can describe a partition ρ through an assignment matrix Y = [ y 1 , . . . , y K ] T ∈ { 0 , 1 } K × n . Every row y k ∈ { 0 , 1 } 1 × n is a multi-hot vector, where y ki = 1 assigns element i to subset S k .

Within the scope of our work, we view a partition of a set of n elements 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.

Probability distribution over subset sizes The multivariate non-central hypergeometric distribution (MVHG) describes sampling without replacement and allows to skew the importance of groups with an additional importance parameter ω (Fisher, 1935; Wallenius, 1963; Chesson, 1976). The MVHG is an urn model and is described by the number of different groups K ∈ N , the number of elements in the urn of every group m = [ m 1 , . . . , m K ] ∈ N K , the total number of elements in the urn ∑ K k =1 m k ∈ N , the number of samples to draw from the urn n ∈ N 0 , and the importance factor for every group ω = [ ω 1 , . . . , ω K ] ∈ R K 0+ (Johnson, 1987). Then, the probability of sampling n = { n 1 , . . . , n K } , where n k describes the number of elements drawn from group K is

where P 0 is a normalization constant. Hence, the MVHG p ( n ; ω , m ) allows us to model dependencies between different elements of a 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 ∀ m k ∈ m : m k = n . We thus use the shorthand p ( n ; ω ) to denote the density of the MVHG. We refer to Appendix A.1 for more details.

Probability distribution over Permutation Matrices Let p ( π ) denote a distribution over permutation matrices π ∈ { 0 , 1 } n × n . 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 π ij = 1 means that element j is ranked at position i in the imposed order. In this work, we assume p ( π ) to be parameterized by scores s ∈ R n + , where each score s i corresponds to an element i . The order given by sorting s in decreasing order corresponds to the most likely permutation in p ( π ; s ) . Sampling from p ( π ; s ) can be achieved by resampling the scores as ˜ s i = β log s i + g i where g i ∼ Gumbel (0 , β ) for fixed scale β , and sorting them in decreasing order. Hence, resampling scores s enables the resampling of permutation matrices π . The probability over orderings p ( π ; s ) is then given by (Thurstone, 1927; Luce, 1959; Plackett, 1975; Yellott, 1977)

where π is a permutation matrix and Z = ∑ n i =1 s i . The resulting distribution is a Plackett-Luce (PL) distribution (Luce, 1959; Plackett, 1975) if and only if the scores s 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

Wepropose the DRPM p ( Y ; ω , s ) , a differentiable and reparameterizable two-stage Random Partition Model (RPM). The proposed formulation separately infers the number of elements i per subset

n ∈ N K 0 , where ∑ K k =1 n k = n , and the assignment of elements to subsets S k by inducing an order on the n elements and filling S 1 , ..., S K sequentially in this order. To model the order of the elements, we use a permutation matrix π = [ π 1 , . . . , π n ] T ∈ { 0 , 1 } n × n , 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 n k -hot vector y k and therefore serves as a subset selection vector, i.e.

such that Y = [ y 1 , . . . , y K ] T . 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 S k to be the empty set ∅ (Mansour and Schork, 2016). We base the following Proposition 4.1 on the MVHG distribution p ( n ; ω ) for the subset sizes n and the PL distribution p ( π ; s ) for assigning the elements to subsets. However, the proposed two-stage approach to RPMs is not restricted to these two classes of probability distributions.

Proposition 4.1 (Two-stage Random Partition Model) . Given a probability distribution over subset sizes p ( n ; ω ) with n ∈ N K 0 and distribution parameters ω ∈ R K + and a PL probability distribution over random orderings p ( π ; s ) with π ∈ { 0 , 1 } n × n and distribution parameters s ∈ R n + , the probability mass function p ( Y ; ω , s ) of the two-stage RPM is given by

where Π Y = { π : y k = ∑ ν k + n k i = ν k +1 π i , k = 1 , . . . , K } , and y k and ν k as 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 ; ω , s ) as a probability of subsets p ( y 1 , . . . , y K ; ω , s ) , which we compute sequentially over subsets, i.e.

where y <k = [ y 1 , . . . , y k -1 ] and

where Π y k in Equation (7) is the set of all subset permutations of elements i ∈ S k . A subset permutation matrix ¯ π represents an ordering over only n k out of the total n elements. The probability p ( y k | y <k ; ω , s ) describes the probability of a subset of a given size n k by marginalizing over the probabilities of all subset permutations p (¯ π | n k , y <k ; s ) . Hence, the sum over all p (¯ π | n k , y <k ; s ) makes p ( y k | y <k ; ω , s ) invariant to the ordering of elements i ∈ S k (Xie and Ermon, 2019). Note that in a slight abuse of notation, we use p (¯ π | n k , y <k ; ω , s ) as the probability of a subset permutation ¯ π given that there are n k elements in S k and thus ¯ π ∈ { 0 , 1 } n k × n .

The probability of a subset permutation matrix p (¯ π | n k , y <k ; s ) describes the probability of drawing the elements i ∈ S k in the order defined by the subset permutation matrix ¯ π given that the elements in S <k are already determined. Hence, we condition on the subsets y <k . This property follows from Luce's choice axiom (LCA, Luce, 1959). Additionally, we condition on n k , the size of the subset S k . The probability of a subset permutation is given by

In contrast to the distribution over permutations matrices p ( π ; s ) in Equation (3), we compute the product over n k terms and have a different normalization constant Z k , which is the sum over the scores s i of all elements i ∈ S k . Although we induce an ordering over all elements i by using a permutation matrix π , the probability p ( y k | y <k ; ω , s ) is invariant to intra-subset orderings of elements i ∈ S k . Finally, we arrive at Equation (5) by substituting Equation (7) into Equation (6),

and applying the definition of the conditional probability p ( n ; ω ) = ∏ K k =1 p ( n k | n <k ; ω ) and by reshuffling indices ∑ π ∈ Π Y p ( π ; s ) = ∏ K k =1 ∑ ¯ π ∈ Π y k p (¯ π | n k , y <k ; s ) .

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 s ∈ R n + and the group importance parameters ω ∈ R K + .

Finally, to sample from the two-stage RPM of Proposition 4.1 we apply the following procedure: First sample π ∼ p ( π ; s ) and n ∼ p ( n ; ω ) . From π and n , compute partition Y by summing the rows of π according to n as described in Equation (4) and illustrated in Figure 1.

4.1 Approximating the Probability Mass Function

The number of permutations per subset | Π y k | scales factorially with the subset size n k , i.e. | Π y k | = n k ! . Consequently, the number of valid permutation matrices | Π Y | is given as a function of n , i.e.

Although Proposition 4.1 describes a well-defined distribution for p ( Y ; ω , s ) , it is in general computationally intractable due to Equation (9). In practice, we thus approximate p ( Y ; ω , s ) using the following Lemma.

Lemma 4.2. p ( Y ; ω , s ) can be upper and lower bounded as follows

We provide the proof in Appendix B. Note that from Equation (3) we see that max ˜ π p (˜ π ; s ) = p ( π s ; s ) , where π s is the permutation that results from sorting the unperturbed scores s .

4.2 The Differentiable Random Partition Model

To incorporate our two-stage RPM into gradient-based optimization frameworks, we require that efficient computation of gradients is 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-toend fashion:

Lemma 4.3 (DRPM) . A two-stage RPM is differentiable and reparameterizable if the distribution over subset sizes p ( n ; ω ) and the distribution over orderings p ( π ; s ) are differentiable and reparameterizable.

We provide the proof in Appendix B. Note that Lemma 4.3 enables us to learn variational posterior approximations and priors using Stochastic Gradient Variational Bayes (SGVB, Kingma and Welling, 2014). In our experiments, we apply Lemma 4.3 using the recently proposed differentiable formulations of the MVHG (Sutter et al., 2023) and the PL distribution (Grover et al., 2019), though other choices would also be valid.

5 Experiments

We demonstrate the versatility and effectiveness of the proposed DRPM in three different experiments. First, we propose a 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 2 .

0123456789

Cluster

0

1

2

3

4

5

7

8

6

10

11

12

13

1415

16

17

18

19

20

21222324252627

28

29

30

31

GeneratedSample

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 x of a dataset X is generated by a latent vector z ∈ R l , where l ∈ N is the latent space size. Traditional V AEs would then assume that all latent vectors z are generated by a single Gaussian prior distribution N ( 0 , I l ) . Instead, we assume every z to be sampled from one of K different latent Gaussian distributions N ( µ k , diag( σ k )) , k = 1 , . . . , K , with µ k , σ k ∈ R l . Further, note that similar to an urn model (Section 3), if we draw a batch from a given finite dataset with samples from different clusters, the cluster assignments within that batch are not entirely independent. Since there is only a finite number of samples per cluster, drawing a sample from a specific cluster decreases the chance of drawing a sample from that cluster again, and the distribution of the number of samples drawn per cluster will follow an MVHG distribution. Previous work on variational clustering proposes to model the cluster assignment y ∈ { 0 , 1 } K of each sample x through independent categorical distributions (Jiang et al., 2016), which might thus be over-restrictive and not correctly reflect reality. Instead, we propose explicitly modeling the dependency between the y of different samples by assuming they are drawn from an RPM. Hence, the generative process leading to X can be summarized as follows: First, the cluster assignments are represented as a partition matrix Y and sampled from our DRPM, i.e., Y ∼ p ( Y ; ω , s ) . Given an assignment y from Y , we can sample the respective latent variable z , where z ∼ N ( µ y , diag( σ y )) , z ∈ R l . Note that we use the notational shorthand µ y := µ arg max( y ) . Like in vanilla VAEs, we infer x by independently passing the corresponding z through a decoder model. Assuming this generative process, we derive the following evidence lower bound (ELBO) for p ( X ) :

Note that computing KL [ q ( Y | X ) || p ( Y )] directly is computationally intractable, and we need to upper bound it according to Lemma 4.2. For an illustration of the generative assumptions and more details on the ELBO, we refer to Appendix C.2.

To assess the clustering performance, we train our model on two different datasets, namely MNIST (LeCun et al., 1998) and Fashion-MNIST (FMNIST, Xiao et al., 2017), and compare it to three baselines. Two of the baselines are based on a Gaussian Mixture Model, where one is directly trained on the original data space (GMM), whereas the other takes the embeddings from a pretrained encoder as input (Latent GMM). The third baseline is Variational Deep Embedding (VaDE, Jiang et al., 2016), which is similar to the DRPM-VC but assumes i.i.d. categorical cluster assignments. For all methods except GMM, we use the weights of a pretrained encoder to initialize the models and priors at the start of training. We present the results of these experiments in Table 1. As can be seen, we outperform all baselines, indicating that modeling the inherent dependencies implied by finite datasets benefits the performance of variational clustering. While achieving decent clustering performance, another benefit of variational clustering methods is that their reconstruction-based nature intrinsically allows unsupervised conditional generation. In Figure 2, we present the result of sampling a partition and the corresponding generations from the respective clusters after training the DRPM-VC on FMNIST. The model produces coherent generations despite not having access to labels, allowing us to investigate the structures learned by the model more closely. We refer to Appendix C.2 for more illustrations of the learned clusters, details on the training procedure, and ablation studies.

5.2 Variational Partitioning of Generative Factors

Data modalities not collected as i.i.d. samples, such as consecutive frames in a video, provide a weak-supervision signal for generative models and representation learning (Sutter et al., 2023). Here, on top of learning meaningful representations of the data samples, we are also interested in discovering the relationship between coupled samples. If we assume that the data is generated from underlying generative factors, weak supervision comes from the fact that we know that certain factors are shared between coupled pairs while others are independent. The supervision is weak because we neither know the underlying generative factors nor the number of shared and independent factors. In such a setting, we can use the DRPM to learn a partition of the generative factors and assign them to be either shared or independent.

In this experiment, we use paired frames X = [ x 1 , x 2 ] from the mpi3d dataset (Gondal et al., 2019). Every pair of frames shares a subset of its seven generative factors. We introduce the DRPM-VAE, which models the division of the latent space into shared and independent latent factors as RPM. We add a posterior approximation q ( Y | X ) and additionally a prior distribution of the form p ( Y ) . The model maximizes the following ELBO on the

marginal log-likelihood of images through a VAE (Kingma and Welling, 2014):

Similar to the ELBO for variational clustering in Section 5.1, computing KL [ q ( Y | X ) || p ( Y )] directly is intractable, and we need to bound it according to Lemma 4.2.

We compare the proposed DRPM-VAE to three methods, which only differ in how they infer shared and latent dimensions. While the Label-VAE (Bouchacourt et al., 2018; Hosoya, 2018) assumes that the number of independent factors is known, the Ada-VAE (Locatello et al., 2020) relies on a heuristic-based approach to infer shared and independent latent factors. Like in Locatello et al. (2020) and Sutter et al. (2023), we assume a single known factor for Label-V AE in all experiments. HG-V AE (Sutter et al., 2023) also relies on the MVHG to model the number of shared and independent factors. Unlike the proposed DRPM-VAE approach, HG-VAE must rely on a heuristic to assign latent dimensions to shared factors, as the MVHG only allows to model the number of shared and independent factors but not their position in the latent vector. We use the code from Locatello et al. (2020) and follow the evaluation in Sutter et al. (2023). We refer to Appendix C.3 for details on the ELBO, the setup of the experiment, the implementation, and an illustration of the generative assumptions.

We evaluate all methods according to their ability to estimate the number of shared generative factors (Figure 3) and how well they partition the latent representations into shared and independent factors (Table 2). Because we have access to the data-generating process, we can control the number of shared n s and independent n i factors. We compare the methods on four different datasets with n s ∈ { 0 , 1 , 3 , 5 } . In Figure 3, we demonstrate that the DRPM-VAE accurately estimates the true number of shared generative factors. It matches the performance of HG-VAE and outperforms the other two baselines, which consistently overestimate the true number of shared factors. In Table 2, we see a considerable performance improvement compared to previous work when assessing the learned latent representations. We attribute this to our ability to not only estimate the subset sizes of latent and shared factors like HG-VAE but also learn to assign specific latent dimensions to the corresponding shared or independent representations. Thus, the DRPM-VAE dynamically learns more meaningful representations and can better separate and infer the shared and independent subspaces for all dataset versions.

The DRPM-VAE provides empirical evidence of how RPMs can leverage weak supervision signals by learning to maximize the data likelihood while also inferring representations that capture the relationship between coupled data samples. Additionally, we can explicitly model the data-generating process in a theoretically grounded fashion instead of relying on heuristics.

5.3 Multitask Learning

Many ML applications aim to solve specific tasks, where we optimize for a single objective while ignoring potentially helpful information from related tasks. Multitask learning (MTL) aims to improve the generalization across all tasks, including the original one, by sharing representations between related tasks (Caruana, 1993; Caruana and de Sa, 1996) Recent works (Kurin et al., 2022; Xin et al., 2022) show that it is difficult to outperform a convex combination of task losses if the task losses are appropriately scaled. I.e., in case of equal difficulty of the two tasks, a classifier with equal weighting of the two classification losses serves as an upper bound in terms of performance. However, finding suitable task weights is a tedious and inefficient approach to MTL. A more automated way of weighting multiple tasks would thus be vastly appreciated.

In this experiment, we demonstrate how the DRPM can learn task difficulty by partitioning a network layer. Intuitively, a task that requires many neurons is more complex than a task that can be solved using a single neuron. Based on this observation, we propose the DRPM-MTL. The DRPM-MTL learns to partition the neurons of the last shared layer such that only a subset of the neurons are used for every task. In contrast to the other experiments (Sections 5.1 and 5.2), we use the DRPM without resampling and infer the partition Y as a deterministic function. This can be done by applying the two-step procedure of Proposition 4.1 but skipping the resampling step of the MVHG and PL distributions. We compare the DRPM-MTL to the unitary loss scaling method (ULS, Kurin et al.,

2022), which has a fixed architecture and scales task losses equally. Both DRPM-MTL and ULS use a network with shared architecture up to some layer, after which the network branches into two task-specific layers that perform the classifications. Note the difference between the methods. While the task-specific branches of the ULS method access all neurons of the last shared layer, the task-specific branches of the DRPM-MTL access only the subset of neurons reserved for the respective task.

We perform experiments on MultiMNIST (Sabour et al., 2017), which overlaps two MNIST digits in one image, and we want to classify both numbers from a single sample. Hence, the two tasks, classification of the left and the right digit (see Appendix C.4 for an example), are approximately equal in difficulty by default. To increase the difficulty of one of the two tasks, we introduce the noisyMultiMNIST dataset. There, we control task difficulty by adding salt and pepper noise to one of the two digits, subsequently increasing the difficulty of that task with increasing noise ratios. Varying the noise, we evaluate how our DRPM-MTL adapts to imbalanced difficulties, where one usually has to tediously search for optimal loss weights to reach good performance. We base our pipeline on (Sener and Koltun, 2018). For more details and additional CelebA MTL experiments we refer to Appendix C.4.

We evaluate the DRPM-MTL concerning its classification accuracy on the two tasks and compare the inferred subset sizes per task for different noise ratios α ∈ { 0 . 0 , . . . , 0 . 9 } of the noisyMultiMNIST dataset (see Figure 4). The DRPM-MTL achieves the same or better accuracy on both tasks for most noise levels (upper part of Figure 4). It is interesting to see that, the more we increase α , the more the DRPM-MTL

1.0

......

0.8

0.6

ULS Left

0.4+

ULS Right

50

DRPM Left (Ours)

DRPM Right (Ours)

Sizes

40

30

Partition

20

10

0+

0.0

0.2

0.4

0.6

0.8

tries to overcome the increased difficulty of the right task by assigning more dimensions to it (lower part of Figure 4, noise ratio α 0 . 6 -0 . 8 ). Note that for the maximum noise ratio of α = 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 on minimizing the earth mover's distance (Monge, 1781) to learn subset assignments could be an alternative to the ordering-based approach in our DRPM and pose an interesting direction for future work. Finally, note that we compute the probability mass function (PMF) p ( Y ; ω , s ) by approximating it with the bounds in Lemma 4.2. While the upper bound is tight when all scores have similar magnitude, the bound loosens if scores differ a lot, leading Equation (10) to overestimate the value of the PMF. In practice, we thus reweight the respective terms in the loss function, but in the future, we will investigate better estimates for the PMF.

Ultimately, we are interested in exploring how to apply the DRPM to multimodal learning under weak supervision, for instance, in medical applications. Section 5.2 demonstrated the potential of learning from coupled samples, but further research is needed to ensure fairness concerning underlying, hidden attributes when working with sensitive data.

Conclusion

In this work, we proposed the differentiable random partition model, a novel approach to random partition models. Our two-stage method enables learning partitions end-to-end by separately controlling subset sizes and how elements are assigned to subsets. This new approach to partition learning enables the integration of random partition models into probabilistic and deterministic gradient-based optimization frameworks. We show the versatility of the proposed differentiable random partition model by applying it to three vastly different experiments. We demonstrate how learning partitions enables us to explore the modes of the data distribution, infer shared and independent generative factors from coupled samples, and learn task-specific sub-networks in applications where we want to solve multiple tasks on a single data point.

References

Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. J. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Józefowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. G. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, I. Sutskever, K. Talwar, P. A. Tucker, V . Vanhoucke, V . Vasudevan, F. B. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and X. Zheng. TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems. CoRR , abs/1603.0, 2016.

P. Adams and R. S. Zemel. Ranking via sinkhorn propagation. arXiv preprint arXiv:1106.1925 , 2011.

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.

Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

M. Bishop and M. Svensen. Robust Bayesian Mixture Modelling. To appear in the Proceedings of ESANN , page 1, 2004. Publisher: Citeseer.

Bouchacourt, R. Tomioka, and S. Nowozin. Multi-level variational autoencoder: Learning disentangled representations from grouped observations. In Thirty-Second AAAI Conference on Artificial Intelligence , 2018.

Cai, J. Fan, W. Guo, S. Wang, Y . Zhang, and Z. Zhang. Efficient deep embedded subspace clustering. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1-10, 2022.

Caruana. Multitask learning: A knowledge-based source of inductive bias. In Machine learning: Proceedings of the tenth international conference , pages 41-48, 1993.

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.

Coates, A. Ng, and H. Lee. An Analysis of Single-Layer Networks in Unsupervised Feature Learning. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics , pages 215-223. JMLR Workshop and Conference Proceedings, June 2011. ISSN: 1938-7228.

Cuturi, O. Teboul, and J.-P. Vert. Differentiable Ranks and Sorting using Optimal Transport, Nov. 2019. arXiv:1905.11885 [cs, stat].

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-9254. doi: 10.1111/j.1467-9876.2007.00569.x.

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. Ieee, 2009.

Dilokthanakul, P. A. M. Mediano, M. Garnelo, M. C. H. Lee, H. Salimbeni, K. Arulkumaran, and M. Shanahan. Deep unsupervised clustering with Gaussian mixture variational autoencoders. arXiv preprint arXiv:1611.02648 , 2016.

A. Fisher. The logic of inductive inference. Journal of the royal statistical society , 98(1):39-82, 1935. Publisher: JSTOR.

Fog. Calculation methods for Wallenius' noncentral hypergeometric distribution. Communications in Statistics-Simulation and Computation® , 37(2):258-273, 2008. Publisher: Taylor & Francis.

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: a New Disentanglement Dataset. In Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

L. Graham, D. E. Knuth, O. Patashnik, S. L.-C. i. Physics, and undefined 1989. Concrete mathematics: a foundation for computer science. aip.scitation.org , 3(5):165, 1989. doi: 10.1063/1. 4822863. Publisher: AIP Publishing.

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.1080/03610929008830345. Publisher: Taylor & Francis.

He, X. Zhang, S. Ren, and J. Sun. Deep Residual Learning for Image Recognition, Dec. 2015.

Higgins, L. Matthey, A. Pal, C. Burgess, and X. Glorot. beta-vae: Learning basic visual concepts with a constrained variational framework. 2016.

Hosoya. A simple probabilistic deep generative model for learning generalizable disentangled representations from grouped data. CoRR , abs/1809.0, 2018.

Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144 , 2016.

Jiang, Y. Zheng, H. Tan, B. Tang, and H. Zhou. Variational deep embedding: An unsupervised and generative approach to clustering. arXiv preprint arXiv:1611.05148 , 2016.

E. Johnson. Multivariate statistical simulation: A guide to selecting and generating continuous multivariate distributions , volume 192. John Wiley & Sons, 1987.

P. Kingma and M. Welling. Auto-Encoding Variational Bayes. In International Conference on Learning Representations , 2014.

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.04122, 2022.

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.

J. Lee and H. Sang. Why the Rich Get Richer? On the Balancedness of Random Partition Models. arXiv preprint arXiv:2201.12697 , 2022.

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-2221, Sept. 2022. ISSN 0920-5691, 1573-1405. doi: 10.1007/s11263-022-01639-z. arXiv:2210.11680 [cs].

W. Linderman, G. E. Mena, H. J. Cooper, L. Paninski, and J. P. 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.

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, B. Schölkopf, O. Bachem, and M. Tschannen. Weakly-supervised disentanglement without compromises. In International Conference on Machine Learning , pages 6348-6359. PMLR, 2020.

Loshchilov and F. Hutter. Decoupled Weight Decay Regularization, Jan. 2019. arXiv:1711.05101 [cs, math].

D. Luce. Individual choice behavior: A theoretical analysis . Courier Corporation, 1959.

MacQueen. Classification and analysis of multivariate observations. In 5th Berkeley Symp. Math. Statist. Probability , pages 281-297, 1967.

Maddison, A. Mnih, and Y. Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations , 2017.

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.

Mansour and M. Schork. Commutation relations, normal ordering, and Stirling numbers . CRC Press Boca Raton, 2016.

Marcus. 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.

E. Mena, D. Belanger, S. W. Linderman, and J. Snoek. Learning Latent Permutations with GumbelSinkhorn Networks. In 6th International Conference on Learning Representations, {ICLR} 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018.

Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci. , pages 666-704, 1781.

Niepert, P. Minervini, and L. Franceschi. Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions. In Advances in Neural Information Processing Systems , volume 34, pages 14567-14579. Curran Associates, Inc., 2021.

Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Z. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. CoRR , abs/1912.0, 2019.

Petersen, C. Borgelt, H. Kuehne, and O. Deussen. Monotonic Differentiable Sorting Networks. In International Conference on Learning Representations , 2021.

Pitman. Some Developments of the Blackwell-Macqueen URN Scheme. Lecture Notes-Monograph Series , 30:245-267, 1996. ISSN 07492170. Publisher: Institute of Mathematical Statistics.

L. Plackett. The analysis of permutations. Journal of the Royal Statistical Society: Series C (Applied Statistics) , 24(2):193-202, 1975. Publisher: Wiley Online Library.

V. Poganˇ ci´ c, 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/2312585. Publisher: JSTOR.

Rubner, C. Tomasi, and L. J. Guibas. The earth mover's distance as a metric for image retrieval. International journal of computer vision , 40(2):99, 2000. Publisher: Springer Nature BV.

Sabour, N. Frosst, and G. E. Hinton. Dynamic Routing Between Capsules. In I. Guyon, U. V. 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.

Santa Cruz, B. Fernando, A. Cherian, and S. Gould. Deeppermnet: Visual permutation learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 3949-3957, 2017.

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 8934-8943, 2019.

Sener and V. Koltun. Multi-task learning as multi-objective optimization. Advances in neural information processing systems , 31, 2018.

Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matrices. The annals of mathematical statistics , 35(2):876-879, 1964. Publisher: JSTOR.

M. Sutter, L. Manduchi, A. Ryser, and J. E. Vogt. Learning Group Importance using the Differentiable Hypergeometric Distribution. In The Eleventh International Conference on Learning Representations , 2023.

L. Thurstone. A law of comparative judgment. In Scaling , pages 81-92. Routledge, 1927.

T. Wallenius. Biased sampling; the noncentral hypergeometric probability distribution. Technical report, Stanford Univ Ca Applied Mathematics And Statistics Labs, 1963.

Xiao, K. Rasul, and R. Vollgraf. Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms, Sept. 2017. arXiv:1708.07747 [cs, stat].

M. Xie and S. Ermon. Reparameterizable subset sampling via continuous relaxations. In Proceedings of the 28th International Joint Conference on Artificial Intelligence , IJCAI'19, pages 3919-3925, Macao, China, Aug. 2019. AAAI Press. ISBN 978-0-9992411-4-1.

Xin, B. Ghorbani, A. Garg, O. Firat, and J. Gilmer. Do Current Multi-Task Optimization Methods in Deep Learning Even Help? arXiv preprint arXiv:2209.11379 , 2022.

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.

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.

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 ∈ N be the number of different classes or groups (e.g. marble colors in the urn), m = [ m 1 , . . . , m K ] ∈ N K describe the number of elements per class (e.g. marbles per color), N = ∑ K k =1 m K be the total number of elements (e.g. all marbles in the urn) and n ∈ { 0 , . . . , N } be the number of elements (e.g. marbles) to draw. Then, the multivariate hypergeometric distribution describes the probability of drawing n = [ n 1 , . . . , n K ] ∈ N K marbles by sampling without replacement such that ∑ K k =1 n k = n , where n k is the number of drawn marbles of class k .

In the literature, two different versions of the noncentral hypergeometric distribution exist, Fisher's (Fisher, 1935) and Wallenius' (Wallenius, 1963; Chesson, 1976) distribution. Sutter et al. (2023) restrict themselves to Fisher's noncentral hypergeometric distribution due to limitations of the latter (Fog, 2008). Hence, we will also talk solely about Fisher's noncentral hypergeometric distribution.

Definition A.1 (Multivariate Fisher's Noncentral Hypergeometric Distribution (Fisher, 1935)) . A random vector X follows Fisher's noncentral multivariate distribution, if its joint probability mass function is given by

The support S of the PMF is given by S = { n ∈ N K : ∀ k n k ≤ m k , ∑ K k =1 n k = n } and ( n k ) = n ! k !( n -k )! .

The class importance ω is a crucial modeling parameter in applying the noncentral hypergeometric distribution (see (Chesson, 1976)).

A.1.1 Differentiable MVHG

Their reparameterizable sampling for the differentiable MVHG consists of three parts:

Reformulate the multivariate distribution as a sequence of interdependent and conditional univariate hypergeometric distributions.

Calculate the probability mass function of the respective univariate distributions.

Sample from the conditional distributions utilizing the Gumbel-Softmax trick.

Following the chain rule of probability, the MVHG distribution allows for sequential sampling over classes k . Every step includes a merging operation, which leads to biased samples compared to groundtruth non-differentiable sampling with equal class weights ω . Given that we intend to use the differentiable MVHG in settings where we want to learn the unknown class weights, we do not expect a negative effect from this sampling procedure. For details on how to merge the MVHG into a sequence of unimodal distributions, we refer to Sutter et al. (2023).

The probability mass function calculation is based on unnormalized log-weights, which are interpreted as unnormalized log-weights of a categorical distribution. The interpretation of the class-conditional unimodal hypergeometric distributions as categorical distributions allows applying the GumbelSoftmax trick (Jang et al., 2016; Maddison et al., 2017). Following the use of the Gumbel-Softmax trick, the class-conditional version of the hypergeometric distribution is differentiable and reparameterizable. Hence, the MVHG has been made differentiable and reparameterizable as well. Again, for details we refer to the original paper (Sutter et al., 2023).

A.2 Distribution over Random Orderings

Yellott (1977) show that the distribution over permutation matrices p ( π ; s ) follows a PlackettLuce (PL) distribution (Plackett, 1975; Luce, 1959), if and only of the perturbed scores ˜ s are sampled independently from Gumbel distributions with identical scales. For each item i , sample g i ∼ Gumbel (0 , β ) independently with zero mean and and fixed scale β . Let ˜ s be the vector of Gumbel perturbed log-weights such that ˜ s i = β log s i + g i . Hence,

We refer to Yellott (1977) or Grover et al. (2019) for the proof. However, Grover et al. (2019) provide only an adapted proof sketch from Yellott (1977). The probability of sampling element i first is given by its score s i divided by the sum of all weights in the set

For z i = log s i , the right hand side of Equation (15) is equal to the softmax distribution softmax ( z i ) = exp( z i ) / ∑ j exp( z j ) as already described in (Xie and Ermon, 2019). Hence, Equation (15) directly leads to the Gumbel-Softmax trick (Jang et al., 2016; Maddison et al., 2017).

A.2.1 Differentiable Sorting

In the main text of the paper we rely on a differentiable function f π ( ˜ s ) , which sorts the resampled version of the scores s

Here, we summarise the findings from Grover et al. (2019) on how to construct such a differentiable sorting operator. As already mentioned in Section 2, there are multiple works on the topic (Prillo and Eisenschlos, 2020; Petersen et al., 2021; Mena et al., 2018), but we restrict ourselves to the work of Grover et al. (2019) as we see the differentiable generation of permutation matrices as a tool in our pipeline.

Corollary A.2 (Permutation Matrix (Grover et al., 2019)) . Let s = [ s 1 , . . . , s n ] T be a real-valued vector of length n . Let A s denote the matrix of absolute pairwise differences of the elements of s such that A s [ i, j ] = | s i -s j | . The permutation matrix π corresponding to sort ( s ) is given by:

where ✶ denotes the column vector of all ones. As we know, the arg max operator is non-differentiable 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 a continuous relaxation π ( τ ) similar to the GS trick (Jang et al., 2016; Maddison et al., 2017). In particular, the i th row of π ( τ ) is given by:

π ( τ )[ i, :] = softmax[( n +1 -2 i ) s -A s ✶ /τ ] (18) where τ > 0 is a temperature parameter. We adapted this section from Grover et al. (2019) and we also refer to their original work for more details on how to generate differentiable permutation matrices.

In this, work we remove the temperature parameter τ to reduce clutter in the notation. Hence, we only write π instead of π ( τ ) , although it is still needed for the generation of the matrix π . For details on how we 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 [ n ] = { 1 , . . . , n } into K subsets {S 1 , . . . , S K } where K is a priori unknown.

Definition B.1 (Partition) . A partition ρ of a set of elements [ n ] = { 1 , . . . , n } is a collection of subsets ( S 1 , . . . , S K ) such that

̸

Put differently, every element i has to be assigned to precisely one subset S k . We denote the size of the k -th subset S k as n k = |S k | . Alternatively, we describe a partition ρ as an assignment matrix Y = [ y 1 , . . . , y K ] T ∈ { 0 , 1 } K × n . Every row y k ∈ { 0 , 1 } 1 × n is a multi-hot vector, where y ki = 1 assigns element i to subset S k .

In this work, we propose a new two-stage procedure to learn partitions. The proposed formulation separately infers the number of elements per subset n k and the assignment of elements to subsets S k by inducing an order on the n elements and filling S 1 , ..., S K sequentially in this order. See Figure 1 for an example.

Definition B.2 (Two-stage partition model) . Let n = [ n 1 , . . . , n K ] ∈ N K 0 be the subset sizes in ρ , with N 0 the set of natural numbers including 0 and ∑ K k =1 n k = n , where n is the total number of elements. Let π ∈ { 0 , 1 } n × n 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 = [ y 1 , . . . , y K ] T ∈ { 0 , 1 } K × n with

Note that in contrast to previous work on partition models (Mansour and Schork, 2016), we allow S k to be the empty set ∅ . 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 a permutation matrix π = [ π 1 , . . . , π n ] T ∈ { 0 , 1 } n × n which is a square matrix where every row and column sums to 1 . This doubly-stochastic property of all permutation matrices π (Marcus, 1960) thus ensures that the columns of Y remain one-hot vectors. At the same time, its rows correspond to n k -hot vectors y k in Definition B.2 and therefore serve as subset assignment vectors.

Corollary B.3. A two-stage partition model Y , which follows Definition B.2, is a valid partition satisfying Definition B.1.

̸

Proof. By definition, every row π i and column π j of π is a one-hot vector, hence every ∑ ν k + n k i = ν k +1 π i results in different, non-overlapping n k -hot encodings, ensuring S i ∩ S j = ∅ ∀ i, j and i = j . Further, since n k -hot encodings have exactly n k entries with 1 , we have ∑ ν k + n k i = ν k +1 ∑ n j =1 π ij = n k . Hence, since ∑ K k =1 n k = n , every element i is assigned to a y k , ensuring S 1 ∪ · · · ∪ S K = [ n ] .

B.2 Two-Stage Random Partition Models

An RPM p ( Y ) 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 a two-stage RPM. To derive the two-stage RPM's probability distribution p ( Y ) , we need to model distributions over n and π . We choose the MVHG distribution p ( n ; ω ) and the PL distribution p ( π ; s ) (see Section 3).

We calculate the probability p ( Y ; ω , s ) sequentially over the probabilities of subsets p y k := p ( y k | y <k ; ω , s ) . p y k itself depends on the probability over subset permutations p ¯ π k := p (¯ π | n k , y <k ; s ) , where a subset permutation matrix ¯ π represents an ordering over n k out of n elements.

Definition B.4 (Subset permutation matrix ¯ π ) . A subset permutation matrix ¯ π ∈ { 0 , 1 } n k × n , where n k ≤ n , must fulfill

We describe the probability distribution over subset permutation matrices p ¯ π k using Definition B.4 and Equation (3).

Lemma B.5 (Probability over subset permutations p ¯ π k ) . The probability p ¯ π k of any subset permutation matrix ¯ π = [ ¯ π 1 , . . . , ¯ π n k ] T ∈ { 0 , 1 } n k × n is given by

where y <k = { y 1 , ..., y k -1 } , Z k = Z -∑ j ∈S <k s j and S <k = ⋃ k -1 j =1 S j .

Proof. We provide the proof for p ¯ π 1 , but it is equivalent for all other subsets. Without loss of generality, we assume that there are n 1 elements in S 1 . Following Equation (3), the probability of a permutation matrix p ( π ; s ) is given by

At the moment, we are only interested in the ordering of the first n 1 elements. The probability of the first n 1 is given by marginalizing over the remaining n -n 1 elements:

where Π 1 is the set of permutation matrices such that the top n 1 rows select the elements in a specific ordering ¯ π ∈ { 0 , 1 } n 1 × n , i.e. Π 1 = { π : [ π 1 , . . . , π n 1 ] T = ¯ π } . It follows

where Z 1 = Z -∑ n 1 j =1 (¯ π s ) j . It follows

Lemma B.5 describes the probability of drawing the elements i ∈ S k in the order described by the subset permutation matrix ¯ π given that the elements in S <k are already determined. Note that in a slight abuse of notation, we use p (¯ π | n k , y <k ; ω , s ) as the probability of a subset permutation ¯ π given that there are n k elements in S k and thus ¯ π ∈ { 0 , 1 } n k × n . Additionally, we condition on the subsets y <k and n k , the size of subset S k . In contrast to the distribution over permutations matrices p ( π ; s ) in Equation (3), we take the product over n k terms and have a different normalization constant Z k . Although we induce an ordering over all elements i in Definition B.2, the probability p y k is invariant to intra-subset orderings of elements i ∈ S k .

Lemma B.6 (Probability distribution p y k ) . The probability distribution over subset assignments p y k is given by

where Π y k = { ¯ π ∈ { 0 , 1 } n k × n : y k = ∑ n k i =1 ¯ π i } and p (¯ π | n k , y <k ; s ) as in Lemma B.5.

Proof. We can proof the statement of Lemma B.6 as follows:

Equation (29) holds by marginalization, where n ′ k denotes the random variable that stands for the size of subset S k . By Bayes' rule, we can then derive Equation (30). The next derivations stem from the fact that we can compute n <k if y <k is given, as the assignments y <k hold information on the size of subsets S <k . More explicitly, n i = ∑ n j =1 y ij . Further, y k is independent of ω if the size n ′ k of subset S k is given, leading to Equation (31). We further observe that p ( y k | n ′ k , y <k ; s ) is only non-zero, if n ′ k = ∑ n i =1 y ki = n k . Dropping all zero terms from the sum in Equation (31) thus results in Equation (32). Finally, by Definition B.2, we know that y k = ∑ ν k + n k i = ν k +1 π i , where ν k = ∑ k -1 ι =1 n ι and π ∈ { 0 , 1 } n × n a permutation matrix. Hence, in order to get y k given y <k , we need to marginalize over all permutations of the elements of y k given that the elements in y <k are already ordered, which corresponds exactly to marginalizing over all subset permutation matrices ¯ π such that y k = ∑ n k i =1 ¯ π i , resulting in Equation (33).

In Lemma B.6, we describe the set of all subset permutations ¯ π of elements i ∈ S k by Π y k . Put differently, we make p ( y k | y <k ; ω , s ) invariant to the ordering of elements i ∈ S k by marginalizing over the probabilities of subset permutations p ¯ π k (Xie and Ermon, 2019).

Using Lemmas B.5 and B.6, we propose the two-stage random partition p ( Y ; ω , s ) . Since Y = [ y 1 , . . . , y K ] T , we calculate p ( Y ; ω , s ) , the PMF of the two-stage RPM, sequentially using Lemmas B.5 and B.6, where we leverage the PL distribution for permutation matrices p ( π ; s ) to describe the probability distribution over subsets p ( y k | y <k ; ω , s ) .

Proposition 4.1 (Two-Stage Random Partition Model). Given a probability distribution over subset sizes p ( n ; ω ) with n ∈ N K 0 and distribution parameters ω ∈ R K + and a PL probability distribution over random orderings p ( π ; s ) with π ∈ { 0 , 1 } n × n and distribution parameters s ∈ R n + , the probability mass function p ( Y ; ω , s ) of the two-stage RPM is given by

where Π Y = { π : y k = ∑ ν k + n k i = ν k +1 π i , k = 1 , . . . , K } , and y k and ν k as in Definition B.2.

Proof. Using Lemmas B.5 and B.6, we write

B.3 Approximating the Probability Mass Function

Lemma 4.2. p ( Y ; ω , s ) can be upper and lower bounded as follows

Proof. Since p ( π ; s ) is a probability we know that ∀ π ∈ { 0 , 1 } n × n p ( π ; s ) ≥ 0 . Thus, it follows directly that:

proving the lower bound of Lemma 4.2.

On the other hand, can prove the upper bound in Lemma 4.2 by:

We can compute the maximum probability max π p ( π ; s ) with the probability of the permutation matrix f π ( s ) , which sorts the unperturbed scores in decreasing order.

B.4 The Differentiable Random Partition Model

We propose the DRPM p ( Y ; ω , s ) , a differentiable and reparameterizable two-stage RPM.

Lemma 4.3 (DRPM). A two-stage RPM is differentiable and reparameterizable if the distribution over subset sizes p ( n ; ω ) and the distribution over orderings p ( π ; s ) are differentiable and reparameterizable.

Proof. To prove that our two-stage RPM is differentiable we need to prove that we can compute gradients for the bounds in Lemma 4.2 and to provide a reparameterization scheme for the two-stage approach in Definition B.2.

Gradients for the bounds: Since we assume that p ( n ; ω ) and p ( π ; s ) are differentiable and reparameterizable, we only need to show that we can compute | Π Y | 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),

Hence, | Π Y | can be computed given a reparametrized version n k , which is provided by the reparametrization trick for the MVHG p ( n ; ω ) . Further, from Equation (14), we immediately see that the most probable permutation is given by the order induced by sorting the original, unperturbed scores s from highest to lowest. This implies that max ˜ π p (˜ π ; s ) = p ( π s ; s ) , which we can compute due to p ( π s ; s ) being differentiable according to our assumptions.

Reparametrization of the two-stage approach: Given reparametrized versions of n and π , we compute a partition as follows:

The challenge here is that we need to be able to backpropagate through n k , which appears as an index in the sum. Let α k = { 0 , 1 } n , such that

Given such α k , we can rewrite Equation (41) with

While this solves the problem of propagating through sum indices, it is not clear how to compute α k in a differentiable manner. Similar to other works on continuous relaxations (Jang et al., 2016; Maddison et al., 2017), we can compute a relaxation of α k by introducing a temperature τ . Let us introduce auxiliary function f : N → [0 , 1] n , that maps an integer x to a vector with entries

such that f i ( x ; τ ) ≈ 0 if x -i τ < 0 and f i ( x ; τ ) ≈ 1 if x -i τ ≥ 0 . Note that σ ( · ) is the standard sigmoid function and ϵ << 1 is a small positive constant to break the tie at σ (0) . We then compute an approximation of α k with

˜ α k ( τ ) ∈ [0 , 1] n . Then, for τ → 0 we have ˜ α k ( τ ) → α k . In practice, we cannot set τ = 0 since this would amount to a division by 0 . Instead, we can apply the straight-through estimator (Bengio et al., 2013) to the auxiliary function f ( x ; τ ) in order to get ˜ α k ∈ { 0 , 1 } n 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 n k . This allows a different path for computing α k which circumvents introducing yet another temperature parameter altogether. We refer to our code in the supplement for more details.

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 6 h8 h (Variational Clustering), 4 h6 h (Generative Factor Partitioning), or ∼ 1 h (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.

Code Release The official code can be found under https://github.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_lib (Locatello et al., 2020). However, the library is based on Tensorflow v1 (Abadi et al., 2016), which makes it more and more difficult to maintain and install. Therefore, we decided to re-implement everything in PyTorch (Paszke et al., 2019).

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 is finite for a given number of elements n and subsets K . Therefore, we can estimate the probability of a fixed partition ˜ Y under given ω and s by sampling M partitions from p ( Y ; ω , s ) , counting the occurrences of ˜ Y in the samples, and dividing the count by M . As M approaches infinity, we can obtain the true probability mass function (PMF) p ( Y ; ω , s ) for every partition Y .

In our experiment, we set n = 5 and aim to evaluate the quality of our bounds for all possible subset combinations of 5 elements, we thus set K = n = 5 . To obtain a reliable estimate of the true PMF, we set M = 10 8 . From Lemma 4.2, we know that:

Let us define

In Figure 5, we present the estimated PMF along with the corresponding upper bounds ( p U ) and lower bounds ( p L ) for four different combinations of RPM parameters ω and s . We observe that when all scores s are equal, as in the priors of the experiments in Sections 5.1 and 5.2, p U ( Y ; ω , s ) approximates p ( Y ; ω , s ) well and serves as a reliable estimate of the PMF. However, when the scores vary, the upper bound becomes looser, particularly for lower probability partitions, as it is dominated by the term max ˜ π p (˜ π ; s ) . Although p L ( Y ; ω , s ) appears looser than p U ( Y ; ω , s ) for certain configurations of ω and s , it provides more consistent results across all hyperparameter combinations.

(a)

ω

=

✶

,

s

=

✶

(b)

ω

∼

U

([0

,

1])

,

s

=

✶

,

s

∼

U

([0

,

1])

(d)

ω

∼

U

([0

,

1])

,

s

∼

U

([0

,

1])

9-

6

-8

-8

log(p(Y; w, s))

log(p(Y; w, s))

-10

10

-12

12

14

14

-16

-16

UpperBound

UpperBound

-18

Estimated

-18

Estimated

LowerBound

Lower Bound

-20

-20

Y~p(Y; w,s)

Y~p(Y; w,s)

-6

-8

-8

log(p(Y; w, s))

((S'

-10

log(p(Y; W, s

-10

-12

12

-14

-14

-16

-16

-18

Estimated

-18

Estimated

LowerBound

LowerBound

20

Y~p(Y; w,s)

-20

Y~p(Y; w, s)

C.2 Variational Clustering with Random Partition Models

C.2.1 Loss Function

As mentioned in Section 5.1, for a 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 can then be summarized as follows: First, we sample the cluster assignments Y from an RPM, i.e., Y ∼ P ( Y ; ω , s ) . Given Y , we can sample the latent variables Z , where for each y we have z ∼ N ( µ y , σ T y I l ) , z ∈ R l . Finally, we sample X by passing each z through a decoder like in vanilla VAEs. Using Bayes rule and Jensen's inequality, we can then derive the following evidence lower bound (ELBO):

We then assume that we can factorize the approximate posterior as follows:

Note that while we do assume conditional independence between z given its corresponding x , we model q ( Y | X ) with the DRPM and do not have to assume conditional independence between

X

p(X|2)

X1

X2

XN-1

XN

q(21|x1)

q(Z2|x2)

q(ZN-1|xN-1)↑

q(ZN|XN)

p(x1|21)

p(x2|z2)

p(xN-1|ZN-1)

p(xN|ZN)

Z

p(Z|Y)

Z1

ZN-1

ZN

p(21y)

p(ZN-1yN-1)

P(ZN|yN)

q(Y|2)

Y

y1

y2

yN-1

YN

different cluster assignments. This allows us to leverage dependencies between samples from the dataset. Hence, we can rewrite the ELBO as follows:

See Figure 6 for an illustration of the generative process and the assumed inference model. Since computing P ( Y ) and q ( Y | X ) is intractable, we further apply Lemma 4.2 to approximate the KLDivergence term in L ELBO , leading to the following lower bound:

where π Y is the permutation that lead to Y during the two-stage resampling process. Further, we want to control the regularization strength of the KL divergences similar to the β -VAE (Higgins et al., 2016). Since the different terms have different regularizing effects, we rewrite Equations (45)

PriorParameters

C

μc

Oc

3

w(X)

~

Y

s(X)

E

D

μ(x)

o(x)

and (46) and weight the individual terms as follows, leading to our final loss:

C.2.2 Architecture

The model for our clustering experiments is a relatively simple, fully-connected autoencoder with a structure as 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 a linear layer and mapping to the respective parameter dimension in the last layer. In our experiments, we use a latent dimension size of l = 10 for MNIST and l = 20 for FMNIST, such that µ ( x ) , σ ( x ) ∈ R l . To understand the architecture choice for the DRPM parameters, let us first take a closer look at Equation (48). For each sample x , this term minimizes the expected KL divergence between its approximate posterior q ( z | x ) = N ( µ ( x ) , diag( σ ( x ))) and the prior at index y given by the partition Y sampled from the DRPM q ( Y | X ; s , ω ) , i.e., N ( µ y , diag( σ y )) . Ideally, the most likely partition should assign the approximate posterior to the prior that minimizes this KL divergence. We can compute such s ( X ) and ω ( X ) given the parameters of the approximate posterior and priors as follows:

where u is a scaling constant that controls the probability of sampling the most likely partition. Note that ω and s minimize Equation (48) if defined this way when given the distribution parameters of the approximate posterior and the priors. The only thing that is left unclear is how much u should scale the scores s . Ultimately, we leave u as a learnable parameter but detach the rest of the computation of s and ω from the computational graph to improve stability during training. Finally, once we resample z ∼ N ( µ ( x ) , σ ( x )) , we pass it through a fully connected decoder D with four layers mapping z to 2000 , 500 , and 500 neurons in the first three layers and then finally back to the input dimension in the last layer to end up with the reconstructed sample ˆ x .

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 q ( z | x ) to sample M samples z ( i ) ∼ q ( z | x ) and compute their reconstruction error to estimate Equation (47). Similarly, we can sample from q ( Y | X ) L times to estimate the terms in Equations (48) to (50), such that we minimize

In our experiments, we set M = 1 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 β = 1 for MNIST and β = 0 . 1 for FMNIST, and otherwise γ = 1 , and δ = 0 . 01 for all experiments.

To resample n and π we need to apply temperature annealing (Grover et al., 2019; Sutter et al., 2023). To do this, we applied the exponential schedule that was originally proposed together with the Gumbel-Softmax trick (Jang et al., 2016; Maddison et al., 2017), i.e., τ = max( τ final , exp ( -rt )) , where t is the current training step and r is the annealing rate. For our experiments, we choose r = log( τ final ) -log( τ init ) 100000 in order to annealing over 100000 training step. Like Jang et al. (2016), we set τ init = 1 and τ final = 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 a reconstruction loss to initialize the decoder properly. To initialize the prior parameters, we fit a GMMto the pretrained embeddings of the training set and took the resulting Gaussian parameters to initialize our priors. Note that we 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 AdamW (Loshchilov and Hutter, 2019) optimizer with a learning rate of 0 . 0001 with a batch size of 256 for 1024 epochs. During initial experiments with the DRPM-VC, we realized that the pretrained weights of the encoder would often lose the learned structure in the first couple of training epochs. We suspect this to be an 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 us 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 on FMNIST, we often observe a 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

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 to 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 as 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 as with the DRPM-VC and freeze the encoder backbone. The results of this additional experiment can be found in 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 easily allow us to generate new samples from each class.

(a) NMI when varying

K

(b) ARI when varying

K

0.7

0.7

DRPM-VC

0.6

0.6

VaDE

LatentGMM

0.5

0.5

NMI

ARI

0.4

DRPM-VC

0.4

VaDE

0.3

0.3

0.2

6

8

10

12

14

16

18

20

0.2

6

8

10

1214

16

18

20

#Clusters

#Clusters

0.7

0.6

0.5

0.4-

DRPM-VC

VaDE

0.3-

LatentGMM

0.2

6

8

10

12

14

16

18

20

#Clusters

(c) Accuracy when varying

K

C.2.7 Samples per cluster

In addition to sampling partitions and then generating samples according to the sampled cluster assignments, we can also directly sample from each of the learned priors. We show some examples of this for both MNIST and FMNIST at the end of the appendix in Figures 17 and 18. We can again see that the DRPM-VC learns accurate cluster representations since each of the samples seems to correspond to 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, of course, not true in practice. We thus also want to 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 ∈ { 6 , 8 , 10 , 12 , 14 , 16 , 18 , 20 } across 5 different seeds on FMNIST. DRPM-VC clearly outperforms the two baselines for all K , except for the extreme case of 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 we are far from the true number of clusters, we additionally generate samples from each prior of the DRPM-VC in Figure 9. Interestingly, we can see that DRPM-VC 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 / 15 ) and without (Priors 4 / 9 ) heels or the cluster T-shirt into clothing of lighter (Prior 0 ) and darker (Prior 3 ) color. Thus, DRPM-VC allows us to investigate clusters in datasets hierarchically, where low values of K detect more coarse and higher values of K more fine-grained patterns in the data.

C.2.9 Clustering of STL-10

We include an additional variational clustering ablation on the STL10 dataset (Coates et al., 2011) that follows the experimental setup of Jiang et al. (VaDE, 2016). As noted by Jiang et al. (2016),

3

５６

8

Prior

9

10

11

12

13

14

15

16

17

18

19

Sample

variational clustering algorithms have difficulties clustering in raw pixel space for natural images, which is why we apply DRPM-VC to representations extracted using an Imagenet (Deng et al., 2009) pretrained ResNet50 (He et al., 2015) as done in Jiang et al. (VaDE, 2016). Note that these results are hard to interpret, as STL10 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 of Generative Factors

We assume that we have access to multiple instances or views of the same event, where only a subset of generative factors changes between views. The knowledge about the data collection process provides a form of weak supervision. For example, we 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 a 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 a 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

E

z

v

1

=









4

7

1

2











z

v

2

=









4

3

8

2











E

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 us reason about the group structure without requiring explicit knowledge in the form of 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 et al., 2020)).

C.3.1 Generative Model

We assume the following generative model for DRPM-VAE

where z = { z s , z 1 , z 2 } . The two frames share an unknown number n s of generative latent factors z s , and an unknown number, n 1 and n 2 , of independent factors z 1 and z 2 . The RPM infers n k and z k using Y . Hence, the generative model extends to

Y

x

1

x

2

X

z

1

z

2

z

s

p

(

z

s

,

z

1

,

z

2

|

Y

)

p

(

x

1

|

z

s

,

z

1

)

p

(

x

2

|

z

s

,

z

2

)

Y

x

1

x

2

X

z

1

z

2

z

s

q

(

z

1

|

Y,

x

1

)

q

(

Y

|

x

1

, x

2

)

q

(

z

2

|

Y,

x

2

)

q

(

z

s

|

Y,

x

1

,

x

2

)

(b) Inference Model

C.3.2 DRPM ELBO

We derive the following ELBO using the posterior approximation q ( z , Y | X )

Following Lemma 4.2, we are able to optimize DRPM-VAE using the following ELBO L ELBO ( X ) :

where π Y is the permutation that lead to Y during the two-stage resampling process. Further, we want to control the regularization strength of the KL divergences similar to the β -VAE (Higgins et al., 2016). The ELBO L ( X ) to be optimized can be written as

where s ( X ) and ω ( X ) denote distribution parameters, which are inferred from X (similar to the Gaussian parameters in the vanilla V AE).

E

D

x

1

x

2

ˆ

x

1

ˆ

x

2

View

Aggregation

q

ϕ

(

z

s

,

z

1

|

x

1

)

q

ϕ

(

z

s

,

z

2

|

x

2

)

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 L samples z ( l ) ∼ q ( z | Y, X ) and compute their reconstruction error to estimate Equation (58). Similarly, we can sample from q ( Y | X ) L times. We use L = 1 to estimate all expectations in L ELBO .

C.3.3 Implementation and Hyperparameters

In this experiment, we use the disentanglement_lib from Locatello et al. (2020). We use the same architectures proposed in the original paper for all methods we compare to. The baseline algorithms, LabelVAE (Bouchacourt et al., 2018; Hosoya, 2018) and AdaVAE (Locatello et al., 2020) are already implemented in disentanglement_lib . For details on the implementation of these methods we refer to the original paper from Locatello et al. (2020). HGV AE is implemented in Sutter et al. (2023). We did not change any hyperparameters or network details. All experiments were performed using β = 1 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 5 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 et al. (2023). It 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

where avg is the averaging function of choice (Locatello et al., 2020; Sutter et al., 2023) and j ∈ { 1 , 2 } . The methods used (i. e. Label-VAE, Ada-VAE, HG-VAE, DRPM-VAE) differ in how to select the subset S.

For DRPM-VAE, we infer ω from the pairwise KL-divergences KL pw between the latent vectors of the two views.

where q ( z j | x j ) are the encoder outputs of the respective images. We do not average or sum across dimensions in the computation of KL pw ( · ) such that the KL pw ( · ) is d -dimensional, where d is the latent space size. The encoder E in Figure 12 maps to µ ( x j ) and σ ( x j ) of a Gaussian distribution. Hence, we can compute the KL divergences above in closed form. Afterwards, we feed the pairwise KL divergence KL pw to a single fully-connected layer, which maps from d to K values

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 KL pw . Instead of using another fully-connected layer, we directly use the log-values of the pairwise KL divergence

Similar to the original works, we also anneal the temperature parameter for p ( n ; ω ) and p ( π ; s ) (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.

0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1

B

D

b

2

7

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 noisyMultiMNIST dataset.

The noisyMultiMNIST dataset modifies the MultiMNIST dataset (Sabour et al., 2017) as follows. In the right image, we set each pixel value to zero with probability α ∈ [0 , 1] . This is done before merging the left and right image in order to only affect the difficulty of the right task. Note that for α = 0 noisyMultiMNIST is equivalent to MultiMNIST and for α = 1 the right task can no longer be solved. This allows us 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

where w L and w L are the loss weights, and L L and L R are the individual loss terms for the respective tasks L and R . In our experiments, we set the task weights to be equal for all dataset versions, i.e. w L = w R = 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 a 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

where C L = C R = 10 for MultiMNIST, gt is a one-hot encoded label vector and p is a categorical vector of estimated class assignments probabilities, i.e. ∑ c p c = 1 .

The predictions for the individual tasks p t are given as

for a sample x ∈ 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-MTL and ULS use the same network enc θ ( · ) 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 a subset of z . DRPM-MTL uses the following prediction scheme

The DRPM-MTL encoder first predicts a latent representation z ← enc θ ( x ) , where x is the input image. Using the same encoder architecture but different parameters φ , we predict a partitioning encoding z ′ ← enc φ ( x ) . With a single linear layer per DRPM log-parameter log ω and log s are computed. Next we infer the partition masks y L , y R ∼ p ( y L , y R ; ω , s ) . We then feed the masked latent representations z L ← z ⊙ y L and z R ← z ⊙ y R into the task specific classification networks h θ L ( z L ) and h θ R ( z R ) respectively to obtain the task specific predictions. Since the two tasks in the MultiMNIST dataset are of similar nature, the task-specifc networks h θ L and h θ R share the same architecture, but have different parameters.

Task L

Task R

h

θ

L

h

θ

R

enc

θ

enc

φ

ω,

s

z

L

z

R

p

(

y

L

,

y

R

;

ω,

s

)

z

C.4.3 Training

For both the ULS and the DRPM-MTL model, we use the Adam optimizer with learning rate 0 . 0005 and train them for 200 epochs with a batch size of 256. We again choose an exponential schedule for the temperature τ and anneal it over the training time, as is explained in Appendix C.2.3.

In our ablation we use α ∈ { 0 , 0 . 1 , 0 . 2 , . . . , 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.

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 of the CelebA dataset serves as an individual task. Hence, using CelebA for MTL results is a 40 task learning problem making the scaling of different task losses more difficult compared to 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 it to the ULS model. We use the same pipeline as for MultiMNIST 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 a ResNet-based encoder to map an image to a representation of d = 64 dimensions. For architectural details, we refer to Sener and Koltun (2018) and https://github.com/isl-org/MultiObjectiveOptimization .

Again, ULS inputs all d = 64 dimensions to the task-specific sub-networks whereas DRPM-MTL partitions the intermediate representations into n T different subsets, which are then fed to the respective task networks. n T is the number of tasks.

Compared to the MultiMNIST experiment (see Appendix C.4.2), we introduce an additional regularization for the DRPM-MTL method. The additional regularization is based on the upper bound in Lemma 4.2 and is penalizing size of | Π Y | for a given n . Hence, the loss function changes to

For all versions of the experiment (i.e. n T ∈ 10 , 20 , 40 ), we set λ = 0 . 015 ≈ 1 64 , which is the number of elements we want to partition. The task losses L t are 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 a MTL experiments using the first 10 attributes out of the 40 attributes. Second, we increase the number of different tasks to 20 . Because we sort the attributes alphabetically in both cases, the first 10 tasks are

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 sub-experiment of the CelebA experiment. In Table 7a, we form a MTL experiment with 10 different tasks. In Table 7b, we form a MTL experiment with 20 different tasks where the first 10 tasks are the same as in the 10 tasks experiment. And in Table 6c, a MTL experiment with all 40 tasks from the dataset. We train both methods for 50 epochs using a learning rate of 0 . 0001 and a 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 (%) as well as the average task accuracy in the bottow row of both subtables.

shared between the two experiment versions. And third, we set n T = 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 outperforms the ULS method on most tasks for the 20 -tasks experiment even though it has only access to d/n T = 64 / 20 = 3 . 2 dimensions on average. And even more extrem for n T = 40 , DRPM-MTL on average has only access to d/n T = 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 of elements is in the case of classification. There, for a given batch X containing B samples, we have Y := ( y 1 , . . . , y B ) where y i ∈ R K is the one-hot encoding of the label of the i -th sample in the batch. In this ablation, we infer ˆ Y := DRPM ( ω θ 1 ( X ) , s θ 2 ( X )) , where we compute ω ( 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

parameters θ 1 and θ 2 by minimizing the following loss:

where L CE denotes the standard cross-entropy loss, n denotes the output of the MVHG leading to ˆ Y , and L 2 ensures that n i matches the number of appearances of label i in the current batch. Using this simple training scheme, we achieve an f1-score of 96 . 43 ± 0 . 02 and 82 . 71 ± 0 . 03 on MNIST and FMNIST, respectively, further demonstrating the applicability and versatility of the DRPM to a number of different problems.

0123456789

Prior

145904h966019061056h02

Sample

0123456789

Prior

32048596793106352898323477691313

2345

Sample

0123456789

Prior

1 2687820827330742409470832990491

Sample

01234567.89

Prior

0

4

5

6

Sample

0123456789

Prior

42950h8226199411880h019520

Sample

0123456789

Prior

98121248947108224610466940/9852

0

Sample

0123456789

Prior

93355739350177872826666541639971

0

Sample

0123456789

Prior

Sample

0123456789

Prior

5

5678

Sample

0123456789

Prior

0

Sample

0123456789

Prior

2345678910111213141516171819202122232425262728293031

Sample

0123456789

Prior

Sample

012m456789

Prior

12345678

Sample

0123456789

Prior

U

4

5

6

8

Sample

2

Prior

9

9

9

9

9

9

6

9

9

5

8

8

8

8

8

X

8

8

8

8

8

6

S

S

5

S

S

5

5

S

5

8

7

7777

7

77

77

77777777

333333

3333

EEEEEEEEEEEEEEEEEEEEEEES

Sample

h

h

4

4444444444444444444444444444

4

h

1

9

6

6

6

6

6

6

6

6

66

666

6

G

6

6

6

6

G

6

6

66666

6

6

O

00o00o0

0

ooo00o

a

2

2022222222

222

222221

2

22222228

&

Prior

6

9

9

6

6

6

9

9

5

8

8

8

8

8

8

8

8

8

5

5

5

5

5

5

S

5

B

7

7

7

7

7

7

77

77

7777

9

9

3

33333

33

3

33

333333333333

33333

333333

33

Sample

0

+

4

44444444444444444444444444444

4

4

1

6

G

6

6

6

G

6

6

6

6

6

6

6

6

6

6

6

6

6

6

6

6

66

6

6

0

oo

ooooooooooo

o

2

2222

又222

222222222

2

Prior

9

9

9

9

9

9

9

9

9

8

8

8

8

8

8

S

5

5

S

5

S

5

5

5

5

7

777

7

7

77777

9

33333333333333333333333

3333333333

Sample

4

4

44444444444444444444444444444

44

6

G

6

6

6

6

6

G

G

6

6

6

6

6

6

6

6

6

6

6

6

G

6

6

6

6

999

6

6

6

oo

o

o

0000

oo

o

。

oo00o

2

2

2

②2

a

22222222

2

2

2

Prior

9

9

9

6

9

6

6

b

5

8

8

8

8

8

8

8

8

5

S

5

5

5

S

5

S

7

7

7

7

7

7

7

7

7

7

7

77777777

3

3

3

33

3

3333333333333333

3

3

3

333333333

Sample

Prior

Prior

Prior

Prior

2

3

8

Sample

3

4

5

0

8

9

Sample

Q

7

3

8

9

Sample

0

2

4

8

Sample

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe motivation for the multivariate hypergeometric distribution could benefit from more elaboration.\n\nQUESTIONS:\n* 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.?\n\nLIMITATIONS:\nYes. 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

comments: SUMMARY:
The paper proposes an interesting approach for differentiable partitioning.

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. 
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.

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 .

QUESTIONS:
-

LIMITATIONS:
Limitations are addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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\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.\n\nQUESTIONS:\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\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).\n\nLIMITATIONS:\nAdequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\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\n- Why you experimented with fixed encoder and different initialization? (in the supplementary)\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}