Learning Selection Bias and Group Importance: Differentiable Reparameterization for the Hypergeometric Distribution

Partitioning a set of elements into a given number of groups of a priori unknown sizes is a critical task in many applications. It can be characterized by a hypergeometric distribution, which describes biased sampling without replacement based on the relative importance between classes of samples. Due to hard constraints, this discrete distribution is not differentiable in its standard formulation, prohibiting its use in modern machine learning frameworks. Hence, previous works mostly fall back on suboptimal heuristics or simplified assumptions. In this work, we propose a differentiable reparameterization trick for the multivariate noncentral hypergeometric distribution. We introduce reparameterizable gradients to enable learning of the importance or the selection bias between groups. We highlight the applicability and usability of the proposed formulation in two different experiments: weakly-supervised learning and clustering.