Differentiable Set Partitioning

Thomas M. Sutter * 1 Alain Ryser * 1 Joram Liebeskind 1 Julia E. Vogt 1

Abstract

sarily assume i.i.d. set elements. On the other hand, mostexisting 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.

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 adataset 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 gradientswith 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 two different challenging experiments:
multitask learning and inference of shared andindependent generative factors under weak supervision.

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 numberof 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 ourapproach in two experiments, demonstrating the proposed
DRPM’s versatility and advantages. First, we perform multitask learning (MTL) by using the DRPM as a building blockin a deterministic pipeline. We show how the DRPM learns
to assign subsets of network neurons to specific tasks andinfers 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). In our second experiment, we demonstrate how to retrieve sets of shared andindependent generative factors of paired images using the
proposed DRPM and how the reparametrizable sampling ofpartitions allows us to learn a novel kind of Variational Autoencoder. In contrast to previous works (Bouchacourt et al.,2018; Hosoya, 2018; Locatello et al., 2020), which rely onstrong assumptions or heuristics, the DRPM enables end-toend inference of generative factors based on theoreticallymotivated modeling assumptions.

1. Introduction
Partitioning a set of elements into subsets is a classicalmathematical problem that attracted much interest over thelast 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, suchas clustering (Bishop & Svensen, 2004) or classification(De la Cruz-Mesı́a et al., 2007). Random partition models(RPM, Hartigan, 1990) define a probability distribution overthe space of partitions. RPMs can explicitly leverage therelationship between elements of a set, as they do not neces-

2. A two-stage Approach to Random Partition
ModelsWe propose the DRPM p(Y ; ω, s), a differentiable and reparameterizable two-stage Random Partition Model (RPM).
Suppose we want to partition a set S with n elements into
K subsets (see Appendix B.1). The proposed formulation
then separately infers theof elements per subset
PnumberKn=n, and the assignment
through n ∈ NK
,
where
0
k=1 k
of elements to subsets Sk by inducing an order on the nelements and filling S1 , ..., SK sequentially in this order.
To model the order of the elements, we use a permutation*
Equal contribution 1 Department of Computer Science,ETH Zurich, Switzerland.
Correspondence to: Thomas
Alain Ryser
M. Sutter <thomas.sutter@inf.ethz.ch>,<alain.ryser@inf.ethz.ch>.

Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1Differentiable Set Partitioning

0
0

0

0

1
0


0
1

e. g. 0
π = 
0

0
0

1
0
0
0
0
0

0
0
0
1
0
0

0
0
0
0
0
1


0
1

0

0

0
0

0
0
0
0
0
1

0
0
1
0
0
0


0
0
 = π̄200
 = π̄3
1
= π̄5
0

1
0
0
0
0
0

0
0
0
1
0
0

0
0
1
0
0
0

create an RPM for n elements: the score parameters s ∈ Rn+
and the group importance parameters ω ∈ RK
+ . Note thatto sample from the two-stage RPM of Proposition 2.1 weapply 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 illustrated in
Figure 1.

 
 
0
n1
3
n2 
  e. g.  

 
n =
n3  = 2
0
n4 
n5
1

Approximating the Probability Mass Function The
number of permutations per subset |Πyk | scales factoriallywith the subset size nk , i.e. |Πyk | = nk !. Consequently,
the number of valid permutation matrices |ΠY | is given as afunction of n, i.e.


1 1 0 0 1 0
0 0 1 0 0 1
0 0 0 1 0 0
|
{z
}
Y

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

|ΠY | =

k=1

|Πyk | =

K
Y

nk !

(1)

k=1

Although Proposition 2.1 describes a well-defined distribution for p(Y ; ω, s), it is in general computationally intractable due to Equation (1). In practice, we thus approximate p(Y ; ω, s) using the following Lemma (we provide a
proof in Appendix C.3).
Lemma 2.2. p(Y ; ω, s) can be upper and lower bounded
as follows

matrix π = [π1 , . . . , πn ]T ∈ {0, 1}n×n , from which weinfer 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 toan nk -hot vector yk andserves as a subset
selecPtherefore
Pk−1
νk +nktion vector, i.e. yk = i=ν
π
,
where
ν
=
nι .
i
k
+1
ι=1
k
such that Y = [y1 , . . . , yK ]T . See Figure 1 for an illustrative example. Note that K defines the maximum number ofpossible subsets, and not the effective number of non-empty
subsets, because we allow Sk to be the empty set ∅ (Mansour& Schork, 2016). We base the following Proposition 2.1 on
the hypergeometric distribution p(n; ω) (MVHG, Fisher,1935) for the subset sizes n and the Plackett-Luce distribution p(π; s) (PL, Luce, 1959; Plackett, 1975) for assigning
the elements to subsets (see Appendix B for more details).
However, the proposed two-stage approach to RPMs is not
restricted to these two classes of probability distributions.

∀π ∈ ΠY : p(Y ; ω, s) ≥ p(n; ω)p(π; s)(2)p(Y ; ω, s) ≤ |ΠY |p(n; ω) max p(π̃; s) (3)π̃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-to-end
fashion:Lemma 2.3 (DRPM). A two-stage RPM is differentiable
and reparameterizable if the distribution over subset sizesp(n; ω) and the distribution over orderings p(π; s) are differentiable and reparameterizable.

Proposition 2.1 (Two-stage Random Partition Model).
Given a probability distribution over subset sizes p(n; ω)Kwith n ∈ NK
0 and distribution parameters ω ∈ R+ and a
PL probability distribution over random orderings p(π; s)with π ∈ {0, 1}n×n and distribution parameters s ∈ Rn+ ,the probability mass function p(Y ; ω, s) of the two-stage
RPM is given byX
p(Y ; ω, s) = p(y1 , . . . , yK ; ω, s) = p(n; ω)p(π; s)We provide the proof in Appendix C.4. Note that Lemma 2.3
enables us to learn variational posterior approximations andpriors using Stochastic Gradient Variational Bayes (SGVB,
Kingma & Welling, 2014). In our experiments, we apply Lemma 2.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.

π∈ΠY

where ΠY = {π : yk =yk and νk as above.

KYPνk +nki=νk +1 πi , k = 1, . . . , K}, and

3. Experiments

We provide the proof and more details in Appendix C.2. In
contrast to previous RPMs, which often need exponentiallymany distribution parameters (Plackett, 1975), the proposed
two-stage approach only requires (n + K) parameters to

We demonstrate the versatility and effectiveness of the proposed DRPM in two different experiments: multitask learning and partitioning of generative factors.
2

Differentiable Set Partitioning

Figure 3. The mean squared errors between the estimated numberof shared factors n̂s and the true number of shared factors ns
across five seeds for the Label-VAE, Ada-VAE, HG-VAE, and
DRPM-VAE.

branches of the ULS method access all neurons of the lastshared 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 D.1 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 totediously search for optimal loss weights to reach good performance. We base our pipeline on (Sener & Koltun, 2018).
For more details and additional CelebA MTL experiments
we refer to Appendix D.1.

Figure 2. ULS and the DRPM-MTL task accuracies on noisyMultiMNNIST (top) and learned number of neurons per task given
different noise ratios (bottom). DRPM-MTL can reach higher
accuracy for most of the different noise ratios α while assigning
the number of dimensions per task according to their difficulty.

3.1. Multitask Learning
Many ML applications aim to solve specific tasks, where
we optimize for a single objective while ignoring potentiallyhelpful 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 & 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 caseof equal difficulty of the two tasks, a classifier with equal
weighting of the two classification losses serves as an upperbound in terms of performance. However, finding suitabletask weights is a tedious and inefficient approach to MTL.
A more automated way of weighting multiple tasks would
thus be vastly appreciated.

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 2). The DRPMMTL achieves the same or better accuracy on both tasks formost noise levels (upper part of Figure 2). It is interesting to
see that, the more we increase α, the more the DRPM-MTL
tries to overcome the increased difficulty of the right task
by assigning more dimensions to it (lower part of Figure 2,
noise ratio α 0.6-0.8). Note that for the maximum noise
ratio of α = 0.9, it seems that the DRPM-MTL basicallysurrenders and starts neglecting the right task, instead focusing on getting good performance on the left task, which
impacts the average accuracy.

In this experiment, we demonstrate how the DRPM can learntask difficulty by partitioning a network layer. Intuitively, atask that requires many neurons is more complex than a taskthat 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 suchthat only a subset of the neurons are used for every task.
Here, we use the DRPM without resampling and infer thepartition Y as a deterministic function. This can be done
by applying the two-step procedure of Proposition 2.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 fixedarchitecture and scales task losses equally. Both DRPMMTL and ULS use a network with shared architecture up
to some layer, after which the network branches into twotask-specific layers that perform the classifications. Notethe difference between the methods. While the task-specific3.2. Variational Partitioning of Generative Factors
Data modalities not collected as i.i.d. samples, such asconsecutive 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
3Differentiable Set Partitioning

Table 1. We evaluate the learned latent representations of the four methods (Label-VAE, Ada-VAE, HG-VAE, DRPM-VAE) with respect
to the shared (S) and independent (I) generative factors. We train linear classifiers on the shared and independent latent dimensions
separately to predict the respective generative factors and report the results in adjusted balanced accuracy across five seeds.
ns = 0ns = 1

ns = 3ns = 5

I

S

I

S

I

S

I

L ABEL0.14±0.01

0.19±0.03

0.16±0.01

0.10±0.00

0.23±0.01

0.34±0.00

0.00±0.00

A DA

0.12±0.01

0.19±0.01

0.15±0.01

0.10±0.03

0.22±0.02

0.33±0.03

0.00±0.00

HG

0.18±0.01

0.22±0.05

0.19±0.01

0.08±0.02

0.28±0.01

0.28±0.01

0.01±0.00

DRPM

0.26±0.02

0.39±0.07

0.2±0.01

0.15±0.01

0.29±0.02

0.42±0.03

0.01±0.00

in discovering the relationship between coupled samples.
If we assume that the data is generated from underlying
generative factors, weak supervision comes from the factthat we know that certain factors are shared between coupledpairs while others are independent. The supervision is weakbecause we neither know the underlying generative factors
nor the number of shared and independent factors. In sucha 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 = [x1 , x2 ] 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 thedivision of the latent space into shared and independent
latent factors as RPM. We add a posterior approximationq(Y | X) and additionally a prior distribution of the formp(Y ). The model maximizes the following ELBO on themarginal log-likelihood of images through a VAE (Kingma& Welling, 2014):(4)

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 D.2 for details on the ELBO, the setupof the experiment, the implementation, and an illustrationof the generative assumptions. We evaluate all methods
according to how well they partition the latent representations according to shared and independent factors (Table 1).
Because we have access to the data-generating process, we
can control the number of shared ns and independent ni
factors. We compare the methods on four different datasetswith ns ∈ {0, 1, 3, 5}. In Figure 3, we demonstrate that the
DRPM-VAE accurately estimates the true number of sharedgenerative 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 1,
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 thecorresponding shared or independent representations.

where we use the short notation z = {zs , z1 , z2 }. Computing KL [q(Y | X)|| p(Y )] directly is intractable, and weneed to bound it according to Lemma 2.2.

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-generatingprocess in a theoretically grounded fashion instead of relying on heuristics.

LELBO =

2
X

Eq(zs ,zj ,Y |X) [log p(xj | zs , zj )]j=1− Eq(Y |X) [KL [q(z | Y, X)||p(z)]]− KL [q(Y | X)|| p(Y )] ,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 ona heuristic-based approach to infer shared and independent
latent factors. Like in Locatello et al. (2020) and Sutter et al.
(2023), we assume a single known factor for Label-VAE inall experiments. HG-VAE (Sutter et al., 2023) also relies on
the MVHG to model the number of shared and independent
factors. Unlike the proposed DRPM-VAE approach, HGVAE must rely on a heuristic to assign latent dimensions toshared factors, as the MVHG only allows to model the number of shared and independent factors but not their position

Conclusion
In this work, we proposed the DRPM, a novel approach torandom partition models. Our two-stage method enables
learning partitions end-to-end by separately controlling subset sizes and how elements are assigned to subsets. The
DRPM integrates random partition models into probabilisticand deterministic gradient-based optimization frameworks.
In two different experiments, we demonstrate how learning
partitions enables us to infer task-specific sub-networks andshared and independent generative factors from coupledsamples.

4Differentiable Set Partitioning

References

Schölkopf, B., and Bauer, S. 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.

Adams, R. P. and Zemel, R. S. Ranking via sinkhorn propagation. arXiv preprint arXiv:1106.1925, 2011.
Bengio, Y., Léonard, N., and Courville, A. Estimating or
propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432,
2013.

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

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

Grover, A., Wang, E., Zweig, A., and Ermon, S. Stochastic Optimization of Sorting Networks via Continuous
Relaxations. In International Conference on Learning
Representations, 2019. URL https://openreview.
net/forum?id=H1eSS3CcKX.

Bouchacourt, D., Tomioka, R., and Nowozin, S. Multi-levelvariational autoencoder: Learning disentangled representations from grouped observations. In Thirty-Second
AAAI Conference on Artificial Intelligence, 2018.
Caruana, R. Multitask learning: A knowledge-based sourceof inductive bias. In Machine learning: Proceedings ofthe tenth international conference, pp. 41–48, 1993.

Hartigan, J. A. Partition models. Communications in Statistics - Theory and Methods, 19(8):2745–2756, 1990. doi:
10.1080/03610929008830345. Publisher: Taylor & Francis.

Caruana, R. and de Sa, V. R. Promoting poor features to supervisors: Some inputs work better as outputs. Advancesin Neural Information Processing Systems, 9, 1996.

Higgins, I., Matthey, L., Pal, A., Burgess, C., and Glorot,
X. beta-vae: Learning basic visual concepts with a constrained variational framework. 2016. URL https://openreview.net/forum?id=Sy2fzU9gl.

Chesson, J. 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.

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

De la Cruz-Mesı́a, R., Quintana, F. A., and Müller, P.
Semiparametric Bayesian Classification with longitudinal Markers. Journal of the Royal Statistical Society Series C: Applied Statistics, 56(2):119–137, March
2007. ISSN 0035-9254. doi: 10.1111/j.1467-9876.
2007.00569.x. URL https://doi.org/10.1111/
j.1467-9876.2007.00569.x.

Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. arXiv preprint
arXiv:1611.01144, 2016.
Jiang, Z., Zheng, Y., Tan, H., Tang, B., and Zhou, H. Variational deep embedding: An unsupervised and generative
approach to clustering. arXiv preprint arXiv:1611.05148,
2016.

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

Kingma, D. P. and Welling, M. Auto-Encoding Variational
Bayes. In International Conference on Learning Representations, 2014.
Kurin, V., Palma, A. D., Kostrikov, I., Whiteson, S., andKumar, M. P. In Defense of the Unitary Scalarization for
Deep Multi-Task Learning. CoRR, abs/2201.04122, 2022.
URL https://arxiv.org/abs/2201.04122.

Fisher, R. A. The logic of inductive inference. Journal of
the royal statistical society, 98(1):39–82, 1935. Publisher:
JSTOR.
Fog, A. Calculation methods for Wallenius’ noncentral hypergeometric distribution. Communications in
Statistics—Simulation and Computation®, 37(2):258–
273, 2008. Publisher: Taylor & Francis.

LeCun, Y., Cortes, C., and Burges, C. J. Gradient-basedlearning applied to document recognition. In Proceedingsof the IEEE, volume 86, pp. 2278–2324, 1998. Issue: 11.
Lee, C. J. and Sang, H. Why the Rich Get Richer? On
the Balancedness of Random Partition Models. arXiv
preprint arXiv:2201.12697, 2022.

Gondal, M. W., Wuthrich, M., Miladinovic, D., Locatello,F., Breidt, M., Volchkov, V., Akpo, J., Bachem, O.,
5

Differentiable Set Partitioning

Li, Y., Yang, M., Peng, D., Li, T., Huang, J., and Peng,
X. Twin Contrastive Learning for Online Clustering.
International Journal of Computer Vision, 130(9):2205–
2221, September 2022. ISSN 0920-5691, 1573-1405. doi:
10.1007/s11263-022-01639-z. URL http://arxiv.
org/abs/2210.11680. arXiv:2210.11680 [cs].

Mena, G. E., Belanger, D., Linderman, S. W., and Snoek,
J. Learning Latent Permutations with Gumbel-Sinkhorn
Networks. In 6th International Conference on Learning
Representations, {ICLR} 2018, Vancouver, BC, Canada,
April 30 - May 3, 2018, Conference Track Proceedings.
OpenReview.net, 2018. URL https://openreview.
net/forum?id=Byt3oJ-0W.

Linderman, S. W., Mena, G. E., Cooper, H. J., Paninski, L., and Cunningham, J. P. Reparameterizing the
Birkhoff Polytope for Variational Permutation Inference.
In Storkey, A. J. and Pérez-Cruz, F. (eds.), International Conference on Artificial Intelligence and Statistics,{AISTATS} 2018, 9-11 April 2018, Playa Blanca, Lanzarote, Canary Islands, Spain, volume 84 of Proceedingsof Machine Learning Research, pp. 1618–1627. PMLR,
2018. URL http://proceedings.mlr.press/
v84/linderman18a.html.

Monge, G. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pp. 666–704,
1781.
Petersen, F., Borgelt, C., Kuehne, H., and Deussen, O.
Monotonic Differentiable Sorting Networks. In International Conference on Learning Representations, 2021.
Pitman, J. Some Developments of the Blackwell-Macqueen
URN Scheme. Lecture Notes-Monograph Series, 30:
245–267, 1996. ISSN 07492170. URL http://www.
jstor.org/stable/4355949. Publisher: Instituteof Mathematical Statistics.

Liu, Z., Luo, P., Wang, X., and Tang, X. Deep Learning Face
Attributes in the Wild, September 2015. URL http://arxiv.org/abs/1411.7766. arXiv:1411.7766[cs].

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

Locatello, F., Poole, B., Rätsch, G., Schölkopf, B., Bachem,O., and Tschannen, M. Weakly-supervised disentanglement without compromises. In International Conference
on Machine Learning, pp. 6348–6359. PMLR, 2020.

Luce, R. D. Individual choice behavior: A theoreticalanalysis. Courier Corporation, 1959.

Prillo, S. and Eisenschlos, J. 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 ofProceedings of Machine Learning Research, pp. 7793–
7802. PMLR, 2020. URL http://proceedings.
mlr.press/v119/prillo20a.html.

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

Rota, G.-C. The Number of Partitions of a Set. The American Mathematical Monthly, 71(5):498, May 1964. ISSN
00029890. doi: 10.2307/2312585. Publisher: JSTOR.

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

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

Manduchi, L., Chin-Cheong, K., Michel, H., Wellmann,S., and Vogt, J. E. Deep Conditional Gaussian Mixture
Model for Constrained Clustering. ArXiv, abs/2106.0,
2021.

Sabour, S., Frosst, N., and Hinton, G. E. Dynamic
Routing Between Capsules. In Guyon, I., Luxburg,
U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural
Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.
neurips.cc/paper/2017/file/2cad8fa47bbef282badbb8de5374b894-Paper.
pdf.

Loshchilov, I. and Hutter, F. Decoupled Weight Decay Regularization, January 2019. URL http://arxiv.org/
abs/1711.05101. arXiv:1711.05101 [cs, math].

Mansour, T. and Schork, M. Commutation relations, normal
ordering, and Stirling numbers. CRC Press Boca Raton,
2016.
Marcus, M.
Some Properties and Applications ofDoubly Stochastic Matrices. The American Mathematical Monthly, 67(3):215–221, 1960.
ISSN
00029890, 19300972. URL http://www.jstor.
org/stable/2309679. Publisher: Mathematical Association of America.

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

Differentiable Set Partitioning

Sener, O. and Koltun, V. Multi-task learning as multiobjective optimization. Advances in neural information
processing systems, 31, 2018.
Sinkhorn, R. A relationship between arbitrary positive matrices and doubly stochastic matrices. The annals ofmathematical statistics, 35(2):876–879, 1964. Publisher:
JSTOR.
Sutter, T. M., Manduchi, L., Ryser, A., and Vogt, J. E.
Learning Group Importance using the Differentiable Hypergeometric Distribution. In International Conference
on Learning Representations, 2023.
Thurstone, L. L. A law of comparative judgment. In Scaling,pp. 81–92. Routledge, 1927.
Wallenius, K. T. Biased sampling; the noncentral hypergeometric probability distribution. Technical report, Stanford
Univ Ca Applied Mathematics And Statistics Labs, 1963.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-MNIST:a Novel Image Dataset for Benchmarking Machine
Learning Algorithms, September 2017. URL http://
arxiv.org/abs/1708.07747. arXiv:1708.07747[cs, stat].
Xie, S. M. and Ermon, S. Reparameterizable subset
sampling via continuous relaxations. arXiv preprint
arXiv:1901.10517, 2019.
Xin, D., Ghorbani, B., Garg, A., Firat, O., and Gilmer, J.
Do Current Multi-Task Optimization Methods in Deep
Learning Even Help? arXiv preprint arXiv:2209.11379,
2022.
Yellott, J. I. 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.

7

Differentiable Set Partitioning

ACKNOWLEDGMENTS
Thomas Sutter is supported by the grant #2021-911 of the Strategic Focal Area “Personalized Health and Related Technologies (PHRT)” of the ETH Domain (Swiss Federal Institutes of Technology). Alain Ryser is supported by the StimuLoop
grant #1-007811-002 and the Vontobel Foundation.

A. 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 & Svensen, 2004). Further, Lee & 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. 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 ina variational fashion by making i.i.d. assumptions about the samples in the dataset and imposing soft assignments of theclusters 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 costfunction 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 & Zemel, 2011) or permutation learning (Santa Cruz et al., 2017; Menaet al., 2018).
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 gradient-based optimization pipelines. Based on the same trick, Xie & Ermon (2019) describe
a top-k elements selection procedure, and 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 & 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 & Eisenschlos (2020) proposed a computationally simpler variant of Grover et al.
(2019).

B. Preliminaries
B.1. Set Partitions
A partiton ρ = (S1 , . . . , SK ) of a set [n] = {1, . . . , n} with n elements is a collection of K subsets Sk ⊆ [n] where K is apriori unknown (Mansour & Schork, 2016). For a partition ρ to be valid, it must hold thatS1 ∪ · · · ∪ SK = [n] and ∀k ̸= l : Sk ∩ Sl = ∅(5)In other words, every element i ∈ [n] has to be assigned to precisely one subset Sk . We denote the size of the k-th subset Sk
as nk = |Sk |. Alternatively, we can describe a partition ρ through an assignment matrix Y = [y1 , . . . , yK ]T ∈ {0, 1}K×n .
Every row yk ∈ {0, 1}1×n is a multi-hot vector, where yki = 1 assigns element i to subset Sk .
B.2. 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 = [m1 , . . . , mK ] ∈ NK describe the number of elements per class (e.g. marbles per color),PK
N = k=1 mK 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
PK describes the probability of drawingn = [n1 , . . . , nK ] ∈ NK marbles by sampling without replacement such that k=1 nk = n, where nk is the number ofdrawn marbles of class k.
In the literature, two different versions of the noncentral hypergeometric distribution exist, Fisher’s (Fisher, 1935) and8

Differentiable Set Partitioning

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 B.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K 
1 Y mk n kωkP (N = n; ω) = p(n; ω) =
nk
P0
k=1

K 
Y
X
mk ηk
where P0 =
ωk
ηk

(6)

(7)

(η1 ,...,ηK )∈S k=1

The support S of the PMF is given by S = {n ∈ NK : ∀kn k ≤ mk ,PK

k=1 nk = 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)).
The MVHG p(n; ω, m) allows us to model dependencies between different elements of a set since drawing one elementfrom the urn influences the probability of drawing one of the remaining elements, creating interdependence between them.
Note that in this paper, we assume ∀ mk ∈ m : mk = n. We thus use the shorthand p(n; ω) to denote the density of the
MVHG.
B.2.1. D IFFERENTIABLE MVHG
Their reparameterizable sampling for the differentiable MVHG consists of three parts:
1. Reformulate the multivariate distribution as a sequence of interdependent and conditional univariate hypergeometric
distributions.
2. Calculate the probability mass function of the respective univariate distributions.
3. Sample from the conditional distributions utilizing the Gumbel-Softmax trick.
Following the chain rule of probability, the MVHG distribution allows for sequential sampling over classes k. Every step
includes a merging operation, which leads to biased samples compared to groundtruth non-differentiable sampling withequal class weights ω. Given that we intend to use the differentiable MVHG in settings where we want to learn the unknownclass weights, we do not expect a negative effect from this sampling procedure. For details on how to merge the MVHG into
a sequence of unimodal distributions, we refer to Sutter et al. (2023).
The probability mass function calculation is based on unnormalized log-weights, which are interpreted as unnormalized
log-weights of a categorical distribution. The interpretation of the class-conditional unimodal hypergeometric distributions
as categorical distributions allows applying the Gumbel-Softmax trick (Jang et al., 2016; Maddison et al., 2017). Following
the use of the Gumbel-Softmax trick, the class-conditional version of the hypergeometric distribution is differentiable andreparameterizable. Hence, the MVHG has been made differentiable and reparameterizable as well. Again, for details werefer to the original paper (Sutter et al., 2023).
B.3. Distribution over Random Orderings
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 ∈ Rn+ , where each score si corresponds to an element i. The order given bysorting s in decreasing order corresponds to the most likely permutation in p(π; s). Sampling from p(π; s) can be achievedby resampling the scores as s̃i = β log si + gi where gi ∼ Gumbel(0, β) for fixed scale β, and sorting them in decreasingorder. Hence, resampling scores s enables the resampling of permutation matrices π. The probability over orderings p(π; s)9Differentiable Set Partitioningis then given by (Thurstone, 1927; Luce, 1959; Plackett, 1975; Yellott, 1977)
p(π; s) = p((π s̃)1 ≥ · · · ≥ (π s̃)n )=where π is a permutation matrix and Z =Pn(πs)1 (πs)2
(πs)n
···
Pn−1
Z Z − (πs)1Z − j=1 (πs)j

(8)
(9)

i=1 si .

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). The probability of samplingelement i first is given by its score si divided by the sum of all weights in the setq(s̃i ) =

si
Z

(10)

P
For zi = log si , the right hand side of Equation (10) is equal to the softmax distribution softmax(zi ) = exp(zi )/ j exp(zj )as already described in (Xie & Ermon, 2019). Hence, Equation (10) directly leads to the Gumbel-Softmax trick (Jang et al.,2016; Maddison et al., 2017).
B.3.1. D IFFERENTIABLE S ORTING
In the main text of the paper we rely on a differentiable function fπ (s̃), which sorts the resampled version of the scores s
π = fπ (s̃) = sort(s̃)(11)

Here, we summarise the findings from Grover et al. (2019) on how to construct such a differentiable sorting operator. As
already mentioned in Appendix A, there are multiple works on the topic (Prillo & 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 B.2 (Permutation Matrix (Grover et al., 2019)). Let s = [s1 , . . . , sn ]T be a real-valued vector of length n. Let
As denote the matrix of absolute pairwise differences of the elements of s such that As [i, j] = |si − sj |. The permutation
matrix π corresponding to sort(s) is given by:(
1 if j = arg max[(n + 1 − 2i)s − As 1](12)π=
0 otherwise
where 1 denotes the column vector of all ones.
As we know, the arg max operator is non-differentiable which prohibits the direct use of Corollary B.2 for gradientcomputation. 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 ith row of π(τ ) is givenby:π(τ )[i, :] = softmax[(n + 1 − 2i)s − As 1/τ ](13)

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

C. Detailed Derivation of the Differentiable Two-Stage Random Partition Model
C.1. Two-Stage Partition Model
We want to partition n elements [n] = {1, . . . , n} into K subsets {S1 , . . . , SK } where K is a priori unknown.
10Differentiable Set Partitioning

Definition C.1 (Partition). A partition ρ of a set of elements [n] = {1, . . . , n} is a collection of subsets (S1 , . . . , SK ) suchthat
S1 ∪ · · · ∪ SK = [n] and ∀i ̸= j : Si ∩ Sj = ∅(14)

Put differently, every element i has to be assigned to precisely one subset Sk . We denote the size of the k-th subset Sk as
nk = |Sk |. Alternatively, we describe a partition ρ as an assignment matrix Y = [y1 , . . . , yK ]T ∈ {0, 1}K×n . Every row
yk ∈ {0, 1}1×n is a multi-hot vector, where yki = 1 assigns element i to subset Sk .
In this work, we propose a new two-stage procedure to learn partitions. The proposed formulation separately infers thenumber of elements per subset nk and the assignment of elements to subsets Sk by inducing an order on the n elements andfilling S1 , ..., SK sequentially in this order. See Figure 1 for an example.
Definition C.2 (Two-stage partition model). Let n = [n1 , . . . , nK ] ∈ NK
0 be the subset sizes in ρ, with N0 the set
PKof natural numbers including 0 and k=1 nk = 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 = [y1 , . . . , yK ]T ∈ {0, 1}K×n with
yk =

νkX
+nk

πi ,

where νk =i=νk +1

k−1
X

nι

(15)

ι=1

T
such that Y = [{yk | nk > 0}K
k=1 ] .

Note that in contrast to previous work on partition models (Mansour & Schork, 2016), we allow Sk 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 nk -hot vectors yk in
Definition C.2 and therefore serve as subset assignment vectors.
Corollary C.3. A two-stage partition model Y , which follows Definition C.2, is a valid partition satisfying Definition C.1.
Pνk +nkπi results in different,
Proof. By definition, every row πi and column πj of π is a one-hot vector, hence every i=ν
k +1
non-overlapping nk -hot encodings, ensuring Si ∩ Sj = ∅ ∀ i, j and i ̸= j. Further, since nk -hot encodings have exactlyPνk +nk Pn
PK
nk entries with 1, we have i=ν
k=1 nk = n, every element i is assigned to a yk ,j=1 πij = nk . Hence, since
k +1
ensuring S1 ∪ · · · ∪ SK = [n].
C.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 C.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 Appendix B).
We calculate the probability p(Y ; ω, s) sequentially over the probabilities of subsets pyk := p(yk | y<k ; ω, s). pyk
itself depends on the probability over subset permutations pπ̄k := p(π̄ | nk , y<k ; s), where a subset permutation matrix π̄represents an ordering over nk out of n elements.
Definition C.4 (Subset permutation matrix π̄). A subset permutation matrix π̄ ∈ {0, 1}nk ×n , where nk ≤ n, must fulfill∀i ≤ nk :nXπ̄ij = 1 and ∀j ≤ n :nkXπ̄ij ≤ 1.

i=1

j=1

We describe the probability distribution over subset permutation matrices pπ̄k using Definition C.4 and Equation (8).
11Differentiable Set Partitioning

Lemma C.5 (Probability over subset permutations pπ̄k ). The probability pπ̄k of any subset permutation matrixπ̄ = [π̄1 , . . . , π̄nk ]T ∈ {0, 1}nk ×n is given bypπ̄k := p(π̄ | nk , y<k ; s) =where y<k = {y1 , ..., yk−1 }, Zk = Z −

P

nk
Y

i=1 Zk −j∈S<k sj and S<k =Sk−1(π̄s)iPi−1

j=1 (π̄s)j(16)

j=1 Sj .

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 n1 elements in S1 . Following Equation (8), the probability of a permutation matrix p(π; s) is given byp(π; s) =(πs)n(πs)1 (πs)2···Pn−1
Z Z − (πs)1Z − j=1 (πs)j(17)

At the moment, we are only interested in the ordering of the first n1 elements. The probability of the first n1 is given by
marginalizing over the remaining n − n1 elements:Xp(π | s)(18)p(π̄ | n1 ; ω) =π∈Π1where Π1 is the set of permutation matrices such that the top n1 rows select the elements in a specific ordering π̄ ∈{0, 1}n1 ×n , i.e. Π1 = {π : [π1 , . . . , πn1 ]T = π̄}. It follows
Xp(π | s)(19)p(π̄ | n1 ; ω) =
π∈Π1

=

n
X Y

(πs)i
Pi−1
j=1 (πs)jπ∈Π1 i=1 Z −

(20)

n1
Y

Y1
X n−n
(π̄s)i(πs)n1 +i
=
Pi−1
Pi−1
P n1
j=1 (π̄s)j −j=1 (π̄s)jj=1 (π̄s)j π∈Π1 i=1 Z −i=1 Z −=
where Z1 = Z −

P n1

X n−n
Y1
(π̄s)i(πs)n1 +i
Pi−1
Pi−1
j=1 (π̄s)j π∈Π1 i=1 Z1 −
j=1 (π̄s)ji=1 Z −
n1
Y

(21)

(22)

j=1 (π̄s)j . It followsp(π̄ | n1 ; ω) =n1
Yi=1 Z −(π̄s)iPi−1
j=1 (π̄s)j(23)Lemma C.5 describes the probability of drawing the elements i ∈ Sk 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(π̄ | nk , y<k ; ω, s)
as the probability of a subset permutation π̄ given that there are nk elements in Sk and thus π̄ ∈ {0, 1}nk ×n . Additionally,
we condition on the subsets y<k and nk , the size of subset Sk . In contrast to the distribution over permutations matrices
p(π; s) in Equation (8), we take the product over nk terms and have a different normalization constant Zk . Although we
induce an ordering over all elements i in Definition C.2, the probability pyk is invariant to intra-subset orderings of elementsi ∈ Sk .
Lemma C.6 (Probability distribution pyk ). The probability distribution over subset assignments pyk is given byX
pyk := p(yk | y<k ; ω, s) = p(nk | n<k ; ω)p(π̄ | nk , y<k ; s)π̄∈Πykwhere Πyk = {π̄ ∈ {0, 1}nk ×n : yk =P nk

i=1 π̄i } and p(π̄ | nk , y<k ; s) as in Lemma C.5.

12

Differentiable Set Partitioning

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

pyk = p(yk | y<k ; ω, s)Xp(yk , n′k | y<k ; ω, s)
=

(24)

n′k

X

p(n′k | y<k ; ω, s)p(yk | n′k , y<k ; ω, s)(25)

p(n′k | n<k ; ω, s)p(yk | n′k , y<k ; s)(26)

= p(nk | n<k ; ω, s)p(yk | nk , y<k ; s)X
= p(nk | n<k ; ω)p(π̄ | nk , y<k ; s)

(27)

=

n′k

=

X
n′k

(28)

π̄∈ΠykEquation (24) holds by marginalization, where n′k denotes the random variable that stands for the size of subset Sk . By
Bayes’ rule, we can then derive Equation (25). The next derivations stem from the fact that we canPcompute n<k if y<k isngiven, as the assignments y<k hold information on the size of subsets S<k . More explicitly, ni = j=1 yij . Further, yk is′
independent of ω if the P
size nk of subset Sk is given, leading to Equation (26). We further observe that p(yk | n′k , y<k ; s) isn′only non-zero, if nk = i=1 yki = nk . Dropping all zero terms from the sum in Equation (26) thus results in Equation (27).
Pνk +nk
Pk−1
Finally, by Definition C.2, we know that yk = i=ν
πi , where νk = ι=1 nι and π ∈ {0, 1}n×n a permutation matrix.
k +1
Hence, in order to get yk given y<k , we need to marginalize over all permutations of the elements of yk given that theelements in y<kPare already ordered, which corresponds exactly to marginalizing over all subset permutation matrices π̄nksuch that yk = i=1
π̄i , resulting in Equation (28).

In Lemma C.6, we describe the set of all subset permutations π̄ of elements i ∈ Sk by Πyk . Put differently, we makep(yk | y<k ; ω, s) invariant to the ordering of elements i ∈ Sk by marginalizing over the probabilities of subset permutationspπ̄k (Xie & Ermon, 2019).
Using Lemmas C.5 and C.6, we propose the two-stage random partition p(Y ; ω, s). Since Y = [y1 , . . . , yK ]T , we calculatep(Y ; ω, s), the PMF of the two-stage RPM, sequentially using Lemmas C.5 and C.6, where we leverage the PL distribution
for permutation matrices p(π; s) to describe the probability distribution over subsets p(yk | y<k ; ω, s).

Proposition 2.1 (Two-Stage Random Partition Model). Given a probability distribution over subset sizes p(n; ω) withKn ∈ NK
0 and distribution parameters ω ∈ R+ and a PL probability distribution over random orderings p(π; s) withn×n
π ∈ {0, 1}
and distribution parameters s ∈ Rn+ , the probability mass function p(Y ; ω, s) of the two-stage RPM is givenbyp(Y ; ω, s) = p(y1 , . . . , yK ; ω, s) = p(n; ω)Xp(π; s)π∈ΠYwhere ΠY = {π : yk =Pνk +nki=νk +1 πi , k = 1, . . . , K}, and yk and νk as in Definition C.2.

13

(29)

Differentiable Set Partitioning

Proof. Using Lemmas C.5 and C.6, we writep(Y ) =p(y1 , . . . , yK ; ω, s) = p(y1 ; ω, s) · · · p(yK | {yj }j<K ; ω, s)Xp(π̄1 | n1 ; s)
= p(n1 ; ω)π̄1 ∈Πy1

(30)=p(n1 ; ω) · · · p(nK | {nK }j<K ; ω)XXp(π̄1 | n1 ; s) · · ·p(π̄K | {nj }j≤K ; s)
·

(31)



· · · p(nK | {nj }j<K ; ω)=p(n; ω) =p(n; ω)π̄K ∈ΠyKp(π̄K | {nj }j≤K ; s)

πK ∈ΠyKπ̄1 ∈Πy1XXπ̄1 ∈Πy1

X

···

X

πK ∈ΠyKp(π | n; s)p(π̄1 | n1 ; s) · · · p(π̄K | {nj }j≤K ; s)

(32)
(33)

π∈ΠY

=p(n; ω)X

p(π; s)

(34)

π∈ΠY

C.3. Approximating the Probability Mass Function
Lemma 2.2.

p(Y ; ω, s) can be upper and lower bounded as follows∀π ∈ ΠY : p(n; ω)p(π; s) ≤ p(Y ; ω, s) ≤ |ΠY |p(n; ω) max p(π̃; s)
π̃

(35)

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

proving the lower bound of Lemma 2.2.
On the other hand, can prove the upper bound in Lemma 2.2 by:p(Y ; ω, s) = p(n; ω)X

π ′ ∈Π

≤p(n; ω)X

π ′ ∈ΠYp(π ′ ; s)Ymax p(π; s)π∈ΠY=p(n; ω) max p(π; s)
π∈ΠY

X

1

π ′ ∈ΠY=|ΠY | · p(n; ω) max p(π; s)π∈ΠY≤|ΠY | · p(n; ω) max p(π; s)π

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.
C.4. The Differentiable Random Partition Model
We propose the DRPM p(Y ; ω, s), a differentiable and reparameterizable two-stage RPM.
14

Differentiable Set Partitioning

Lemma 2.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 inLemma 2.2 and to provide a reparameterization scheme for the two-stage approach in Definition C.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 2.2
are differentiable. By definition (see Section 2),|ΠY | =KY

|Πyk | =

k=1

K
Y

nk !.

k=1

Hence, |ΠY | can be computed given a reparametrized version nk , which is provided by the reparametrization trick for the
MVHG p(n; ω). Further, from Equation (8) we immediately see that the most probable permutation is given by the orderinduced 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:
yk =

νkX
+nk

πi ,

where νk =i=νk +1

k−1
X

nι

(36)

ι=1

The challenge here is that we need to be able to backpropagate through nk , which appears as an index in the sum. Let
αk = {0, 1}n , such that(
1 if νk < i ≤ νk+1
(αk )i =0 otherwise
Given such αk , we can rewrite Equation (36) with
yk =

n
X

(αk )i πi .

(37)

i=1

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


x−i+ϵ
,
fi (x; τ ) = στx−isuch that fi (x; τ ) ≈ 0 if x−i
τ < 0 and fi (x; τ ) ≈ 1 if τ ≥ 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 (τ ) = f (νk ; τ ) − f (νk−1 ; τ ),α̃k (τ ) ∈ [0, 1]n . Then, for τ → 0 we have α̃k (τ ) → αk . In practice, we cannot set τ = 0 since this would amount to adivision by 0. Instead, we can apply the straight-through estimator (Bengio et al., 2013) to the auxiliary function f (x; τ ) inorder to get α̃k ∈ {0, 1}n and use it to compute Equation (37).
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 nk . This allows a different path for computing αk which circumvents introducing yet anothertemperature parameter altogether. We refer to our code in the supplement for more details.
15

Differentiable Set Partitioning

Table 2. Total GPU hours per experiment. We report the cumulative training and testing hours to generate the results shown in the mainpart of this manuscript. We relied on our internal cluster infrastructure equipped with RTX2080Ti GPUs. Hence, we report the number ofcompute hours for this GPU-type.

Experiment

Computation Time (h)MTL (Section 3.1)
Partitioning of Generative Factors (Section 3.2)
Clustering (Appendix D.3)

100
480
100

Figure 4. Samples from the noisyMultiMNIST dataset with increasing noise ratio in the right task.

D. Experiments
In the following, we describe each of our experiments in more detail and provide additional ablations. All our experiments
were run on RTX2080Ti GPUs. Each run took 6h-8h (Variational Clustering), 4h-6h (Generative Factor Partitioning), or∼ 1h (Multitask Learning) respectively. We report the training and test time per model. Please note that we can only report
the numbers to generate the final results but not the development time.
D.1. Multitask Learning
D.1.1. M ULTI MNIST 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 theright 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 4.
D.1.2. I MPLEMENTATION & A RCHITECTURE
The multitask loss function for the MultiMNIST dataset is
L = w L LL + w R LR

(38)

where wL and wL are the loss weights, and LL and LR 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. wL = wR = 0.5. We use these loss weights forthe 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 aredefined as cross-entropy losses
Lt = −CtX

gtc log pc = −gtT log p

(39)

c=1

where CL = CR = 10 for MultiMNIST,P gt is a one-hot encoded label vector and p is a categorical vector of estimatedclass assignments probabilities, i.e. c pc = 1.
16

Differentiable Set Partitioning

ω, s

encϕ

Task LhθR

Task R

zL
zR

z

encθ

hθL

p(yL , yR ; ω, s)Figure 5. Overview of the multitask learning pipeline of the DRPM-MTL method.

The predictions for the individual tasks pt are given as
pt = hθt (z), wherez = encθ (x)

(40)
(41)

for a sample x ∈ X (see also Figure 5). We use an adaptation of the LeNet-5 architecture (LeCun et al., 1998) to themultitask learning problem (Sener & Koltun, 2018). Both DRPM-MTL and ULS use the same network encθ (·) with sharedarchitecture 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 thedigit using only a subset of z. DRPM-MTL uses the following prediction scheme
pt = hθt (zt ), where(42)zt = z ⊙ y t
yt = DRPM(ω, s)t = DRPM(encϕ (x))t

(43)
(44)

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 yL , yR ∼ p(yL , yR ; ω, s).
We then feed the masked latent representations zL ← z ⊙ yL and zR ← z ⊙ yR into the task specific classification networkshθL (zL ) and hθR (zR ) 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.
D.1.3. T RAINING
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 D.3.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.
D.1.4. C ELEBA FOR MTL
In addition to the experiment shown in Section 3.1, 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 toMultiMNIST (see Section 3.1) 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 D.1.2 and D.1.3). We use
the pipeline of Sener & Koltun (2018) with a ResNet-based encoder to map an image to a representation of d = 6417

Differentiable Set Partitioning

Table 3. 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 3a, we form a MTL experiment with 10
different tasks. In Table 3b, we form a MTL experiment with 20 different tasks where the first 10 tasks are the same as in the 10 tasks
experiment. We train both methods for 50 methods using a learning rate of 0.0001 and a batch size of 128. The temperature annealingschedule 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.
(b) 20 Tasks(a) 10 Tasks

T0
T1
T2
T3
T4
T5
T6
T7
T8
T9
avg(Tasks)

ULS

DRPM

92.0±0.583.8±0.480.2±0.581.9±0.8
98.5±0.2
95.2±0.2
80.0±1.482.0±0.3
89.7±0.7
94.6±0.587.8±0.3

92.4±0.583.7±0.280.2±0.482.2±0.6
98.5±0.1
95.3±0.282.4±0.482.2±0.2
90.7±0.2
95.0±0.2
88.3±0.1

T0
T1
T2
T3
T4
T5
T6
T7
T8
T9
T10
T11
T12
T13
T14
T15
T16
T17
T18
T19
avg(Tasks)

ULS

DRPM

92.4±0.7
83.7±0.679.9±0.682.4±0.5
98.6±0.1
95.2±0.1
82.0±1.382.5±0.1
90.1±0.994.7±0.2
95.9±0.1
84.9±0.1
91.0±0.4
94.7±0.1
95.4±0.3
99.2±0.0
95.8±0.3
97.3±0.3
91.2±0.3
87.0±0.3
90.7±0.2

93.0±0.2
83.9±0.780.1±0.4
83.0±0.7
98.6±0.1
95.5±0.0
84.4±0.482.8±0.2
91.0±0.4
95.1±0.1
95.9±0.1
84.6±0.3
91.6±0.2
94.9±0.1
96.0±0.1
99.2±0.1
96.0±0.1
97.5±0.2
91.2±0.187.3±0.2
91.1±0.1

dimensions. For architectural details, we refer to Sener & 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 nT different subsets, which are then fed to the respective task networks. nT is the number of tasks.
Compared to the MultiMNIST experiment (see Appendix D.1.2), we introduce an additional regularization for the DRPMMTL method. The additional regularization is based on the upper bound in Lemma 2.2 and is penalizing size of |ΠY | for agiven n. Hence, the loss function changes tonT
1 X
Lt + λ · LregnT t=1!
nT
nT
X
Y
log Γ(nt + 1)
nt ! =
where Lreg = log

L=

t=1

(45)

(46)

t=1

1
For both versions of the experiment (i.e. nT = 10 and nT = 20), we set λ = 0.015 ≈ 64
, which is the number of elements
we want to partition. The task losses Lt 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 thefirst 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 shared between the two experiment versions.
18Differentiable Set Partitioning

Table 3 shows the results of 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 eventhough it has only access to d/nT = 64/20 = 3.2 dimensions on average. On the other hand, the ULS method can accessthe full set of 64 dimensions for every single task.
D.2. Variational Partitioning of Generative Factors

E

E

 
4


 

7
zv 1 =
1


 

2
 
4


 

3
zv2 =
8


 

2

Figure 6. Motivation for the Partitioning of Generative Factors under weak supervision. The knowledge about the data collection
process provides a weak supervision signal. We have access to a dataset of pairs of images of the same robot arm with a subset ofshared generative factors (in red). We want to learn the shared and independent generative factors in addition to learning from thedata. The images of the robot arms are taken from Locatello et al. (2020) but originate from the mpi3d toy dataset (see https://github.com/rr-learning/disentanglement_dataset). The image is from Sutter et al. (2023) and their ICLR 2023
presentation video (see https://iclr.cc/virtual/2023/poster/10707).

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 subsetof 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 twolatent factors, depicted in red, out of a total of four factors. Please note that we chose four generative in Figure 6 only forillustrative reason as there are seven generative factors in the mpi3d toy dataset. Hence, the goal of learning under weaksupervision is not 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 groupstructure holds promise for learning more generalizable and disentangled representations (see (e.g., Locatello et al., 2020)).
D.2.1. G ENERATIVE M ODELWe assume the following generative model for DRPM-VAE
p(X) ==Z

Zz

p(X, z)dz(47)p(X | z)p(z)dz

(48)

z

where z = {zs , z1 , z2 }. The two frames share an unknown number ns of generative latent factors zs , and an unknown
number, n1 and n2 , of independent factors z1 and z2 . The RPM infers nk and zk using Y . Hence, the generative model
19

Differentiable Set Partitioning

extends top(X) =
=
=

Z

Z

Z

p(X | z)zXp(z | Y )p(Y )dz

Yp(x1 , x2 | zs , z1 , z2 )zXp(z | Y )p(Y )dz

Yp(x1 | zs , z1 )p(x2 | zs , z2 )zs ,z1 ,z2Xp(zs , z1 , z2 | Y )p(Y )dzs dz1 dz2

(49)

Y

Figure 7 shows the generative and inference models assumptions in a graphical model.
X
x2

x1

x1

p(x1 | zs , z1 )x2

Xp(x2 | zs , z2 )q(Y | x1 , x2 )
Y

z1

zs

q(z2 | Y, x2 )q(z1 | Y, x1 )

z2

z2

z1

p(zs , z1 , z2 | Y )q(zs | Y, x1 , x2 )Yzs(a) Generative Model(b) Inference Model

Figure 7. Graphical Models for DRPM-VAE models in the weakly-supervised experiment.

D.2.2. DRPM ELBO
We derive the following ELBO using the posterior approximation q(z, Y | X)q(z, Y | X)
LELBO (X) = Eq(z,Y |X) log p(X | z, Y ) − logp(z, Y )q(z | Y, X)q(Y | X)
= Eq(z,Y |X) log p(X | z) − log
p(z)p(Y )q(z | Y, X)
q(Y | X)
= Eq(z,Y |X) log p(x1 , x2 | z) − log
− log
p(z)
p(Y )

= Eq(z,Y |X) [log p(x1 | zs , z1 )] − Eq(z,Y |X) [log p(x2 | zs , z2 )]q(zs , z1 , z2 | Y, X)
q(Y | X)
− Eq(z,Y |X) log− Eq(z,Y |X) log
p(zs , z1 , z2 )
p(Y )

(50)
(51)
(52)

(53)

Following Lemma 2.2, we are able to optimize DRPM-VAE using the following ELBO LELBO (X):LELBO ≥ Eq(z,Y |X) [log p(x1 | zs , z1 )] − Eq(z,Y |X) [log p(x2 | zs , z2 )]q(zs , z1 , z2 | Y, X)− Eq(z,Y |X) log
p(zs , z1 , z2 ) |ΠY | · q(n | X; ω)− Eq(Y |X) log
p(n; ωp )p(πY ; sp )− log max q(π̃ | X; s) ,
π̃

20

(54)
(55)
(56)
(57)

Differentiable Set Partitioning

where πY is the permutation that lead to Y during the two-stage resampling process. Further, we want to control theregularization strength of the KL divergences similar to the β-VAE (Higgins et al., 2016). The ELBO L(X) to be optimized
can be written as
LELBO = Eq(z,Y |X) [log p(x1 | zs , z1 )] + Eq(z,Y |X) [log p(x2 | zs , z2 )]q(zs , z1 , z2 | Y, X)− β · Eq(z,Y |X) log
p(zs , z1 , z2 ) |ΠY | · q(n; ω(X))− γ · Eq(Y |X) log
p(n; ωp ) maxπ̃ q(π̃; s(X))− δ · Eq(Y |X) log
p(πY ; sp )

(58)
(59)
(60)
(61)

where s(X) and ω(X) denote distribution parameters, which are inferred from X (similar to the Gaussian parameters in
the vanilla VAE).
As in vanilla VAEs, we can estimate the reconstruction term in Equation (54) with MCMC by applying the reparametrization
trick (Kingma & 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 (54). Similarly, we can sample from q(Y | X) L times. We use L = 1 to estimate all expectations
in LELBO .
D.2.3. I MPLEMENTATION AND H YPERPARAMETERS
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). HGVAE 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.
x1

E

x2

qφ (zs , z1 | x1 )View
Aggregationx̂1Dqφ (zs , z2 | x2 )x̂2Figure 8. Setup for the weakly-supervised experiment. The three methods differ only in the View Aggregation module.

We adapted Figure 8 from Sutter et al. (2023). It shows the baseline architecture, which is used for all methods. As alreadystated in the main part of the paper, the methods only differ in the View Aggregation module, which determines theshared and independent latent factors. Given a subset S of shared latent factors, we haveqφ (zi | xj ) = avg(qφ (zi | x1 ), qφ (zi | x2 ))qφ (zi | xj ) = qφ (zi | xj )

∀ i∈S

(62)

else

(63)

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 KLpw between the latent vectors of the two views.
KLpw (x1 , x2 ) =1
1KL[q(z1 | x1 )||q(z2 | x2 )] + KL[q(z2 | x2 )||q(z1 | x1 )]
2
2

(64)

where q(zj | xj ) are the encoder outputs of the respective images. We do not average or sum across dimensions in thecomputation of KLpw (·) such that the KLpw (·) is d-dimensional, where d is the latent space size. The encoder E in
Figure 8 maps to µ(xj ) and σ(xj ) of a Gaussian distribution. Hence, we can compute the KL divergences above in closed
21

Differentiable Set Partitioning

form. Afterwards, we feed the pairwise KL divergence KLpw to a single fully-connected layer, which maps from d to K
valueslog ω = F C(KLpw (x1 , x2 ))(65)

where d = 10 and K = 2 in this experiment. d is the total number of latent dimensions and K is the number of groups inthe latent space. To infer the scores s(X) we again rely on the pairwise KL divergence KLpw . Instead of using anotherfully-connected layer, we directly use the log-values of the pairwise KL divergencelog s = log KLpw (x1 , x2 )(66)

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 D.3). We anneal thetemperature τ from 1.0 to 0.5 over the complete training time.
D.3. Variational Clustering with Random Partition Models
The proposed DRPM cannot only be used to partition neurons in specific network layers but also to partition entire samples,e.g., for clustering. In this additional experiment, we introduce a new version of a Variational Autoencoder (VAE, Kingma &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 alatent vector z ∈ Rl , where l ∈ N is the latent space size. Traditional VAEs would then assume that all latent vectors z
are generated by a single Gaussian prior distribution N (0, Il ). 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 ∈ Rl . Further, note that similar to an
urn model (Appendix B.2), 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 thenumber 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 thedependency 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 andsampled from our DRPM, i.e., Y ∼ p(Y ; ω, s). Given an assignment y from Y , we can sample the respective latent variablez, where z ∼ N (µy , diag(σy )), z ∈ Rl . 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):X
X
Eq(Y |X) [KL[q(z|x)||p(z|Y )]] − KL[q(Y |X)||p(Y )]Eq(z|x) [log p(x|z)] −LELBO =x∈X

x∈X

Note that computing KL[q(Y |X)||p(Y )] directly is computationally intractable, and we need to upper bound it according
to Lemma 2.2. We provide the detailed derivation of the ELBO in Appendix D.3.1.
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 theembeddings 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 4. As can be seen, we outperform all baselines, indicating that modelingthe 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 Figures 11 and 12, we present the result of sampling a partition
and the corresponding generations from the respective clusters after training the DRPM-VC on MNIST and FMNIST. 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. To further investigate the learned structures, we can also directly sample from each of the22

Differentiable Set Partitioning

Table 4. We compare the clustering performance of the DRPM-VC on test sets of MNIST and FMNIST between Gaussian Mixture
Models (GMM), GMM in latent space (Latent GMM), and Variational Deep Embedding (VADE). We measure performance in terms of
the Normalized Mutual Information (NMI), Adjusted Rand Index (ARI), and cluster accuracy (ACC) over five seeds and put the bestmodel in bold.

MNIST
GMM
L ATENT GMM
VADE
DRPM-VCFMNIST

NMI

ARI

ACC

NMI

ARI

ACC

0.32±0.01
0.86±0.02
0.84±0.01
0.89±0.01

0.22±0.02
0.83±0.06
0.76±0.05
0.88±0.03

0.41±0.01
0.88±0.07
0.82±0.04
0.94±0.02

0.49±0.01
0.60±0.00
0.56±0.02
0.64±0.00

0.33±0.00
0.47±0.01
0.40±0.04
0.51±0.01

0.44±0.01
0.62±0.01
0.56±0.03
0.65±0.00

learned priors without first sampling a partition. We show some examples of this for both MNIST and FMNIST in Figures 13
and 14. We can again see that the DRPM-VC learns accurate cluster representations since each of the samples seems tocorrespond to one of the classes in the datasets. Further, the clusters also seem to capture the diversity in each cluster, as wesee a lot of variety across the generated samples.
D.3.1. L OSS F UNCTION
As mentioned in above, 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 wehave z ∼ N (µy , σyT Il ), z ∈ Rl . 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):!
Z X
log(p(X)) = log
p(X, Y, Z)dZY≥ Eq(Z,Y |X) log:= LELBO (X)p(X|Z)p(Z|Y )p(Y )q(Z, Y |X)

We then assume that we can factorize the approximate posterior as follows:Yq(z|x)q(Z, Y |X) = q(Y |X)x∈X

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 different cluster assignments. This allows us toleverage dependencies between samples from the dataset. Hence, we can rewrite the ELBO as follows:
LELBO (X) =Eq(Z|X) [log(p(X|Z))]− Eq(Y |X) [KL[q(Z|X)||p(Z|Y )]]− KL[q(Y |X)||p(Y )]X
Eq(z|x) [log p(x|z)]
=
x∈X

−

X

Eq(Y |X) [KL[q(z|x)||p(z|Y )]]x∈X− KL[q(Y |X)||p(Y )]See Figure 9 for an illustration of the generative process and the assumed inference model. Since computing P (Y ) andq(Y |X) is intractable, we further apply Lemma 2.2 to approximate the KL-Divergence term in LELBO , leading to the23Differentiable Set Partitioning

following lower bound:
LELBO ≥

X

Eq(z|x) [log p(x|z)]

(67)

x∈X

−

X

Eq(Y |X) [KL[q(z|x)||p(z|Y )]]

(68)



(69)

x∈X



|ΠY | · q(n; ω(X))− Eq(Y |X) log
p(n; ω)p(πY ; s)
− log max q(π̃; s(X)) ,
π̃

(70)

where πY is the permutation that lead to Y during the two-stage resampling process. Further, we want to control theregularization strength of the KL divergences similar to the β-VAE (Higgins et al., 2016). Since the different terms havedifferent regularizing effects, we rewrite Equations (69) and (70) and weight the individual terms as follows, leading to our
final loss:
L := −

X

Eq(z|x) [log p(x|z)]

(71)

x∈X

+β·

X

Eq(Y |X) [KL[q(z|x)||p(z|Y )]]

(72)

x∈X


 
|ΠY | · q(n; ω(X))+ γ · Eq(Y |X) log
p(n; ω) maxπ̃ q(π̃; s(X))+ δ · Eq(Y |X) log
p(πY ; s)

(73)
(74)

Figure 9. Generative model of the DRPM clustering model. Generative paths are marked with thin arrows, whereas inference is in bold.

24

Differentiable Set Partitioning

Figure 10. Autoencoder architecture of the DRPM-VC model.

D.3.2. A RCHITECTURE
The model for our clustering experiments is a relatively simple, fully-connected autoencoder with a structure as seen in
Figure 10. 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 therespective 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) ∈ Rl . To understand the architecture choice for the DRPM parameters, let
us first take a closer look at Equation (72). 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:∀xi ∈ X : si (xi ) = u · (K − arg min (KL[N (µ(xi ), diag(σ(xi ))||N (µk , diag(σk ))]))kN
1 X
ω(X) =|X|x∈X(

N (x|µk , diag(σk ))PK

k′ =1 N (x|µk , diag(σk ))
′

′

)K

,

k=1

where u is a scaling constant that controls the probability of sampling the most likely partition. Note that ω and s minimize
Equation (72) 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 thereconstructed sample x̂.
D.3.3. T RAINING
As in vanilla VAEs, we can estimate the reconstruction term in Equation (71) with MCMC by applying the reparametrization
trick (Kingma & Welling, 2014) to q(z|x) to sample M samples z (i) ∼ q(z|x) and compute their reconstruction error toestimate Equation (71). Similarly, we can sample from q(Y |X) L times to estimate the terms in Equations (72) to (74), such
25

Differentiable Set Partitioningthat we minimizeL̃ := −M
X 1 X
log p(x|z (i) )
M i=1

x∈X

L

β XX
KL[q(z|x)||p(z|Y (i) )]·Lx∈X i=1
Lγ X|ΠY (i) | · q(n(i) ; ω(X))
log
+ ·
L i=1
p(n(i) ; ω)


L
δ X
maxπ̃ q(π̃; s(X))
log
+ ·
L i=1
p(πY (i) ; s)+In our experiments, we set M = 1 and L = 100 since the MVHG and PL distributions are not concentrated around theirmean 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, weapplied the exponential schedule that was originally proposed together with the Gumbel-Softmax trick (Jang et al., 2016;
Maddison et al., 2017), i.e., τ = max(τf inal , exp(−rt)), where t is the current training step and r is the annealing rate.
log(τf inal )−log(τinit )For our experiments, we choose r =in order to annealing over 100000 training step. Like Jang et al.
100000(2016), we set τinit = 1 and τf inal = 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
GMM to 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.
To optimize the DRPM-VC in our experiments, we used the AdamW (Loshchilov & 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 thefirst three layers of the encoder when training the DRPM-VC, giving us much better results.
Finally, when training the VADE baseline and the DRPM-VC on FMNIST, we often observe a local optimum where theprior 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.

26Differentiable Set Partitioning

Figure 11. Additional partition samples from the DRPM-VC trained on MNIST. The different sets of each partition match each of thedigits very well, even after repeatedly sampling from the model.

27

Differentiable Set Partitioning

Figure 12. Additional partition samples from the DRPM-VC trained on FMNIST. Most clusters accurately represent one of the clothing
categories and generate new samples very well. The only problem is with the handbag class, where the DRPM-VC learns two different
clusters for different kinds of handbags (cluster 5 and 6).

28

Differentiable Set Partitioning

Figure 13. Various samples from each of the generative priors. Each prior learns to represent one of the digits. Further, we see a lot ofvariation between the different samples, suggesting that the clusters of the DRPM-VC manage to capture some of the diversity present in29
the dataset.

Differentiable Set Partitioning

Figure 14. Various samples from each of the generative priors. Each prior learns to represent one of the digits. The DRPM-VC learns nicerepresentations that provide coherent generations of most classes. For high-heels (cluster 4), generating new samples seems difficult due
30
to the heterogeneity within that class.