VAE

keywords encoder, decoder

Concepts

Variational autoencoder(VAE) is a latent model $p(x,z|\theta )$.​

probabilistic encoder-decoder:

• probabilistic decoder/likelihood/generating proba: $p(x|z,\theta )$​, the distribution of the decoded variable given the encoded one,

• probabilistic encoder/posterior/predicting(discriminant, respondant) proba/encoding: $p(z|x,\theta )$​, the distribution of the encoded variable given the decoded one.

graph: $z\to x(\leftarrow \theta )$​​​​​ 👀

Task: to find $q(z)\approx p(z|x)$​​​​​​​​​​, or denoted as $q(z|x)$​​ to emphsize the dependence on $x$. (variational inference)

loss: $D_{KL}(q(z)\Vert p(z|x))$​

likelihood: $l(\theta )={\sum }_x\ln {\int }_{z}p(x|z,\theta )p(z|\theta )dz$

Variational Lower Bound/Evidence Lower Bound(ELBO)
$$\begin{gathered} L_x:=E_{z\sim q}(\log p(x,z)-\log q(z)) \\ =Q(p,q)+H(q)(=F(p,q)) \end{gathered}$$
where $q$ alwayes depends on $x$​.

Remark $L_x$ is also called free energy.

ELBO = expected likelihood + entropy

(EM algo. / MM algo.)

Facts

Identity

variational inequality
$$\log p(x)=D_{KL}(q(z)\Vert p(z|x))+L_x\ge L_x$$

remark likelihood = divergence + ELBO(free energy)

$$L_x=E_{z\sim q}(\log p(x|z))-D_{KL}(q(z)\Vert p_Z(z)),$$

ELBO = Reconstruction loss + Regularization term

Reconstruction loss: $-E_{z\sim q}(\log p(x|z))$​​

Regularization term: $D_{KL}(q\Vert p_Z)$

For samples $D$,
$$L_D=\sum _{x\in D}{E}_{z\sim q}\log p(x|z)-D_{KL}(q\Vert p_Z)$$

remark For fixed $p$​, ${\min }_{p,q}{D}_{KL}⟺{\max }_{p,q}{L}_x$​​​​.

diff. between EM and VAE

• EM algorithm: solve $\max L_x$ by coordinate ascent
• VAE: by SGD

Parameter form
$$\begin{gathered} L_D(\theta ,\varphi )=\sum _{x\in D}(E_{z\sim q_{x,\varphi }}\log p(x|z,\theta )-D_{KL}(q(z|x,\varphi )\Vert p(z|\theta ))) \\ =\sum _{x\in D}(E_{z\sim q_{x,\varphi }}\log p(x|z,\theta )-D_{KL}(q(z|x,\varphi )\Vert p(z)))(if\theta \to x\leftarrow z) \end{gathered}$$
where variational parameters:$\varphi$, generative parameters: $\theta$​.

Assumptions

Gaussian assumption: $p(z),p(x|z)$: Gaussian distr.
$$\begin{gathered} Z\sim N(0,1), \\ X|Z=z\sim N(f(z),c),f\in F,c>0. \end{gathered}$$
It is intractable to compute $p(z|x)$ by Bayesian formula.

variational inference(VI)

VI is a technique to approximate complex distributions

Continue from (4)

We are going to approximate $p(z|x)$​ by a Gaussian distribution $q_x(z)$​​ whose mean and covariance are defined by two functions, $g$​ and $h$​, of the parameter $x$​.
$$q(Z|x)\sim N(g(x),h(x)),g\in G,h\in H,$$
where variational parameter $\varphi =(g,h)\in G\times H$​.

Fixed $f$ (hence $\log p(x)$ is a constant), solve the following optimialization problem,

$$\begin{gathered} \max L_x⟺\underset{g,h}{\min }{D}_{KL}(q_x(z)\Vert p(z|x)) \\ \simeq \underset{g,h}{\min }-E\log p(x|z)+D_{KL}(q_x\Vert p_Z) \\ =\underset{g,h}{\min }\frac{1}{2c}E\Vert x-f(z){\Vert }^2+D_{KL}(q_x\Vert p_Z). \end{gathered}$$

Then find optimal $f$​​​​,
$$\underset{f}{\max }{E}_{z\sim N(g(x),h(x))}\log p(x|z)=\underset{f}{\min }{E}_z\Vert x-f(z){\Vert }^2$$

===>
$$=\frac{1}{2c}\underset{f,g,h}{\min }{E}_z\Vert x-f(z){\Vert }^2+D_{KL}(q_x\Vert p_Z).$$

variational autoencoder(AVE) architecture

Based on Gaussian assumption.

architecture

encoder: $x\to z=g(x)+h(x)\zeta \sim N(g(x),h(x))$​​​​​​​​ and $\zeta \sim N(0,1)$​​​​​​ (reparameterization), as an approximation of $p(z|x)$​​

decoder: $z\to x\sim N(f(z),c)$​​

reparametrisation trick
$$L(x)=C{E}_z\Vert x-f(z){\Vert }^2+D_{KL}(N(g(x),h(x))\Vert N(0,1))$$

where $D_{KL}$​ of the diagonal normal distr. and standard normal distr. is $D_{KL}(DN,SN):=\frac{1}{2}{\sum }_{i=1}^k({\sigma }_i^2+{\mu }_i^2-1-\ln {\sigma }_i^2)$​​​​​​​, where ${\mu }_i=g_i(x),{\sigma }_i^2=h_i(x)$.

see Wiki of KL divergence.

Bayes-by-BP

Algo

1. initalize $\varphi \leftarrow {\varphi }_0$
2. loop from i=0 to N
   1. $z\sim q(z|x,\varphi )$;
   2. $\theta$​ maximizes $\log p(x|z,\theta )$;
   3. calculate ELBO $L(\theta ,\varphi )$;
   4. update $\varphi$ by GD;

General form of Algorithm

opt $F(\theta )=E_{z\sim q_{\theta }}f(\theta ,z)$​​​​​​:

1. guess $\theta$, generate $z\sim q_{\theta }$;
2. Do GD for $f(\theta ,z)$​​​​ to update $\theta$​​;

PCA/SVD as a VAE

General form encoder-decoder model
$$\underset{d,e}{\min }|x-d(e(x))|$$

SVD, dim of data space = p, laten space = q
$$\begin{gathered} \underset{V}{\min }|X-X{V}_q{V}_q^{\prime }| \\ V:O(p) \end{gathered}$$

xi ~ N(0,1)
   \
D -- > Z --> X = ZV'
1
2
3

where $Z\sim N(0,D),X=Z{V}^{\prime }$​​​.

stat. model of encoder-decoder

• encoder: $p(x|z)$
• decoder: $p(z|x)$

as the conditional proba. of $p(x,z)$

generating rask: $z\sim N(0,1)\Rightarrow x\sim p(x|z)$

Basic Model

graph: c -> z -> x (<-theta), where $c$: context.

Likelihood: $p(D|\theta )={\prod }_{x\in D}\int p(x|z,\theta )p(z|c,\theta )dz$​​

variational lower bound:
$$L_D=E_{c\sim q_D}(L_{D|c})-D_{KL}(q_D\Vert p_c)$$

Full Model

graph: c -> z1,...,zT -> x (<-theta), where $c$: context.

References

understanding-variational-autoencoders

Ming Ding. The road from MLE to EM to VAE: A brief tutorial,2022.
H Edwards, A. Storkey. Towards a neural statistician,2017.

Codes

code on line:
VAE-keras

my code:

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
VAE
"""

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers


class Sampling(layers.Layer):
    """Sampling Layer for reparameterization
    Uses (z_mean, z_log_var) to sample z, the vector encoding a digit.
    """

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch, dim = tf.shape(z_mean)[0], tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon


class VAE(keras.Model):
    # architechture of VAE
    def __init__(self, latent_dim, **kwargs):
        super().__init__(**kwargs)
        self.latent_dim = latent_dim
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
        ]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    keras.losses.binary_crossentropy(data, reconstruction), axis=(1, 2)
                )
            )  # Hb(x_i, f(z_i)) z_i ~ N(g(x_i),h(x_i))
            kl_loss = 0.5 * tf.reduce_sum(tf.square(z_mean) + tf.exp(z_log_var) - 1 - z_log_var, axis=1)
            kl_loss = tf.reduce_mean(kl_loss)    # KL-div of N and SN
            total_loss = reconstruction_loss + kl_loss
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)

        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result()}

    @classmethod
    def make_from_data(cls, X, latent_dim, *args, **kwargs):
        shape = X.shape[1:3]
        return cls.make(shape, latent_dim, *args, **kwargs)

    @classmethod
    def make(cls, shape, latent_dim, *args, **kwargs):
        assert len(shape)>=2, "Sorry, this VAE works only for images; the ndim of each sample>=2!"
        model = cls(latent_dim, *args, **kwargs)
        model.encoder = cls.make_encoder(shape, latent_dim)
        model.decoder = cls.make_decoder(shape, latent_dim)
        model.compile(optimizer=keras.optimizers.Adam())
        return model

    @classmethod
    def make_encoder(cls, shape, latent_dim):
        # shape: the shape of input data
        if len(shape)==2:
            height, width = shape
            n_channels = 1
        elif len(shape)==3:
            height, width, n_channels = shape

        encoder_inputs = layers.Input(shape=(height, width, n_channels))
        x = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")(encoder_inputs)
        x = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
        x = layers.Flatten()(x)
        x = layers.Dense(16, activation="relu")(x)
        z_mean = layers.Dense(latent_dim, name="z_mean")(x)       # g(x)
        z_log_var = layers.Dense(latent_dim, name="z_log_var")(x) # h(x)
        z = Sampling()([z_mean, z_log_var])                       # z ~ N(g(x),h(x))
        return keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")

    @classmethod
    def make_decoder(cls, shape, latent_dim):
        # x = f(z), shape: the shape of the output of the decoder
        if len(shape)==2:
            height, width = shape
            n_channels = 1
        elif len(shape)==3:
            height, width, n_channels = shape
        small_height, small_width = height // 4, width // 4
        decoder = keras.Sequential(name="decoder")     # f(z)
        latent_inputs = keras.Input(shape=(latent_dim,))
        decoder.add(latent_inputs)
        decoder.add(layers.Dense(small_height * small_width * 64, activation="relu"))
        decoder.add(layers.Reshape((small_height, small_width, 64)))
        decoder.add(layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same"))
        decoder.add(layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same"))
        decoder.add(layers.Conv2DTranspose(1, 3, activation="sigmoid", padding="same"))
        return decoder

    def trainsform(self, X):
        M, _, _ = self.encoder(X)
        return M

    def inverse_trainsform(self, Z):
        return self.decoder.predict(Z)


def plot_latent_space(vae, n=11, *args, **kwargs):

    import itertools
    import matplotlib.pyplot as plt

    # display a n*n 2D manifold of digits
    scale = 1.0

    # linearly spaced coordinates corresponding to the 2D plot
    # of digit classes in the latent space
    grid_x = np.linspace(-scale, scale, n)
    grid_y = np.linspace(-scale, scale, n)[::-1]
    z_sample = np.hstack((list(itertools.product(grid_y, grid_x)), np.random.normal(size=(n**2, vae.latent_dim-2))))
    x_decoded = vae.inverse_trainsform(z_sample)

    figure = np.block([[x_decoded[i*n+j].reshape((width, height)) for j in range(n)] for i in range(n)])

    plt.figure(*args, **kwargs)
    start_range_x, start_range_y = width // 2, height // 2
    end_range_x = n * width + start_range_x
    end_range_y = n * height + start_range_y
    pixel_range_x = np.arange(start_range_x, end_range_x, width)
    pixel_range_y = np.arange(start_range_y, end_range_y, height)
    sample_range_x = np.round(grid_x, 1)
    sample_range_y = np.round(grid_y, 1)
    plt.xticks(pixel_range_x, sample_range_x)
    plt.yticks(pixel_range_y, sample_range_y)
    plt.xlabel("$z_0$")
    plt.ylabel("$z_1$")
    plt.imshow(figure, cmap="Greys_r")
    plt.show()

if __name__ == '__main__':
    
    # from get_hanzi import X_train
    # input your data here, images in size of 4m X 4n

    X_train /= 255
    X_train = X_train>0.5  # binarize

    X_train = np.expand_dims(X_train, -1).astype("float32")
    vae = VAE.make_from_data(X_train, latent_dim=15)
    vae.fit(X_train, epochs=200, batch_size=16, verbose=False)
    plot_latent_space(vae)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170