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