<p>I was trying to implement VAE on the CelebA dataset inspired by the Tensorflow implementation of MNIST. I have tried varying batch size but there seems to be no effect from that. The image formed is mostly grey all the time. Ideally, we want KL divergence and Reconstruction loss both to be close to zero but in my case, both are increasing exponentially.
<a href="https://i.sstatic.net/CbmTfuar.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/CbmTfuar.png" alt="Here's the loss curve I am getting" /></a>
Here is the Loss function definition block:</p>
<pre><code>optimizer = tf.keras.optimizers.Adam(1e-4)
def log_normal_pdf(sample, mean, logvar, raxis=1):
  log2pi = tf.math.log(2. * np.pi)
  return tf.reduce_sum(
      -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
      axis=raxis)
  
def compute_loss(model, x):
  mean, logvar = model.encode(x)
  z = model.reparameterize(mean, logvar)
  x_logit = model.decode(z)
  cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
  #logpx_z = tf.reduce_mean(tf.square(x - x_logit), axis=[1, 2, 3])
  logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
  logpz = log_normal_pdf(z, 0., 0.)
  logqz_x = log_normal_pdf(z, mean, logvar)
  return -tf.reduce_mean(logpx_z + logpz - logqz_x), logpx_z, logqz_x-logpz
</code></pre>
<p>My latent dimension is 16 and batch size is 500. Also, my input is only 500 images.</p>
<p>I have already tried changing the size of input but there seems to be no impact from that. Here is the model definition:</p>
<pre><code>class CVAE(tf.keras.Model):
    def __init__(self, latent_dim):
        super(CVAE, self).__init__()
        self.latent_dim = latent_dim
        self.encoder = tf.keras.Sequential(
            [
                tf.keras.layers.InputLayer(input_shape=(64, 64, 3)),
                tf.keras.layers.Conv2D(
                    filters=32, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=64, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=128, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=256, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=512, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Flatten(),
                # No activation
                tf.keras.layers.Dense(latent_dim + latent_dim),
            ]
        )

        self.decoder = tf.keras.Sequential(
            [
                tf.keras.layers.InputLayer(input_shape=(latent_dim,)),
                tf.keras.layers.Dense(units=4*4*256, activation=tf.nn.relu),
                tf.keras.layers.Reshape(target_shape=(4, 4, 256)),
                tf.keras.layers.Conv2DTranspose(
                    filters=128, kernel_size=3, strides=2, padding='same',
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=64, kernel_size=3, strides=2, padding='same',
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=32, kernel_size=3, strides=2, padding='same',
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=3, kernel_size=3, strides=2, padding='same'),
            ]
        )
    @tf.function
    def sample(self, eps=None):
        if eps is None:
            eps = tf.random.normal(shape=(100, self.latent_dim))
        return self.decode(eps, apply_sigmoid=True)

    def encode(self, x):
        mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)
        return mean, logvar

    def reparameterize(self, mean, logvar):
        eps = tf.random.normal(shape=mean.shape)
        return eps * tf.exp(logvar * .5) + mean

    def decode(self, z, apply_sigmoid=False):
        logits = self.decoder(z)
        if apply_sigmoid:
            probs = tf.sigmoid(logits)
            return probs
        return logits

</code></pre>
<p><a href="https://colab.research.google.com/gist/theFulminatedHuman/0e4cc065ac624501f3858a2f72a2f431/vae-on-celebsa.ipynb" rel="nofollow noreferrer">Here's is the link to the colab notebook</a></p>

CC BY-SA 4.0
2024-07-01T18:51:13.407Z
the ELBO should be positive, as it’s the sum of the reconstruction and the KL
Alberto
CC BY-SA 4.0
2024-07-01T20:01:51.833Z
@Alberto Ya I just reversed the sign. ELBO in my code is the sum of negative log-likelihood and KL Divergence.
Vedant Bhardwaj
CC BY-SA 4.0
2024-07-02T10:12:24.750Z
please post the new loss
Alberto
<convolutional-neural-networks><tensorflow><generative-model><variational-autoencoder><kl-divergence>
How to solve the exploding gradient problem in VAE training?
Vedant Bhardwaj
# How to solve the exploding gradient problem in VAE training?

(Asked by: username_0 on 2024-07-01)

I was trying to implement VAE on the CelebA dataset inspired by the Tensorflow implementation of MNIST. I have tried varying batch size but there seems to be no effect from that. The image formed is mostly grey all the time. Ideally, we want KL divergence and Reconstruction loss both to be close to zero but in my case, both are increasing exponentially.

Here is the Loss function definition block:
optimizer = tf.keras.optimizers.Adam(1e-4)
def log_normal_pdf(sample, mean, logvar, raxis=1):
  log2pi = tf.math.log(2. * np.pi)
  return tf.reduce_sum(
      -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
      axis=raxis)
  
def compute_loss(model, x):
  mean, logvar = model.encode(x)
  z = model.reparameterize(mean, logvar)
  x_logit = model.decode(z)
  cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
  #logpx_z = tf.reduce_mean(tf.square(x - x_logit), axis=[1, 2, 3])
  logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
  logpz = log_normal_pdf(z, 0., 0.)
  logqz_x = log_normal_pdf(z, mean, logvar)
  return -tf.reduce_mean(logpx_z + logpz - logqz_x), logpx_z, logqz_x-logpz

My latent dimension is 16 and batch size is 500. Also, my input is only 500 images.
I have already tried changing the size of input but there seems to be no impact from that. Here is the model definition:
class CVAE(tf.keras.Model):
    def __init__(self, latent_dim):
        super(CVAE, self).__init__()
        self.latent_dim = latent_dim
        self.encoder = tf.keras.Sequential(
            [
                tf.keras.layers.InputLayer(input_shape=(64, 64, 3)),
                tf.keras.layers.Conv2D(
                    filters=32, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=64, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=128, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=256, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=512, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Flatten(),
                # No activation
                tf.keras.layers.Dense(latent_dim + latent_dim),
            ]
        )

        self.decoder = tf.keras.Sequential(
            [
                tf.keras.layers.InputLayer(input_shape=(latent_dim,)),
                tf.keras.layers.Dense(units=4*4*256, activation=tf.nn.relu),
                tf.keras.layers.Reshape(target_shape=(4, 4, 256)),
                tf.keras.layers.Conv2DTranspose(
                    filters=128, kernel_size=3, strides=2, padding='same',
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=64, kernel_size=3, strides=2, padding='same',
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=32, kernel_size=3, strides=2, padding='same',
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=3, kernel_size=3, strides=2, padding='same'),
            ]
        )
    @tf.function
    def sample(self, eps=None):
        if eps is None:
            eps = tf.random.normal(shape=(100, self.latent_dim))
        return self.decode(eps, apply_sigmoid=True)

    def encode(self, x):
        mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)
        return mean, logvar

    def reparameterize(self, mean, logvar):
        eps = tf.random.normal(shape=mean.shape)
        return eps * tf.exp(logvar * .5) + mean

    def decode(self, z, apply_sigmoid=False):
        logits = self.decoder(z)
        if apply_sigmoid:
            probs = tf.sigmoid(logits)
            return probs
        return logits


Here's is the link to the colab notebook

**Tags:** convolutional-neural-networks, tensorflow, generative-model, variational-autoencoder, kl-divergence

> **Comments:**
> **username_1**: the ELBO should be positive, as it’s the sum of the reconstruction and the KL  
> **username_0**: @username_1 Ya I just reversed the sign. ELBO in my code is the sum of negative log-likelihood and KL Divergence.  
> **username_1**: please post the new loss  

---

## 0 answers

