# Differentiable Random Partition Models

arXiv: https://arxiv.org/abs/2305.16841 · HTML: https://mlsanity.com/paper/2305.16841/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 76 references that are in the map.
- First posted on arXiv 26 May 2023.

## Builds on (most cited first, 30 of 36)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [Dynamic Routing Between Capsules](https://mlsanity.com/paper/1710.09829/) (2017)
- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [Weakly-Supervised Disentanglement Without Compromises](https://mlsanity.com/paper/2002.02886/) (2020)
- [Differentiation of Blackbox Combinatorial Solvers](https://mlsanity.com/paper/1912.02175/) (2019)
- [Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders](https://mlsanity.com/paper/1611.02648/) (2016)
- [Variational Deep Embedding: An Unsupervised and Generative Approach to Clustering](https://mlsanity.com/paper/1611.05148/) (2016)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [Multi-Level Variational Autoencoder: Learning Disentangled Representations from Grouped Observations](https://mlsanity.com/paper/1705.08841/) (2017)
- [On the Transfer of Inductive Bias from Simulation to the Real World: a New Disentanglement Dataset](https://mlsanity.com/paper/1906.03292/) (2019)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [Efficient Parameter-free Clustering Using First Neighbor Relations](https://mlsanity.com/paper/1902.11266/) (2019)
- [Deep Spectral Clustering using Dual Autoencoder Network](https://mlsanity.com/paper/1904.13113/) (2019)
- [Reparameterizable Subset Sampling via Continuous Relaxations](https://mlsanity.com/paper/1901.10517/) (2019)
- [In Defense of the Unitary Scalarization for Deep Multi-Task Learning](https://mlsanity.com/paper/2201.04122/) (2022)
- [SoftSort: A Continuous Relaxation for the argsort Operator](https://mlsanity.com/paper/2006.16038/) (2020)
- [Do Current Multi-Task Optimization Methods in Deep Learning Even Help?](https://mlsanity.com/paper/2209.11379/) (2022)
- [Twin Contrastive Learning for Online Clustering](https://mlsanity.com/paper/2210.11680/) (2022)
- [SIMPLE: A Gradient Estimator for k-Subset Sampling](https://mlsanity.com/paper/2210.01941/) (2022)
- [DeepPermNet: Visual Permutation Learning](https://mlsanity.com/paper/1704.02729/) (2017)

## Built on by (most cited first, 3 of 3)

- [Unity by Diversity: Improved Representation Learning in Multimodal VAEs](https://mlsanity.com/paper/2403.05300/) (2024)
- [Edge Probability Graph Models Beyond Edge Independency: Concepts, Analyses, and Algorithms](https://mlsanity.com/paper/2405.16726/) (2024)
- [You Only Train Once: Differentiable Subset Selection for Omics Data](https://mlsanity.com/paper/2512.17678/) (2025)

## Abstract (arXiv)

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.