# Sampling From Autoencoders’ Latent Space via Quantization And Probability
Mass Function Concepts

Aymene Mohammed Bouayed ${}^{\dagger\,\ddagger}$, Adrian Iaccovelli ‡, David
Naccache †
†DIÉNS, ÉNS, CNRS, PSL University, Paris, France
45 rue d’Ulm, 75230, Paris cedex 05, France
<EMAIL_ADDRESS>
‡Be-Ys Reseach, France
46 rue du ressort, Clermont-Ferrand, France
<EMAIL_ADDRESS>
###### Abstract

In this study, we focus on sampling from the latent space of generative models
built upon autoencoders so as the reconstructed samples are lifelike images.
To do to, we introduce a novel post-training sampling algorithm rooted in the
concept of probability mass functions, coupled with a quantization process.
Our proposed algorithm establishes a vicinity around each latent vector from
the input data and then proceeds to draw samples from these defined
neighborhoods. This strategic approach ensures that the sampled latent vectors
predominantly inhabit high-probability regions, which, in turn, can be
effectively transformed into authentic real-world images. A noteworthy point
of comparison for our sampling algorithm is the sampling technique based on
Gaussian mixture models (GMM), owing to its inherent capability to represent
clusters. Remarkably, we manage to improve the time complexity from the
previous $\mathcal{O}(n\times d\times k\times i)$ associated with GMM sampling
to a much more streamlined $\mathcal{O}(n\times d)$, thereby resulting in
substantial speedup during runtime. Moreover, our experimental results, gauged
through the Fréchet inception distance (FID) for image generation, underscore
the superior performance of our sampling algorithm across a diverse range of
models and datasets. On the MNIST benchmark dataset, our approach outperforms
GMM sampling by yielding a noteworthy improvement of up to $0.89$ in FID
value. Furthermore, when it comes to generating images of faces and ocular
images, our approach showcases substantial enhancements with FID improvements
of $1.69$ and $0.87$ respectively, as compared to GMM sampling, as evidenced
on the CelebA and MOBIUS datasets. Lastly, we substantiate our methodology’s
efficacy in estimating latent space distributions in contrast to GMM sampling,
particularly through the lens of the Wasserstein distance.

##  1 Introduction

The realm of image generation in biometric imaging has undergone a notable
surge in recent times, with a variety of generative models rooted in
Generative Adversarial Networks (GANs) [12], as in [23, 36, 47]. In contrast,
our approach charts a distinct path by employing autoencoder-based generative
models for image synthesis. This decision is underpinned by the intricate and
potentially unstable learning phase inherent to GANs, where improper tuning
can lead to complications such as mode collapse [32]. As a result, our focus
in this study centers on the nuanced task of sampling within the latent space
of autoencoders.

Existing literature presents three primary strategies for latent space
sampling within autoencoders. The first approach involves the incorporation of
a prior distribution onto the latent space through a regularization term. To
acquire a vector sample, we draw from this prior distribution. Variational
Autoencoders (VAEs) [18] are a representative model category following this
strategy. However, a limitation of these methods emerges after training, where
the latent space might not precisely adhere to the predefined prior
distribution, leading to the emergence of clusters within the latent space
[10]. Consequently, samples drawn from the prior distribution might deviate
from the actual ground truth distribution. Consequently, the decoder model
struggles to faithfully reconstruct these samples into tangible real-world
data points [10].

The second avenue revolves around modeling clusters within the latent space
through a Gaussian mixture model (GMM)-based sampling approach. This method
entails fitting a GMM to the latent vectors from the training data, then
generating new samples from this learned mixture model [10]. Generally, this
technique yields superior samples compared to sampling from a prior. However,
due to the expansive nature of the Gaussian distribution across the entire
Euclidean space, there are instances where generated samples fall outside the
decoder network’s capacity to transform them into authentic real-world samples
[4].

Lastly, an alternative strategy involves the random sampling of latent
vectors, followed by their manipulation through either discrete or continuous
normalizing flows [2, 43], aiming to guide them towards high-density regions.
Various methodologies to identify such high-density regions have been
formulated in the literature [2, 4, 43, 46]. Nonetheless, these techniques
entail significant computational overhead and are applicable only to specific
types of autoencoder-based models necessitating tailored model architectures
[26]. For instance, the approach presented in [4] introduces supplementary
learnable layers to estimate the metric of the latent space, which assumes the
form of a Riemannian manifold.

Within the scope of this study, we craft a novel technique for sampling from
latent spaces, a methodology that can be seamlessly integrated with any
autoencoder model. Our approach effectively sidesteps the challenge of drawing
samples from regions that cannot be faithfully reconstructed into authentic
real-world images, a shortcoming often encountered in GMM sampling. This bears
significance in biometrics, where superior latent samples can be adeptly
transformed into highly realistic synthetic images. As a solution, we present
a post-training density estimation algorithm designed to operate on the latent
space of any autoencoder-based generative model. This algorithm employs both a
quantization step and the concept of probability mass function. The salient
attribute of our algorithm lies in its impressive time complexity of
$\mathcal{O}(n\times d)$, where $n$ denotes the dataset size and $d$
symbolizes the latent space dimension. This translates into a significant
acceleration in runtime compared to GMM sampling. Moreover, we validate the
quality of the generated images through our methodology both qualitatively and
quantitatively. This validation is performed on three distinct image
generation benchmark datasets (including two in the biometric domain) and five
distinct autoencoder-based models. In terms of image quality gauged by the
Fréchet Inception Distance (FID) metric, our approach demonstrates a
considerable enhancement compared to GMM sampling. Specifically, on the MNIST
[20], CelebA [21], and MOBIUS [30, 31, 40, 41, 42] datasets, we achieve
improvements of up to $0.89$, $1.69$, and $0.87$ respectively when compared to
GMM sampling.

To delve into the specifics, Section 2 presents related works addressing the
challenge of latent space sampling within autoencoders. In Section 3, we
introduce essential background definitions and notations. The mechanics of our
sampling methodology, anchored in quantization and probability mass function
concepts, are presented in Section 4. The efficacy of our method in image
generation is rigorously tested and verified across diverse datasets in
Section 5. Lastly, we conclude and outline avenues for future exploration in
Section 6.

##  2 Related works

###  2.1 Sampling from a prior

Numerous studies have introduced a regularization term to encourage the
distribution of latent vectors to approximate a predefined prior distribution,
often taking the form of a normal distribution. Within this category of
methods, notables include variational autoencoders (VAEs) [18], and
$\beta$-VAE [16]. These approaches model the latent space as a distribution
and employ the Kullback-Leibler divergence [6] to steer each distribution
towards a normal distribution. Furthermore, the landscape includes Wasserstein
autoencoders (WAE) [37], which retain the deterministic essence of
autoencoders while harnessing the Wasserstein distance [29] to drive the
overall distribution of latent vectors towards a Gaussian distribution.

Following the minimization of these regularization terms, the process of
latent vector sampling often involves drawing from the prior distribution and
subsequently reconstructing the sampled vector into a data point.
Consequently, due to the inherent non-zero nature of the resulting loss after
enforcing such constraints, clusters may materialize within the latent space.
Additionally, the latent vectors might not precisely adhere to the predefined
prior distribution [10]. To counteract this challenge, $\beta$-VAEs adopt an
elevated weighting for the regularization term, striving to align the
distribution of latent vectors with the prior distribution. However, this
intensified weighting negatively impacts the autoencoder model’s
reconstruction prowess and significantly hampers its generative potential
[16].

###  2.2 Sampling from a Gaussian mixture model

Gaussian Mixture Models (GMMs) assume the presence of subpopulations in the
data and have demonstrated high performance in multiple works [5, 10].
Therefore, GMMs model each subpopulation $i$ via a Gaussian distribution
$\mathcal{N}_{i}(\cdot|\mu_{i},\Sigma_{i})$ with a mean parameter $\mu_{i}$
and covariance parameter $\Sigma_{i}$. To express the probability distribution
of the entire data set, a weighted sum of the probability distributions
$\mathcal{N}_{i}(\cdot|\mu_{i},\Sigma_{i})$ by weights $\alpha_{i}$ is taken,
where $\sum_{i}\alpha_{i}=1$. This way, the probability of any data point in
the dataset can be written as:

$p(x)=\sum_{i}\alpha_{i}\mathcal{N}_{i}(x|\mu_{i},\Sigma_{i}).$

The parameters of the GMM can be estimated by using the Expectation-
Maximization (EM) algorithm [24]. Nonetheless, this sampling method presents
two drawbacks. 1) The large value of $\mathcal{O}(n\times d\times k\times i)$
computational complexity of the EM algorithm, where $n,d,k$ and $i$ represent
respectively the size of the data set, the dimension of the latent space, the
number of distributions in the GMM model and the number of iterations. In
fact, we could have an infinite time complexity for this algorithm, if not for
the number of iterations. 2) With low probability, outliers can be sampled,
since the model is based on distributions defined on the whole latent space.

###  2.3 Sampling using normalizing flows

Recent works like in [2, 3, 4, 45, 46] propose to use normalizing flows [19,
26] to sample latent vectors from the latent space of autoencoders. These
methods select vectors randomly as sample vectors, then move them to high
probability regions. Normalizing flow methods can be divided into two
families, notably discrete and continuous normalizing flows [26]. The
difference comes from the definition of the flow process $T(z)$ as discrete
steps via the composition of functions $\\{T_{i}\\}^{k}_{i=1}$ as in [45, 46]:

$T(z)=T_{k}\circ\dots\circ T_{1}(z)$

or continuous via an integral over an interval as in [3]:

$T(z)=z+\int_{t=t_{0}}^{t_{1}}g_{\phi}(t,z_{t})dt.$

Recent works such as [2, 3, 4] focus mostly on continuous normalizing flows
due to their better modeling of the flow dynamics. These works propose to see
the Euclidean latent space of VAEs as a Riemannian manifold [33] with the sum
of all covariance matrices weighted by the distance to the mean as the
manifold’s metrics [3]. This metric was further improved in [4] by learning a
covariance matrix from the data. Then, using the determinant of this metric,
the probability density at any point can be estimated. Finally, via a
Hamiltonian Monte Carlo process, the drawn samples are moved to regions of
high probability. These regions can be reconstructed into realistic data
points [4]. Yet, these methods do not represent a general off-the-shelf method
to sample from the latent space because they require a dedicated model
architecture. Also, as noted in [26], using Riemannian metrics is only
applicable to topologies homeomorphic to Euclidean space $\mathbb{R}^{d}$
which is not applicable, for example, to spherical latent spaces such as in
Hyperspherical Variational Auto-Encoders [8].

In this work, our proposed sampling method performs density estimation over
the latent space via probability mass function and quantization processes.
Consequently, it does not require a special regularization term (as in
sampling from a prior distribution) nor a special autoencoder architecture (as
in sampling using normalizing flows). Furthermore, in addition to not sampling
outliers, training our sampling method improves the time complexity of
Gaussian mixture model sampling.

##  3 Background

In this section, we introduce necessary background definitions and notations.

###  3.1 Autoencoders

A Vanilla Autoencoder [11] is a neural network architecture comprised of two
main components: an encoder $E_{\phi}$ and a decoder $D_{\theta}$. The role of
the encoder is to map input data $x$ from the space $\mathcal{X}$ to a latent
representation $z$ within a lower-dimensional space $\mathcal{Z}$, where
$\mbox{dim}(\mathcal{Z})\ll\mbox{dim}(\mathcal{X})$. Conversely, the decoder’s
function is to reconstruct the input data $\hat{x}$ using the latent
representation $z$. Both the encoder and decoder are equipped with parameters
$\phi$ and $\theta$ correspondingly. The optimization of the weights $\phi$
and $\theta$ is achieved through gradient descent, aimed at minimizing the
reconstruction loss. This loss is generally defined as the squared $\ell_{2}$
norm of the difference between the input data $x$ and its corresponding
reconstruction $\hat{x}$, and can be mathematically articulated as
$\mathcal{L}_{AE}(x,\hat{x})=||x-\hat{x}||_{2}^{2}$.

###  3.2 Quantization

Quantization involves the process of discretizing an input, often derived from
a continuous or extensive set of values, into a discrete integer set. In the
case of a vector, each dimension is commonly treated independently. Machine
learning and autoencoder studies on quantization can be divided into two main
categories: quantization-aware training and post-training quantization. In the
quantization-aware training approach, quantization is applied during the
learning phase. This approach yields a model that has learned to handle
quantized vectors, making the resulting sampling techniques model-specific.
Notable among such methods is VQ-VAE [39], which develops a dictionary of
quantized vectors and maps each latent vector of input data to the nearest
quantized vector. On the other hand, the second approach of post-training
quantization applies quantization subsequent to the model’s training. This
method’s advantage lies in its applicability to various models without
altering the learning phase nor necessitating extra layers for fine-tuning.
Within the realm of post-training quantization methods, examples include
lattice quantization [1, 22] and clustering-based non-uniform quantization
[13, 34]. While these methods offer superior quantization quality, they tend
to be computationally intensive. In this study, we adopt a post-training
quantization strategy to leverage its broad applicability. However, given our
focus on simplicity and reduced computational complexity, we introduce a
quantization method in Section 4 with a time complexity of just
$\mathcal{O}(d)$, where $d$ represents the length of the vector.

###  3.3 Probability density function

Given a continuous random variable $x$ taking on an innumerable infinite
number of possible values with support in $\mathcal{C}$. A probability density
function $PDF(\cdot)$ estimates the probability of $a<x<b$, i.e. $P(a<x<b)$
and satisfies the following conditions :

  * •

$PDF(x)\geq 0\quad\forall x\in\mathcal{C}$.

  * •

$\int_{\mathcal{C}}PDF(x)dx=1$.

  * •

$P(a<x<b)=\int_{a}^{b}PDF(x)dx$.

  * •

$P(x=a)=0\quad\forall a\in\mathcal{C}$.

###  3.4 Probability mass function

A _Probability mass function_ (PMF) is an adaptation of the probability
density function to the case of discrete random variables. Given a discrete
random variable $x$ taking a finite or countably infinite number of possible
values with support in $\mathcal{D}$, a PMF estimates the probability
$P(x=a)\quad\forall a\in\mathcal{D}$ and satisfies the following conditions:

  * •

$P(x=a)\geq 0\quad\forall a\in\mathcal{D}$.

  * •

$\sum_{a\in\mathcal{D}}P(x=a)=1$.

###  3.5 Fréchet Inception Distance

Fréchet Inception Distance (FID) [15] is a metric used for evaluating the
quality of synthetic images generated by generative models such as GANs and
AEs. It compares the distribution of generated images to that of a set of
real-world images, with a score of 0 indicating a perfect match between the
distribution of synthetic and real images. The FID metric has been employed in
various works targeting biometric image generation, including IrisGAN [23].

To compute the FID, first, an encoding of real and synthetic images is
calculated using the Inception V3 network [35] without the classification
layer. Subsequently, under the assumption that the encodings follow Gaussian
distributions, mean values $\mu_{r}$ and $\mu_{s}$ as well as covariance
matrices $\Sigma_{r}$ and $\Sigma_{s}$ are estimated for real and synthetic
images, respectively. Finally, the FID is evaluated using the following
formula:

$FID=\left\|\mu_{r}-\mu_{s}\right\|_{2}^{2}+\text{tr}\left(\Sigma_{r}+\Sigma_{s}-2\left(\Sigma_{r}\Sigma_{s}\right)^{\frac{1}{2}}\right).$

###  3.6 Wasserstein Distance

The Wasserstein distance [29] is a metric used for calculating the distance
between probability distributions. It can be computed between two
distributions, $p$ and $q$, as follows:

$\mathcal{W}_{\xi}(p,q)=\left(\inf_{\gamma\in\mathcal{P}(p(x),q(x^{\prime}))}\mathbb{E}_{\gamma(x,x^{\prime})}[d^{\xi}(x,x^{\prime})]\right)^{\frac{1}{\xi}}$

where $d(x,x^{\prime})$ represents the distance between $x$ and $x^{\prime}$,
$\xi$ is a positive integer, and $\mathcal{P}(p(x),q(x^{\prime}))$ represents
the joint distribution between $p$ and $q$. Given that computing the
infimum111The infimum of a subset $S$ of a partially ordered set $P$ is the
greatest $p\in P$ that is less than or equal to each element of $S$, if such
an element exists. is often computationally challenging, entropy-regularized
optimal transport was introduced in [7] alongside the Sinkhorn algorithm.

##  4 Probability mass function sampling

In this work, we introduce the Probability Mass Function Sampling (PMFS)
method, an innovative approach for sampling from the latent space of any
autoencoder model. This post-training method employs discrete density
estimation through probability mass functions and quantization procedures.
PMFS presents a solution to the challenge of continuous density estimation,
which becomes complex due to the infinite nature of the space.

Formally, given a set of latent vectors
$\mathcal{Z}=\\{z_{i}\in\mathbb{R}^{d}\\}_{i=1}^{n}$ and a parameter $k$
indicating the number of uniform partitions or bins per dimension222The
hyperparameter $k$ is chosen via a tuning process to minimize the FID metric
(See Figure 2).. Initially, we identify the maximum $max_{j}$ and minimum
$min_{j}$ values per dimension $j\in\\{1,\dots,d\\}$. Subsequently, we compute
the width of each partition $w_{j}=\frac{max_{j}-min_{j}}{k}$ for dimension
$j$. The quantization step is then performed for each vector $z_{i}$ in the
set $\mathcal{Z}$. This process assigns one of the $k^{d}$ global partitions
in the space to each latent vector $z_{i}$. The quantized values $Q_{z_{i,j}}$
corresponding to the latent vector $z_{i}$ and each dimension $j$ are
determined using the formula:

$Q_{z_{i,j}}=\left\lfloor\frac{z_{i,j}-min_{j}}{w_{j}}\right\rfloor=\left\lfloor\frac{k(z_{i,j}-min_{j})}{max_{j}-min_{j}}\right\rfloor.$
(1)

In Appendix A, we provide an illustrative example of the proposed quantization
step.

Lastly, we define our probability mass function as follows333This definition
can be easily verified to represent a probability mass function.:

$P(x=p)=\frac{\text{\\#vectors in the global partition $p$}}{n}$ (2)

where $p$ denotes a global partition. This formulation of the probability mass
function can also be interpreted as a weight assigned to each global
partition. Hence, to sample a latent vector, a global partition $p$ is sampled
according to its weight calculated using Equation 2. Subsequently, given that
each partition has upper and lower bounds $p_{min}$ and $p_{max}$, a vector
$z$ is uniformly sampled from the partition, i.e.,
$z\sim\mathcal{U}_{[p_{min},p_{max}]}(\cdot)$. The proposed PMFS sampling
method can be visualized as defining volumes with controllable dimensions
around the known latent vectors, followed by uniform sampling from these
volumes. PMFS sampling ensures that the generated samples can be accurately
reconstructed into real-world data points, given their proximity to actual
latent vectors. We provide a summary of the PMFS sampling model and its per-
step time-complexity in Algorithm 1. Additionally, the related Python and
Mathematica code is available in Appendix B.

Algorithm 1 Probability mass function sampling

Input : $\mathcal{Z}=\\{z_{i}\in\mathbb{R}^{d}\\}_{i=1}^{n}$ set of latent
vectors, $k$ number of partition per dimension.

Output : partitions and their weight;

1:for each $j\in\\{1\dots d\\}$ do $\triangleright$ $\mathcal{O}(n\times d)$

2: $max_{j}$ = $\max_{i}\\{z_{i,j}\\}_{i=1}^{n}$; $\triangleright$
$\mathcal{O}(n)$

3: $min_{j}$ = $\min_{i}\\{z_{i,j}\\}_{i=1}^{n}$; $\triangleright$
$\mathcal{O}(n)$

4:end for each

5:$P=0$ $\triangleright$ hashmap of partition weights.

6:for each $j\in\\{1\dots n\\}$ do $\triangleright$ $\mathcal{O}(n\times d)$
$\triangleright$ Quantize vector $z_{i}$.

7: $Q_{z_{i}}$ = $Q(z_{i},k,\\{min_{j}\\}_{j=1}^{d},\\{max_{j}\\}_{j=1}^{d})$
$\triangleright$ $\mathcal{O}(d)$ $\triangleright$ Increment the weight of the
partition $Q_{z_{i}}$.

8: $P[X=Q_{z_{i}}]=\frac{P[X=Q_{z_{i}}]\times n+1}{n}$ $\triangleright$
$\mathcal{O}(1)$

9:end for each

Analyzing and comparing the PMFS sampling algorithm to the Gaussian Mixture
Model (GMM) sampling algorithm, we observe that in terms of time complexity,
we transition from $\mathcal{O}(n\times d\times k\times i)$ to
$\mathcal{O}(n\times d)$, which constitutes a significant improvement.
Moreover, we shift from relying on the infinite Expectation-Maximization (EM)
algorithm [24] with the number of iterations $i$ as a stopping criterion to a
finite algorithm that depends solely on the number of latent vectors $n$ and
their dimension $d$.

Furthermore, for performance optimization, we retain only bins containing data
points by utilizing a hashmap to store the weights. This is a strategy to
avoid an exponential growth in the number of partitions as the latent
dimension $d$ increases, given the $k^{d}$ global partitions generated.
Consequently, we only maintain in memory the weights of existing partitions.
This approach is equivalent to preserving all bins without sacrificing
performance nor generality, as global partitions lacking data samples possess
a weight of zero and therefore never get selected.

To summarize, the PMFS method facilitates density estimation and the sampling
of high-probability vectors. When these vectors are reconstructed by the
decoder network, they yield synthetic data points that closely resemble real-
world data (See to Section 5 and Appendix G). Also, the estimation process in
our model operates with a time complexity of $\mathcal{O}(n\times d)$. This
stands as a significant improvement over Gaussian mixture model sampling,
which entails a time complexity of $\mathcal{O}(n\times d\times k\times i)$.
Finally, we address the challenge of an exponentially growing number of
partitions by utilizing a hash map to retain only pertinent partitions.

##  5 Experiments

This section presents experimental results validating the proposed PMFS
sampling method described in Section 4, compared to GMM sampling. In Section
5.1 we present an overview of the autoencoder model architectures employed in
this work. Appendix C and D present our experimental setup and a description
of the used datasets. Moreover, in Appendix F, we provide samples of image
reconstructions using different autoencoder models. Furthermore, we lay out a
qualitative comparison on the impact of the sampling method on the synthetic
images in Appendix G.

###  5.1 Model architectures

The autoencoder models used is this work contain an encoder network with
convolutional layers and a linear layer which allows operating on images. The
decoder network comprises a non-linear layer and a set of transpose
convolutional layers to reconstruct the input images (See Figure 1, ample
network architecture details are provided in Appendix E).

Figure 1: Overview of the models’ architecture used in this work.

###  5.2 Training parameters

The MOBIUS [30, 31, 40, 41, 42], MNIST [20], and CelebA [21] datasets used in
this work were maintained at their original resolutions. The different models
were trained for $50$ epochs using the Adam optimizer with a learning rate of
$10^{-3}$ and $\beta_{1}=0.9,\beta_{2}=0.999$, and a batch size of $64$.
However, for the CelebA dataset, we trained our models for $100$ epochs with a
learning rate of $10^{-4}$. The rest of the parameters remained unchanged. We
ensured the convergence of all models. In terms of latent dimensionality, we
employed sizes of $32$, $256$, and $256$ for the MNIST, CelebA, and MOBIUS
datasets, respectively, for all models. Moreover, we used the latent vectors
of the validation set images to fit the GMM and PMFS models. The models
utilized in this work encompass the Vanilla Autoencoder [11], Variational
Autoencoder [18], $\beta$-VAE [16] with $\beta=2$, Wasserstein Autoencoder
[37] with $\beta=100$, and InfoVAE [44] with $\alpha=0,\lambda=1000$. The
hyperparameters of the networks are those yielding the best results in the
original papers introducing each model, except for $\beta$-VAE, where we set
$\beta=2$ to better minimize the reconstruction loss.

###  5.3 Results and discussion

Following the training of various autoencoder-based generative models, notably
the Vanilla Autoencoder [11], Variational Autoencoder [18], $\beta$-VAE [16],
Wasserstein Autoencoder [37], and InfoVAE [44] on the MNIST [20], CelebA [21],
and MOBIUS [30, 31, 40, 41, 42] datasets, we computed the Fréchet Inception
Distance FID metric. The obtained results are presented in Table 1.
Furthermore, Figure 2 illustrates the progression of the FID metric in
relation to the number of distributions for GMM sampling and the number of
partitions for PMFS sampling. It’s important to note that for calculating the
FID metric, we generated a quantity of synthetic images equivalent to the
number of images in the respective test set for each dataset. Lastly, Figure 3
displays the distribution of the latent space of the validation set, providing
a comparison with samples drawn from GMM and PMFS models.

Table 1: A comparison between GMM sampling and PMFS based on the lowest FID,
Wasserstein distance (denoted by $\mathcal{W}$), and model fitting time in
seconds (denoted by $T$). The comparison is performed across the MNIST,
CelebA, and MOBIUS datasets. The columns labeled #dist and $k$ respectively
signify the number of distributions utilized in GMM sampling and the number of
partitions per dimension in the PMFS sampling method, aiming to minimize the
FID metric. In the table, the most favorable values are highlighted in bold.

| Model | GMM | PMFS
---|---|---|---
| FID $\downarrow$ | $\mathcal{W}\downarrow$ | T $\downarrow$ | #dists | FID $\downarrow$ | $\mathcal{W}\downarrow$ | T $\downarrow$ | $k$
MNIST | AE |  04.66 | 839.77 | 25.64 | 20 |  04.19 | 551.81 | 0.02 |  08
VAE |  08.66 | 13.63 | 20.59 | 12 |  08.57 | 14.16 | 0.07 | 10
$\beta$-VAE | 11.98 | 13.50 | 24.64 | 12 | 11.74 | 53.71 | 0.06 |  02
WAE |  04.67 | 880.54 | 29.24 | 18 |  04.24 | 570.55 | 0.04 |  08
InfoVAE |  08.84 | 13.03 | 18.81 | 10 |  07.95 | 49.18 | 0.08 |  02
CelebA | AE | 10.63 | $4.6\times 10^{5}$ | 95.56 |  08 |  08.94 | $\mathbf{3.3\times 10^{5}}$ | 0.11 | 20
VAE | 10.49 | 190.93 | 65.31 | 16 |  09.25 | 159.14 | 0.18 | 14
$\beta$-VAE | 11.10 | 196.48 | 50.01 | 20 |  09.91 | 170.28 | 0.16 | 20
WAE | 10.58 | $4.9\times 10^{5}$ | 52.86 | 20 |  08.89 | $\mathbf{3.8\times 10^{5}}$ | 0.15 | 12
InfoVAE | 10.56 | 193.85 | 48.90 | 20 |  09.33 | 159.26 | 0.18 | 12
MOBIUS | AE | 24.61 | $\mathbf{6.4\times 10^{3}}$ |  01.59 |  02 | 24.33 | $2\times 10^{4}$ | 0.02 |  02
VAE | 24.55 | 528.51 |  01.48 |  06 | 23.68 | 1765.73 | 0.06 |  02
$\beta$-VAE | 24.74 | 303.68 |  01.69 |  02 | 24.05 | 836.28 | 0.02 |  02
WAE | 23.60 | $\mathbf{5.9\times 10^{3}}$ |  01.57 |  02 | 23.66 | $1.8\times 10^{4}$ | 0.03 |  02
InfoVAE | 24.18 | 387.27 |  01.78 |  02 | 23.41 | 1156.14 | 0.05 |  02

(a) MNIST data set.

(b) CelebA data set.

(c) MOBIUS data set.

Figure 2: The evolution of the FID metric in relation to the number of
partitions in PMFS or the number of distributions in a GMM. Dashed and solid
lines respectively illustrate the FID evolution using GMM and PMFS sampling
methods. The colors green, red, magenta, yellow, and black correspond to the
Vanilla Autoencoder, Variational Autoencoder (VAE), $\beta$-VAE, Wasserstein
Autoencoder (WAE), and InfoVAE models, respectively. This figure is best
viewed in color.

(a) Latent space distribution on MNIST.

(b) Latent space distribution on CelebA.

(c) Latent space distribution on MOBIUS.

Figure 3: An illustration of the distribution of the latent vectors from the
validation set of the Vanilla Autoencoder, along with samples generated using
GMM and PMFS sampling. These figures portray a dimensionality reduction
through PCA into the two-dimensional space $\mathbb{R}^{2}$, encompassing the
first hundred images from the validation set as well as a hundred samples
extracted from the GMM and PMFS sampling methods. This figure is best viewed
in color.

(a) MNIST GMM samples.

(b) MNIST PMFS samples.

(c) CelebA GMM samples.

(d) CelebA PMFS samples.

(e) MOBIUS GMM samples.

(f) MOBIUS PMFS samples.

Figure 4: Sample synthetic images from a Vanilla Autoencoder model on the
MNIST, CelebA and MOBIUS data sets using GMM and PMFS sampling. Figure better
viewed in color.

From Table 1 and Figure 2, it becomes evident that, in terms of GMM sampling,
the optimal choice for the number of distributions (#dist) is rarely the same
as the number of classes in the dataset, even for relatively simple datasets
like MNIST. As a result, relying on the heuristic of setting the number of
classes as the number of distributions in a GMM, as is commonly done [10], is
likely to yield suboptimal outcomes. Furthermore, a comparison of the lowest
FID values achieved by GMM sampling with those obtained through PMFS sampling
underscores the enhancements introduced by PMFS sampling across all models and
datasets. This improvement can be attributed to the strategy of defining a
neighborhood around each latent vector and sampling from within this
neighborhood, leading to the acquisition of superior, high-probability samples
that can be effectively reconstructed by the decoder network into realistic
data points. This contrasts with GMM sampling, which operates on distributions
encompassing the entire space and might yield samples lying outside the
regions reconstructable by the decoder network. Such a characteristic is
crucial in biometric image synthesis, where the goal is to generate synthetic
images that closely resemble real-world counterparts.

Furthermore, Table 1 reveals substantial gains in terms of the FID metric when
applying our proposed PMFS sampling method to deterministic models,
particularly the Vanilla Autoencoder [11] and the Wasserstein Autoencoder
[37]. This trend can be attributed to the stochastic learning phase in these
models, which capitalizes on the vicinity of each mean latent vector
corresponding to an image. Consequently, our method, rooted in the concept of
neighborhoods, yields slightly reduced efficacy for this model type. However,
it still outperforms the GMM sampling method in terms of FID. Consequently,
when employing the PMFS sampling method on biometric data, the outcome is a
collection of synthetic images resembling more closely the actual instances.
Additionally, a comparison of the model fitting time between PMFS and GMM
sampling reveals a substantial improvement, quantifiable in terms of orders of
magnitude, thus corroborating the anticipated theoretical disparity in time
complexity.

Moreover, we conducted an investigation into the distribution of drawn samples
via GMM sampling and PMFS, comparing them to the underlying distribution of
latent vectors generated by the validation dataset. This investigation was
executed quantitatively using the Wasserstein distance [29] (see the
$\mathcal{W}$ column in Table 1) and qualitatively (see Figure 3). From the
$\mathcal{W}$ column in Table 1, it is evident that, particularly for large
datasets like MNIST and CelebA, PMFS has the capacity to generate
distributions which closely approximate the real ones, demonstrating an
improvement over GMM sampling. This holds true, especially when the value of
$k$ is high. Consequently, for smaller datasets like MOBIUS, GMM sampling
exhibits superior performance over PMFS due to the low number of partitions
$k$ employed in PMFS, aiming to maximize FID while compensating for the
distribution discrepancy between the training and testing sets. These findings
are visually supported by Figure 3 and Appendix H.

Upon analyzing Figure 2, it becomes evident that, for all the datasets, there
exists a specific number of partitions $k$ for PMFS sampling that yields the
lowest FID outcome, irrespective of the number of distributions employed in
GMM sampling. Particularly on the CelebA dataset, known for face generation, a
substantial advantage of PMFS sampling over GMM sampling is apparent. Here, a
notable margin separates the images generated using GMM sampling from those
produced via PMFS sampling. Additionally, it’s noticeable that the optimal
value for the number of partitions $k$ in PMFS sampling varies depending on
the dataset.

Figure 4 showcases samples of synthetic images generated through GMM and PMFS
sampling using a Vanilla Autoencoder model (Additional images are provided in
Appendix G). From this figure, we notice that the images generated with PMFS
sampling possess a heightened realism. Notably, when PMFS sampling is applied
to the CelebA dataset. For example, the first image depicts the actor Daniel
Radcliffe (Harry Potter) without glasses. This observation underscores how
PMFS sampling introduces data augmentation and modifications to images present
in the training dataset. This enhancement emphasises the significance of our
method in biometric image synthesis and data augmentation on biometric
datasets.

Furthermore, in Table 2, we present a face image quality assessment on the
synthetic images generated using PMFS and GMM sampling via the SDD-FIQA metric
[25]. Since the SDD-FIQA metric uses a model which was not trained on the
CelebA data set [21], we start by establishing a baseline by calculating the
value of SDD-FIQA metric on the original untouched images of the CelebA data
set. Then we generate $10,000$ synthetic images per sampling strategy, notably
GMM sampling and PMFS, and per autoencoder model. The obtained results are
presented in Table 2. From this table we observe improvements in terms of
synthetic face image quality on all models when samples are generated through
PMFS which emphasizes the performance gains of PMFS over GMM sampling.

Table 2: Face image quality assessment on the CelebA data set. In the table,
the most favorable values are highlighted in bold.

Maximum obtainable value on CelebA data set | 41.55
---|---
Model | GMM | PMFS
Vanilla Autoencoder | 39.28 | 39.85
VAE | 38.78 | 39.43
$\beta$-VAE | 38.71 | 38.79
WAE | 37.60 | 38.85
InfoVAE | 38.68 | 39.39

In this section, we have presented a quantitative comparison through the FID
metric and a qualitative comparison via synthetic images between the GMM
sampling method and our proposed method. These comparisons were conducted
across five distinct autoencoder-based generative models and three diverse
datasets, including two biometric datasets. Throughout these experiments, PMFS
sampling consistently outperformed GMM sampling, showcasing the lowest FID
values, achieving fitting times several orders of magnitude faster, producing
more lifelike synthetic images across all scenarios, and on large datasets
with a high number of partitions generating samples with distributions that
closely resemble real-world ones.

##  6 Conclusion

In this study, our focus revolves around the task of sampling from the latent
space of autoencoder-based generative models. Our primary objective is to
extract samples from regions characterized by high density, ensuring that
these samples can be subsequently reconstructed into images that possess a
realistic quality. To accomplish this goal, we introduce a novel sampling
algorithm that leverages the concept of probability mass functions coupled
with a quantization process. Notably, this algorithm boasts a time complexity
of $\mathcal{O}(n\times d)$, resulting in a substantial acceleration compared
to Gaussian mixture model (GMM) sampling. The core principle underlying our
algorithm entails the definition of neighborhoods surrounding each latent
encoding corresponding to input data points. Subsequently, the sampling
process targets these neighborhoods, guaranteeing that the sampled latent
vectors fall within regions of elevated probability. Consequently, these
vectors can be seamlessly transformed into genuine images. Our experimentation
covers a diverse array of scenarios, spanning three distinct datasets—MNIST
for digit generation, CelebA for face generation, and MOBIUS for ocular
generation—as well as five autoencoder-based generative models. The empirical
results across our experiments consistently showcase the superiority of the
PMFS sampling method in comparison to Gaussian mixture model (GMM) sampling.
Impressively, this superiority is especially pronounced on the biometric
datasets—CelebA and MOBIUS—where PMFS exhibits a remarkable enhancement of
$1.69$ and $0.87$ in terms of the Fréchet inception distance (FID),
respectively, when compared to GMM sampling. These findings underscore the
exceptional efficacy of our novel sampling algorithm, particularly in
generating high-quality images in diverse contexts, with a marked emphasis on
biometric images. Moreover, when considering the Wasserstein distance, PMFS
outperforms GMMs in approximating the true latent distribution. Prospective
extensions of our work could encompass enhancements aimed at improving the
image quality produced by autoencoder models by mitigating the blur introduced
by the $\ell_{2}$ norm [11]. Additionally, we aspire to differentiate between
whether the generated synthetic images serve as data augmentations over the
training set or constitute an entirely new dataset, thereby necessitating the
quantification of data leakage. This perspective is particularly salient given
the sensitive nature of biometric data, which warrants stringent privacy
considerations. In summary, these insights lay the groundwork for further
exploration and refinement of the PMFS method.

## Acknowledgement

The authors extend their gratitude to Stéphane Ayache, Laurent Galathée, and
Adnan Ben Mansour for their invaluable input and constructive feedback on the
content of this paper. Furthermore, the authors acknowledge the generous
access granted to them for the utilization of High-Performance Computing (HPC)
resources available through MesoPSL. This valuable resource is funded by the
Région Île-de-France and is a part of the Equip@Meso project (reference
ANR-10-EQPX-29-01). The Equip@Meso project is a constituent of the programme
investissements d’avenir, which is overseen by France’s Agence nationale pour
la recherche. The MOBIUS Data set images employed in this research have been
generously provided by the University of Ljubljana, Slovenia [30, 31, 40, 41,
42].

## References

  * [1] E. Agrell and B. Allen.  On the best lattice quantizers.  IEEE Transactions on Information Theory, pages 1–1, 2023.
  * [2] A. L. Caterini, A. Doucet, and D. Sejdinovic.  Hamiltonian variational auto-encoder.  In Advances in Neural Information Processing Systems, volume 31, 2018.
  * [3] C. Chadebec and S. Allassonnière.  A geometric perspective on variational autoencoders.  arXiv preprint arXiv:2209.07370, 2022.
  * [4] C. Chadebec, E. Thibeau-Sutre, N. Burgos, and S. Allassonnière.  Data Augmentation in High Dimensional Low Sample Size Setting Using a Geometry-Based Variational Autoencoder.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022\.
  * [5] C. Chadebec, L. J. Vincent, and S. Allassonniere.  Pythae: Unifying Generative Autoencoders in Python - A Benchmarking Use Case.  In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
  * [6] I. Csiszar.  $I$-Divergence Geometry of Probability Distributions and Minimization Problems.  The Annals of Probability, 3(1):146 – 158, 1975.
  * [7] M. Cuturi.  Sinkhorn distances: Lightspeed computation of optimal transport.  In Advances in Neural Information Processing Systems, 2013.
  * [8] T. R. Davidson, L. Falorsi, N. De Cao, T. Kipf, and J. M. Tomczak.  Hyperspherical Variational Auto-Encoders.  34th Conference on Uncertainty in Artificial Intelligence (UAI-18), 2018.
  * [9] W. Falcon and The PyTorch Lightning team.  PyTorch Lightning, 3 2019.
  * [10] P. Ghosh, M. S. M. Sajjadi, A. Vergari, M. Black, and B. Scholkopf.  From variational to deterministic autoencoders.  In International Conference on Learning Representations, 2020.
  * [11] I. Goodfellow, Y. Bengio, and A. Courville.  Deep Learning.  MIT Press, 2016.
  * [12] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio.  Generative adversarial nets.  In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
  * [13] S. Han, H. Mao, and W. J. Dally.  Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding, 2016.
  * [14] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. Fernández del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant.  Array programming with NumPy.  Nature, 585:357–362, 2020.
  * [15] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter.  GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium.  In Advances in Neural Information Processing Systems, NIPS, volume 30, 2017.  arXiv:1706.08500.
  * [16] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner.  Beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework.  In International Conference on Learning Representations, 2017\.
  * [17] J. D. Hunter.  Matplotlib: A 2D Graphics Environment.  Computing in Science & Engineering, 9(3):90–95, 2007.
  * [18] D. P. Kingma and M. Welling.  Auto-Encoding Variational Bayes, 2013.  arXiv:1312.6114.
  * [19] I. Kobyzev, S. J. Prince, and M. A. Brubaker.  Normalizing flows: An introduction and review of current methods.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021\.
  * [20] Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner.  Gradient-Based Learning Applied to Document Recognition.  Proceedings of the IEEE, 1998.
  * [21] Z. Liu, P. Luo, X. Wang, and X. Tang.  Deep Learning Face Attributes in the Wild.  In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
  * [22] C. Metz, T. Allenet, J. Thiele, A. Dupret, and O. Bichler.  Lattice quantization.  In 2023 Design, Automation & Test in Europe Conference & Exhibition (DATE), pages 1–2, 2023.
  * [23] S. Minaee and A. Abdolrashidi.  Iris-gan: Learning to generate realistic iris images using convolutional gan, 2018.
  * [24] T. Moon.  The Expectation-Maximization Algorithm.  IEEE Signal Processing Magazine, 13(6):47–60, 1996.
  * [25] F.-Z. Ou, X. Chen, R. Zhang, Y. Huang, S. Li, J. Li, Y. Li, L. Cao, and Y.-G. Wang.  Sdd-fiqa: Unsupervised face image quality assessment with similarity distribution distance.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 7670–7679, June 2021.
  * [26] G. Papamakarios, E. Nalisnick, D. Jimenez Rezende, S. Mohamed, and B. Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference.  Journal of Machine Learning Research, 12 2019.
  * [27] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala.  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  In Advances in Neural Information Processing Systems 32, 2019\.
  * [28] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay.  Scikit-learn: Machine Learning in Python.  Journal of Machine Learning Research, 12:2825–2830, 2011.
  * [29] G. Peyré and M. Cuturi.  Computational Optimal Transport: With Applications to Data Science.  Foundations and Trends in Machine Learning, 11(5-6):355–607, 2019\.
  * [30] P. Rot, v. Emeršič, V. Štruc, and P. Peer.  Deep multi-class eye segmentation for ocular biometrics.  In IEEE International Work Conference on Bioinspired Intelligence (IWOBI), pages 1–8, 07 2018.
  * [31] P. Rot, M. Vitek, K. Grm, v. Emeršič, P. Peer, and V. Štruc.  Deep sclera segmentation and recognition.  In A. Uhl, C. Busch, S. Marcel, and R. N. J. Veldhuis, editors, Handbook of Vascular Biometrics (HVB), pages 395–432. Springer, 2020.
  * [32] T. Salimans, I. Goodfellow, W. Zaremba, V. Cheung, A. Radford, X. Chen, and X. Chen.  Improved techniques for training gans.  In Advances in Neural Information Processing Systems, 2016.
  * [33] S. Sommer, T. Fletcher, and X. Pennec.  Introduction to differential and riemannian geometry.  In Riemannian Geometric Statistics in Medical Image Analysis. Medical Image Analysis, 2020.
  * [34] P. Stock, A. Joulin, R. Gribonval, B. Graham, and H. Jégou.  And the bit goes down: Revisiting the quantization of neural networks.  In International Conference on Learning Representations, 2020.
  * [35] C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna.  Rethinking the inception architecture for computer vision.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
  * [36] P. Tinsley, A. Czajka, and P. J. Flynn.  Haven’t i seen you before? assessing identity leakage in synthetic irises.  In 2022 IEEE International Joint Conference on Biometrics (IJCB), pages 1–9, 2022.
  * [37] I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schoelkopf.  Wasserstein Auto-Encoders, 2017.  arXiv:1711.01558.
  * [38] TorchVision, maintainers & contributors.  TorchVision: PyTorch’s Computer Vision library.  https://github.com/pytorch/vision, 2016.
  * [39] A. van den Oord, O. Vinyals, and K. Kavukcuoglu.  Neural discrete representation learning, 2018.
  * [40] M. Vitek, A. Das, D. R. Lucio, L. A. Zanlorensi, D. Menotti, J. N. Khiarak, M. A. Shahpar, M. Asgari-Chenaghlu, F. Jaryani, J. E. Tapia, A. Valenzuela, C. Wang, Y. Wang, Z. He, Z. Sun, F. Boutros, N. Damer, J. H. Grebe, A. Kuijper, K. Raja, G. Gupta, G. Zampoukis, L. Tsochatzidis, I. Pratikakis, S. Aruna Kumar, B. Harish, U. Pal, P. Peer, and V. Štruc.  Exploring bias in sclera segmentation models: A group evaluation approach.  IEEE Transactions on Information Forensics and Security (TIFS), 18:190–205, 2023.
  * [41] M. Vitek, A. Das, Y. Pourcenoux, A. Missler, C. Paumier, S. Das, I. De Ghosh, D. R. Lucio, L. A. Zanlorensi Jr., D. Menotti, F. Boutros, N. Damer, J. H. Grebe, A. Kuijper, J. Hu, Y. He, C. Wang, H. Liu, Y. Wang, Z. Sun, D. Osorio-Roig, C. Rathgeb, C. Busch, J. Tapia Farias, A. Valenzuela, G. Zampoukis, L. Tsochatzidis, I. Pratikakis, S. Nathan, R. Suganya, V. Mehta, A. Dhall, K. Raja, G. Gupta, J. N. Khiarak, M. Akbari-Shahper, F. Jaryani, M. Asgari-Chenaghlu, R. Vyas, S. Dakshit, S. Dakshit, P. Peer, U. Pal, and V. Štruc.  SSBC 2020: Sclera segmentation benchmarking competition in the mobile environment.  In IEEE International Joint Conference on Biometrics (IJCB), pages 1–10, 10 2020.
  * [42] M. Vitek, P. Rot, V. Štruc, and P. Peer.  A comprehensive investigation into sclera biometrics: A novel dataset and performance study.  Neural Computing & Applications (NCAA), pages 17941–17955, 2020\.
  * [43] Z. Xiao, Q. Yan, and Y. Amit.  Generative Latent Flow, 2019.  arXiv:1905.10485.
  * [44] S. Zhao, J. Song, and S. Ermon.  Infovae: Balancing learning and inference in variational autoencoders.  Proceedings of the AAAI Conference on Artificial Intelligence, 2019\.
  * [45] Y. Zhao, Q. Ding, and X. Zhang.  AE-FLOW: Autoencoders with normalizing flows for medical images anomaly detection.  In The Eleventh International Conference on Learning Representations, 2023.
  * [46] Z. Ziegler and A. Rush.  Latent normalizing flows for discrete sequences.  In International Conference on Machine Learning, 2019.
  * [47] H. Zou, H. Zhang, X. Li, J. Liu, and Z. He.  Generation textured contact lenses iris images based on 4dcycle-gan.  In 2018 24th International Conference on Pattern Recognition (ICPR), pages 3561–3566, 2018.

Supplementary material

##  Appendix A Quantization example

Given the dimension $d=3$ latent vector $z_{i}=[1.5;2.6;8]$ and the two
vectors $[-19;-5;0]$, $[5.7;3;20]$ representing the minimums and maximums for
each dimension, respectively, and $k=10$. The quantized components of the
vector $Q_{z_{i}}$ corresponding to the latent vector $z_{i}$ are calculated
as follows :

$\displaystyle Q_{z_{i,1}}$
$\displaystyle=\left\lfloor\frac{10(1.5+19)}{5.7+19}\right\rfloor=\left\lfloor\frac{10\times
20.5}{24.7}\right\rfloor=\left\lfloor 8.29\right\rfloor=\mathbf{8}$
$\displaystyle Q_{z_{i,2}}$
$\displaystyle=\left\lfloor\frac{10(2.6+5)}{3+5}\right\rfloor=\left\lfloor\frac{10\times
7.6}{8}\right\rfloor=\left\lfloor 9.5\right\rfloor=\mathbf{9}$ $\displaystyle
Q_{z_{i,3}}$
$\displaystyle=\left\lfloor\frac{10(8-0)}{20-0}\right\rfloor=\left\lfloor\frac{10\times
8}{20}\right\rfloor=\left\lfloor 4\right\rfloor=\mathbf{4}.$

Thus, the quantized vector $Q_{z_{i}}$ is $[8;9;4]$. This means the vector
$z_{i}$ belongs to the global partition $[8;9;4]$ (partition $8$ in the first
dimension, partition $9$ in the second dimension and partition $4$ in the last
dimension). We point out that we start the numbering of partitions from zero.

##  Appendix B Probability mass function sampling code

Python code:

⬇

def PMFS_sampling(z, k):

max_, min_ = np.max(z, axis=0, keepdims=True), np.min(z, axis=0,
keepdims=True)

partition_sizes = (max_ - min_)/k

z = ((z-min_)/partition_sizes).astype(np.int8)

unique, counts = np.unique(z, axis=0,return_counts=True)

bin_edges = (min_, max_, num=k)

return unique, counts, bin_edges

Mathematica code:

⬇

MFSSampling[z_List, k_Integer] := Module[{max, min, partitionSizes,
quantizedZ, unique, counts,

binEdges}, max = Max[z]; min = Min[z];

partitionSizes = (max - min)/k;

quantizedZ = Floor[(z - min)/partitionSizes];

{unique, counts} = Transpose[Tally[quantizedZ]];

binEdges = {min, max, k};

{unique, counts, binEdges}]

##  Appendix C Experimental setup

To train the models for this work, we used the Pytorch [27] and Pytorch
Lightning [9] libraries in conjunction with Torchvision [38] to obtain the
training benchmark data sets. We conducted experiments on Nvidia Tesla V100
GPUs with $32$ GB of VRAM, available in the MesoPSL computing cluster. With
these GPUs, training on the MNIST, CelebA and MOBIUS data sets took
approximately $1$ minute, $14$ minutes and $25$ minutes per epoch
respectively, while computing the FID required approximately $30$ minutes per
data set.

The Numpy [14] and the Scikit-learn [28] libraries are respectively used to
implement the PMFS and GMM sampling strategies. The Matplotlib library [17] is
used to generate the figures in this work.

##  Appendix D Datasets

###  D.1 MNIST

The MNIST dataset [20], comprises ten classes of grayscale images, each with a
size of $28\times 28$, which we employ for training and testing our models.
The dataset is divided into : a training set with $50,000$ images, a
validation set with $10,000$ images, and a test set with $10,000$ images. This
dataset serves as a standard benchmark for computer vision tasks, particularly
for evaluating models on simple, low-resolution grayscale images.

###  D.2 CelebA

The Celebrity Faces dataset [21], encompasses colored images of celebrity
faces with a size of $178\times 218$, which we employ for training and testing
our models. It comprises a total of $200,000$ images. In our experiments, we
utilize the original train, validation, and test dataset split provided with
the dataset. This dataset serves as a standard benchmark for computer vision
applications, particularly in the domain of face generation, rendering it
highly relevant to the field of face biometrics.

###  D.3 MOBIUS

The MOBIUS dataset [30, 31, 40, 41, 42] is a compilation of $16,717$ colored
ocular images with dimensions $1,700\times 3,000$, which we utilize for
training and testing our models. During our experiments, we removed blurry
images, resulting in a total of $14,331$ images. We partitioned the dataset
into training, validation, and test subsets, consisting of $9,172$, $2,293$,
and $2,866$ images, respectively. The images from the MOBIUS dataset employed
in this work have been generously provided by the University of Ljubljana,
Slovenia, representing a typical biometric dataset suitable for image
generation tasks.

##  Appendix E Model architecture details

The models we trained in this work have the architectures details in Tables 3,
4 and 5. In all tables, $\star$ stands for a duplicated layer in stochastic
models to calculate the mean and the variance, and BN stands for batch
normalization.

Table 3: Autoencoders model architecture for the MNIST data set.  Layers | Encoder | Decoder
---|---|---
Layer 1 | $\mbox{Conv}(128,(4,4),s=2,p=1)$ | $\mbox{Linear}(32,16384)$
| BN, ReLU | $\mbox{Reshape}(1024,4,4)$
Layer 2 | $\mbox{Conv}(256,(4,4),s=2,p=1)$ | $\mbox{ConvT}(512,(4,4),s=2,p=1)$
| BN, ReLU | BN, ReLU
Layer 3 | $\mbox{Conv}(512,(4,4),s=2,p=1)$ | $\mbox{ConvT}(256,(4,4),s=2,p=2)$
| BN, ReLU | BN, ReLU
Layer 4 | $\mbox{Conv}(1024,(4,4),s=2,p=1)$ | $\mbox{ConvT}(256,(4,4),s=2,p=1)$
| BN, ReLU | Sigmoid
Layer 5 | $\mbox{Linear}(1024,32)\star$ | -
Table 4: Autoencoders model architecture for the CelebA data set.  Layers | Encoder | Decoder
---|---|---
Layer 1 | $\mbox{Conv}(128,(4,4),s=2,p=1)$ | $\mbox{Linear}(256,146432)$
| BN, ReLU | BN, ReLU, $\mbox{Reshape}(1024,4,4)$
Layer 2 | $\mbox{Conv}(256,(4,4),s=2,p=1)$ | $\mbox{ConvT}(512,(3,4),s=2,p=(0,1))$
| BN, ReLU | BN, ReLU
Layer 3 | $\mbox{Conv}(512,(4,4),s=2,p=1)$ | $\mbox{ConvT}(256,(3,4),s=2,p=(0,1))$
| BN, ReLU | BN, ReLU
Layer 4 | $\mbox{Conv}(1024,(4,4),s=2,p=1)$ | $\mbox{ConvT}(128,(3,4),s=2,p=1)$
| BN, ReLU | BN, ReLU
Layer 5 | $\mbox{Linear}(146432,256)\star$ | $\mbox{ConvT}(3,(2,4),s=2,p=0)$
|  | Sigmoid
Table 5: Autoencoders model architecture for the MOBIUS data set. Layers | Encoder | Decoder
---|---|---
Layer 1 | $\mbox{Conv}(32,(7,7),s=3,p=0)$ | $\mbox{Linear}(256,3072)$
| BN, ReLU | BN, ReLU, $\mbox{Reshape}(1024,1,3)$
Layer 2 | $\mbox{Conv}(64,(7,7),s=3,p=0)$ | $\mbox{ConvT}(512,(7,7),s=2,p=(0,0))$
| BN, ReLU | BN, ReLU
Layer 3 | $\mbox{Conv}(128,(7,7),s=3,p=0)$ | $\mbox{ConvT}(256,(5,7),s=3,p=(1,0))$
| BN, ReLU | BN, ReLU
Layer 4 | $\mbox{Conv}(256,(7,7),s=3,p=0)$ | $\mbox{ConvT}(128,(5,7),s=3,p=(1,2))$
| BN, ReLU | BN, ReLU
Layer 5 | $\mbox{Conv}(512,(7,7),s=3,p=0)$ | $\mbox{ConvT}(64,(5,7),s=3,p=(1,2))$
| BN, ReLU | BN, ReLU
Layer 6 | $\mbox{Conv}(1024,(5,5),s=2,p=0)$ | $\mbox{ConvT}(32,(5,7),s=3,p=(1,2))$
| BN, ReLU | BN, ReLU
Layer 7 | $\mbox{Linear}(3072,256)\star$ | $\mbox{ConvT}(3,(4,6),s=3,p=(1,0))$
|  | Sigmoid

##  Appendix F Image reconstruction

###  F.1 MNIST

(a) Vanilla Autoencoder.

(b) VAE.

(c) $\beta$-VAE.

(d) WAE.

(e) InfoVAE.

Figure 5: Reconstruction of MNIST images via different autoencoder based
models.

###  F.2 CelebA

(a) Vanilla Autoencoder.

(b) VAE.

(c) $\beta$-VAE.

(d) WAE.

(e) InfoVAE.

Figure 6: Reconstruction of CelebA images via different autoencoder based
models.

###  F.3 MOBIUS

(a) Vanilla Autoencoder.

(b) VAE.

(c) $\beta$-VAE.

(d) WAE.

(e) InfoVAE.

Figure 7: Reconstruction of MOBIUS images via different autoencoder based
models.

##  Appendix G Synthetic images

###  G.1 MNIST

(a) GMM Vanilla Autoencoder.

(b) PMFS Vanilla Autoencoder.

(c) GMM VAE.

(d) PMFS VAE.

(e) GMM $\beta$-VAE.

(a) PMFS $\beta$-VAE.

(b) GMM WAE.

(c) PMFS WAE.

(d) GMM InfoVAE.

(e) PMFS InfoVAE.

Figure 9: Synthetic images generated using GMM and PMFS sampling on the MNIST
data set via different autoencoder based models.

###  G.2 CelebA

(a) GMM Vanilla Autoencoder.

(b) PMFS Vanilla Autoencoder.

(c) GMM VAE.

(d) PMFS VAE.

(e) GMM $\beta$-VAE.

(a) PMFS $\beta$-VAE.

(b) GMM WAE.

(c) PMFS WAE.

(d) GMM InfoVAE.

(e) PMFS InfoVAE.

Figure 11: Synthetic images generated using GMM and PMFS sampling on the
CelebA data set via different autoencoder based models.

###  G.3 MOBIUS

(a) GMM Vanilla Autoencoder.

(b) PMFS Vanilla Autoencoder.

(c) GMM VAE.

(d) PMFS VAE.

(e) GMM $\beta$-VAE.

(a) PMFS $\beta$-VAE.

(b) GMM WAE.

(c) PMFS WAE.

(d) GMM InfoVAE.

(e) PMFS InfoVAE.

Figure 13: Synthetic images generated using GMM and PMFS sampling on the
MOBIUS data set via different autoencoder based models.

##  Appendix H Distribution plots

###  H.1 MNIST

(a) Latent space distribution of Vanilla Autoencoder.

(b) Latent space distribution of VAE.

(c) Latent space distribution of $\beta$-VAE.

(d) Latent space distribution of WAE.

(e) Latent space distribution of InfoVAE.

Figure 14: Distribution of the validation set’s latent vectors of different
Autoencoder models and samples generated via GMM and PMFS sampling on the
MNIST data set. These figures represents a dimensionality reduction using PCA
into the two dimensional $\mathbb{R}^{2}$ space of the first hundred images
from the validation set, as well as a hundred samples from GMM and PMFS
sampling methods.

###  H.2 CelebA

(a) Latent space distribution of Vanilla Autoencoder.

(b) Latent space distribution of VAE.

(c) Latent space distribution of $\beta$-VAE.

(d) Latent space distribution of WAE.

(e) Latent space distribution of InfoVAE.

Figure 15: Distribution of the validation set’s latent vectors of different
Autoencoder models and samples generated via GMM and PMFS sampling on the
CelebA data set. These figures represents a dimensionality reduction using PCA
into the two dimensional $\mathbb{R}^{2}$ space of the first hundred images
from the validation set, as well as a hundred samples from GMM and PMFS
sampling methods.

###  H.3 MOBIUS

(a) Latent space distribution of Vanilla Autoencoder.

(b) Latent space distribution of VAE.

(c) Latent space distribution of $\beta$-VAE.

(d) Latent space distribution of WAE.

(e) Latent space distribution of InfoVAE.

Figure 16: Distribution of the validation set’s latent vectors of different
Autoencoder models and samples generated via GMM and PMFS sampling on the
MOBIUS data set. These figures represents a dimensionality reduction using PCA
into the two dimensional $\mathbb{R}^{2}$ space of the first hundred images
from the validation set, as well as a hundred samples from GMM and PMFS
sampling methods.