Self-Supervised Variational Auto-Encoders

Abstract

Density estimation, compression, and data generation are crucial tasks in artificial intelligence. Variational Auto-Encoders (VAEs) constitute a single framework to achieve these goals. Here, we present a novel class of generative models, called self-supervised Variational Auto-Encoder (selfVAE), which utilizes deterministic and discrete transformations of data. This class of models allows both conditional and unconditional sampling while simplifying the objective function. First, we use a single self-supervised transformation as a latent variable, where the transformation is either downscaling or edge detection. Next, we consider a hierarchical architecture, i.e., multiple transformations, and we show its benefits compared to the VAE. The flexibility of selfVAE in data reconstruction finds a particularly interesting use case in data compression tasks, where we can trade-off memory for better data quality and vice-versa. We present the performance of our approach on three benchmark image data (Cifar10, Imagenette64, and CelebA).

1. Introduction

The framework of variational autoencoders (VAEs) provides a principled approach for learning latent-variable models. As it utilizes a meaningful low-dimensional latent space with density estimation capabilities, it forms an attractive solution for generative modeling tasks. However, its performance in terms of the test log-likelihood and quality of generated samples is often dissatisfying, and thus many modifications were proposed. In general, one can obtain a tighter lower bound and, thus, a more powerful and flexible model, by advancing over the following three components: the encoder [1,2,3,4], the prior (or marginal over latents) [5,6,7,8,9], and the decoder [10]. Recent studies have shown that, by employing deep hierarchical architectures and by carefully designing the building blocks of the neural networks, VAEs can successfully model high-dimensional data and reach state-of-the-art test likelihoods [11,12,13].

In this work, we present a novel class of VAEs, called self-supervised Variational Auto-Encoders, where we introduce additional variables to VAEs that result from the discrete and deterministic transformations of observed images. Since the transformations are deterministic, and they provide a specific aspect of images (e.g., contextual information through detecting edges or downscaling), we refer to them as self-supervised representations.

The introduction of the discrete and deterministic variables allows training deep hierarchical models efficiently by decomposing the task of learning a highly complex distribution into training smaller and conditional distributions. In this way, the model allows integrating prior knowledge about the data but still enables the synthesis of unconditional samples. Furthermore, the discrete and deterministic variables could be used to conditionally reconstruct data, which could be of great use in data compression and super-resolution tasks.

We make the following contributions: (i) We propose an extension of the VAE framework by incorporating self-supervised representations of the data. (ii) We analyze the impact of modeling natural images with different data transformations as self-supervised representations. (iii) This new type of generative model (self-supervised Variational Auto-Encoders), which can perform both conditional and unconditional sampling, demonstrated improved quantitative performance in terms of density estimation and the generative capabilities on image benchmarks.

2. Background

2.1. Variational Auto-Encoders

Let $\mathbf{x}\in {\mathcal{X}}^{\mathrm{D}}$ be a vector of observable variables, where $\mathcal{X}\subseteq \mathbb{R}$ or $\mathcal{X}\subseteq \mathbb{Z}$ and $\mathbf{z}\in {\mathbb{R}}^M$ denote a vector of latent variables. We are interested in a latent variable model of the following form:

$$p_{\vartheta }(\mathbf{x},\mathbf{z})=p_{\theta }\left(\mathbf{x}\right|\mathbf{z})p_{\lambda }\left(\mathbf{z}\right).$$

(1)

For training, we require marginalizing $\mathbf{z}$ out; however, calculating the marginal likelihood $p_{\vartheta }\left(\mathbf{x}\right)=\int p_{\vartheta }(\mathbf{x},\mathbf{z})\mathrm{d}\mathbf{z}$ is computationally intractable for non-linear stochastic dependencies. As a potential solution, a variational family of distributions could be used for approximate inference. Then, the logarithm of the marginal likelihood could be lower-bounded and the following objective function could be derived, namely, the evidence lower bound (ELBO) [14]:

$$\begin{array}{c}\hfill ln{p}_{\vartheta }\left(\mathbf{x}\right)\ge \underset{q_{\varphi }\left(\mathbf{z}\right|\mathbf{x})}{\mathbb{E}}\left[ln{p}_{\theta }\left(\mathbf{x}\right|\mathbf{z})+ln{p}_{\lambda }\left(\mathbf{z}\right)-ln{q}_{\varphi }\left(\mathbf{z}\right|\mathbf{x})\right]\end{array}$$

(2)

Where $q_{\varphi }\left(\mathbf{z}\right|\mathbf{x})$ is the variational posterior (or the encoder), $p_{\theta }\left(\mathbf{x}\right|\mathbf{z})$ is the conditional likelihood function (or the decoder), $p_{\lambda }\left(\mathbf{z}\right)$ is the prior (or marginal), and $\varphi$, $\theta$, and $\lambda$ denote parameters. The expectation is approximated by Monte Carlo sampling while exploiting the reparameterization trick to obtain unbiased gradient estimators [15]. The models are parameterized by neural networks. This generative framework is known as Variational Auto-Encoder (VAE) [1,15].

2.2. VAEs with Bijective Priors

Even though the lower-bound suggests that the prior plays a crucial role in improving the variational bounds, usually a fixed distribution is used, e.g., a standard multivariate Gaussian. While being relatively simple and computationally cheap, the fixed prior is known to result in over-regularized models that tend to ignore most of the latent dimensions [9,16,17]. Moreover, even with powerful encoders, VAEs may still fail to match the variational posterior to a unit Gaussian prior [18].

However, it is possible to obtain a rich, multi-modal prior distribution $p\left(\mathbf{z}\right)$ by using a bijective (or flow-based) model [19]. Formally, given a latent code $\mathbf{z}$, a base distribution $p_V\left(\mathbf{v}\right)$ over latent variables $\mathbf{v}\in {\mathbb{R}}^M$, and $f:{\mathbb{R}}^M\stackrel{}{\to }{\mathbb{R}}^M$, consisting of a sequence of L diffeomorphic transformations (i.e., invertible and differentiable transformations), where $f_i\left({\mathbf{v}}_{i-1}\right)={\mathbf{v}}_i$, ${\mathbf{v}}_0=\mathbf{v}$, and ${\mathbf{v}}_L=\mathbf{z}$, the change of variable can be used sequentially to express the distribution of $\mathbf{z}$ as a function of $\mathbf{v}$ as follows:

$$logp\left(\mathbf{z}\right)=log{p}_V\left(\mathbf{v}\right)-\sum _{i=1}^{L}log\left|\frac{\partial f_i\left({\mathbf{v}}_{i-1}\right)}{\partial {\mathbf{v}}_{i-1}}\right|,$$

(3)

where $\left|\frac{\partial f_i\left({\mathbf{v}}_{i-1}\right)}{\partial {\mathbf{v}}_{i-1}}\right|$ is the Jacobian-determinant of the $i$th transformation. Thus, using the bijective prior yields the following lower-bound:

$$lnp\left(\mathbf{x}\right)\ge {\mathbb{E}}_{q_{\varphi }\left(\mathbf{z}\right|\mathbf{x})}\left[log{p}_{\theta }\left(\mathbf{x}\right|\mathbf{z})-log{q}_{\varphi }\left(\mathbf{z}\right|\mathbf{x})+log{p}_V\left({\mathbf{v}}_0\right)+\sum _{i=1}^{L}log\left|\frac{\partial f_i^{-1}\left({\mathbf{v}}_i\right)}{\partial {\mathbf{v}}_i}\right|\right].$$

(4)

In this work, we utilize Real-Valued Non-Volume Preserving (RealNVP) transformations [19] as the prior; however, any other flow-based model could be used [3,20]. In the RealNVP, invertible transformations are composed of coupling layers. A single coupling layer, for a given input ${\mathbf{v}}_i=[{\mathbf{v}}_{i,a},{\mathbf{v}}_{i,b}]$, processes only ${\mathbf{v}}_{i,b}$ to obtain ${\mathbf{v}}_{i+1,b}=exp\left(s\left({\mathbf{v}}_{i,a}\right)\right)\odot {\mathbf{v}}_{i,b}+t\left({\mathbf{v}}_{i,a}\right)$, where $s(·)$ and $t(·)$ denote scale and translation neural networks, respectively. Then, the Jacobian-determinant of a single coupling layer can be calculated analytically, namely:

$$\left|\frac{\partial f_i^{-1}\left({\mathbf{v}}_i\right)}{\partial {\mathbf{v}}_i}\right|=\prod _{j=1}^{D-d}exp{\left(s\left({\mathbf{v}}_{i,a}\right)\right)}_j=exp\left(\sum _{j=1}^{D-d}s{\left({\mathbf{x}}_{i,a}\right)}_j\right).$$

(5)

3. Method

3.1. Motivation

The idea of self-supervised learning is about utilizing original unlabeled data to create additional context information. This could be achieved in multiple manners, e.g., by adding noise to data [21] or masking data during training [22]. Self-supervised learning could also be seen as turning an unsupervised model into a supervised model by, e.g., treating predicting the next pixels as a classification task [23,24]. These are only a few examples of a quickly growing research line [25].

Here, we propose to use non-trainable transformations to obtain information about the image data. Our main hypothesis is that, since working with high-quality images is challenging, we could alleviate this problem by additionally considering partial information about them. Fitting a model to images of lower quality and then enhancing them to match the target distribution appears to be an overall easier task [26,27].

By incorporating compressed transformations (i.e., the self-supervised representations) that still contain global information, with the premise that it would be easier to approximate, the process of modeling a high-dimensional complex density can be broken down into simpler tasks. In this way, the expressivity of the model will grow and gradually result in richer, better generations.

Examples of image transformations are presented in Figure 1, such as downscaling through bicubic interpolation, blurring using Gaussian kernels, considering an n-bit image representation, edge detection, or turning an RGB image into grayscale. We notice that, with these transformations, even though they discard a great deal of information, the global structure is preserved. As a result, in practice, the model should have the ability to extract a general concept of the data and to add local information afterward. In this work, we focus on downscaling (Figure 1b–d) and edge detection or sketching (Figure 1i).

A positive effect of the proposed framework is that the model allows us to integrate prior knowledge through the image transformations, without losing its unconditional generative functionality. Overall, we end up with a two-level VAE with three latent variables, where one is a data transformation that can be obtained in a self-supervised fashion. In Figure 2, a schematic representation of the proposed approach with downscaling is presented.

3.2. Model Formulation

In our model, we consider representations that result from deterministic and discrete transformations of an image. Formally, we introduce a transformation $d:{\mathcal{X}}^D\to {\mathcal{X}}^C$ that takes $\mathbf{x}$ and returns an image representation $\mathbf{y}$, e.g., a downscaled image. Since we lose information about the original image, $\mathbf{z}$ could be seen as a variable that compensates for lost details in $\mathbf{x}$. Further, we propose to introduce an additional latent variable, $\mathbf{u}\in {\mathbb{R}}^N$, to model $\mathbf{y}$ and $\mathbf{z}$.

We can define the joint distribution of $\mathbf{x}$ and $\mathbf{y}$ as $p(\mathbf{x},\mathbf{y})=p\left(\mathbf{y}\right|\mathbf{x}\left)p\right(\mathbf{x})$, where $p\left(\mathbf{y}\right|\mathbf{x})$ becomes deterministic due to the deterministic transformation $d(·)$. Formally, $p\left(\mathbf{y}\right|\mathbf{x})=\delta (\mathbf{y}-d\left(\mathbf{x}\right))$, where $\delta (·)$ denotes the Kronecker delta (for $\mathcal{X}=\mathbb{Z}$) or the Dirac delta (for $\mathcal{X}=\mathbb{R}$). Once we define the deterministic function d, for given data $\mathcal{D}=\{{\mathbf{x}}_1,\dots ,{\mathbf{x}}_N\}$, we can obtain corresponding $\mathbf{y}$’s. At first glance, it might be not apparent what is a potential advantage of our approach except for knowing how to generate pairs of data $\{{({\mathbf{x}}_n,{\mathbf{y}}_n\}}_{n=1}^N$. In order to obtain insight into that, we first consider the joint entropy of $\mathbf{x}$ and $\mathbf{y}$, which is expressed as follows:

$$\begin{array}{cc}\hfill \mathbb{H}[\mathbf{x},\mathbf{y}]& =\mathbb{H}\left[\mathbf{y}\right|\mathbf{x}]+\mathbb{H}[\mathbf{x}]\hfill \end{array}$$

(6)

$$\begin{array}{cc}\hfill & =\mathbb{H}\left[\mathbf{x}\right|\mathbf{y}]+\mathbb{H}[\mathbf{y}].\hfill \end{array}$$

(7)

Immediately, we can notice that, since $\mathbf{y}$ is a result of a deterministic function of $\mathbf{x}$, the first component in Equation (6) is 0, i.e., $\mathbb{H}\left[\mathbf{y}\right|\mathbf{x}]=0$. This follows from the fact that $\mathbf{y}$ is a processed version of $\mathbf{x}$ and contains no new information about $\mathbf{x}$. As a result, we can express the entropy of $\mathbf{x}$ as follows:

$$\mathbb{H}\left[\mathbf{x}\right]=\mathbb{H}\left[\mathbf{x}\right|\mathbf{y}]+\mathbb{H}[\mathbf{y}].$$

(8)

In other words, the entropy of the original data is equal to the conditional entropy of the original data and the entropy of the processed data. In general, modeling conditional distributions is an easier task than modeling the original data, and modeling $p\left(\mathbf{y}\right)$ may be also less complicated than modeling $p\left(\mathbf{x}\right)$ because $\mathbf{y}$ corresponds to a representation (a processed version) of $\mathbf{x}$. As a result, we can turn the problem of modeling a complex marginal distribution $p\left(\mathbf{x}\right)$ into a self-supervised problem, where we factorize the joint distribution of $\mathbf{x}$ and $\mathbf{y}$ as $p(\mathbf{x},\mathbf{y})=p\left(\mathbf{x}\right|\mathbf{y}\left)p\right(\mathbf{y})$. Further, we propose to compensate lost information in $\mathbf{y}$ by introducing $\mathbf{z}$, and we also model $\mathbf{y}$ and $\mathbf{z}$ by introducing the additional latent variable $\mathbf{u}$. The joint distribution over all random variables in our model is as follows:

$$p(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{u})=p\left(\mathbf{x}\right|\mathbf{y},\mathbf{z}\left)p\right(\mathbf{y}\left|\mathbf{u}\right)p\left(\mathbf{z}\right|\mathbf{y},\mathbf{u}\left)p\right(\mathbf{u}).$$

(9)

All dependencies in our model are presented in Figure 2.

Eventually, for training, we need to calculate the logarithm of the marginal likelihood, namely, $p(\mathbf{x},\mathbf{y})=\int \int p(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{u})d\mathbf{z}\mathbf{u}$. Calculating these integrals is intractable; therefore, we again utilize the variational inference. For this purpose, we propose to apply the following family of variational distributions over latent variables:

$$Q(\mathbf{u},\mathbf{z}|\mathbf{x},\mathbf{y})=q(\mathbf{u}\left|\mathbf{y}\right)q\left(\mathbf{z}\right|\mathbf{x}).$$

(10)

Notice that we skip a dependency between $\mathbf{u}$ and $\mathbf{x}$. Utilizing the variational inferences gives the following lower-bound on the logarithm of the marginal likelihood:

$$\begin{array}{cc}\hfill lnp(\mathbf{x},\mathbf{y})\ge {\mathbb{E}}_Q[ln{p}_{\theta }\left(\mathbf{x}\right|\mathbf{y},\mathbf{z})+lnp\left(\mathbf{z}\right|\mathbf{u},\mathbf{y})+& lnp\left(\mathbf{y}\right|\mathbf{u})+lnp(\mathbf{u})+\hfill \\ \hfill & -lnq\left(\mathbf{z}\right|\mathbf{x})-lnq(\mathbf{u}\left|\mathbf{y}\right)].\hfill \end{array}$$

(11)

Intuitively, the premise for self-supervised Variational Auto-Encoder (selfVAE) is that the latents $\mathbf{u}$ will capture the global structure of the input data and the latents $\mathbf{z}$ will encode the missing information between $\mathbf{y}$ and $\mathbf{x}$, guiding the model to discover the distribution of the target observations. To highlight the self-supervised part in our model, we refer to it as the self-supervised Variational Auto-Encoder (or selfVAE for short). Further, we propose to choose the following distributions:

$$\begin{array}{cc}\hfill p\left(\mathbf{v}\right)& =\mathcal{N}\left(\mathbf{v}|\mathbf{0},\mathbf{1}\right)\hfill \end{array}$$

(12)

$$\begin{array}{cc}\hfill p_{\lambda }\left(\mathbf{u}\right)& =p\left(\mathbf{v}\right)\prod _{i=1}^F|det\frac{\partial f_i\left({\mathbf{v}}_{i-1}\right)}{\partial {\mathbf{v}}_{i-1}}{|}^{-1}\hfill \end{array}$$

(13)

$$\begin{array}{cc}\hfill p_{{\theta }_1}\left(\mathbf{y}|\mathbf{u}\right)& =\sum _{i=1}^I{\pi }_i\left(\mathbf{u}\right)\mathrm{Dlogistic}\left({\mu }_i\left(\mathbf{u}\right),s_i\left(\mathbf{u}\right)\right)\hfill \end{array}$$

(14)

$$\begin{array}{cc}\hfill p_{{\theta }_2}\left(\mathbf{z}|\mathbf{y},\mathbf{u}\right.)& =\mathcal{N}\left(\mathbf{z}|{\mathit{\mu }}_{{\theta }_2}(\mathbf{y},\mathbf{u}),diag\left({\sigma }_{{\theta }_2}(\mathbf{y},\mathbf{u})\right.)\right)\hfill \end{array}$$

(15)

$$\begin{array}{cc}\hfill p_{{\theta }_3}\left(\mathbf{x}|\mathbf{z},\mathbf{y}\right.)& =\sum _{i=1}^I{\pi }_i(\mathbf{z},\mathbf{y})\mathrm{Dlogistic}\left({\mu }_i(\mathbf{z},\mathbf{y}),s_i(\mathbf{z},\mathbf{y})\right)\hfill \end{array}$$

(16)

$$\begin{array}{cc}\hfill q_{{\varphi }_1}\left(\mathbf{u}|\mathbf{y}\right)& =\mathcal{N}\left(\mathbf{u}|{\mathit{\mu }}_{{\varphi }_1}\left(\mathbf{y}\right),\mathrm{diag}\left({\sigma }_{{\varphi }_1}\left(\mathbf{y}\right)\right)\right)\hfill \end{array}$$

(17)

$$\begin{array}{cc}\hfill q_{{\varphi }_2}\left(\mathbf{z}|\mathbf{x}\right)& =\mathcal{N}\left(\mathbf{z}|{\mathit{\mu }}_{{\varphi }_2}\left(\mathbf{x}\right),\mathrm{diag}\left({\sigma }_{{\varphi }_2}\left(\mathbf{x}\right)\right)\right).\hfill \end{array}$$

(18)

For the images $\mathbf{x}$ and $\mathbf{y}$, we use $\mathrm{Dlogistic}$, which is the discretized logistic distribution [28,29]. This is defined as a difference of two CDFs of the logistic distribution with a mean $\mu$ and a scale $\nu$:

$$\mathrm{Dlogistic}\left(x\right)=\mathrm{sigm}\left(\frac{(x+0.5-\mu )}{\nu }\right)-\mathrm{sigm}\left(\frac{(x-0.5-\mu )}{\nu }\right),$$

(19)

where $\mathrm{sigm}$ is the sigmoid function, and we utilize a flow-based model for $p_{\lambda }\left(\mathbf{u}\right.)$. We use the discretized logistic distribution because the images are represented by values between 0 and 255. For integer-valued random variables, other distributions, such as Gaussian, are inappropriate.

3.3. Generation and Reconstruction in selfVAE

Latent variable generative models, such as VAEs, can be used to synthesize novel content through the following process: $\mathbf{z}\sim p\left(\mathbf{z}\right)\to \mathbf{x}\sim p\left(\mathbf{x}\right|\mathbf{z})$, i.e., we first sample $\mathbf{z}$, which is further fed to the conditional distribution to sample $\mathbf{x}$. Additionally, these models can be utilized in reconstructing a datapoint ${\mathbf{x}}^{*}$ by using the following scheme: $\mathbf{z}\sim q\left(\mathbf{z}\right|{\mathbf{x}}^{*})\to \mathbf{x}\sim p\left(\mathbf{x}\right|\mathbf{z})$. During training, VAEs use the reconstruction process.

Interestingly, our approach allows utilizing more operations regarding data generation and reconstruction. First, analogously to VAEs, the selfVAE can generate data by applying the following hierarchical sampling process (we refer to it as generation, see Figure 3i): $\mathbf{u}\sim p\left(\mathbf{u}\right)\stackrel{}{\to }\mathbf{y}\sim p\left(\mathbf{y}\right|\mathbf{u})\stackrel{}{\to }\mathbf{z}\sim p\left(\mathbf{z}\right|\mathbf{u},\mathbf{y})\stackrel{}{\to }\mathbf{x}\sim p\left(\mathbf{x}\right|\mathbf{y},\mathbf{z}).$ In this procedure, we start with sampling $\mathbf{u}$ and proceed to sample $\mathbf{z}$ and $\mathbf{y}$ given $\mathbf{u}$, and eventually we can obtain $\mathbf{x}$. Alternatively, we can use the ground-truth $\mathbf{y}$ (i.e, ${\mathbf{y}}^{*}=d\left({\mathbf{x}}^{*}\right)$), and sample or infer $\mathbf{z}$.

Then, the generative process is the following (we refer to it as conditional generation, see Figure 3ii): $\mathbf{u}\sim q\left(\mathbf{u}\right|{\mathbf{y}}^{*})\stackrel{}{\to }\mathbf{z}\sim p\left(\mathbf{z}\right|\mathbf{u},{\mathbf{y}}^{*}),\stackrel{}{\to }\mathbf{x}\sim p\left(\mathbf{x}\right|{\mathbf{y}}^{*},\mathbf{z}).$ Knowing ${\mathbf{y}}^{*}$ allows us to use the variational posterior to sample $\mathbf{u}$ instead of applying the prior $p\left(\mathbf{u}\right)$. This scenario is possible if we consider ${\mathbf{y}}^{*}$ as a compressed version of ${\mathbf{x}}^{*}$ that is shared instead of the original content.

Similarly to the conditional generation, we can conditionally reconstruct ${\mathbf{x}}^{*}$ by utilizing ${\mathbf{y}}^{*}$. Then, the procedure is the following (we refer to it as conditional reconstruction, see Figure 3iii): $\mathbf{z}\sim q\left(\mathbf{z}\right|{\mathbf{x}}^{*})\stackrel{}{\to }\mathbf{x}\sim p\left(\mathbf{x}\right|{\mathbf{y}}^{*},\mathbf{z}).$ Here, we skip inferring $\mathbf{u}$, because for the given ${\mathbf{y}}^{*}$ and ${\mathbf{x}}^{*}$, the posterior could be used to sample $\mathbf{z}$ and then obtain $\mathbf{x}$. If $\mathbf{y}$ is a result of a downscaling transformation of the input image, selfVAE is equivalent to the super-resolution as in [27].

Further, we can generate $\mathbf{y}$ instead of using ${\mathbf{y}}^{*}$ and choose to sample or infer $\mathbf{z}$. In this way, we can reconstruct an image in two ways. First, we go from ${\mathbf{x}}^{*}$ to $\mathbf{u}$, and then go back to $\mathbf{x}$ using a purely generative path (we refer to it as reconstruction 1, see Figure 3iv): ${\mathbf{y}}^{*}=d\left({\mathbf{x}}^{*}\right)\stackrel{}{\to }\mathbf{u}\sim q\left(\mathbf{u}\right|{\mathbf{y}}^{*})\stackrel{}{\to }\mathbf{y}\sim p\left(\mathbf{y}\right|\mathbf{u})\stackrel{}{\to }\mathbf{z}\sim p\left(\mathbf{z}\right|\mathbf{u},\mathbf{y})\stackrel{}{\to }\mathbf{x}\sim p\left(\mathbf{x}\right|\mathbf{z},\mathbf{y}).$ Alternatively, we can use the variational posterior to sample $\mathbf{z}$ as in the following procedure (we refer to it as reconstruction 2, see Figure 3v): $({\mathbf{y}}^{*}=d\left({\mathbf{x}}^{*}\right)\stackrel{}{\to }\mathbf{u}\sim q\left(\mathbf{u}\right|{\mathbf{y}}^{*})\stackrel{}{\to }\mathbf{y}\sim p\left(\mathbf{y}\right|\mathbf{u})),\mathrm{then}\mathbf{z}\sim q\left(\mathbf{z}\right|{\mathbf{x}}^{*})\stackrel{}{\to }\mathbf{x}\sim p\left(\mathbf{x}\right|\mathbf{y},\mathbf{z}).$

In both reconstructions, we want to obtain $\mathbf{u}$. To understand the purpose, we can again think of compression. Instead of sending high-dimensional variables, we can share the low-dimensional $\mathbf{u}$ and, then, infer all other variables. However, in the reconstruction 2 procedure, we must send $\mathbf{z}$ as well. As we will see in the experiments, this trade-off of the number of sent variables allows us to play with the compression ratio and the quality of the reconstruction.

The presented versions of generating and reconstructing images could be useful in the compression task. As we will see in the experiments, each option creates a different ratio of the reconstruction quality against the memory that we need to allocate to send information. However, every inferred variable needs to be sent; thus, more sampling corresponds to lower memory requirements. We want to highlight that, during training, we sample from variational posteriors, ${\mathbf{z}}^{*}\sim q\left(\mathbf{z}\right|{\mathbf{x}}^{*}),{\mathbf{u}}^{*}\sim p\left(\mathbf{u}\right|{\mathbf{y}}^{*})$, and ${\mathbf{y}}^{*}=d\left({\mathbf{x}}^{*}\right)$, and then we use these samples for the generative part and in the ELBO eventually. None of the procedures presented in this subsection are used for training. Here, we outline that the selfVAE allows various options for generation and reconstruction, with potential use for compression.

3.4. Hierarchical Self-Supervised VAE

The proposed approach can be further extended and generalized by introducing multiple transformations. For instance, we can achieve this by applying downscaling multiple times. As a result, by incorporating multiple self-supervised representations of the data, the process of modeling a high-dimensional complex density breaks down into K simpler modeling tasks. The resulting model is a K-level VAE architecture, where the overall expressivity of the model grows even further and gradually leads to generations of higher quality. It is important to realize that some transformations cannot be applied multiple times, e.g., edge detection; however, others could be used sequentially, e.g., downscaling.

Formally, we consider K self-supervised data transformations $d_k(·)$ that result in K representations denoted by ${\mathbf{y}}_{1:K}=[{\mathbf{y}}_1,\dots ,{\mathbf{y}}_K]$, where ${\mathbf{y}}_k=d_k\left({\mathbf{y}}_{k-1}\right)$ and ${\mathbf{y}}_0=\mathbf{x}$. Thus, the joint distribution is factorized as follows:

$$p(\mathbf{x},{\mathbf{y}}_{1:K})=p\left(\mathbf{x}\right)\prod _{k=1}^K\delta \left({\mathbf{y}}_k-d_k\left({\mathbf{y}}_{k-1}\right)\right).$$

(20)

However, analogously to the model with a single $\mathbf{y}$, we can introduce additional latent variables for each level to compensate for lost information, ${\mathbf{z}}_k$, and another latent variable on top, $\mathbf{u}$. As a result, our model can be expressed in the following form:

$$\begin{array}{cc}\hfill p(\mathbf{x},{\mathbf{y}}_{1:K},{\mathbf{z}}_{1:K},\mathbf{u})=& p\left(\mathbf{x}\right|{\mathbf{y}}_1,{\mathbf{z}}_1)\prod _{k=1}^{K-1}\{p\left({\mathbf{z}}_k\right|{\mathbf{y}}_k,{\mathbf{z}}_{k+1})\times p\left({\mathbf{y}}_k\right|{\mathbf{y}}_{k+1},{\mathbf{z}}_{k+1})\}\times \hfill \\ \hfill & \times p\left({\mathbf{z}}_K\right|\mathbf{u},{\mathbf{y}}_K)p\left({\mathbf{y}}_K\right|\mathbf{u})p\left(\mathbf{u}\right).\hfill \end{array}$$

(21)

Further, due to the intractability of the marginal likelihood function, we utilize the following family of variational distributions:

$$Q(\mathbf{u},\mathbf{z}|\mathbf{x},{\mathbf{y}}_{1:K})=q\left(\mathbf{u}\right|{\mathbf{y}}_K)q\left({\mathbf{z}}_1\right|\mathbf{x})\prod _{k=1}^{K-1}q\left({\mathbf{z}}_{k+1}\right|{\mathbf{y}}_k).$$

(22)

After applying the variational inference, we obtain the following objective:

$$\begin{array}{cc}\hfill lnp(\mathbf{x},{\mathbf{y}}_{1:K})\ge & {\mathbb{E}}_Q[ln{p}_{\theta }\left(\mathbf{x}\right|{\mathbf{y}}_1,{\mathbf{z}}_1)+\sum _{k=1}^{K-1}\{lnp\left({\mathbf{z}}_k\right|{\mathbf{y}}_k,{\mathbf{z}}_{k+1})+\hfill \\ \hfill & +lnp\left({\mathbf{y}}_k\right|{\mathbf{y}}_{k+1},{\mathbf{z}}_{k+1})\}+lnp\left({\mathbf{z}}_K\right|\mathbf{u},{\mathbf{y}}_K)+\hfill \\ \hfill & +lnp\left({\mathbf{y}}_K\right|\mathbf{u})+lnp\left(\mathbf{u}\right)-lnq\left(\mathbf{u}\right|{\mathbf{y}}_K)+\hfill \\ \hfill & -lnq\left({\mathbf{z}}_1\right|\mathbf{x})-\sum _{k=1}^{K-1}lnq\left({\mathbf{z}}_{k+1}\right|{\mathbf{y}}_k)].\hfill \end{array}$$

(23)

4. Experiments

4.1. Experimental Setup

• Datasets We evaluated the proposed model on the CIFAR-10, Imagenette64, and CelebA datasets:

  -

  CIFAR-10 The CIFAR-10 (https://www.cs.toronto.edu/~kriz/cifar.html (accessed on 11 June 2021)) dataset is a well-known image benchmark data containing 60,000 training natural images and 10,000 validation natural images. Each image is of size 32 px × 32 px. From the training data, we put aside $15\%$ randomly selected images as the test set. We augmented the training data by using random horizontal flips and random affine transformations and normalized the data uniformly in the range (0, 1).

  -

  Imagenette64 Imagenette64 (https://github.com/fastai/imagenette (accessed on 11 June 2021)) is a subset of 10 easily classified classes from Imagenet https://www.image-net.org/ (accessed on 11 June 2021 ) (tench, English springer, cassette player, chain saw, church, French horn, garbage truck, gas pump, golf ball, and parachute). We downscaled the dataset to 64 px × 64 px images. Similarly to CIFAR-10, we put aside $15\%$ randomly selected training images as the test set. We used the same data augmentation as in CIFAR-10. Please note that Imagenette64 is not equivalent to Imagenet64.

  -

  CelebA The Large-scale CelebFaces Attributes (CelebA)(http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html (accessed on 11 June 2021)) dataset consists of $202,599$ images of celebrities. We cropped the original images on the 40 vertical and 15 horizontal components of the top left corner of the crop box, where the height and width were cropped to 148. In addition to the uniform normalization of the image, no other augmentation was applied.

• Architectures The encoders and decoders consisted of building blocks composed of DenseNets [30], channel-wise attention [31], and ELUs [32] as activation functions. The dimensionality of all the latent variables was kept at $8\times 8\times 16=1024$, and all models were trained using AdaMax [33] with data-dependent initialization [34]. Regarding the selfVAEs, in CIFAR-10, we used an architecture with a single downscaled transformation (selfVAE-downscale), while, on the remaining two datasets (CelebA and Imagenette64), we used a hierarchical three-leveled selfVAE with downscaling, and a selfVAE with sketching. Note that using sketching multiple times would result in no new data representation.
  All models were employed with the bijective prior (RealNVP). All models were comparable in terms of the number of parameters (the range of the weights of all models was from 32 million to 42 million). Due to our limited computational resources, these were the largest models that we were able to train. A fair comparison to the current SOTA models that use over 100 million weights [12,13] is, therefore, slightly skewed. For more details on the model architectures, please refer to the Appendix A.1.
• Evaluation We approximated the negative log-likelihood using 512 IW-samples [16] and express the scores in bits per dimension (bpd). Additionally, for CIFAR-10, we used the Fréchet Inception Distance (FID) [35]. We used the FID score indicating that sometimes the likelihood-based models achieved very low bpd; however, their perceptual quality was lower.

4.2. Quantitative Results

We present the results of the experiments on the benchmark datasets in

Table 1. Quantitative comparison on the test sets from CIFAR-10, CelebA, and Imagenette64. If available, we provide a score measured on the training set (depicted by *). Best results are in bold.

Dataset	Model	bpd↓	FID ↓
CIFAR-10	PixelCNN [36]	3.14	65.93
	GLOW [20]	3.35	65.93
	ResidualFlow [37]	3.28	46.37
	BIVA [12]	3.08	-
	NVAE [13]	2.91	-
	Very deep VAEs [38]	2.87
	DDPM [39]	3.75	5.24 (3.17 *)
	VAE (ours)	3.51	41.36 (37.25 *)
	selfVAE-downscale	3.65	34.71 (29.95 *)
CelebA	RealNVP [19]	3.02	-
	VAE (ours)	3.12	-
	selfVAE-sketch	3.24	-
	selfVAE-downscale-3lvl	2.88	-
Imagenette64	VAE (ours)	3.85	-
	selfVAE-downscale-3lvl	3.70	-

. First, we notice that, on CIFAR-10, our implementation of the VAE was still lagging behind the other generative models in terms of the bpd; however, it was better or comparable in terms of the FID. Surprisingly, the selfVAE-downscale achieved worse bpd compared with the VAE with the bijective prior. A possible explanation may lie in the small image size ($32\times 32$), as the benefits of breaking down the learning process into two or more steps are not obvious given the small target dimensional space.

Nevertheless, the selfVAE-downscale achieved significantly better FID scores over the other generative models, except for the diffusion-based generative model [39]. This result could follow from the fact that downscaling allows maintaining context information about the original image, and, as a result, the general coherence is of higher quality. We want to highlight that the diffusion-based generative model proposed in [39] consists of 1000 layers; thus, comparing it to our approach is not entirely fair.

Interestingly, on the two other datasets, CelebA and Imagenette64, a three-level selfVAE-downscale achieved better bpd scores compared with the VAE with the bijective prior. This indicates the benefit of employing a multi-leveled self-supervised framework against the VAE in higher-dimensional data, where the plain model fails to scale efficiently. It appears that the hierarchical structure of self-supervised random variables allows encoding the missing information more efficiently in ${\mathbf{z}}_k$, in contrast to the vanilla VAE, where all information about images must be coded in $\mathbf{z}$. This result is promising and indicates that the proposed approach is of great potential for generative modeling.

Lastly, comparing the selfVAE-sketch against the hierarchical selfVAE-downscale on CelebA, we clearly see that, in terms of the bpd, the hierarchical selfVAE-downscale was significantly better. The selfVAE-sketch achieved an even worse bpd than the VAE and RealNVP. However, the generations and reconstructions (see Figure 4 and Figure A4) indicate that learning with $\mathbf{y}$ allowed obtaining crisp images. The only issue was learning a proper color tone (e.g., see Figure A4).

4.3. Qualitative Results

We present generations on CelebA in Figure 4 and on CIFAR-10 and Imagenette64 in Figure A2 (see the Appendix), and reconstructions on CIFAR-10 & CelebA in Figure 5.

We first notice that the generations from selfVAE seem to be more coherent in contrast with these from VAE, which produced, overall, more contextless and distorted generations. This result appears to be in line with the FID scores. Especially for CelebA, we observed impressive synthesis quality, great sampling diversity, and coherent generations (Figure 4). On the Imagenette64 dataset, we also observed crisper generations for our method compared to the VAE (see Figure A2 in the Appendix A.3).

Furthermore, the hierarchical selfVAE seems to be of great potential for compression purposes. In contrast to the VAE, which is restricted to using a single way of reconstructing an image, the selfVAE allows four various options with different quality/memory ratios (Figure 5). In the selfVAE-sketch, we can retrieve an image with high accuracy by using only 16% of the original data, as it manages to encode all the texture of the image to $\mathbf{z}$ (Figure A3). This shows the advantage of choosing prior knowledge in the learning process. Lastly, the latent variables added extra information, which defined the result, and we can alter the details of an image, such as the facial expressions (Figure 6ii, see the Appendix A.3).

In Figure 6, we visualize (i) interpolations through two ground-truth images through the latent code $\mathbf{u}$ and (ii) image reconstructions, where we keep the latent code $\mathbf{u}$ but vary all the others (${\mathbf{z}}_1$ and ${\mathbf{z}}_2$) of the three-leveled selfVAE architecture. As with the previous cases, in the first case, we see that the model incorporated a rich latent space $\mathbf{u}$, which is responsible for the generation and construction of the global structure of the image.

Moving from one latent code $\mathbf{u}$ of a given image to another, we obtained meaningful modifications of the image that resulted in images that share characteristics from both of them. However, in the latter case, we see that we can alter only the high-level features of the image when we keep the values of $\mathbf{u}$ but vary the others: ${\mathbf{z}}_1$ and ${\mathbf{z}}_2$. Interestingly, we see that, given the ground-truth image that is illustrated on the very left, we can sample different expressions and characteristics of the same person, as the latent variable $\mathbf{u}$ is kept constant.

5. Conclusions

In this paper, we showed that taking deterministic and discrete transformations resulted in coherent generations of high visual quality, and allowed integrating the prior knowledge without losing its unconditional generative functionality. The experimental results confirm that the hierarchical architectures performed better and allowed obtaining better bpd scores as well as better generations and reconstructions.

In the experiments, we considered two classes of image transformations, namely, downscaling and edge detection (sketching). However, there are many possible other transformations (see Figure 1), and we leave investigating them for future work. Moreover, we find the proposed approach interesting for the compression task. A similar approach with a multi-scale auto-encoder for image compression was proposed, e.g., by [40,41]. However, we still used a probabilistic framework and indicate that various non-trainable image transformations (not only multiple scales) could be of great potential.