C ONTINUOUS R ELAXATION FOR THE M ULTIVARIATE N ONCENTRAL H YPERGEOMETRIC D ISTRIBUTION
A BSTRACT
Partitioning a set of elements into a given number of groups of a priori unknown sizes is an essential task in many applications.
Due to hard constraints, it is a non-differentiable problem that prohibits its direct use in modern machine learning frameworks.
Hence, previous works mostly fall back on suboptimal heuristics or simplified assumptions.
The multivariate hypergeometric distribution offers a probabilistic formulation of sampling a given number of elements from multiple groups.
Unfortunately, as a discrete probability distribution, it neither is differentiable.
We propose a continuous relaxation for the multivariate non-central hypergeometric distribution.
We introduce an efficient and numerically stable sampling procedure that enables reparameterized gradients for the hypergeometric distribution and its integration into automatic differentiation frameworks.
We additionally highlight its advantages on a weakly-supervised learning task.
I NTRODUCTION
A variety of machine learning approaches, such as deep generative models, rely on differentiable sampling procedures, most notably the reparameterization trick for Gaussian distributions (Kingma & Welling, 2014; Rezende et al., 2014).
As part of either generative models with discrete latent variables (Kingma et al., 2014) or attention mechanisms (Xu et al., 2015), the categorical distribution is another example.
The concrete distribution (Maddison et al., 2017) or Gumbel-Softmax trick (Jang et al., 2016) has recently boosted its usage 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.
Following the Gumbel-Softmax trick, reparameterization procedures combined with continuous relaxations for discrete distributions and other non-differentiable algorithms gained attraction in machine learning research (Grover et al., 2019; Xie & Ermon, 2019; Paulus et al., 2020).
Despite enormous progress in recent years, the extension to more complex probability distributions is still missing or comes with a trade-off regarding differentiability or computational speed (Huijben et al., 2021).
The hypergeometric distribution is a discrete probability distribution that describes the probability of x successes out of n trials from a total population of size N with m members belonging to the success class (Upton & Cook, 2014).
The most common application is sampling without replacement, for which an urn model is the standard example.
In the univariate case, the urn consists of marbles in two different colours, whereas in the multivariate case, there are marbles in c different colours.
The hypergeometric distribution plays an essential role in various applications ranging from biology and gene mutations to computer science and recommender systems or social sciences and the analysis of networks (Becchetti et al., 2011; Lodato et al., 2015; Casiraghi et al., 2016).
It is essential wherever the choice of a single element influences the distribution over classes of the remaining elements.
Previous work uses the hypergeometric distribution implicitly to model assumptions or as a tool to prove theorems. However - to the best of our knowledge - it is not yet integrated into the learning processes themselves.
In this work, we propose a continuous relaxation for the multivariate non-central hypergeometric distribution that enables reparameterized gradients with respect to the distribution parameters.
A differentiable formulation and reparameterizable gradients enable the integration of the hypergeo- metric distribution into stochastic networks and modern learning frameworks.
The hypergeometric distribution introduces dependencies between elements or samples where previous works had to assume independence.
We demonstrate the potential of the hypergeometric distribution for a weakly-supervised learning task where we assign latent dimensions to groups of unknown size.
We perform a second experiment described in the appendix, where we highlight the advantage of the proposed model in a clustering algorithm.
We model the number of samples per cluster using an adaptive hypergeometric distribution prior, overcome the simplified i.i.d.
assumption of previous work, and establish a dependency structure between dataset samples.
M ETHOD
Definition 2.1 (Multivariate Fisher’s Non-Central Hypergeometric Distribution) .
A random vector
X follows Fisher’s non-central multivariate distribution, if its joint probability mass function is given by (Fisher, 1935) [Equation] where [Equation]
where (cid:0) nk (cid:1) = n !
k !
( n − k )! , c ∈ N + is the number of different classes (e.g. marble colours in the urn), m =
[ m 1 , . .
, m c ] ∈ N c describes the number of elements per class (e.g. marbles per colour),
N = (cid:80) ci =
1 m i is the total number of elements (e.g. all marbles in the urn) and n ∈ { 0 , . . .
, N } is the number of elements (e.g. marbles) to draw.
{+ω ∈ R c 
+}
+ describe the class weights, which introduce selection bias between classes.
The support of the non-central hypergeometric distribution is equal to the support S of the central distribution (Equation (10)).
For general information on the hypergeometric distribution, we refer the interested reader to Appendix A.
2. S AMPLING FROM THE N ON -C ENTRAL H YPERGEOMETRIC D ISTRIBUTION
There exist different sampling procedures for the non-central hypergeometric distribution in the literature.
Each of these comes with its own advantages and disadvantages (for more details see (Fog, 2008b)).
For the purpose of finding an efficient and numerically stable sampling procedure that is differentiable and batch-wise computable, we choose the conditional sampling procedure.
The first class X 1 is sampled from the marginal distribution,
then the following classes X i with i > 1 are sampled from the conditional distribution given all previously sampled classes { X 1 = x 1 , . . .
, X i −
= x i − 1 } (Johnson, 1987).
As all classes are sampled sequentially, this algorithm scales linearly with the number of classes, and not with the number of possible outcomes of a random draw as in previous works (Liao & Rosen, 2001).
Broadly speaking, the sampling algorithm consists ofparts:
Reformulate the multivariate distribution as a sequence of interdependent and conditional univariate distributions.
Calculate the PMF of the respective univariate distributions.
Sample from this conditional distribution utilizing the Gumbel-Softmax trick.
The sampling procedure is described in Algorithm 1 using pseudo-code, while the details of the individual steps are explained and discussed in the following sections of this work.
2. S EQUENTIAL S AMPLING U SING C ONDITIONAL D ISTRIBUTIONS
As stated above, we reformulate the problem using a conditional sampling scheme by sequentially sampling from a conditional univariate distribution instead of a multivariate one.
In the first step, class i = 1 is seen as the success class.
The remaining classes i > 1 are merged together into a single, negative class.
We sample x 1 elements from class i = 1 out of n elements where n is the total number of sampled elements.
In the second step, class i = 2 is seen as success class and, again, the remaining classes i > 2 are seen as the negative class.
We sample x 2 elements from class i = 2 .
Because we already sampled x 1 elements in the first step, the number of elements to sample are n − x 1 elements at this stage.
From step 1 to step 2 , the number of elements to sample from (e.g. the total number of marbles in the urn) changes from N 1 = (cid:80) ci =
1 m i to N 2 = N 1 − m 1 .
We continue with this procedure until no class is left to be sampled from, i.e. i =
c
(Fog, 2008b).
More formally, we define two classes as L and R with the respective parameters that are given by (Fog, 2008b)
[Equation]
where ω l and ω r are the respective class weights.
In the first step, we select a single class i as set L , and the remaining classes as set R .
Then, for each remaining steps, we select a single class j from set R as our new set L .
Although we follow a simple approach, there exist various algorithms for creating univariate distributions out of subsets of classes.
We leave the exploration of different and more sophisticated subset selection strategies for future work.
For non-central distributions, this method produces biased samples as merging multiple classes into one is only approximately equal to the univariate non-central distribution.
The selection of which classes are merged together might influence the quality of the approximation (Fog, 2008b).
2. C ONTINUOUS R ELAXATION
After simplifying the problem of sampling from the joint distribution into sequentially sampling from conditional distributions, we present continuous relaxations for these.
Continuous relaxations describe procedures to make discrete distributions differentiable with respect to their parameters (Huijben et al., 2021).
Lemma 2.2.
The Gumbel-Softmax trick can be applied to the conditional distribution P ( X i =
x i |{ X k = x k } i − 1 k =1 ) of class i given the already sampled classes
k < i .
The proof for Lemma 2.2 can be found in Appendix B.2.
The Gumbel-Softmax trick enables the reparameterization of categorical distributions and their gradients with respect to its parameters.
We make use of the Gumbel-Softmax trick to reparameterize the conditional distributions, which results in an efficient and reliable sampling procedure.
The reparameterization scheme for the univariate non-central hypergeometric distribution goes as follows:
[Equation]
where u ∈ [0 , 1] m i +1 is a random sample from an i.i.d.
uniform distribution U .
As a result, g i is i.i.d.
gumbel noise.
ˆ r
i is the perturbed conditional probability for class i given the class conditional unnormalized log-weights α i .
We use the softmax function to generate ( m i + 1) -dimensional sample vectors from the perturbed unnormalized weights ˆ r
i /τ , where τ is the temperature parameter.
Note that softmax is the continuous and differentiable approximation to the argmax function.
We refer the reader to Jang et al.
(2016) or Maddison et al.
(2017) for more details on the Gumbel-Softmax trick itself.
This procedure relates to the contRelaxSample function in Algorithm 1.
For the multivariate non-central hypergeometric distribution, the learnable parameters of the distribution are the relative class weights ω .
It is important to mention that the assignment of ω is independent of the sampling procedure and therefore ω are not affected by the sequential sampling of classes.
This is essential for the reparameterization of distributions (for more details see Huijben et al.
The remaining parameters of the multivariate distribution are fixed by design of experiments, m , or do not influence the ordering of elements, n .
[Figure]
2. C ALCULATE P ROBABILITY M ASS F UNCTION
In Sections 2.2 and 2.3, we derive a sequential sampling procedure.
In this sequential procedure we repeatedly simulate a univariate distribution to simplify sampling.
As a result, the PMF of the joint multivariate distribution is not needed in the sampling procedure.
Instead we can restrict the computation of the PMF to the univariate case.
For the multivariate extension, see Appendix B.1.
The
PMF P X L ( x ) for the hypergeometric distribution of two classes L and R defined by m L , m R , ω L , ω R and n is given as
[Equation]
P 0 is defined as in Equation (2).
Lemma 2.3.
The unnormalized log-probabilities
[Equation]
define the unnormalized weights of a categorical distribution that follows Fisher’s non-central hypergeometric distribution.
ψ F ( x ) is defined as follows
[Equation]
The proof for Lemma 2.3 can be found in Appendix B.3.
Additionally, calculations in the log -domain increase numerical stability and also preserve the ordering between different values x (Huijben et al., 2021).
This subsection relates to the calcLogPMF function in Algorithm 1.
Equation (7) and more general Equation (6) can be calculated efficiently and reliably as common automatic differentiation frameworks have numerically stable implementations of log Γ( x ) , e.g. Tensorflow (Abadi et al., 2016) or PyTorch (Paszke et al., 2019).
Additionally, all necessary steps can be calculated batch-wise.
Using the multivariate form of Lemma 2.3 (see Appendix B.1), it is possible to directly calculate the categorical weights for the multivariate states.
A multivariate state is defined as valid combination of X = x .
It would result in a significant speed-up for large number of classes c compared to the conditional sampling procedure.
However, the number of multivariate states is (cid:81) ci =1 m i , which quickly results in unfeasible memory requirements.
Therefore, we would be restricted to settings with no practical relevance.
E XPERIMENT ON W EAKLY -S UPERVISED L EARNING
In this experiment, we induce weak-supervision through pairs of images that may share part of their generating factors.
In general, neither the true number of generative factors nor the number of shared latent factors is known.
We use the setting and code from Locatello et al.
and compare the models on the mpi3d toy dataset (Gondal et al., 2019).
All models are trained as a variational autoencoder (Kingma & Welling, 2014, VAE) to maximize an evidence lower bound (ELBO) on the marginal log-likelihood of the pair of images.
We perform the same experiment for multiple dataset versions with different numbers of independent generative factors k .
By independent we mean generative factors that are not shared between pairs of images.
We compare three methods that aim at inferring shared and independent latent factors.
All methods aggregate shared factors that replace the individual factors.
The difference among methods is only in the selection process of shared factors.
The first one assumes that the number of independent factors is known (Bouchacourt et al., 2018; Hosoya, 2018, LabelVAE).
Similar to the experiments in
Locatello et al.
, we assume this to be 1 for all experiments.
The second method relies on an adaptive heuristic to infer the number of shared latent factors (Locatello et al., 2020, AdaptiveVAE).
First, they calculate the Kullback-Leibler (KL) divergence between pairs (across images) of latent factors.
Then, their heuristic defines a threshold to determine whether a pair of factors encodes the same information, i.e. is shared, or not.
In our approach (HypergeometricVAE), we model the number of shared and independent latent factors of a pair of images as a multivariate hypergeometric distribution with unknown ω ∈ R 20+ .
For d being the number of latent factors, we have m i =
d where i ∈ { 1 , 2 } and the number of elements to draw n = d .
Given class weights ω , we are able to sample estimates for the k independent and d − k shared factors.
The proposed formulation allows to infer such ω and simultaneously learn the latent representation in a fully differentiable setting.
Like in AdaptiveVAE, we utilize the KL divergences between latent factors.
They are passed to a single dense layer which outputs ω .
We make use of the stochastic sorting algorithm (Grover et al., 2019) for sorting the latent factors by KL divergence.
Next, we define the top- k latent factors to be independent, and the remaining d − k to be shared ones where k and d − k are samples from the hypergeometric distribution.
For a more detailed description of HypegeometricVAE, the baseline models and the dataset, see Appendix C.2.
To assess the performance of the three methods, we evaluate the trained models on two different tasks and over four mpi3d based training sets with the true number of independent factors k set to 2 , 4 , 6 and uniformly random between 1 and 7 ( k = − 1) .
We measure the mean squared error of predicting the number of shared latent factors (Figure 1a) and quantify downstream performance of the learned latent representations on a classification task (Figure 1b).
As can be seen, previous methods are not able to accurately estimate the number of shared factors.
In fact, both baseline methods estimate the number of shared factors to be approximately constant for all experimentsindependent of the underlying ground truth number of shared factors.
For the first model, this does not come as a surprise, whereas this was quite unexpected for the second approach given their in theory adaptive heuristic.
On the other hand, the proposed approach is able to accurately estimate the number of shared factors which is reflected in the low mean squared error (MSE) for all experiments, including the most difficult experiment with random number of independent factors.
To assess the quality of the learned latent representation, we evaluate it with respect to all ground truth generative factors of the dataset.
We train classifiers for all factors and calculate their accuracy.
The reported accuracy is the average over the factor-specific accuracies.
Given the general nature of the method, the positive results of the proposed method are quite surprising.
Different to previous works, it was not explicitly designed for weakly-supervised learning but is nevertheless able to achieve results that are more than comparable to field-specific models.
Additionally, the proposed method is able to give accurate estimates on the latent space structure for different experimental settings.
C ONCLUSION
We propose a continuous relaxation for the multivariate noncentral hypergeometric distribution in this work.
In combination with the Gumbel-Softmax trick, this new formulation enables reparameterized gradients with respect to the class weights ω of the hypergeometric distribution.
We show the potential of the hypergeometric distribution in machine learning by applying it to a weakly-supervised learning task.
In this application, the method making use of the hypergeometric distribution reaches state-ofthe-art performance.
We believe this work is an essential step toward integrating the hypergeometric distribution into more machine learning algorithms.
Applications in biology and social sciences represent potential directions for future work.
A P RELIMINARIES
The hypergeometric distribution is a discrete probability distribution that describes the probability of x successes in n draws without replacement from a finite population of size N with m elements that are part of the success class.
This is different to the binomial distribution which describes the probability distribution of x successes in n draws with replacement.
Definition A.1 (Hypergeometric Distribution (Gonin, 1936) 1 ) .
A random variable X follows the hypergeometric distribution, if its probability mass function (PMF) is given by
[Equation]
Urn models are classical examples of the hypergeometric probability distribution.
If we think of an urn with marbles in two different colours, e.g. green and purple, we can label as success the drawing of a green marble.
Then N defines the total number of marbles and m the number of green marbles in the urn.
x is the number of green marbles and n − x the number of purple marbles that are drawn.
The multivariate hypergeometric distribution describes an urn with more than two colours, e.g. green, purple and yellow in the simplest case with three colours.
As described in Johnson (1987), the definition is given by:
Definition A.2 (Multivariate Hypergeometric Distribution) .
A random vector X follows the multivariate hypergeometric distribution, if its joint probability mass function is given by
[Equation]
where c ∈ N + is the number of different classes (e.g. marble colours in the urn), m =
[ m 1 , . .
, m c ] ∈ N c describes the number of elements per class (e.g. marbles per colour),
N = (cid:80) ci =
1 m i is the total number of elements (e.g. all marbles in the urn) and n ∈ { 0 , . . .
, N } is the number of elements (e.g. marbles) to draw.
The support S of the PMF is given by
[Equation]
Input: m ∈ Z c 0+ , ω ∈ R c 0+ , n ∈ N , τ ∈ R +
Output: x ∈ Z c 0+ , { α i ∈
R m i } ci =1 , { ˆ r i ∈
R m i } ci = for i ∈ { 1 , . . .
, c } do
L ← i, R ← { (cid:83) cj
= i +1 j } # Formulate the multivariate as a m → m L , m R ∈
Z 0+ , ω → ω L , ω R ∈ R 0+ # univariate distribution (Section 2.2) x L , α L , ˆ r L ← sampleUNCHG( m L , m R , ω L , ω R , n, τ )
# Sample from univariate distribution n ← n − x L , m ← m \ m L , ω
← ω \ ω L # re-assign classes for next step x
i ← x L , α i ← α L , ˆ r i ← ˆ r L # assign values for class i end for return x , { α i } ci =1 , { ˆ r i } ci = function SAMPLE UNCHG
( m i , m j , ω
i , ω j , n, τ ) α i ← calcLogPMF( m i , m j , ω i , ω j , n ) # Section 2.
x i , ˆ r i ← contRelaxSample
( α i , τ ) )
# Section 2.
return x i , α i , ˆ r i end function
Definitions A.1 and A.2 assume that the only property to distinguish between marbles is their colour.
In other words, there is no property that leads to a selection bias between the different classes.
Every marble is to be picked equally likely and the number of selected class samples is proportional to the ratio between class elements and total number of elements in the urn.
Often this assumption is too restrictive and selection bias is a desired model property.
Biased generalizations, which make certain colours more likely to be picked, exist for the uni- and the multivariate case.
We focus on the multivariate distribution as the univariate distribution is a special case thereof.
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), we will refer to Fisher’s version of the noncentral hypergeometric distribution in the remaining part of this work.
B M ETHODS B. PMF FOR THE MULTIVARIATE F ISHER ’ S NONCENTRAL DISTRIBUTION
In this section, we give a detailed derivation for the calculation of the log-probabilities of the multivariate Fisher’s noncentral hypergeometric distribution.
We end up with a formulation that is proportional to the actual log-probabilities.
Because the ordering of categories is not influenced by scaling with a constant factor (addition/subtraction in log domain), these are unnormalized log-probabilities of the multivariate Fisher’s noncentral hypergeometric distribution.
[Equation]
[Equation]
where P 0 is defined as in Equation (2).
From there it follows
[Equation]
Constants factor can be removed as the argmax is invariant to scaling with a constant factor which equals addition or subtraction in log -space.
It follows
[Equation]
In the last line we used the relation Γ( k + 1)
=
k ! .
[Equation]
[Equation]
where ψ F ( x ) = − (cid:80) ci =1 log (Γ( x i + 1)Γ
( m i − x i + 1)) .
The Gamma function is defined as (Whittaker & Watson, 1996)
[Equation]
B. P ROOF FOR L EMMA 2.
Proof.
When sampling class i , we draw x i samples from class i where x i ≤ m i .
The conditional distribution P ( X i = x
i |{ X k = x k } i − 1 k =1 ) for class i given the already sampled classes
k < i simultaneously defines the weights of a categorical distribution.
Sampling x
i elements from class i can be seen as selecting the x i th category from the distribution defined by the weights P
( X i =
x i |{ X k = x k } i − 1 k =1 ) .
Therefore, (cid:80) x i ≤ m i P X i ( x i ) = 1 , which allows us to apply the Gumbel-Max trick and, respectively the Gumbel-Softmax trick.
B. P ROOF FOR L EMMA 2.
Proof.
Factors that are constant for all x do not change the relative ordering between different values of x .
Hence, removing them preserves the ordering of values x (Barrett, 2017).
[Equation]
[Figure]
m i = 200 ∀ i and n = 180 .
The class weights ω 1 and ω 3 for classes 1 and 3 are set to 1 . 0 , ω 2 is increased from 1 . 0 to 10 .
with a step size of 1 . 0 (defined as w2 in the figure legend).
Figure 2a shows histograms of the number of elements for class 2 .
We see that the behaviour across different class weights ω 2 is very similar between the proposed and the reference implementation.
This behaviour is also reflected in values of the KS test and their respective p -values in Figure 2b.
Using the definition of the binomial coefficient (see Appendix A) and its relation to the Gamma function 2 Γ( k + 1)
=
k ! , it follows log P X L ∝ x · log ω L + ( n − x ) · log ω R (24) − log Γ( x + 1) − log Γ( m L − x + 1) − log Γ( n − x + 1) − log Γ( m R − n + x + 1)
With ψ F ( x ) as defined in Equation (7) it follows log P X L ∝ x · log ω L +( n − x ) · log ω R + ψ F ( x ) .
C E XPERIMENTS C. K OLMOGOROV -S MIRNOV T EST
To assess the accuracy of the proposed method, we evaluate it against a reference distribution using the Kolmogorov-Smirnov test (Kolmogorov, 1933; Smirnov, 1939, KS).
It is a nonparametric test to evaluate the equality of two distributions.
It can be used for continuous as well as discrete 1dimensional probability distributions.
We compare the class conditional hypergeometric distributions against a reference distribution from SciPy (Virtanen et al., 2020).
As described in Section 2 class conditional distributions are used to sample from the multivariate distribution.
For this experiment, we use a multivariate hypergeometric distribution of three classes.
We perform a sensitivity analysis with respect to the class weights ω .
We keep ω 1 and ω 3 fixed at 1 . 0 , ω 2 is increased from 1 . 0 to 10 .in steps of 1 . 0 .
For every value of ω 2 , we sample 50000 i.i.d.
random vectors.
The KS test quantifies a distance between the empirical distributions of two groups of samples.
Hence, a small test value implies that the underlying distributions are similar.
The null distribution of this test is calculated under the null hypothesis that the two groups of samples are drawn from the same distribution.
In our setting, we would like the test to fail to reject the null hypothesis.
In this case, the two groups of samples were generated by the same distribution, i.e. the two underlying distributions are equal.
To reject the null hypothesis, the p -value should be below the significance threshold t = 0 .
And vice versa, p > 0 .
05 implies that the null hypothesis cannot be rejected which is desirable for our application.
Figure 2 shows the histogram of values of class 2 for all values of ω 2 (Figure 2a) and the results of the KS test for all classes (Figure 2b).
The histograms for class 1 and 3 can be found in the Appendix (Figure 3).
We see that the histograms of the proposed and the reference distribution are visually similar.
Additionally, the calculated distances of the KS-test are small and the corrected p > 0 . 05 .
We correct the p -values for false discovery rate of multiple comparisons as we are performing c =tests per joint distribution.
We use the Benjamini-Hochberg correction (Benjamini & Hochberg, 1995).
In Figure 2b, we see that the test is clearly significant in 30 out 30 cases.
For all tests p -values [Figure] and are well above the threshold and many are close to 1.0.
The results of the KS test give strong evidence that the proposed continuous relaxation and sampling procedure follow a noncentral hypergeometric distribution.
C. W EAKLY -S UPERVISED L EARNING C.2. M ETHOD , I MPLEMENTATION AND H YPERPARAMETERS
In this section we give more details on the used methods.
We make use of the disentanglement lib (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 disentanglement 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 β = 1 .
as this is the best performing β according to Locatello et al.
For all experiments we train 5 models with different random seeds.
All experiments are performed using GroupVAE (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.
Different to GroupVAE, their ML-VAE aggregates shared latent factors by using the Product of Experts (i.e. geometric mean).
C.2. H YPERGEOMETRIC VAE ( IN MORE DETAIL )
In our approach (HypergeometricVAE), we model the number of shared and independent latent factors of a pair of images as a hypergeometric distribution with unknown ω .
In reference to the urn model, shared and independent factors each correspond to one colour and the urn contains d marbles of each colour, where d is the dimensionality of the latent space.
Given the correct weights ω and when drawing from the urn d times, the number of each respective colour corresponds, in expectation, to the correct number of independent/shared factors.
The proposed formulation allows to simultaneously infer such ω and learn the latent representation in a fully differentiable setting within the weakly-supervised pipeline by Locatello et al.
To integrate the procedure described above in this framework, we need two additional building blocks.
First, we introduce a function that returns ω .
To achieve this, we use a single dense layer followed by a tangens hyperbolicus activation function which returns logits ˜ ω .
The input to this layer is a vector γ containing the symmetric version of the KL divergences between pairs of latent distributions, i.e. for latent P and Q , the vector contains 12 ( KL ( P || Q ) + KL ( Q || P )) .
An estimation ˆ ω is than given by exp (˜ ω ) .
Second, sampling from the hypergeometric distribution with these weights leads to an estimate ˆ k .
Consequently, we need a method to select ˆ k factors out of the d that are given.
Similar to the original paper, we select the factors achieving the highest symmetric KL-divergence.
In order to do this, we sort γ 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 − ˆ k shared latent factors.
Like AdaptiveVAE, we substitute the shared factors by the mean value of the original latent code before continuing the VAE forward pass in the usual fashion.
C.2. D ATA
The mpi3d dataset (Gondal et al., 2019) consists of frames displaying a robot arm and is based ongenerative factors.
• object colour • object shape • object size • camera height • background colour • horizontal axis • vertical axis
[Table]
For more details on the dataset and in general [Table] we refer to https://github.com/ rr-learning/disentanglement_dataset .
C.2. D OWNSTREAM TASK ON THE LEARNED LATENT REPRESENTATIONS
For the downstream task we sample randomly 10000 samples from the training set and 5000 samples from the test set.
For every generative factor of the dataset, an individual classifier is trained on the 10000 training samples.
Afterwards, every classifier evaluates the latent representations of the 5000 test samples.
The classification accuracies of the individual classifiers is averaged into a single average accuracy.
The average accuracy is the one reported.
C. D EEP V ARIATIONAL C LUSTERING
We investigate the use of the multivariate noncentral hypergeometric distribution in a deep clustering task.
Several techniques have been proposed in the literature to combine long-established clustering algorithms, such as K-means or 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 et 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.
This permits a probabilistic approach to clustering where a clear generative assumption of the data is defined and optimised within the framework of stochastic gradient variational Bayes (Kingma & Welling, 2014; Rezende et al., 2014).
A major drawback of the above models is that the samples are either assumed to be i.i.d.
or they require pairwise side information, which limits their applicability in real-world scenarios.
The multivariate noncentral hypergeometric distribution can be easily integrated in VAE-based clustering algorithms to overcome limitations of current approaches.
Given a dataset X = { x i } Ni = that we wish to cluster into K groups, we consider the following generative assumptions: [Equation]
[Equation]
where c = { c i } Ni =1 are the cluster assignments, z
i ∈ R D are the latent embeddings of a VAE and x i is assumed to be binary for simplicity.
In other words, we assume the data is generated from a random process where the cluster assignments are first drawn from a prior probability p ( c ; π ) , then each latent embedding z i is sampled from a Gaussian distribution, whose mean and variance [Figure] depend on the selected cluster c i .
Finally the sample x i is generated from a Bernoulli distribution whose parameter µ x i is the output of a neural network parameterized by θ , as in the classical VAE.
With these assumptions, the latent embeddings z i follow a mixture of Gaussian distributions, whose means and variances, { µ i , σ 2 i } Ki =1 , 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.1 for a complete description of the optimisation procedure).
Previous work (Jiang et al., 2016) modeled the prior distribution as p ( c ; π ) =
(cid:81)
i p ( c i ) =
(cid:81) i Cat ( c i | π ) with either tunable or fixed parameters π .
In this task, we instead replace this prior with the multivariate noncentral hypergeometetric distribution with weights π and K classes where every class relates to a cluster.
Hence, we sample the number of samples per cluster (or cluster size) following Definition 2.1.
The hypergeometric distribution permits to create a dependence between samples.
The prior probability of a sample to be assigned to a given cluster is not independent of the remaining samples anymore, allowing us to loosen the over-restrictive i.i.d. assumption.
We explore the effect of three different prior probabilities in Equation 25, namely (i) the categorical distribution, by setting p ( c ; π ) =
(cid:81) i Cat ( c i | π ) ; (ii) the uniform distribution, by fixing π i
= 1 /K
for i = 1 . . .
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 a 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 1 we evaluateevaluate the methods with respect to their clustering accuracy (Acc), normalized mutual information (NMI) and adjusted rand index (ARI).
As can be seen the hypergeometric prior distribution shows fairly good clustering performance in all datasets.
It is also assuring that the model is able to learn the weights, π , which reflect the subsampling rates of each cluster (see Figure 4).
Although the uniform distribution performs reasonably good, it assumes the clusters have equal importance, hence it might fail in more complex scenarios.
The categorical distribution, on the other hand, 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 hypergeometric prior, but rather boost its clustering performance, especially in the imbalanced dataset.
C.3. M ODEL
We follow a deep variational clustering approach as described by Jiang et al.
In particular, assuming the generative process described in Equations 25, 26 and 27, we can write the join probability of the data, also known as the likelihood function, as
[Equation]
Different from Jiang et al.
(2016), the prior probability p ( c ; π ) cannot be factorized as p ( c i ; π ) for i = 1 , . .
, K are not independent.
By using a variational distribution q ϕ ( Z, c | X ) , we have the following evidence lower bound
[Equation]
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):
[Equation]
From where it follows
[Equation]
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
[Equation]
where we use the SGVB estimator and the Gumbel-Softmax trick (Jang et al., 2016) to sample from the variational distributions q ϕ
( z i | x i ) and q ϕ ( c i | x i ) respectively.
The latter is set to a categorical distributions with weights given by:
[Equation]
L is the number of Monte Carlo samples and it is set to 1 in all experiments.
C.3. I MPLEMENTATION D ETAILS
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 VAE is pretrained using the same layer-wise pretraining procedure used by Jang et al.
Each data set is divided into training and test sets, and all the reported results are computed on the latter.
We employed the same hyper-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, we used an annealing schedule for the temperature of the Gumbel-Softmax trick.