Differentiable Random Partition Models

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.

Paper

Similar papers

Peer review

Reviewer Ynzj7/10 · confidence 2/52023-07-03

Summary

This paper proposes a probabilistic model for partitioning a non-i.i.d. set of points into a fixed number of subsets, such that the probabilistic model can be fit to the data by variational inference, using the reparameterization trick to enable computation of gradients. The probabilistic model for partitioning is composed of two processes: one random process, using a multivariate hypergeometric distribution, for determining the sizes of the subsets, and a another random process, which assigns a random ordering to the points according to the Plackett-Luce distribution over permutations, which determines the assignment of points to subsets. To implement the inference, an evidence lower bound for the marginal probability p(X) is derived. Three additional methods utilizing the random partitioning model are presented. One is a method for clustering the data, one is a method for inferring shared and unique latent factors for a pair of data points, and another is a method for partitioning final-layer neurons by task in a multi-task learning setup.

Strengths

The paper shows how random partitioning models have many applications in machine learning, and explains the advantages of this particular approach. A large number of technical optimizations, such as lower bounds of the loss function, are developed, which seem essential for making the method usable.

Weaknesses

The motivation for the multivariate hypergeometric distribution could benefit from more elaboration.

Questions

* Rather than MVHG, wouldn't it be simpler to use a multinomial distribution instead? * Under what data generation conditions should the multivariate hypergeometric distribution be advantageous to the multinomial? * 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.?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Yes. The authors mention the memory inefficiency and the possibility of improving the approximations used in the ELBO, which do seem be the most important limitations of the method.

Reviewer 6PXU7/10 · confidence 4/52023-07-06

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

-

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations are addressed.

Reviewer jtds6/10 · confidence 4/52023-07-07

Summary

The authors tackle the problem of partitioning a set of elements into a variable number of disjoint subsets in a differentiable manner. The method starts by first inferring the number of elements per subset. Consequently, they fill these subsets in a learned order.

Strengths

- The problem of partitioning a set of elements into variable number of disjoint sets while remaining amenable to back-prop is indeed an important and timely problem. The authors spend the time to attempt to develop an approach top the problem starting from first principles, which I appreciate.

Weaknesses

- I found the writing to be quite a bit confusing and very notation heavy at times, with very little by way of intuition or meaning of the equations. A running example, e.g. an urn with marbles, would've greatly helped with the exposition. - After reading the paper, I'm still not really sure how the authors overcome the non-differentiability of sampling. - 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?) - The related work can be further improved as it's missing a lot of work regarding differentiable discrete distributions e.g. [1], [2], [3] References: [1] Marin Vlastelica, P., Anselm Paulus, Vít Musil, Georg Martius and Michal Rolinek. “Differentiation of Blackbox Combinatorial Solvers.” ArXiv abs/1912.02175 (2019) [2] Mathias Niepert, Pasquale Minervini and Luca Franceschi. “Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions.” Neural Information Processing Systems (2021). [3] Kareem Ahmed, Zhe Zeng, Mathias Niepert and Guy Van den Broeck. “SIMPLE: A Gradient Estimator for k-Subset Sampling.” ICLR 2023.

Questions

- In equation 2, what purpose does the superscript in $\omega_k^{n_k}$ serve? It is my understanding that the weights are specific to each class, not each instance? Also, am I correct in my understanding that $P_0$ here is intractable? - 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? - If I'm not mistaken, Equation (3) is very similar to weighted reservoir sampling (see e.g. [1])? - It is not clear how you overcome the non-differentiability of sampling. Do you use Gumbel-Softmax when sampling both the permutation matrix as well as an order? - 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$? References: [1] Sang Michael Xie and Stefano Ermon. “Reparameterizable Subset Sampling via Continuous Relaxations.” International Joint Conference on Artificial Intelligence (2019).

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Adequately addressed.

Reviewer WEvf6/10 · confidence 3/52023-07-09

Summary

This paper proposes a two-stage differentiable partitioning method where the first stage outputs number of elements in each cluster and the second stage assigns samples by ordering them through a permutation matrix. The differentiability in their formulation is achieved via gumbel-softmax trick for distributions sampling. The paper shows the performance of their method for 1-a variational clustering task on MNIST nad FMNIST, 2-learn a partition of generative factors on mpi3d dataset, 3-neuron grouping task for multi-task learning on MultiMNIST.

Strengths

- The proposed formulation and its usage in a differentiable setting is novel from my knowledge. Further applying this method for grouping neural networks for MTL and other tasks are novel from my knowledge. However, I am not an expert in this field, so I may not be aware of some of the previous work. - 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. - Although the applicability to a less trivial setting -- with larger datasets and harder tasks -- are not clear, the method shows promise to be improved and has sufficient contribution. Sampling from cluster numbers and permutation matrix distributions in a differentiable way to create partitioning is an interesting idea and application.

Weaknesses

- The experiment results need more clarification for a couple of concerns I have to understand the impact and contribution better: - 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. - 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. - Why didn't you try any dataset other than MNIST for clustering tasks, for instance STL-10 or Reuters datasets at least? - 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]. 1) Yang, Xu, et al. "Deep spectral clustering using dual autoencoder network." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019. 2) Cai, Jinyu, et al. "Efficient deep embedded subspace clustering." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. 3) 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.

Questions

- A quadratic relationship wrt n is mentioned in the limitations. How does your method's computational complexity compare with the baselines? Is this why you did not go beyond MNIST? - Why you experimented with fixed encoder and different initialization? (in the supplementary)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Reviewer jtds2023-08-15

Thanks for you response. [Writing Clarity] I appreciate the technical nature of the topic, requiring sufficient notation to describe the problem. I do not, however, believe that contradicts with the clarity of exposition. For instance, equation (2) is In my opinion, stated without any commentary beyond it corresponding to the probability of sampling $\mathbf{n}$ and $P_0$ being the normalization constant. Another instance of this is your re-statement of the problem as an urn model. Unless I'm mistaken, I do not see the need to denote by $n$ both the number of different colors in the urn as well as the number of marbles per color and then *again* the sample size? [Differentiability] As a paper on *differentiable* subset sampling I feel like perhaps the means by which you render such an operation differentiable should come earlier, with more discussion, instead of simply being relegated to the appendix. Other than that, I believe you've answered my concerns. I'm raising my score to a weak accept.

Authorsrebuttal2023-08-15

Official comment by the authors

Thanks for your reply! We would like to respond to your comments and clarify them. > For instance, equation (2) is In my opinion, stated without any commentary beyond it corresponding to the probability of sampling $\mathbf{n}$ and $P_0$ being the normalization constant. We agree that we keep the description of equation (2) short. However, due to the page limit, we decided to keep the preliminaries relatively compact and refer to the appendix for more details. > Unless I'm mistaken, I do not see the need to denote by n both the number of different colors in the urn as well as the number of marbles per color and then again the sample size? Thanks for your question. We want to clarify that we denote the number of different colors, the number of marbles per color, and the sample size by $n$ because in our version of the urn model (partition model) they all have the same value. What often leads to confusion is the number of subsets. Note that, at most, there can be as many subsets as there are elements in the set we want to partition, namely if each element is placed in a separate subset. Since each color stands for one subset, there are a maximum of $n$ colors. Our paper follows previous work [a], which introduced $K$ as the number of subsets where $K \leq n$. Further, in our preliminaries section of the MVHG distribution and equation (2), we initially follow the notation of [4] and state that there are $m_k$ marbles of color $k$ in the urn. In the case of a partition, each subset can contain, at most, as many elements as there are elements in the set we want to partition. Thus, we set the number of marbles for each color to $n$, i.e., $\forall k\in[K]: m_k=n$. Hence, we decided to keep $n$ as a central part of our notation instead of introducing additional notation that denotes the same value. [a] Mansour, T., & Schork, M. (2015). Commutation relations, normal ordering, and Stirling numbers. CRC Press. > As a paper on differentiable subset sampling I feel like perhaps the means by which you render such an operation differentiable should come earlier, with more discussion, instead of simply being relegated to the appendix. We decided to focus the main part of our paper on the idea of the two-stage procedure, which we view as our main contribution, and keep the technical details brief. In our opinion, this improves the clarity of our paper and helps in understanding the core part of our contribution. However, we see your point, and without space constraints, we would have added more discussion on the differentiability aspect of the proposed formulation.

Reviewer Ynzj2023-08-16

Additional questions

Thank you for the answers to my questions. From reading the other reviews and discussions, I realize that you do not cluster the entire dataset in one step, but rather apply the clustering one batch at a time. I think this fact should be made more apparent in the main paper. So is $\omega$ the only parameter that is learned across batches? And do you think it would make sense to include the same point in multiple batches, and then to aggregate the results somehow (similarly to including the same pixel in multiple overlapping windows for image segmentation)?

Authorsrebuttal2023-08-17

Thanks for your reply and the additional questions! >From reading the other reviews and discussions, I realize that you do not cluster the entire dataset in one step, but rather apply the clustering one batch at a time. I think this fact should be made more apparent in the main paper. Indeed, we cluster the dataset one batch at a time. We briefly mention this at the beginning of Section 5.1 when we make the example with the urn model. However, we acknowledge that this fact may not be entirely apparent and will make sure to add further clarification. >So is $\omega$ the only parameter that is learned across batches? Generally, we treat the parameters of the prior distributions $p(Y)$ and $p(\mathbf{z})$ as global parameters. We set the corresponding $\mathbf{s}=const.$ and learn $\mathbf{\omega}$, and $\mathbf{\mu}_k$, $\mathbf{\sigma}_k$ for $k\in\\{1, \ldots, K\\}$ globally across batches. On the other hand, we learn all parameters of the approximate posterior distributions $q(\cdot|X)$ as functions of the batch $X$. Finally, we’d like to note that there are no constraints on how the DRPM parameters $\mathbf{s}$ and $\mathbf{\omega}$ should be learned. As demonstrated in the clustering experiment, we can learn them both globally and at batch level. Further, if we want to partition a set on instance-level, we can also learn the parameters per sample as presented in the generative factor partitioning and multitask experiments. >And do you think it would make sense to include the same point in multiple batches, and then to aggregate the results somehow (similarly to including the same pixel in multiple overlapping windows for image segmentation)? We thank the reviewer for their suggestion! Indeed, it could be interesting to see whether an adaption with overlapping batches could further improve clustering performance in the variational clustering experiment. However, we don’t think such an extension is feasible for the other experiments. As mentioned above, in the other experiments, we do partitioning at instance-level and aggregating batches would thus not influence the predictions.

Reviewer Ynzj2023-08-18

Thanks for the clarifications

I am pleased that you will clarify how batching works in the main paper. As for my review, I am keeping my score as is.

Area Chair GARg2023-08-19

Thanks for the response

Dear authors, Thanks for your efforts to clarify the concerns raised by the reviewer. Although the reviewer has not responded yet, I'd like to assure you that your response will be incorporated into the reviewers' discussion and my final recommendation. Best wishes, AC

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC