Learning Group Importance using the Differentiable Hypergeometric Distribution
Authors: Thomas M. Sutter
Laura Manduchi
Alain Ryser
Julia E Vogt
Venue: ICLR 2023 notable top 25% (2023)
Keywords: hypergeometric distribution
weakly-supervised learning
reparameterization trick
group importance
variational clustering
gumbel softmax

Abstract: Partitioning a set of elements into subsets of a priori unknown sizes is essential in many applications. These subset sizes are rarely explicitly learned - be it the cluster sizes in clustering applications or the number of shared versus independent generative latent factors in weakly-supervised learning. Probability distributions over correct combinations of subset sizes are non-differentiable due to hard constraints, which prohibit gradient-based optimization. In this work, we propose the differentiable hypergeometric distribution. The hypergeometric distribution models the probability of different group sizes based on their relative importance. We introduce reparameterizable gradients to learn the importance between groups and highlight the advantage of explicitly learning the size of subsets in two typical applications: weakly-supervised learning and clustering. In both applications, we outperform previous approaches, which rely on suboptimal heuristics to model the unknown size of groups.

Published as a conference paper at ICLR 2023
LEARNING GROUP IMPORTANCE USING THE DIFFEREN- TIABLE HYPERGEOMETRIC DISTRIBUTION
Thomas M. Sutter, Laura Manduchi, Alain Ryser, Julia E. Vogt
Department of Computer Science
ETH Zurich
Switzerland
Thomas.sutten. laura manduchinn alain. ryser, julia. vogt] ethz.ch
ABSTRACT
Partitioning a set of elements into subsets of priori unknown sizes essential in many applications. These subset sizes are rarely explicitly Iearned be it the cluster sizes in clustering applications or the number of shared versus independent generative latent factors in weakly-supervisee learning. Probability distributions over correct combinations subset sizes are non- difffeerrentiblee due to hard con- straints, which prohibit gradient-basse optimization. In this work, propose the differentiable hypergeometric distribution. The hypergeometric distribution models the probability of different group sizes based their relative importance. We introduce reparameterizable gradients to learn the importance between groups and highlight the advantage of explicitly learning the size of subsets two typical applications: weakly-sup supeevs-e learning and clustering. In both applications, we outperform previous approaches, which rely on suboptimal heuristics to model the unknown size of groups.
INTRODUCTION
Many machine learning approaches rely on differentiable sampling procedures, from which the reparameterization trick for Gaussian distributions is best known (Kingma & Welling, 2014; Rezende et al., 2014). The non-differentiable nature of discrete distributions has long hindered their use in machine learning pipelines with end-to-end gradient- distribution (Maddison et al., 2017) or Gumbel-Softmax trick (Jang et al., 2) boosted the use of categorical distributions in stochastic networks. Unlike the high-variance gradients of score-based methods such as REINFORCE (Williams, 1992), these works enable reparameterized and low- variance gradients with respect to the categorical weights. Despite enormous progress in recent years, the extension to more complex probability distributions is still missing comes with trade-off regarding differentiability or computational speed (Huijben et al., 2021).
optimization. Only the concrete
The hypergeometric distribution plays vital role in various areas of science, such social and computer science and biology. The range of applications goes from modeling gene mutations and recommender systems to analyzing social networks (Becchetti al., 2011: Casiraghi al., 2016; Lodato et al., 2015). The hypergeometric distribution describes sampling without replacement and, therefore, models the number of samples per group given limited number of total samples. Hence, is essential wherever the choice of single group element influences the probability of the remaining elements being drawn. Previous work mainly uses the hypergeometric distribution implicitly model assumptions or as a tool to prove theorems. However, its hard constraints prohibited integrating the hypergeometric distribution into gradient-based optimization processes.
In this work, we propose the differentiable hypergeometric distribution. It enables the reparameteri- zation trick for the hypergeometric distribution and allows its integration into stochastic networks of modern, gradient-based learning frameworks. In turn, we learn the size of groups by modeling their relative importance in an end-to-end fashion. First, we evaluate our approach by performing Saanogorov-Smmmoon test, where we compare the proposed method to a dinedffeerettablee refer- ence implementation. Further, we highlight the advantages of our new formulation two different applications, where previous work failed to learn the size of subgroups of samples explicitly. Our

---

Published as a conference paper at ICLR 2023
first application is weekky-suupeviin learning task where two images share unknown number of generative factors. The differentiable hypergeometric distribution learns the number of shared and independent generative factors between paired views through gradient-baset optimization. In contrast, previous work has to infer these numbers based on heuristics or rely on prior knowledge about the connection between images. Our second application integrates the hypergeometric distribution into a variational clustering algorithm. We model the number of samples per cluster using an adaptive hypergeometric distribution prior. By doing so, we overcome the simplified i.i.d. assumption and establish dependency structure between dataset samples. The contributions of our work are the following: i) we introduce the differentiable hypergeometric distribution, which enables its use for gradient-basee optimization, ii) we demonstrate the accuracy of our approach by evaluating it against a reference implementation, and iii) we show the advantages of explicitly learning the size of groups in two different applications, namely weakly-supervised learning and clustering.
2 RELATED WORK
In recent years, finding continuous relaxations for discrete distributions and non-differentiable algorithms to integrate them into differentiable pipelines gained popularity. Jang et al. (2016) and Maddison et al. (2017) concurrently propose the Gumbel-Softmax gradient estimator. It enables reparameterized gradients with respect to parameters of the categorical distribution and their use in differentiable models. Methods to select & elements instead of only one- are subsequently introduced. Kool et al. (2019; 2020a) implement sequential sampling without replacement using a stochastic beam search. Kool et al. (2020b) extend the sequential sampling procedure to a reparameterizable estimator using REINFORCE. Grover et al. (2019) propose a relaxed version of a sorting procedure, which simultaneouslyy serves differentiable and reparameterizable top-k element selection procedure. Xie & Ermon (2019) propose a relaxed subset selection algorithm to select a given number out of n elements. Paulus al. (2020) generalize stochastic softmax tricks to combinatorial spaces." Unlike Kool et al. (2020b), who also use a sequence of categorical distributions, the proposed method describes a differentiable reparameterization for the more complex but well-defined hypergeometric distribution. Differentiable reparameterizations of complex distributions with learnable parameters enable new applications, as shown in Section
The classical use case for the hypergeometric probability distribution is sampling without replacement, for which urn models serve as the standard example. The hypergeometric distribution has previously been used as modeling distribution in simulations of social evolution (Ono et al., 2003; Paolucci et al., 2006; Lashin et al., 2007), tracking of human neurons and gene mutations (Lodato et al., 2015; 2018), network analysis (Casiraghi et al., 2016), and recommender systems (Becchetti et al., 2011). Further, it is used as a modeling assumption in submodular maximization (Feldman et al., 2017; Harshaw et al., 2019), multimodal VAEs (Sutter & Vogt, 2021). k-means clustering variants (Chien et al., 2018), or random permutation graphs (Bhattacharya & Mukherjee, 2017). Despite not being differentiable, current sampling schemes for the multivariate hypergeometric distribution are a trade-off between numerical stability and computational efficiency (Liao & Rosen, 2001; Fog, 2008a;b).
PRELIMINARIES
Suppose have an ur with marbles in different colors. Let N be the number of different
classes or groups (e.g. marble colors in the urn) m = [m1...., m,,] E No describe the number of elements per class (e.g. marbles per color), N N=>i E; m, be the total number of elements (e.g. all marbles in the urn) and n E {o....,N' . be the number of elements (e.g. marbles) to draw. Then, the multivariate hypergeometric distribution describes the probability of drawing æ = [Z1...., Te] E N6 marbles by sampling without replacement such that Σi 24 n. where I; is the number of drawn marbles of class i.
Using the central hypergeometric distribution, every marble is picked with equal probability. The number of selected elements per class is then proportional to the ratio between number of elements per class and the total number of elements in the urn. This assumption is often too restrictive, and we
HHijhenn et al. (2021) provide a great review article of the Gumbel-Max trick and its extensions describing recent algorithmic developments and applications.
2

---

Published as a conference paper at ICLR 2023
want an additional modeling parameter for the importance of a class. Generalizations, which make certain classes more likely to be picked, are called noncentral hypergeometric distributions.
In the literature, two different versions of the noncentral hypergeometric distribution exist, Fisher's (Fisher, 1935) and Wallenius' (Wallenius, 1963; Chesson, 1976) distribution. Due to limitations of the latter (Fog, 2008a), will refer Fisher's version of the noncentral hypergeometric distribution in the remaining part of this work.
Definition 3.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
P(X =x (x; 1
II
w
Po
where Po = EyEs II (m)) w The support S of the PMF is given by = fa E No Vi mi; Ei=1 23 = n) and (+) RJ!
24 S
The total number of samples per class m, the number of samples to draw n, and the class importance parameterize the multivariate distribution. Here, we assume m and be constant per experiment and are mainly interested in the class importance w. Consequently, we only use as the distribution parameter Equation (1) and the remaining part of this work.
The class importance w is crucial modeling parameter in applying the noncentral hypergeometric distribution (see (Chesson, 1976).. resembles latent factors like the importance, fitness, or adaptation capabilities of a class of elements, which are often more challenging to measure in field experiments than the sizes different populations. Introducing a differentiable and reparameterizable formulation enables the learning of class importance from data (see Section 5). We provide more detailed introduction to the hypergeometric distribution in Appendix A.
4 METHOD
The reparameterizable sampling for the proposed differentiable hypergeometric distribution consists of three parts:
1. Reformulate the multivariate distribution sequence of interdependent and conditional univariate hypergcometric distributions.
2. Calculate the probability mass function of the respective univariate distributions.
3. Sample from the conditional distributions utilizing the Gumbel-Softmax trick.
We explain all steps in the following Sections 4.1 to 4.3. Additionally, Algorithm and Algorithm (see Appendix B.5) describe the full reparameterizable sampling method using pseudo-code and Figures 5 and 6 in the appendix illustrate it graphically.
4.1 SEQUENTIAL SAMPLING USING CONDITIONAL DISTRIBUTIONS
Because it scales linearly with the number of classes and not with the size of the support S (see Definition 3.1), we use the conditional sampling algorithm (Liao Rosen, 2001; Fog, 2008b). Following the chain rule of probability, we sample from the following sequence of conditional probability distributions
Px (x: px, (I1:W) Il
px. px. (zil (aall w))
Following Equation (2), every Px, describes the probability of a single class i of samples given the already sampled classes
In the conditional sampling method, model every conditional distribution Px, (-) as a univariate hypergeometric distribution with two classes L and R: for = 1..c, define class L {i) as the left class and class R (i iAj j c) the right class.
3

---

Published as a conference paper at ICLR 2023
Algorithm 1 Sampling from the differentiable hypergeometric distribution. The different blocks are explained in more detail Sections 4.1 to 4.3 and Algorithm 2.
Input: m E N,, w E 184,
N. T E Ro+ Output: 2 E No. {ar E 1m Ji=1 {r, E RRD. )
for E {I..... c) do
Formulate the multivariate univariate # distribution (Section 4.1) IL Q. PL + -aamleeUNCHGmm mR. WL: wR, n, T) # Sample from univariate distribution - L1
+ i,
UU 1
m m,, m& é Zo+, - WR, E Ro-
m m., w WL
Re-assign classes for next step Assign values for class
I; - It, a; t a fi t FT.
end for
return æ; (ox (ox);=n {A:)
function SAMPLEUNCHG(M.
0₁ f- calcLogPMMP(m my, wi, Wj, n) F. <- contRelxxSamllec;;.77
Section 4.2 # Section 4.3
return Zi, 0;, P.
end function
To sample from the original multivariate hypergeometric distribution, we sequentially sample from the urn with only two classes I. and R, which simplifies to sampling from the univariate noncentral hypergeometric distribution given by the following parameters (Fog, 2008b):
ΣiEL Wi
WR ΣPER
Σ mL,
Σ my,
mL
M.
WL
mL
mR
JEL
We leave the exploration of different and more sophisticated subset selection strategies for future work.
Samples drawn using this algorithm are only approximately equal samples from the joint noncentral multivariate distribution with equal W. Because of the merging operation in Equation (3), the approximation error is only equal to zero for the central hypergeometric distribution. One way to reduce this approximation eтror independent of the learned w is different subset selection algorithm (Fog, 2008b). Note that the proposed method introduces an approximation error compared non-differentiable reference implementation with the same w (see Section 5.1), but not the underlying and desired true class importance. We can still recover the true class importance because a different overcomes the approximation error introduced by merging needed for the conditional sampling procedure.
4.2 CALCULATE PROBABILITY MASS FUNCTION
In Section 4.1. derive sequential sampling procedure for the hypergeometric distribution, in which we repeatedly draw from a univariate distribution to simplify sampling. Therefore, we only need to compute the PMF for the univariate distribution. See Appendix B.1, for the multivariate extension. The PMF px, (Z1; w) for the hypergeometric distribution of two classes and R defined by m,, mR W1, Wp and is given
pxx(1 Po
mL
WT.
mR
R-HL R
JT.
FT.
Po is defined as Equation (1), Wr,, WR and their derivation from W, and my, p. n as in Equa- tion (3). The large exponent of and the combinatorial terms can lead to numerical instabilities making the direct calculation of the PMF in Equation (4) infeasible. Calculations in log-domain increase numerical stability for such large domains, while keeping the relative ordering.
Lemma 4.1. The unnormalized log- pogpaobbasilittes
log Px (ZL:W) = log wt + (n xt) log wi + (bF(zL) + C define the unnormalized weights of a categorical distribution that follows Fisher's noncentral hyper- geometric distribution. C is a constant and &F (æ is defined as
(5)
(x -log (T( + 1)T(n FL. + +)) log [T(m
Ft. 1)T(mR n xL + 1)) (6)
4

---

Published as a conference paper at ICLR 2023
We provide the proof to Lemma 4.1 in Appendix B.2. Common automatic differentiation frameworks" have numerically stable implementations of log [()) Therefore, Equation (6) and more importantly Equation (5) can be caleulated efficiently and reliably. Lemma 4.1 relates to the caIcLogPMP
function in Algorithm 1, and Algorithm 2 describes ca cLogPMF in more detail.
Using the multivariate form of Lemma 4.1 (see Appendix B.1), it is possible directly calculate the categorical weights for the full support S. Compared to the conditional sampling procedure, this would result in computational speed-up for large number of classes c. However, the size of the support S is [T== my, quickly resulting in unfeasible memory requirements. Therefore, we would restrict ourselves to settings with no practical relevance.
4.3 CONTINUOUS RELAXATION FOR THE CONDITIONAL DISTRIBUTION
Continuous relaxations describe procedures to make discrete distributions differentiable with respect to their parameters (Huijben et al., 2021). Following Lemma 4.1, we make use of the Gumbel- Softmax trick to reparameterize the hypergeometric distribution via its conditional distributions Px, (-). The Gumbel-Softmax trick enables a reparameterization of categorical distributions that allows the computation of gradients with respect to the distribution parameters. We state Lemma 4.2, and provide a proof in Appendix B.3.
Lemma 4.2. The Gumbel- Soofeel-offmax trick can be applied to the conditional distribution Px, (=x))= of class i given the already sampled classes k < i.
Lemma 4.2 connects the Gumbel-Softmal trick to the hypergeometric distribution. Hence, reparame- terizing enables gradients with respect to the parameter of the hypergeometric distribution:
U0,
gi log( log(u)),
A [w) + gi
(7)
where u E [0, + is a random sample from an i.i.d. uniform distribution U. gi is therefore i.i.d. gumbel noise. F are the perturbed conditional probabilities for class i given the class conditional unnormalized log-weights (w))
(w) log px, w) C [log Px; (0:w), ... log px, (m;w))
(8)
We use the softmax function to generate (m; + 1) didmennioo sample vectors from the perturbed unnormalized weights f where T is the temperature parameter. Due Lemma 4.2, we do not need to calculate the constant C' in Equations (5) and (8). We refer to the original works (Jang et al., 2016; Maddison et al., 2017) or Appendix A.2 for more details on the Gumbel-Softmal trick itself. Lemma 4.2 corresponds to the contRelaxSamg more details).
function in Algorithm (see Algorithm 2 for
Note the difference between the categorical and the hypergeometric distribution concerning the Gumbel-Softmax trick. Whereas the (unnormalized) category weights are also the distribution param- eters for the former, the log-weigg weights a of class pre-defined æ; = (0...., mi] for the latter. It follows that for a sequence of categorical distributions, would have Σi=1 m; learnable parameters, whereas for the differentiable hypergeometric distri- bution we only have learnable parameters. In Appendix B.4, we discuss further difficulties in using sequence of unconstrained differentiable categorical distributions.
function of the class importance and the
EXPERIMENTS
We perform three experiments that empirically validate the proposed method and highlight the versatility and applicability of the differentiable hypergeometric distribution to different important areas of machine learning. We first test the generated samples of the proposed differentiable formu- lation procedure against a non non-differentiable reference implementation. Second, we present how the hypergeometric distribution helps detecting shared generative factors of paired samples setting. Our third experiment demonstrates the hypergeometric distribution as prior in variational clustering algorithms. In our experiments we focus on applications in the field of distribution inference where make use of the reparameterizable gradients. Nevertheless, the
FF. g. Tensorflow (Abadi al., 2016) Py (Paszke al., 2019)
5

---

Published as a conference paper at ICLR 2023
W2-1
W2+2
25
30
75
100
125
150
175
200
W2-5
10 dass
W2-6
dass
10
0.05
W2+8
8.05
w2-9
w2+10
1871
25
30
75
100
125
150
175
200
2
7
9
10
Number of elements
(a) Histograms
(b) Kolmogorov -Siinoo Test
Figure 1: Comparing random variables from the proposed differentiable formulation to a non- differentiable reference implementation. We draw samples from a multivariate noncentral hypergeo- metric distribution consisting of three classes. m, = 200 vi and n = 180. The class weights wi and 1.0, Wa increased from 1.0 to 10.0 with a step size of 1.0 (w2 in
wy for classes 1 and 3 are the figure). Figure la shows histograms of the number of elements for class 2. Figure lb represents the calculated distance values of the KS test between the reference and proposed implementation (upper plot) and their respective p- valuue (lower plot).
proposed method is applicable to any application with gradient-based optimization, in which the underlying process models sampling without replacement
5.1 KOLMOGOROV--MIRNGE TEST
To assess the accuracy of the proposed method, evaluate it against reference implementation using the Kolmogorov-Smingg test to estimate the equality of two distributions by quantifying the distance between the empirical distributions of their samples. The null distribution of this test is calculated under the null hypothesis that the two groups of samples are drawn from the same distribution. If the test fails to reject the null hypothesis, the same distribution generated the two groups of samples, i.e., the two underlying distributions are equal. As described in Section 4, we use class conditional distributions to sample from the differentiable hypergeometric distribution. We compare samples from our differentiable formulation to samples from non-differentiable reference implementation (Virtanen al., 2020, SciPy). For this experiment, we use multivariate hypergeometric distribution of three classes. We perform sensitivity analysis with respect to the class weights w. We keep w1 and w3 fixed at 1.0, and W2 is increased from 1.0 to 10.0 in steps of 1.0. For every value of w2, sample 50000 i...d. random vectors. We use the Benjamini- Senjamini-Hochberg correction (Benjamini Hochberg, 1995) to adjust the v-values for false discovery rate of multiple comparisons joint distribution. Given a significance threshold of / 0.05, p > 0.05 implies that we cannot reject the null hypothesis, which is desirable for our application.
test (Kolmogorov, 1933; Smirnov, 1939, KS). It is a nonparametric
are performing = 3 tests per
Figure la shows the histogram of class 2 samples for all values of w2, and Figure 1b the results the KS test for all classes. The histograms for classes and 3 are in the Appendix (Figure 4). We that the calculated distances of the KS-test are small, and the corrected p-values well above the threshold. Many are even close to 1.0. Hence, the test clearly fails reject the null hypothesis in 30 out of 30 cases. Additionally. the proposed and the reference implementation histograms are visually similar. The results of the KS test strongly imply that the proposed differentiable formulation effectively follows noncentral hypergeometric distribution. We provide more analyses and results from KS test experiments in Appendix C.
5.2 WEAKLY-SUPERVED LEARNING
Many data modalities, such as consecutive frames in video, are not observed as i.i.d. samples, which provides a weak-supervision signal for representation learning and generative models. Hence, we are only interested in learning meaningful representations and approximating the data distribution but also in the detailed relation between frames. Assuming underlying factors generate such coupled
"The code he found here:

---

Published as a conference paper at ICLR 2023
Table 1: To compare the three methods (LabeIVAE, AdaVAE, HGVAE) in the weakly-supervised experiment, we evaluate their learned latent representations with respect to shared (S) and independent (I) generative factors. To assess the amount of shared and independent information in the latent representation, we train linear classifiers on the respective latent dimensions only. We report the adjusted balanced classification accuracy, such that the random classifier achieves score O.
8
5
S
S
S
LABEL ADA
0.14
0.19
0.16#0.01
0.10=0.00
0.23+0.01
0.34
0.00++.00
0.12
0.19
0.1520.01
0.10=0.03
0.22:0.02
0.33=0.03
0.00+0.00
HG (OURS)
0. 0.18
0.22:10.0
0.19
0.08
0.28:10.01
0.2820.01
0.01=0.00
data, a subset of factors should be shared among frames to describe the underlying concept leading to the coupling. Consequently, differing concepts between coupled frames stem from independent factors exclusive to each frame. The differentiable hypergeometric distribution provides a principled approach to learning the number of shared and independent factors in an end-to-end fashion.
In this experiment, we look at pairs of images from the synthetic mpi3D toy dataset (Gondal et al., 2019). We generate coupled dataset by pairing images, which share certain number of their seven generative factors. We train all models as variational autoencoders (Kingma & Welling, 2014, VAE) to maximize an evidence lower bound (ELBO) on the marginal log-likelihood of images using the setting and code from Locatello et al. (2020). We compare three methods, which sequentially encode both images to some latent representations using a single encoder. Based on the two representations. the subset of shared latent dimensions is aggregated into single representation. Finally. a single decoder independently computes reconstructions for both samples based on the imputed shared and the remaining independent latent factors provided by the respective encoder. The methods only differ in how they infer the subset of shared latent factors. We refer to Appendix C.2 or Locatello al. (2020) for more details on the setup of the weakly-supervised experiment.
The LabelVAB assumes that the number of independent factors is known (Bouchacourt et al., 201 8; Hosoya, 2018, 8, LabelVAE). Like in Locatello et al. (2020), we assume known factor for all experiments. The Ada relies on heuristic to infer the shared factors (Locatello et al., 2020, AdaVAE). Based on the Kullback-Leibler (KL) divergence between corresponding latent factors across images, the decision between shared and independent factors is based on a hand-designed threshold. The proposed HGVAL uses the differentiable hypergeometric distribution model the number of shared and independent latent factors. We infer the unknown group importance w E RR with single dense layer, which uses the KI. divergences between cor- responding latent factors as input (similar to VAE). Based on and with d being the number of latent di- mensions, we sample random variables to estimate the & independent and == d k shared factors by reparameter- izing the hypergeometric distribution, providing with & and == d k. The proposed differentiable formulation allows us to infer such w and simultaneously Iearn the latent representation in a fully differentiable setting. After sorting the latent factors by KL divergence (Grover et al., 2019), we define the top-k latent factors as independent, and the remaining as shared ones. For a more detailed description of HGVAF, the baseline models, and the dataset, see Appendix C.2.
80
Label
Ada
70-
116 (Ours)
60-
MSE(s.5)
50
40
30
20
10-
0
Number of Shared Factors 5
Figure 2: We report the mean squared error MSB(s, 8) between true & and es- timated & number of shared factors to assess the models' performance.
To evaluate the methods, we compare their performance on two different tasks. We measure the mean squared error MSE(s, &) between the actual number of shared latent factors and the estimated number 3 (Figure 2) and the classification accuracy of predicting the generative factors on the shared and independent subsets of the learned representations (Table 1). We train classifiers for all factors on the shared and independent part of the latent representation and calculate their balanced accuracy. The reported scores are the average over the factor-specific balanced accuracies. Because the number of different classes differs between the discrete generative factors, we report the adjusted balanced accuracy as a classification metric. These two tasks challenge the methods regarding their estimate
7

---

Published as a conference paper at ICLR 2023
of the relationship between images. Generating the dataset controls the number of independent and shared factors, k and 8, which allows us to evaluate the methods on different versions of the same underlying data regarding the number of shared and independent generative factors. We generate four weakly-supervised datasets with s = {0. 1, 1,3, 5). On purpose, we also evaluate the edge case of 0. which is equal to the two views not sharing any generative factors.
Figure 2 shows that previous methods cannot accurately estimate the number of shared factors. Both baseline methods estimate the same number of shared factors & independent of the underlying ground truth number of shared factors 8. What is not surprising for the first model is unexpected for the second approach, given their- in theory- adaptive heuristic. On the other hand, the low mean squared error (MSE) reflects that the proposed HGVAE can dynamically estimate the number of shared factors for every number of shared factors 8. These results suggest that the differentiable hypergeometric distribution is able to learn the relative importance of shared and independent factors in the latent representations. Previous works' assumptions, though, do not reflect the data's generative process (LabeIVAE), or the designed heuristics are oversimplified (AdaVAE). We also see the effect of these oversimplified assumptions evaluating the learned latent representation. Table I shows that the estimation of large number of shared factors & leads to an inferior latent representation of the independent factors, which is reflected in the lower accuracy scores of previous work compared to the proposed method. More surprisingly, for the shared latent representation, our HGVAE reaches the same performance on the shared latent representation despite being more flexible.
Given the general nature of the method, the positive results of the proposed method are encouraging. Unlike previous works, it is not explicitly designed for weakly-supervised learning but achieves results that are more than comparable to field-specific models. Additionally, the proposed method accurately estimates the latent space structure for different experimental settings.
5.3 DEEP VARIATIONAL. CLUSTERING
We investigate the use of the differentiable hypergeometric distribution in a deep clustering task. Several techniques have been proposed in the literature to combine lotg-established clustering algorithms, such K-means Gaussian Mixture Models, with the flexibility of deep neural networks to learn better representations of high-dimensional data (Min et al., 2018). Among those, Jiang al. (2016), Dilokthanakul et al. (2016), and Manduchi et al. (2021) include a trainable Gaussian Mixture prior distribution in the latent space of a VAE. A Gaussian Mixture model permits a probabilistic approach clustering where an explicit generative assumption of the data is defined. All methods are optimized using stochastic gra- dient variational Bayes (Kingma & Welling, 2014; Rezende et al., 2014). A major drawback of the above models is that the samples are ei- ther assumed to be i.i.d. or they require pairwise
Uniform
Welght
0.30
0.05
0.00
Subsampling 80%
side information, which limits their applicability in
Welgit
0.30
scenarios.
0.05
The differentiable hypergcometric distribution can be easily integrated in VAE-based clustering algorithms to overcome limitations of current approaches. We want to cluster given dataset X == [x into K subgroups. Like previous work (Jiang et al., 2016), we assume the data is generated from a random process where the cluster assignments are first drawn from prior probability p(c;T) then each latent embedding z; is sampled from a Gaussian distribution, whose mean and variance depend on the selected cluster cs. Finally the sample x, is generated from a Bernoulli distribution whose parameter Mrx is the output of neural network parameterized by 0, as in the classical VAE. With these assumptions, the latent embeddings z; follow a mixture of Gaussian distributions, whose means and variances, {14 of Je are tunable parameters. The above generative model can then be optimised by maximising the ELBO using the stochastic gradient variational Bayes estimator (we refer to Appendix C.3 for complete description of the optimisation procedure). Previous work (Jiang et al., 2016) modeled the prior distribution as p(c; πT) = I[. p(c;) =
Subsampling 60%
Weight
0.10
0.05
Cluster
Figure 3: True class vs. learned hypergeomet- ric cluster weights ग.
II. Cal(c: π)) with either tunable fixed parameters In this task, we instead replace this prior
8

---

Published as a conference paper at ICLR 2023
Table 2: Evaluation of the clustering experiment on the MNIIS datasets. We compare the methods on 3 different dataset versions, namely i) uniform class distribution ii) subsampling with 80% of samples and iii) subsampling with only 60% of samples. We subsample half of the classes. Accuracy (Acc), normalized mutual information (NMI), and adjusted rand index (ARI) are used as evaluation metries. Higher is better for all metrics. Mean and standard deviations are computed across runs. For fair comparison with the baselines all methods use the pretraining weights provided by Jiang et al. (2016).
DATASET VERSION
UNIFORM
CATEGORICAL
HYPERGEOMETRIC
Acc (%)
92.0 + 3.0
87.2 5.0
91.4 + 5.0
UNIFORM
NMT (%)
84.8 + 2.2
81.8 + 1.9
85.6 + 2.0
ARI (%)
84.2 4.3
78.3 + 4.6
84.8 + 4.6
ACC (%)
90.8 + 4.0
87.4 = 4.7
92.5 + 0.5
SUBSAMPLING 80%
NMI (%)
84.122.2
81.8 = 2.3
84.6 0.8
ARI (%)
83.2 3.6
78.2 = 5.0
84.4 1.0
ACC (%)
83.5 3.9
86.5 = 4.9
89.7 + 4.3
SUBSAMPLING 60%
NMI (%)
10.7:14
81.3 = 2.9
82.9 + 2.2
ARI (%)
2.6
77.7 6.3
81.5 + 3.9
with the multivariate noncentral hypergeometric distribution with weights π and K classes where every class relates to a cluster. Hence, we sample the number of elements per cluster (or cluster size) following Definition 3.1 and Algorithm 1. The hypergeometric distribution permits create dependence between samples. The prior probability of sample to be assigned is not independent of the remaining samples anymore, allowing us to loosen the over-restrictive i.dd assumption. We explore the effect of three different prior probabilities in Equation (54), namely (i) the categorical distribution, by setting p(c; T) = ПI, Cat(c; T): (ii) the uniform distribution, by fixing म. = 1/K V i {I,.... K}: and (iii) the multivariate noncentral hypergeometric distribution. We compare them on three different MNIST versions (LeCun & Cortes, 2010). The first version is the standard dataset which has balanced class distribution. For the second and third dataset version we explore different ratios of subsampling for half of the classes. The subsampling rates are 80% in the moderate and 60% in the severe case. In Table 2 we evaluate the methods with respect to their clustering accuracy (Acc), normalized mutual information (NMI) and adjusted rand index (ARI).
given cluster
The hypergeometric prior distribution shows fairly good clustering performance in all datasets. Although the uniform distribution performs reasonably well, it assumes the clusters have equal importance. Hence it might fail more complex scenarios. On the other hand, the categorical distribution has subpar performance compared to the uniform distribution, even in the moderate subsampling setting. This might be due to the additional complexity given by the learnable cluster weights, which results in unstable results. On the contrary, the additional complexity does not seem to affect the performance of the proposed hypergeometrie prior but instead boosts its clustering performance, especially in the imbalanced dataset. In Figure 3, we show that the model is able to learn the weights, ग, which reflect the subsampling rates each cluster, which is not directly possible using the uniform prior model.
6 CONCLUSION
We propose the differentiable hypergeometric distribution in this work. In combination with the Gumbel-Softtaa trick, this new formulation enables reparametrized gradients with respect to the class weights w of the hypergeometric distribution. We show the various possibilities of the hypergeometric distribution in machine learning by applying it to two common areas, clustering, and weakly- supervised learning. In both applications, methods using the hypergeometric distribution reach at least sfate-of-the-all performance. We believe this work is an essential step toward integrating the hypergeometric distribution into more machine learning algorithms. Applications biology and social sciences represent potential directions for future work.
ACKNOWLEDGMENTT
We would like to thank Ričards Marcinkevičs for good discussions and help with the Kolmogorov- Smirnoy-Test. TS is supported by the grant #2021-911 of the Strategic Focal Area "Personalized

---

Published as a conference paper at ICLR 2023
Health and Related Technologies (PHRT)" of the ETH Domain (Swiss Federal Institutes of Tech- nology). LM supported by the SDSC PhD Fellowship #1-001568-007. AR is supported by the StimuLoop grant #1-00781 1--02 and the Vontobel Foundation.
ETHICS STATEMENT
In this work, we propose a general approach to learning the importance of subgroups. In that regard, potential ethical concerns arise with the different applications method could be applied to. We intend to apply our model the medical domain in future work. Being able to correctly model the dependencies of groups of patients is important and offers the potential of correctly identifying underlying causes of diseases on a group level. On the other hand, grouping patients needs to be handled carefully and further research is needed ensure fairness and reliability with respect underlying and hidden attributes of different groups.
REPRODUCIBILITY STATEMENT
For all theoretical statements, provide detailed derivations and state the necessary assumptions. We present empirical support on both synthetic and real data to back our idea of introducing the differentiable hypergeometric distribution. To ensure empirical reproducibility. the results of each experiment and every ablation were averaged over multiple seeds and are reported with standard deviations. All of the used datasets are public or can be generated from publicly available resources using the code that we provide in the supplementary material. Information about implementation details, hyperparameter settings, and evaluation metrics are provided in the main text or the appendix.
REFERENCES
Martin Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Gregory S Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Józefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Gordon Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, llya Sutskever, Kunal Talwar, Paul A Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda B Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems. CoRR, abs/1603.0, 2016.
Bruce Barrett, note exact calculation of the non central hypergeometric distribution. Communi- cations in Statistics Theory and Methods, 46(13):6737-674), 2017.
Luca Becchetti, Ugo Maria Colesanti, Alberto Marchetti- Spacheti-SSpaccamlla and Andrea Vitaletti. Rec- ommending items in pervasive scenarios: models and experimental analysis. Knowledge and information systems, 28(3)5555-578, 2011.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 330..333,, 2013.
Yoav Benjamini and Yosef Hochberg. Controlling the false discovery rate: a practical and powerful approach to multiple testing. Journal of the royal statistical society series b-methodological, 57: 289-300, 1995.
Bhaswar B Bhattacharya and Sumit Mukherjee. Degree sequence of random permutation graphs. The Annals of Applied Probability, 27(1):439 484, 2017. doi: 10. 10.121/11-AAP12000
Diane Bouchacourt, Ryota Tomioka, and Sebastian Nowozin. Multi-level variational autoencoder: Learning disentangled representations from grouped observations. Thirty-Second AAAI Confer- ence on Artificial Intelligence, 2018.
Giona Casiraghi, Vahan Nanumyan, Ingo Scholtes, and Frank Schweitzer. Generalized Hypergeomet- ric Ensembles: Statistical Hypothesis Testing in Complex Networks. 7 2016.
10

---

Published as a conference paper at ICLR 2023
Jean Chesson. A non-central multivariate hypergeometric distribution arising from biased sampling with application to selective predation. Journal of Applied Probability. 13(4):795-797, 1976.
I Chien, Chao Pan, and Olgica Milenkovic. Query k-means clustering and the double dixie cup problem. arXiv preprini 1806. 05938, 2018.
David R Cox. The regression analysis of binary sequences. Journal of the Royal Statistical Society: Series B (Meshoalolggica)) 20(2):215-232, 1958.
M Mediano, M Garnelo, M MCHLe Lee, Salimbeni, K Arulkumaran, and M Shanahan. Deep unsupervised clustering with Gaussian mixture variational autoencoders. arXiv preprint ArXiv: 1.02648, 2016.
Dilokthanakul,
Moran Feldman, Christopher Harshaw, and Amin Karbasi. Greed is good: Near-optimal submodular maximization via greedy optimization. In Conference on Learning Theory, pp. 758-784. PMLR, 2017.
Ronald A Fisher. The logic of inductive inference. Journal of the royal statistical society, 98(1): 39 39-82,11 1935.
Agner Fog. Calculation methods for Wallenius' noncentral hypergeometric distribution. Communica- fions in Statistics-Simulatte and Computation@o, 37(2):258-22) 2733, 2008a.
Agner Fog. Sampling methods for Wallenius' and Fisher's noncentral hypergeometric distributions. and ComputationO, 37(2):2412577, 2008b.
Communications
Muhammad Waleed Gondal, Manuel Wuthrich, Djordje Miladinovic, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Akpo. Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. On the Transfer of Inductive Bias from Simulation to the Real World: New Disentanglement Dataset. In Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
T Gonin. XIV. The use of factorial moments in the treatment of the hypergeometric distribution and in tests for regression. The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science, 21(139):215-226 1936.
Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon. Stochastic Optimization of Sorting Networks via Continuous Relaxations. In International Conference on Learning Representations. OpenReview.net, 2019.
Emil Julius Gumbel. Statistical theory of extreme values and some practical applications: a series of lectures, volume 33. US Government Printing Office, 1954.
Chris Harshaw, Moran Feldman, Justin Ward, and Amin Karbasi. Submodular maximization beyond non-negativity: Guarantees, fast algorithms, and applications. In International Conference on Machine Learning, pp.
PMLR, 2019.
Irina Higgins, Loie Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAFE Learning Basic Visual Concepts with Constrained Variational Framework. ICLR, 2(5):6, 2017.
Haruo Hosoya. A simple probabilistic deep generative model for learning generalizable disentangled representations from grouped data. CoRR, abs/1809/ 2018.
Iris A M Huijben, Wouter Kool, Max B Paulus, and Ruud J G van Sloun. A Review of the Gumbel- max Trick and its Extensions for Discrete Stochasticity in Machine Learning. arXiv preprint arXiv:21 10.01515, 2021.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizatioo with gumbel-softmax. arXiv preprint XXXv:1611.0014. 2016.
Zhuxi Jiang, Yin Zheng, Huachun Tan, Bangsheng Tang, and Hanning Zhou. Variational deep embed- ding: An unsupervised and generative approach to clustering. arXiv preprint arXiv: 161.0.5418.88 2016.
11

---

Published as a conference paper at ICLR 2023
Mark E Johnson. Multivariate statistical simulation: A guide selecting and generating continuous multivariate distributions, volume 192. John Wiley & Sons, 1987.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint urXiv::412.660 2014.
Diederik P Kingma and Max Welling, Auto-Encoding Variational Bayes. International Conference on Learning Representations. OpenReview.net, 2014.
Andrey Kolmogorov. Sulla determinazione empirica di una Igge di distribuzione. Inst. Ital, Attuari, Giorn., 4:83-91, 1933.
Wouter Kool, Herke van Hoof, and Max Welling. Stochastic Beams and Where To Find Them: The Gumbel- -oonbel-Top-kk Trick for Sampling Sequences Without Replacement. In International Conference on Machine Learning, pp. 3499-3508. PMLR, 2019.
Wouter Kool, Herke van Hoof, and Max Welling. Ancestral Gumbel-Top-k Sampling for Sampling Without Replacement. Journal of Machine Learning Research, 21(47): 2020a.
Wouter Kool, Herke van Hoof, and Max Welling. Estimating Gradients for Discrete Random Variables by Sampling without Replacement. In International Conference on Learning Representations. OpenReview.net, 2020b.
Sergey A Lashin, Valentin Suslov, Nikolay A Kolchanov, and Yury Matushkin. Simulation of coevolution in community by using the" Evolutionary Constructor" program. In silico biology, 7 (3):261-275, 2007.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010.
Jason G Liao and Ori Rosen. Fast and stable algorithms for computing and sampling from the noncentral hypergeometric distribution. The American Statistician. 55(4):366 2001.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentan- gled representations. In International Conference on Machine Learning, pp. 4114-4124. PMLR, 2019.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In International Conference on Machine Learning, pp. 6348-6359 PMLR, 2020.
Michael Lodato, Mollie B. Woodworth, Semin Lee, Gilad D. Evrony, Bhaven K. Mchta, Amir Karger, Soohyun Lee, Thomas W. Chittenden, Alissa M. D DaGam, Xuyu Cai, Lovelace J. Luquette, Eunjung Lee, Peter J. Park, and Christopher A. Walsh. Somatic mutation in single human neurons tracks developmental and transcriptional history. Science, 350(6256))94 98, 10 2015.
Michael A. Lodato, Rachel E. Rodin, Craig L. Bohrson, Michael E. Coulter, Alison R. Barton, Minseok Kwon, Maxwell A. Sherman, Carl M. Vitzthum, Lovelace J. Luquette, Chandri N. Yandava, Pengwei Yang, Thomas W. Chittenden, Nicole E. Hatem, Steven C. Ryu, Mollie B. Woodworth, Peter J. Park, and Christopher A. Walsh. Aging and neurodegenerattion are associated with increased mutations in single human neurons. Science, 359(6375):555 559, 2018.
Maddison, A Mnih, and Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations, OpenReview.net, 2017.
Chris J Maddison, Daniel Tarlow, and Tom Minka. A sampling. Advances neural information processing systems, 27, 2014.
Laura Manduchi, Kieran Chin-Cheone Holger Michel, Sven Wellmann, and Julia E Vogt. Deep Conditional Gaussian Mixture Model for Constrained Clustering. ArXiv, abs/2106.0, 2021.
Min, X Guo, Q Liu, G Zhang, Cui, and Long. A Survey of Clustering With Deep Learning: From the Perspective of Network Architecture. IEEE Access, 6:39501-39511 2018.
12

---

Published as a conference paper at ICLR 2023
Seiji Ono, Kazuharu Misawa, and Kazuki Tsuji. Effect of group selection the evolution of altruistic behavior. Journal of theoretical biology, 220(1):55-66) 2003.
Mario Paolucci, Rosaria Conte, and Gennaro Di Tosto. A model of social organization and the evolution of food sharing in vampire bats. Adaptive Behavior, 14(3):223-238 2006.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Z Yang. Zach De Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch An Imperative Style, High- -erforeeffrmance Deep Learning Library. CoRR, abs/1912.0, 2019.
Max B Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris Maddison. Gradient estimation with stochastic softmax tricks. arXiv preprint arXiv:2006. 08063, 2020.
Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, and others. Scikit-learn: Machine learning in Python. Journal of machine learning research, 12(Oct):2825 2830,, 2011.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wienstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, pp. 1278-1286. PMLR, 2014.
Nikolai V Smirnov. On the estimation of the discrepancy between empirical curves of distribution for two independent samples. Bull. Math. Univ. Moscou, 2(2):3-14, 1939.
Thomas M Sutter and Julia E Vogt. Multimodal Relational VAE. 2021.
Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, and others. SciPy 1.0: fundamental algorithms for scientific computing in Python. Nature methods, 17(3):261-722, 2020.
Kenneth Wallenius. Biased sampling: the noncentral hypergcometric probability distribution. Technical report, Stanford Univ Ca Applied Mathematics And Statistics Labs, 1963.
E T Whittaker and G N Watson. A Course of Modern Analysis. Cambridge Mathematical Library. Cambridge University Press, 4 edition, 1996.
Ronald J Williams. Simple statistical frolient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-256, 1992.
Sang Michael Xie and Stefano Ermon. Reparameterizable subset sampling via continuous relaxations. arXiv preprint arXiv: 10517, 2019.
A PRELIMINARIES
A. HYPERGEOMETRIC DISTRIBUTION
The hypergeometric distribution is a discrete probability distribution that describes the probability of æ successes in n draws without replacement from finite population of size N with elements that are part of the success class Unlike the binomial distribution, which describes the probability distribution of successes
draws with replacement.
Definition A.1 (Hypergeometrri Distribution (Gonin, 1936)*), random variable follows the hypergeometric distribution, if its probability mass function (PMF) is given by
=(X=z) = px
17
n-
AAthouugh the distribution itself is older, Gonin (1936) were the first name it hypergeometric distribution
13

---

Published as a conference paper at ICLR 2023
Urn models typical examples of the hypergeometric probability distribution. Suppose we think of an urn with marbles in two different colors, e. g. green and purple, we can label as success the drawing of green marble. Then N defines the total number of marbles and m the number of green marbles in the urn. x is the number green marbles, and nt x is the number of drawn purple marbles.
The multivariate hypergeometric distribution describes an urn with more than two colors, e.g. green, purple and yellow in the simplest case with three colors. As described in Johnson (1987), the definition is given by:
Definition A.2 (Multivariate Hypergeometric Distribution). A random vector X follows the multi- variate hypergeometric distribution, if its joint probability mass function is given by
P(X a:) = 11 (m.) (
(10)
where E N4. is the number of different classes (e.g. marble colors in the urn), m [mi,.... me]
N describes the number of elements per class (e.g. marbles per color), N = - m, is the total number of elements (e.g. all marbles in the um) and {O....,N] N) is the number of elements (e.g. marbles) to draw. The support S of the PMF is given by
S = x E EE
: Vi
H, S m&
(II)
1=1
A.2 GUMBEL-SSFFMMX--ER
Most of the information in this section is from (Jang et al., 2016; Maddison et al., 2017), which concurrently introduced the Humbel-Softmal trick. Gumbel-Softmax is continuous distribution that approximate the categorical distribution, and whose parameter gradients can be easily computed using the reparameterization trick.
Let z be a categorical variable with categorical weights π = [x1,-... mc| such that Ee T 1. Following (Jang al., 2016), assume that categorical samples are encoded one-hot vectors. The Gumbel-Max trick (Gumbel, 1954; Maddison et al., 2014) defines an efficient way to draw samples from categorical distribution with weights æ::
onc_hot(arg max log(TT) + +gkk)
(12)
where g = lg.............................. ... 9c] are i. i. d. samples drawn from Gumbel(0.1), We can efficiently sample g from 10 by drawing sample from a uniform distribution U (0, 1) and applying the transform = log( log u). For more details, we refer to (Gumbel, 1954; Maddison et al., 2014).
(Jang al., 2016; Maddison al., 2017) both use the softmax function continuous and differen- tiable approximation to arg max. The softmax function is defined as
exp((log Th + gx)/T)
for k == -=,,... I....,C
(13)
exp((log n, + 9;)/T)
where is temperature parameter. As goes to zero, the softmax function approximates the argmax function. Hence, the Gumbel--Sofmax distribution approximates the categorical distribution.
B METHODS
B. PMF FOR THE MULTIVARIATE FISHER'S 'S NONCENTRAL DISTRIBUTION
In this section, we give a detailed derivation for the calculation of the log- Dog-brbbabilities of the multivariate Fisher's noncentral hypergeometric distribution. We end up with formulation that proportional to the actual log probabilities. Because the ordering of categories is not influenced by scaling with constant factor (additiion/uubraction in log domain), these are unnormalized
14

---

Published as a conference paper at ICLR 2023
log- oog-pobbabilities of the multivariate Fisher's noncentral hypergeometric distribution.
px
(14)
Po
24
where Po defined as in Equation (1). From there it follows
log lx; = log
(15)
Po ==1
Iz
1
Mi
= log
+
log
(16)
Po
3i
1
mi
= log
+
log
W
(17)
Po
24
k=1
1
(m²)
log
log
+ log (w
(18)
Po
244
1 Po
(m²)
log
log
+ z, log (wi)
(19)
Constant factor can removed as the argmax is invariant to scaling with constant factor which equals addition or subtraction in log-space. It follows
1o xx w)
log
(mi
+ log (wi) + C
(20)
x
+ xs log (wi)) +6
log
(21)
a..[(m: x.,)
log (T(a; + 1)f(mm - H, + 1)) + 2, log (wi) + C
(22)
(23)
In the last line we used the relation [[k + 1) = kd. Setting C Č, it directly follows
log (x: w) Σ
Iy log wi + the (x) + C
(24)
where (x(x)) = Ei -i_plop log (T(2; 1)r(m--)) I; + 1)).
The Gamma function is defined in Whittaker & Watson (1996) as r(=) (e x dx
(25)
B.2 PROOF FOR LEMMA 4.1
Proof. Factors that constant for all x do not change the relative ordering between different values of æ. Hence, removing them preserves the ordering of values x (Barrett, 2017).
log px,, (#L:W) loog
w
m-
WR
n-z1
(26)
Po
x
n
IL
=log
(..))
log
mR
log (w + log (w)=*2) C
(27)
JL
n
FLL
15

---

Published as a conference paper at ICLR 2023
Proposed
0,5
0.05
25
50
300 Reference
125
150
175
200
50
75
100
125
150
200
42
Raference
W2-0
0.5
W0-1
WS-T
60-4
60-2
0.05
17-0
K2+10
W2m10
125
175
200
125
190
200
(a) Histograms class
(b) Histograms class
W2-3
0.00
25
50
75
100
125
150
175
200
60:5
62-4
60-6
W2-1
K2-2
K2-5
0.00
25
50
75
100
125
150
175
200
Number of elements
(c) Histograms class 3
Figure 4: Comparing random variables drawn from the proposed distribution to a reference distribu- tion.
Using the definition of the binomial coefficient (see Section 3) and its relation to the Gamma function I[k + 1) = k., it follows
log (I1: w) =r log wL + (n - TL) log WR
(28)
log (T + 1)1 (n 2g + 1))
log (m(m
+][ - æ7 + 1)) +
With (z) as defined in Equation (6), Equation (5) follows directly.
B.3 PROOF FOR LEMMA 4.2
Proof. When sampling class we draw x; samples from class i where x; m, The conditional distribution Px, ) w) for class given the already sampled classes k < simultaneously defines the weights of categorical distribution. Sampling æ; elements from class can be seen as selecting the æ,th category from the distribution defined by the weights px, (I: (I:l 1i-1 Therefore.
; w.).
Σ px; (arl( =1,
(29)
053,
which allows to apply the Gumbel-Max trick and, respectively, the Gumbel-Softmax trick.
B.4 THOUGHT EXPERIMENT FOR MODELING THE HYPERGEOMETRIC DISTRIBUTION WITH A SEQUENCE OF UNCONSTRAINED CATEGORICAL DISTRIBUTIONS
We quickly touched on the topic of using sequence of categorical distributions in the main text (Section 4.3). To further describe and discuss the problem of using a sequence of categorical distributions, we provide a more detailed explanation and example here. Of course, the same constraints as the ones described in the main paper apply, e. g. we want our method to be differentiable, scalable in the number of random states and computationally efficient. Hence, methods that infer at least all states of single class in parallel, and not in methods that sequentially iterate over all possible random states for every re-sampling operation.
interested in
see Appendix B.l for the definition the Gamma function
16

---

Published as a conference paper at ICLR 2023
In our example, we want to model a hypergeometric distribution with classes and the following specifications
m 10, m2 7, my 8, n 9
(30)
We use one categorical distribution for every class. The categories of the categorical distributions describe the number of elements to sample from this class. Based on the three categorical distributions, we would like to sample vi E No such that Σj zjj n. The sequence of distributions ideally should describe the class conditional distributions of the hypergeometric distribution described in Section 4.1. Or at least model sampling without replacement,
fulfill the necessary constraint.
We then have three vectors T1, T2 and T33 which define the categorical weights of the three categorical distributions. As such Le J 1 V j. Using three categorical distributions, we are not able to explicitly model w. but if we are able to fulfill the constraints, there would be some matching w. To make it differentiable, trick, such that we can use everything in differentiable pipeline. In most differentiable settings, the categorical weights are inferred using some neural network, e. g.
approximate the categorical distributions using the Gumbel-Softmax (GS)
fg.e(-)
(31)
Knowing that our categories correspond to integer values, we can use the straight-throus estimator (Bengio et al., 2013) together with the Gumbel-Softmnn trick and a bit of matrix multiplication to convert an one-hot vector to an integer value:
Yj strraight (Grr
(32)
arg max(GG((T)))
(33)
Please check the provided code for more details on the details of the matrix multiplications.
We can either infer the categorical weights of and sample from the three distributions in parallel or sequentially. We first try to infer the weights of all distributions in parallel. With the constraint Σ, = ni not all combinations of T,, Vj are valid anymore. We might be able to learn the constraints, such that the number of samples fulfilling them is maximized. But it not possible to have the constraints respected due to the formulation of this method. Hence, it is not guaranteed that we sample zj, Vj such that Σ, xz = n, which then results in non-valid samples.
Therefore, we try infer the π, sequentially. The sequential procedure models the following behaviour (similar to our proposed class conditional sampling)
p( J2, 03) p(a1 (p(22 x1 (1(p((s 22)
(34)
Without loss of generality, we assume that we first infer лi and sample Cat(TT)) such that I = 7. follows that not all combination of weights मं and ny are valid anymore. When inferring T2 some weights have to be zero to guarantee that Ej Bi and sample valid random vector. Hence,
2 0 k>2
(35)
If we assign any nonzero probability to T2 k > 2, we are not able to fulfill the constraint Σ, Ir anymore for every generated sample. Additionally, from Equation (12) it follows that we would need to constrain the gumbel noise as well. We are unaware of previous work that proposed constrained Guubbel-Sfftaa trick which would model this behaviour. Also, there is no guarantee that T2.k > 0, < 2, leading to additional heuristics-base solution that we would need to implement. There is unclear effect on the calculation of the gradients, which would need to be investigated. Also, for sampling and inferring of weights π there arise questions of ordering between classes and how valid random samples be guaranteed. We summarize that the implementation and modeling of a hypergeometric distribution using a sequence of unconstrained categorical distributions is far from being trivial, because of the open question of how implement constraints in a general and dynamic way when using the
trick.
Note the important difference Softmax trick to generate random samples, we are able to infer w in parallel, which does not introduce an implicit ordering between classes and the constraint E3 xj n guaranteed. The only constraint with respect w have satisfy, is wy 0, Vj. wj 2 easily be satisfied using ReLU
hypergeometric distribution. Although also use the Gumbel-
17

---

Published as a conference paper at ICLR 2023
Algorithm 2 Subroutines for sampling From Multivariate Noncentral Hypergeometric Distribution.
function SAMPLEUNCHG(m)
Oi f calcLogPMI((m;; myj wi; Wi, n)
Section 4.2 # Section 4.3
3g F. < contRelaxSample(a,,
return 24, ax,, f'
end function
function CALCLoGPMFON my, Wi, wr
for k E {0, } do
31,b + (k 1)
kr.k + [ReLU(n k) 1)
end for
[ < log [x: +) + log I' (ml + 1)
# see Appendix B.5
r <- log [ (a., + 1) log T (m, - Frr + 1)
a log + 2,, log (1+r)
return or end function
function CONTRELAXSAMPLECCN
9 log( log
f) f- ag +
pi < Softmax(Fi/9
T + Count-Index (hrighttTThoogh(()))
Count-Index and Shright-Through:: # see Appendix B.5
return 21, fz end function
activation function. We are then able to calculate the class conditional distributions as function of w such that all constraints are satisfied. We not need to dynamically change categorical weights, e. model parameters, during the sampling procedure, which is the case for the sequence of constrained categorical distributions.
B.5 ALGORITHM AND MINIMAL EXAMPLE
In the main text, we drafted the pseudocode for our proposed algorithm (see Algorithm 1). We only did for the main functionality, but the subroutines described in Sections 4.2 and 4.3). Algorithm 2 describes the subroutines used in Algorithm 1 and explained in Section 4.
The straight-through operator uses hard one-hot embedding in the forward path, but the relaxed vector for the backward path and the calculation of the derivative (Bengio et al., 2013). The Count- Index maps the one-hot vector to an index, which is equal number of selected class elements in our case.
B.5.1 MINIMAL EXAMPLE
We describe minimal example application using a step-by-step procedure to provide intuition and further illustrate the proposed method. Here, we learn the generative model of an urn model using stochastic gradient descent when given samples from an urn model with a priori unknown weights w. Note that given dataset, we could also estimate the weights of the hypergeometric distribution by minimizing the negative log probability of the data given the parameters log px (x; w). In contrast, we generative approach to demonstrate how method allows backpropagation when modeling the generative process of the samples by reparameterizing the hypergeometric distribution. Additionally, illustrate the minimal example with two figures (Figures and 6), which explain the sampling procedure visually.
Let us start with our minimal example. We are given a dataset of i.dd samples X, E KK*. from multivariate noncentral distribution with unknown wyt RR_ K is the number of samples in the
18

---

Published as a conference paper at ICLR 2023
m1 =3
m2 5
m1
N = 12
Figure 5: Illustration the basic setting the multivariate hypergeometric distribution. We have classes of elements (green, orange, and blue) with different and unknown importance we for every class In our urn, the total number of elements N is given by the sum of elements of all classes me- From this urn, draw group of n samples. In this example, often unknown, and difficult to estimate. Our formulation helps to learn we using gradient-biie based optimization when simulating how given samples are drawn from the urn.
5. The group importance w. is
mL 11
mL m2
122.. mi
9
1
ma =
w1. 41
WL
WL
WR ug waggg
R
WR w7
U
X1 ~ (n, my. W1. wR)
Xg (n,
WR)
X3 ~ px v. (n, my, U, wr. 0)
Figure 6: Illustration of the proposed conditional sampling from the multivariate noncentral hyperge- ometric distribution. We the same in Figure 5 with [3, 5, 5, and n 5. As described in Section 4, we sequentially sample random variates of the individual classes. Hence, we start by sampling class 1. For that, merge classes 2 and 3 (illustrated by the half blue and half orange balls) creating the necessary parameters m,, MR WL, for px. (.) (described in the left column). This is also described in Algorithms 1 and 2. Using the univariate distribution Px: (.) we sample the random variable X₁, which is equal to 1 in our example (symbolized by the single green ball). We continue with sampling the class 2, which is described in the middle column. The merge operation simplifies to assigning my
and n is the original minus X₁. We draw X, 3 our example (again illustrated by the 3 orange balls below urn). Because the number drawn balls must sum to n, the last class X3 is fully determined already.
and ma
19

---

Published as a conference paper at ICLR 2023
dataset and n. defines the number of classes. A sample from the dataset XD is denoted as Xp. For every XD E Xp. it holds that Ee X,, = nt. We assume that we know the total number of elements in the um, e.g.
[m,, m2...., mp]-
In our minimal example, we are interested learning the unknown importance weights w with a generative model using stochastic gradient descent (SGD). Hence, we assume data generating distribution px (z; w) such that X ~ px (x; (I; The loss function given as
Σ XnEXn
(xppx(aw)
(x²
(36)
Σ
Ex<px(aw) [XX,,
(37)
XDEXD
where L the loss per sample. Px (æ; is a noncentral multivariate hypergeometric distribution as defined in Definition 3.1 where the class importance is unknown.
To minimize [X)(,),, we want to optimize w. Using SGD, we optimize the parameters in an iterative manner:
d Wi+1 WP 17 E Ex-px (m=w
XX)
(38)
where is the learning rate, and is the step in the optimization process. Unfortunately, we do have a reparameterization estimator - [ X]] because of the jump discontinuities of the arg max function in the categorical distributions.
As described in Sections 4.1 and 4.3, we can rewrite px (x:; w) as sequence of conditional distribu- tions. Every conditional distribution is itself categorical distribution, which prohibits the calculation of dee E[L(X,, (X, X)].
In more detail, we rewrite the joint probability distribution px (æz; w) as
ne (x w) =xx, (w1:w) I px. (ee F1, Ti. Io-l
(39)
where every distribution px, (www is categorical distribution. We sample every X. using Equa- tion (5), i.e.
px. (Xt.ii)))) 1., logwr, + (n. xr.) loguR, + ( C
(40)
Who Wn, m..; m... and ne = Ejsc X, are calculated according to Equation (3) and sequentially for every class.
The expected element-wiss loss Ex xxxxxxxee (mmw) [L(XD, X)] changes to
Ex (x=) [L(D,, = x Xxxpx (anw)
(ND xx)²
(41)
= Ex Exxpx((ww)
Σ L(XD,c: Xc)
(42)
E Ex (#:w [LXXD.,
(43)
Hence,
E[L(X,, x))
Σ
du E Exspx(aw) DL,,, Xc)l
(44)
de
d E [X D.. X.)), we face the problem of not having a reparameterizablle Unfortunately, for every du gradient estimator. We cannot calculate the gradients of the loss directly, but px., (-) being categorical distributions allows us to use the Gumbel-Softme gradient estimator (Jang et al., 2016: Maddison al., 2014; Paulus et al., 2020).
20

---

Published as a conference paper at ICLR 2023
The Gumbel-Softmar trick is differentiable approximation to the Gumbel-Mar trick (Maddison et al., 2014), which provides a simple and efficient way to draw samples from a categorical distribution with weights The Cumbel-Softunn trick uses the softmax function differentiable approximation to the a remax function used in the Gumbel-Max trick. It follows (Jang et al., 2016)
y = 9)/T)
(45)
=softmax, (log π + 9)
(46)
where 91. 91.....9k 99 are i...d samples drawn from Gumbel 1). and is temperature parameter, y is continuous approximation to one-hot vector, i.e. 0 y/ such that >; 1.
Different to the standard Gumbel-Softman trick, we infer the weights ι from the probability density function log Pxx (-) (see Equations (7) and (8)). We write for single conditional class LL. as the procedure is the same for all classes. follows
[ g) softmax, Px, (æL,;W) + g)
(47)
where LL = [0, mL.). Due to the translation invariance of the softmax function, we do not need to calculate the constant C' log pxx. (z).:W)
The Gumbel-Softmax approximation is smooth for 7 > 0, and therefore ElL(X,,,,, has well- defined gradients Ai
We write the loss function to optimize and its gradients
Eg XxW,,)))
Σ Fg [T((D,, X Xx,c(w,)))
(48)
d Fg L(X,,XX( du
The Σ
d
Fyy
L(XD,e, (x.(..9))
(49)
dw
By replacing the categorical distribution in Equation (5) with the Gumbel-Softmax distribution (see Lemma 4.2), we can thus use backpropagation and automatic differentiation frameworks to compute gradients and optimize the parameters (Jang al., 2016).
We implemented our minimal example for n, = 3 classes We set m = [m,, m2, mg] [200, 200, 200] and Wgt and show the performance of the proposed method. From these 1000 samples, we use 800 for training and 200 for validation. Similar to the setting we use for the KS-test (see Section 5.1), choose 10 values for wgt.2, i.e. Wgr,2 [1.0, 2.0, 10.0). The values for Wgt.1 and Wgt.3 are set to 1.0 for all datasets versions.
180. We create 10 datasets X E 1100003 generated from different
As described above the model does not have access to the data generating true Wyt. So, for every dataset XD, we optimize the unknown based on the loss L defined Equation (43).
Figure 7 shows the training and validation losses over the training steps. We train the model for 10 epochs, but we see that the model converges earlier. The losses only differ at the beginning of the training procedure, which is probably an initialization effect, but quickly converge to similar values independent of the We value that generated the dataset. Figure shows the estimation of log w. The a-axxs shows the training step, the y-axis shows the estimated value. Figures 8a to 8b demonstrate that the hypergeometric distribution is invariant to the scale of w. With increasing value of wgt,2, the values of W1 and Wg decrease, although their ground truth values W and Wat.3 do not change. The training and validation loss do not increase though (Figures 7a and 7b), which demonstrates the scale-in valeeinvariance.
C EXPERIMENTS
All our experiments were performed on our internal compute cluster, equipped with NVIDIA RTX 2080 and NVIDIA RTX 1080. Every training and test run used only a single NVIDIA RTX 2080 or NVIDIA RTX 1080. The weakly-supervised runs took approximately 3-4 hours each, where the clustering runs only took about hours each. We report detailed runtimes for the weakly supervisd experiments in Table 3 to highlight the efficiency of the proposed method.
21

---

Published as a conference paper at ICLR 2023
train_lop loss
validation. loss
Showing first 10 runs
Showing first 10
group: w2_6.0
group: w2_8.0
group: w2_6.0
group: w2_8.0
group: w2_9.0
group: w2_1.0
group: w2_9.0
group: w2_1.0
group: w2_5.0
group: w2_4.0
group: w2_5.0
group: w2_4.0
group: w2_7.0
group: w2_3.0
group: w2_7.0
group: w2_3.0
group: w2_10.0
group: w2_2.0
group: w2_18.8
group: w2_2.0
1000
9000
3000
6000
2000
4000
1000
traingrploball
trainosigle www.olloooooooo
2k
4k
6k
2k
3k
4k
5k
7k
(a) Training Loss
(b) Validation Loss
Figure 7: Training and validation losses for different values Wyt our minimal example described in Appendix B.5.1.
Tog, MM
1og
1og
en1i0: #0.1-6
+0.2-9
group:
group:
group:
grou:
67242:
group: N2.2.
gro.p:
19.0
greec: $2.2.8
greapt 18.0
gro4p: 42.2.0
02
45
08
36
(a) Estimation of logwr
(b) Estimation of log W/2
(c) Estimation log wa
Figure 8: The estimated log values over training procedure for different ground truth Wyt values of our minimal example (see Appendix B.5.I). We see nicely that the hypergcometric distribution is invariant the scale of w. With increasing value of w2, the estimated values for w1 and Wy change as well, but the training and validation loss remain low (see Figure 7).
22

---

Published as a conference paper at ICLR 2023
proposed
proposed
m11
mm
1-10
10
50
175
200
0=100
100
n=200
400
100
25
50
75
100
125
150
175
200
25
73
100
125
15c
175
230
Numbar of ahumanta
Namber ulmmmnta
(a) Histograms class
(b) Histograms class 2
proposes
<5 test
175
m:10
0.00
15
P volues
10P
10a1
Number
atemmnts
(c) Histograms class
(d) Distance and p-values
Figure 9: Sensitivity analysis for a varying number of samples to draw n measured with the Golmogorov-Sminoe test. this experiment, we have the following specifications: m {200, 200, 200}, w = {1.0, 5.0, 1.0]. The values of n are given in the plot.
C.I KOLMOGOROV-SMIRING TEST
As described in Section 5.1 in the main text, we report the histograms all classes next to each other here in Figure 4. We that class and class 3 have (at least visually) the same distribution histograms for different values of w2. And they also match their respective reference histograms.
C.I.1 ABLATION STUDIES FOR KOLMOGOROV MIRNOO TEST
To highlight our accurate approximation, we provide more results from KS-tests here in the appendix. We perform additional tests with varying and m2. See Figures 9 and 10 for the detailed results. We see that over all combinations, our approximation of the hypergeometric distribution performs well
and produces samples of approximately the same quality as the reference distribution.
C.2 WEAKLY-SUPERVISED LEARNING
C.2.1 METHOD, IMPLEMENTATION AND HYPERPARAMETERS
In this section we give more details on the used methods. We make use of the 1ibb (Locatello et al., 2019) which is also used in the original paper we compare to (Locatello et al., 2020). The baseline algorithms (Bouchacourt et al., 2018; Hosoya, 2018) are already implemented in ci disentangl ement_lib For details on the implementation of models, we refer to the original paper. We did not change any hyperparameters or network settings. All experiments were performed using 3 1.0 as this is the best performing / according to Locatello et al. (2020). For all experiments we train three models with different random seeds.
All experiments performed using Group VAE (Hosoya, 2018). Using GroupVAE shared latent factors are aggregated using an arithmetic mean. Bouchacourt et al. (2018) assume also knowledge about shared and independent latent factors. In contrast to GroupVAE, their ML-VAE aggregates shared latent factors by using the Product of Experts (i.e. geometric mean).
Figure 12 shows the basic architecture. The different architectures only differ in the View Aggregation ion module. this module, every method selects the latent factors 2; E S, which should be aggregated over different views e and 222. Given subset S of shared latent factors,
23

---

Published as a conference paper at ICLR 2023
poppsed
proposed
0.05
m2-1
00-3
0.00
200
200
m2-100
m2=10
100
200
md-100
m2=10
m2-200
02-200
m2-300
m2-300
0.05
0.00
23
100
125
150
175
200
50
100
220
300
Number of elaments
Number of alamanaa
(a) Histograms class
(b) Histograms class 2
proposed
ICS test
0.05
0.00
0.00
Ief
5 voluss
m2-300
10"
0.05
10
0.00
Number elements
(c) Histograms class
(d) Distance and p-values
Figure 10: Sensitivity analysis for a varying total number of elements m2 measured with the Coamnggorov-Sminee test. In this experiment, we have the following specifications: 200, = {1.0, 5.0, 1.0), my = 200 and m3 200. The values of m2 are given in the plot.
Label Adarivv Pppergeements
58
2.
00
(a) Disentanglement
(b) Learned Representation
Figure 1I: As an additional evaluation, we perform analysis on how disentangled the latent repre- sentations are. And related to that, assess the quality of the learned latent representation using linear classifier. We see that the dynamics over different & seems to related for the disentantlement and downstream performance of the learned latent representations.
24

---

Published as a conference paper at ICLR 2023
31
g0(=
x1)
E
View
Aggregation
22
g0(2 x22)
$2
Figure 12: Setup for the weakly-supervised experiment. The three methods differ only in the View Aggregation module.
follows
g0(2) 2z)) =avg(go(z x1), go(2) x2))
A i
(50)
(0(z; xz)) Q (24 x.))
else
(51)
where avg is the averaging function of choice as described above and j {1,2). The methods used (i. e. LabeIVAE, Ada HGVAE) differ in how to select the subset S.
C.2.2 HYPERGEOMETRICVAE (IN MORE DETAIL)
In our approach (HGVAE), we model the number of shared and independent latent factors of a pair of images as discrete random variables following hypergeometric distribution with unknown w. In reference to the urn model, shared and independent factors each correspond to one color and the urn contains d marbles of each color, where d is the dimensionality of the latent space. Given the correct weights w and when drawing from the um d times, the number of each respective color corresponds, expectation, to the correct number of independent/shared factors. The proposed formulation allows to simultaneously infer such w and learn the latent representation in fully differentiable setting within the supervss pipeline by Locatello al. (2019).
To integrate the procedure described above in this framework, we need two additional building blocks. First, we introduce function that returns log w. achieve this, we use a single dense layer which returns the logits log w. The input to this layer is a vector containing the symmetric version of the KI. divergences between pairs of latent distributions, i.e. for latent and Q, the vector contains ((L((P))) Q) + KL(Q)I Second, sampling from the hypergeometric distribution with these weights leads to an estimate and s. Consequently, we need a method to select l factors out of the that are given. Similar to the original paper, select the factors achieving the highest symmetric Kl.-divergence. In order to do this, we sort y in descending order using the stochastic sorting procedure neuralsort (Grover et al., 2019). This enables us to select the top k independent as well as the bottom å = d shared latent factors. Like Ada VAE, we substitute the shared factors by the mean value of the original latent code before continuing the VAE forward pass the usual fashion.
C.2.3 HYPERPARAMETEA SENSITIVITY
We perform ablations to examine how sensitive HypergeometricWAA is to certain hyperparameters. We find that the temperature and the learning rate have the most influence on training stability and convergence. First, if we set the temperature too low, we often observe vanishing or exploding gradients. This well-known artifact of the Gumbel-Softmax trick can be avoided by introducing temperature annealing. Hence, similar to the original Gumbel-Softma trick (Jang et al., 2016; Maddison et al., 2017) and the neuralsort implementation (Grover et al., 2019), we anneal the temperature using exponential function
Tixit exp(- (rt)
(52)
where / is the current training step, Tinit is the initial temperature, and r is the annealing rate:
log Tinst
log inaal
(53)
Rsteps
Tfinal is the final temperature value and Nateps is the number of annealing steps. As shown in Figure 13, training loss and shared factor estimation then converge almost independently of the final temperature. In our final experiments, identical temperatures for both the differentiable
25

---

Published as a conference paper at ICLR 2023
6500
6500
6490
6490
6480
6480
6170
6470
6460
6460
6450
64500
6440
6440
6430
6430
0.0025
0.01
0.04
0.16
0.52
1e-06
1e:05
0.0001
0.
001
(a) Average train loss when varying Tysmat.
(b) Average train loss when varying the learning rate.
10
8-
6-
6-
4-
2-
0-
0-
-2
-2
0.0025
0.01
0.04
0.16
0.52
1e:06
1e-05
0.0001
0.001
(c) Mean squared error estimating the number shared factors when varying Tyonul-
(d) Mean squared error estimating the number shared factors when varying the learning rate.
Figure 13: Ablation of varying the final temperature Tyinal (left) and the learning rate (right) of the
HGVAE in the weakly supervised experiment. We made ablations for the training loss (top) and mean squared error of estimating the number of shared factors (bottom). The final temperature had minimal impact on stability and convergence, whereas higher learning rates led to some instabilities.
hypergeometric distribution and neuralsort. We set the initial temperature Tener to 10 and the final temperature Tyinal to 0.01, which is annealed over nstepe 50000. Further, find the learning rate to be the most crucial hyperparameter in terms of convergence. Higher learning rates generally seem to lead worse training losses. On the other hand, estimating the number of shared factors seems robust on average, although high standard deviations imply decreasing consistency for higher learning rates. We demonstrate this finding Figure 13. We used an initial learning rate of 10- 'ogether with the Adam optimizer (Kingma & Ba, 2014) for our final experiments. Finally, we also experimented with weighting the KL-divergence with a term like the 3-VAE Higgins et al. (2017), where did not find an influence on stability and convergence, Hence, we left it at default value of in our experiments.
C.2.4 DATA
The mpi3d dataset (Gondal al., 2019) consists of frames displaying robot arm and is based generative factors:
object color, shape and size
camera height
background color
horizontal and vertical axis
26

---

Published as a conference paper at ICLR 2023
Table We report the runtimes for the three methods used in the weakly-supervised experiment, i. labelVAE, adaptive VAE and hypergeometrie We report the runtime mean and standard deviation over runs per experiment. We report the runtimes for the different number of independent factors k - 43.
RUNTIME [s]
METHOD
k:
k:
21907.4 + 273.0
20921.1 + 705.9
21389.3 + 163.8
21761.1 + 567.9
ADAPTIVE HYPERGEOMETRIC
29071.5 + 133.2
28609.8 + 439.9
29479.1 + 487.1
29966.3 + 303.3
21888.9 + 632.9
21299.4 = 293.6
21863.9 + 190.1
22241.0 + 137.8
For details the dataset and in general, refer to https://github.com rr-
C.2.5 DOWNSTREAM TASK ON THE LEARNED LATENT REPRESENTATIONS
For the downstream task sample randomly 10000 samples from the training set and 5000 samples from the test set. For each sample, we extract the predicted shared and the predicted independent parts of both views. Then, for every generative factor of the dataset, three individual classifiers are trained on the respective latent representations of the 10000 training samples. Afterwards, every classifier evaluates its predictive performance on the latent representations of the 5000 test samples. To arrive at the final scores, we extract the prediction of the shared factors the shared representation and compute the balanced accuracy. Similarly, we calculate the balanced accuracy of the independent factors the respective independent representation classifiers and average their balanced accuracy. Because the number of classes differs between generative factors we report the adjusted balanced accuracy. We use the implementation from scikit-learn (Pedregosa al., 2011), I). For details, see pettps:/scikktt-tee @elearn.metricn balanced wwwpaaaceuuc_ceeeeeeeeeeeeeeeeee scoreess__seeeee html.
org/ /tgg/stee stable wttaal//mduulleeseeeeeeeeee /ttm//md///////////////////// generated/ /
For all shared generative factors, we average the accuracies of the individual classifiers into a single average balanced accuracy. We do the same for the independent factors. This allows to report the amount of shared and independent information that is present in the learned latent representation. Consequently, we report these averages in the main text.
To evaluate the latent representation we train linear classifiers. More specifically, in this work we use logistic regression classifiers (Cox, 1958) from scikit-learn (Pedregosa et al., 2011). To train the model, we increased the mar_iter parameter so that all models converged and left everything else on default settings.
C.2.6 RUNTIMES OF DIFFERENT ALGORITHMS
In general, our sampling method scales with the number of classes O(c) and the sampling for single class scales with O(m)), which is calculated parallel and single forward pass. The get better and empirically validated picture of the overhead created by using the hypergeometric distribution, we also report the training times for the three methods compared. All methods were trained for an equal number of epochs and on identical hardware equipped with NVIDIA GeForce GTX 1080. Table 3 reports the runtimes for all methods and averaged over runs.
We see that the overhead of using the hypergeometric distribution is almost negligible. The proposed hypergeometric VAE reaches approximately equal training runtimes as the labeIVAE, which assumes that the number of shared factors is known. It even outperforms the adaptive VAE, which uses very simple heuristic, but needs to result to more sophisticated methods order to avoid cutting gradients in their thresholding function.
27

---

Published as a conference paper at ICLR 2023
C.3 CLUSTERING
In this section, we provide the interested reader with more details on the clustering experiments. In the following we describe the models, the optimisation procedure and the implementation details used in the clustering task.
C.3.1
MODEL
We follow a deep variational clustering approach as described by Jiang et al. (2016).
Given dataset X {x that we wish to cluster into K groups, consider the following generative assumptions:
C~ (--p(c; )
-~~ ~(p(c))) -N (Z Ii
(x(xx Ber(Ua,) (54)
where c= [c are the cluster assignments, Z; E IRD are the latent embeddings of VAE and x;
assumed be binary for simplicity.
In particular, assuming the generative process described in Equation (54), we can write the joint probability of the data, also known as likelihood function, as
p(X) = E 2,
Σ p(c; = Σ plc; π II p(x112)p((z(c) c
(55)
Different from Jiang et al. (2016), the prior probability p(c; π) cannot be factorized p(e; T)) for i = 1...., K are not independent. By using variational distribution gz(Z, clX),, we have the following evidence lower bound
p((c (p(c J) (x
)} - LELBO.
log ((X) 2 (z.((i)))) log
(56)
(2,
For sake of simplicity, we assume the following amortized mean-field variational distribution, as in previous work (Jiang et al., 2016; Dilokthanakul et al., 2016):
2 (Z, c)) yo (o( II go go(22 (=(()(xx)))(((((x))
(57)
From where it follows
CELBO = =Bgg(2 p(g)(()))) log log p(Z(c) - log g0 (Z, c(x))
(58)
= =Bqg(211x)aaeeh [log p(c|Tr) (z))) [lo + (2)x [log p( [|c]
(g x) log 90 (2] (. xg(ee [log (c(x)))
(59)
In the ELBO formulation all terms, except the first one, can be efficiently calculated as in previous work (Jiang et al., 2016). For the remaining term, we rely on the following sampling scheme
(2x(2(1))))111)) [logp(c|<))] Σ / qs (c)c)))))))))))ll
(60)
/ 9 (c/c
(61)
log logp(c' m)),
(62)
where we use the SGVB estimator and the Gumbel--Sfttax trick (Jang et al., 2016) to sample from the variational distributions qs (z/ (x:) and go(cx|x;) respectively. The latter categorical distributions with weights given by:
p(G|ZA))) = N(zi|Me, o Te,
(63)
Nz, (z
L is the number of Monte Carlo samples and it is to 1 in all experiments.
28

---

Published as a conference paper at ICLR 2023
qp(2 xz)
E
D
ge(C x))
Figure 14: General Architecture for the clustering experiments. All methods have the same architec- ture details. They only differ their definition of the prior distribution p(c; πT)
C.3.2 IMPLEMENTATION DETAILS
To implement our model we adopted a feed-forward architecture for both the encoder and decoder of the VAE with four layers of 500, 500, 2000, D units respectively, where D = 10. The VAF is pretrained using the same layer-wise pretraining procedure used by Jiang et al. (2016). Each data set is divided into training and test sets, and all the reported results computed the latter. We employed the same eeper-parameters for all experiments. In particular, the learning rate is set to 0.001, the batch size is set to 128 and the models are trained for 1000 epochs. Additionally, used an annealing schedule for the temperature of the Gumbel-Softmax trick. As the VaDE is rather sensitive initialization, weights have been selected by the baseline to enhance the performance of their method, leading optimistic outcome. If random initialization were used instead, the clustering performance would be lower. Nonetheless, the focus of our work is on comparing methods rather than absolute performance values. Figure 14 displays the general architecture of all methods. The methods only differ in their definition of the prior probability distribution p(c; πT.
used the
pretraining weights provided by Jiang et al. (2016). These
29


Reviews:
[{"review_id": "yaNFlYXKLTk", "reviewer": "Reviewer_yMNr", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an approximation of the hypergeometric distribution that introduces a differentiable hyperparameter. The benefit is that, in machine learning algorithms where data may fall into discrete categories, the number and size of the categories is potentially learnable with gradient methods. The paper illustrates this distribution numerically in two settings: learning the number of dependent components in a weak supervision setting and clustering.", "summary_of_the_review": "Summary\n\n* Potentially a useful methodology for working with distributions having underlying categories in some situations.\n\n", "strength_and_weaknesses": "Strengths\n\n* Seems like a potentially useful tool.\n\n\nWeaknesses\n\n* Methodology could still use further investigation. The Kolmogorov-Smirnov test is one nice step, but it would also be nice to mathematically quantify the difference between the hypergeometric and the proposed distribution with standard measures of distributional similarity (various divergences, TV distance, etc.).\n\n\nQuestions\n\n* It seems like this methodology is most applicable for variational inference, which combines gradient optimization and distributional inference. Are there other methodological applications I should have in mind?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "sR_JYUZoqx", "reviewer": "Reviewer_gdaG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an approximation of the multivariate non-central hypergeometric distribution in which, first the multivariate distribution is expressed through the product rule as a product of conditional distributions, each of which can be treated as a univariate non-central hypergeometric distribution by grouping the remaining (un-accommodated/sampled) classes into one. Finally, these univariate non-central hypergreometric distributions are approximated using the Gumbel-Softmax approximation technique. This makes the differentiation of the joint distribution with respect to the class importance weights feasible.\n\nThe paper concludes with some experiments; first showing the accuracy of the approximation through simulating from the true and approximate distributions, and then in two real applications where it shows superiority to some existing methods in learning numbers of shared and independent latent generative factors from coupled image observations and in clustering, where the proposed distribution is used as a tunable prior for a VAE.", "summary_of_the_review": "The paper covers an interesting and under-explored topic in machine learning, and the proposed method seems to offer options for modelling, estimation and inference which are otherwise quite limited to heuristics. Although the paper is very readable and the details are well explained, as a reader one might feel still somewhat at a loss for some of the higher level practical aspects of the proposed approximation.\n The experiments show that the proposed distribution has realisable benefits in some important and interesting applications over existing alternatives and more naive straightforward formulations.", "strength_and_weaknesses": "Strengths:\n- The work investigates a relatively under-studied topic in machine learning where alternatives perform modelling based on i.i.d. sampling and may, at best, use post-hoc assessment of group sizes or similar to infer importance weights. By building the importance weights into the estimation both inference and estimation can be improved.\n- The empirical results show that the method can offer improvements on existing methods in some interesting and relevant problems.\n\nWeaknesses:\n- The main weakness of the paper is that it expects a lot of knowledge from the reader, and many \"for details see...\" pieces of text which may make the work not stand well along as a single item of literature.\n- There are also potential issues with clarity, where while the individual points are clear it is not immediately forthcoming how the optimisation rendered possible by the differentiability of the mass function is implemented.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "X8TiSlkDeFV", "reviewer": "Reviewer_xbJ8", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a differentiable formulation of the hypergeometric distribution using the Gumbel-Softmax trick. They provide a sampling algorithm and demonstrate its use on two tasks that require modeling of discrete distributions. The first is a weakly-supervised representation learning problem for video frames. The second is a deep variational clustering application. In both applications, the use of the differentiable hypergeometric distribution improves performance. Additionally, the authors demonstrate that for a wide range of parameter values, one cannot distinguish between the a baseline non-differentiable implementation of the hypergeometric distribution and the proposed differentiable implementation using a KS test.\n", "summary_of_the_review": "The authors propose a useful tool for modeling discrete problems in differentiable frameworks.\n", "strength_and_weaknesses": "Modeling categorical data using neural architectures has always been a challenge. The proposed method extends the tools available to model different classes of problems. Additionally, this is accomplished without significantly increasing the number of parameters.\n\nMy only concern with this work is the issue of fitting the importance weights. The authors do not seem to spend any time discussing the implications of this method on the convergence of gradient optimization routines. Can this method cause exploding or vanishing gradients? If so, under what conditions? Some remarks, analysis, or empirical evidence on this topic would be welcome.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "IRm_R6P66D7", "reviewer": "Reviewer_uGmy", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a differentiable approximation of the hypergeometric distribution and shows that it works well.", "summary_of_the_review": "Well written comprehensive treatment of a practical tool that makes building models that rely on the hypergeometric distribution easier.", "strength_and_weaknesses": "Strengths:\n- Very clear contribution in form of a practical tool that can be used by others\n- Potential for high impact, in a similar manner as the concrete/GS distribution had an impact on how categorical variables are used in broad range of models\n- Very well written for a paper this technical\n\nWeaknesses:\n- Nothing really, but a graphical illustration of the hypergeometric distribution and the role of the augmentation would not hurt; it would be especially useful if this becomes a core reference for people not so familiar with the literature", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%