Variational Partitioning

Thomas M. Sutter, Alain Ryser, Joram Liebeskind, Julia E Vogt

Keywords: random partitions, weakly-supervised learning, representation learning

TL;DR: We propose a differentiable random partition model to be integrated into modern gradient-based optimization pipelines.

Abstract:

Partitioning a set of elements into an unknown number of mutually exclusive subsets is essential in many machine-learning problems. However, assigning elements to an unknown and discrete number of subsets is inherently non-differentiable, prohibiting end-to-end gradient-based optimization of parameters. We propose a novel two-step method for learning distributions over partitions, including a reparametrization trick, to allow the inclusion of partitions in variational inference tasks.
Our method works by first inferring the number of elements per subset and then sequentially filling these subsets in an order learned in a second step. We highlight the versatility of our general-purpose approach on two different experiments: multitask learning and unsupervised conditional sampling.