Improving Variational Auto-Encoders using Householder Flow

Variational Auto-Encoder

with respect to parameters. This task could be troublesome due to intractability of the marginal likelihood, e.g., when the model is parameterized by a neural network (NN). To overcome this issue one can introduce an inference model (an encoder) $q(\mathbf{z}|\mathbf{x})$ and optimize the variational lower bound:

where $p(\mathbf{x}|\mathbf{z})$ is called a decoder and $p(\mathbf{z})=\mathcal{N}(\mathbf{z}|\mathbf{0},\mathbf{I})$ is the prior. There are various ways of optimizing this lower bound but for continuous $\mathbf{z}$ this could be done efficiently through a re-parameterization of $q(\mathbf{z}|\mathbf{x})$ . Then the architecture is called a variational auto-encoder (VAE).

Improving posterior flexibility using Normalizing Flows

A (finite) normalizing flow, first formulated by and further developed by , is a powerful framework for building flexible posterior distribution by starting with an initial random variable with a simple distribution for generating $\mathbf{z}^{(0)}$ and then applying a series of invertible transformations $\mathbf{f}^{(t)}$, for $t=1,\ldots,T$. As a result, the last iterate gives a random variable $\mathbf{z}^{(T)}$ that has a more flexible distribution. Once we choose transformations $\mathbf{f}^{(t)}$ for which the Jacobian-determinant can be computed, we aim at optimizing the following objective:

In fact, the normalizing flow can be used to enrich the posterior of the VAE with small or even none modifications in the architecture of the encoder and the decoder.

There are two main kinds of normalizing flows, namely, general normalizing flows and volume preserving flows. The difference between these types of flow is in the manner how the Jacobian-determinant is handled. The general normalizing flows aim at formulating the flow for which the Jacobian-determinant is relatively easy to compute. On the contrary, the volume-preserving flows design series of transformations such that the Jacobian-determinant equals $1$ while still it allows to obtain flexible posterior distributions. The reduced computational complexity is the main reason why volume-preserving flows are so appealing. The question is whether one can propose a series of transformations for which the Jacobian-determinant is equal one, they are cheap to calculate and are general enough to model flexible posteriors. In the next subsection we present a new volume-preserving flow that applies series of Householder transformations that we refer to as the Householder flow.

Householder Flow

In general, any full-covariance matrix $\boldsymbol{\Sigma}$ can be represented by the eigenvalue decomposition using eigenvectors and eigenvalues:

Generally, the task of modelling an orthogonal matrix in a principled manner is rather non-trivial. However, first we notice that any orthogonal matrix can be represented in the following form :

(The Basis-Kernel Representation of Orthogonal Matrices) For any $M\times M$ orthogonal matrix $\mathbf{U}$ there exist a full-rank $M\times K$ matrix $\mathbf{Y}$ (the basis) and a nonsingular (triangular) $K\times K$ matrix $\mathbf{S}$ (the kernel), $K\leq M$, such that:

The value $K$ is called the degree of the orthogonal matrix. Further, it can be shown that any orthogonal matrix of degree $K$ can be expressed using the product of Householder transformations , namely:

Any orthogonal matrix with the basis acting on the $K$-dimensional subspace can be expressed as a product of exactly $K$ Householder transformations:

where $\mathbf{H}_{k}=\mathbf{I}-\mathbf{S}_{kk}\mathbf{Y}_{\cdot k}(\mathbf{Y}_{\cdot k})^{\top}$, for $k=1,\ldots,K$.

Theoretically, Theorem 2 shows that we can model any orthogonal matrix in a principled fashion using $K$ Householder transformations. Moreover, the Householder matrix $\mathbf{H}_{k}$ is orthogonal matrix itself . Therefore, this property and the Theorem 2 put the Householder transformation as a perfect candidate for formulating a volume-preserving flow that allows to approximate (or even capture) the true full-covariance matrix.

2 Definition

where $\mathbf{H}_{t}=\mathbf{I}-2\frac{\mathbf{v}_{t}\mathbf{v}_{t}^{\top}}{||\mathbf{v}_{t}||^{2}}$ is called the Householder matrix.

Related Work

In invertible linear normalizing flows with known Jacobian determinant were proposed. These are easy to calculate, i.e., the determinant of the Jacobian can be analytically computed, however, many such transformations are needed to capture high-dimensional dependencies. A different approach relies on volume-preserving flows, for which the absolute value of Jacobian determinant is equal $1$, such as, Non-linear Independent Components Estimation (NICE) , Hamiltonian Variational Inference (HVI) and linear Inverse Autoregressive Flow (linIAF) or its non-linear version (nlIAF) .

The HF is similar in spirit to the linIAF where the linear transformation is also applied but it is the lower triangular inverse Cholesky matrix with ones on the diagonal instead of the Householder matrix. Nevertheless, the motivation of the linIAF differs from ours completely.

The Householder transformations (reflections) were also exploited in the context of learning recurrent neural nets . They were used for modelling unitary weights instead of more flexible variational posterior, however, these served as a component for representing a matrix of eigenvectors, similarly to our approach.

Experiments

In the first experiment we used the MNIST dataset that contains 60,000 training and 10,000 test images of ten handwritten digits ($28\times 28$ pixels in size). From the training set we put aside 10,000 images for validation and tuning hyper-parameters. We used the dynamically binarized dataset as in .

In order to make more reliable comparison to , we trained the VAE with $40$ stochastic hidden units and the encoder and decoder were parameterized with two-layered neural networks ($300$ hidden units per layer). We used the gating mechanism as the activation function (see Appendix for details). The length of the HF was $T=\{1,10\}$ (VAE+HF ($T$=1/10)).Taking larger values of $T$ did not resulted in better performance. For training we utilized ADAM with the mini-batch size equal $100$ and one example for estimating the expected value. The learning rate was set according to the validation set. The maximum number of epochs was $5000$ and early-stopping with a look-ahead of $100$ epochs was applied. We used the warm-up for first $200$ epochs. We initialized weights according to .

We compared our approach to linear normalizing flow (VAE+NF) , and finite volume-preserving flows: NICE (VAE+NICE) , HVI (VAE+HVI) . When appropriate, the length of a flow $T$ is given. The methods were compared according to the lower bound of marginal log-likelihood measured on the test set.

The code of the proposed approach can be found on: https://github.com/jmtomczak.

The results presented in Table 3 reveal that the proposed flow is competitive comparing to other normalizing flows. This outcome is especially interesting since in our approach sampling from the variational posterior requires application of linear transformations, thus, the HF is reasonably cheap computationally. In our approach we need to store only $T\times M$ more parameters comparing to the vanilla VAE while, e.g., the VAE+NICE and VAE+NF require $\frac{M\times(M-1)}{2}$ and $\mathcal{O}(T\times M)$ more parameters, respectively. All these remarks are in favor of the proposed flow.

In order to get some insight into the behavior of the HF we also examined the values of components of the training objective (see Figure 3). First, the application of the HF helps the VAE to better model data (lower reconstruction error). Second, the HF gives higher flexibility of the posterior because the approximation of the orthogonal matrix allows the initial posterior to model eigenvalues instead of the whole information about the true posterior, which results in smaller Kullback-Leibler penalty (i.e., the variances do not need to model whole information about the true posterior distribution).

2 Histopathology data

In the second experiment we applied the VAE and the VAE+HF to a more challenging problem of grayscale histopathology data. We prepared a dataset basing on histopathological images freely available on-linehttp://www.enjoypath.com/. We selected 16 patientsPatient numbers: 272, 274, 283, 289, 290, 291, 292, 295, 297, 298, 299. and each histopathological image represented a bone marrow biopsy. Diagnoses of the chosen cases were associated with different kinds of cancer (e.g., lymphoma, leukemia) or anemia. All images were taken using HE, $40\times$, and each image was of size $336\times 448$. The original RGB representation was transformed to grayscaleA pixel’s value was determined according to the formula: $0.299\times R+0.587\times G+0.114\times B$.. Further, we divided each image into small patches of size $28\times 28$ (see Figure 4 for exemplary image patches). Eventually, we picked 10 patients for training, 3 patients for validation and 3 patients for testing, which resulted in 6,800 training images, 2,000 validation images and 2,000 test images. The selection of patients was performed in such fashion that each dataset contained representative images with different diagnoses and amount of fat.

We used the same architecture of the VAE and the same training procedure as in the previous experiment. Since the data are grayscale, we modeled data using normal distribution instead of the Bernoulli distribution with means constrained to $$by applying the sigmoid function. We compared the vanilla VAE with the VAE+HF ($T$=1,10,20).

The results presented in Table 1 confirmed findings of the previous experiment that the application of the HF helps in obtaining more flexible posterior. Additionally, in this case application of the series of 20 HF steps provided a slight improvement. We have also noticed that the application of the HF allows to obtain less variable results (smaller standard deviation comparing to the vanilla VAE).

Conclusion

In this paper, we proposed a new volume-preserving flow using Householder transformations. Our idea relies on the observation that the true full-covariance matrix can be decomposed to the diagonal matrix with eigenvalues on the diagonal and the orthogonal matrix of eigenvalues that could be further modeled using a series of Householder transformations. The obtained results in the experimental studies reveal that fully connected VAEs with Householder flows could perform better than other volume-preserving flows. This makes the HF a very promising direction for further investigation. In the future, we aim at using the outlined flow in recently proposed extensions of the variational inference, e.g., importance weighted VAE , Rényi Divergence for VAE and Ladder VAE . Moreover, we believe that the proposed flow can be beneficiary in modelling natural and medical images, where flexibility of the posterior is crucial . We leave investigating this issue for future research.

We would like to give special thanks to Diederik P. Kingma for fruitful discussions and insightful remarks on the paper. The research conducted by Jakub M. Tomczak was funded by the European Commission within the Marie Skłodowska-Curie Individual Fellowship (Grant No. 702666, ”Deep learning and Bayesian inference for medical imaging”).

References

Appendix

Let $\mathbf{x}$ be a $D$-dimensional input and $\mathbf{h}_{1},\ldots,\mathbf{h}_{L}$ be hidden units in a neural network (NN).To keep the notation uncluttered we set $\mathbf{h}_{0}=\mathbf{x}$. Typically, a hidden layer in NN is calculated using ReLU activation function:

where $\otimes$ denotes the element-wise multiplication.

Very recently, a new type of activation function was proposed that applies a gating mechanism :

Calculation of the gating mechanism for a single layer is presented in Figure 5.