Body: <p>I am trying to train a variational auto-encoder where <code>x ≈ f_VAE(x) = x_hat</code>. In my real problem, I have 100-400 dimensional data that I would like to compress to around 30 latent variables. For now I have a toy problem where I have 14 dimensional data that I would like to compress to 7 dimensions. All dimensions are real numbers and have been standardized (zero mean, unit standard deviation). To give some idea of the data here is a scatter plot:
  <a href="https://i.sstatic.net/VhDuzJth.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/VhDuzJth.png" alt="enter image description here" /></a></p>
  <p>I train the VAE by minimizing the following loss <code>loss = MSE + beta * KL_div</code>, where <code>MSE</code> is the mean-squared-error, <code>KL_div</code> is the KL-divergence term and <code>beta</code> is a scaling factor. I keep running into the <strong>posterior collapse problem</strong>. In brief, the <code>MSE ≈ 1.0</code> and does not drop lower no matter how long you train and what model architecture you use. The <code>KL_div ≈ 0.0</code>, thus vanishing and the reconstructed features <code>x_hat</code> is a vector of values around 0 (mean of the standardized data) <em>regardless of input <code>x</code> / latent variables <code>z</code></em>. Posterior collapse has been described in the literature and there are several proposed solutions such as beta-VAE [<a href="https://openreview.net/forum?id=Sy2fzU9gl" rel="nofollow noreferrer">2</a>] and cyclical annealing VAE [<a href="https://arxiv.org/abs/1903.10145" rel="nofollow noreferrer">3</a>]. For my problem, I have tried the following:</p>
  <ul>
  <li>normal VAE (beta = 1)</li>
  <li>beta VAE (beta &lt; 1)</li>
  <li>linear annealing with warmup (<code>beta = 0</code> for 4 epochs and then rises to 1 over the other 16 epochs and is incremented every mini-batch)</li>
  <li>cyclical annealing VAE with warmup (<code>beta = 0</code> for 4 epochs and then goes through 4 cycles between 0 and 1 with increments every mini-batch)</li>
  </ul>
  <p>At the end of the day, none of these strategies work. During warmup, when <code>beta = 0</code>, the MSE goes down to around 0.2, actually yielding decent reconstructions <code>x_hat</code>. As soon as the KL-regularization kicks in, the network stops caring about the <code>MSE</code> reconstruction error and the <code>KL_div</code> vanishes. Using ray Tune, I tried many different VAE architectures with differing numbers of hidden layers, different sizes of latent and hidden layers and other learning parameters. One interesting observation is that training a normal auto-encoder (AE, where <code>beta = 0</code> for 20 epochs) does seem to work with decent reconstructions. The AE even performs somewhat OK on the validation data.</p>
  <p><strong>Questions</strong> in decreasing order of importance:</p>
  <ul>
  <li>What else can I try to make VAEs work for my problem</li>
  <li>Why is is possible to train an AE for this data, but as soon as you introduce some KL-div regularization, the whole task becomes impossible</li>
  <li>What could possible explanations be for the KL regularization leading to posterior collapse</li>
  <li>What other machine learning approaches exist for latent variable learning that I could look into</li>
  </ul>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-12-09T21:23:01.640Z
Comments.Text: On common cause of posterior collapse is decoder is too powerful to overfit the reconstruction, try use dropout, weight decay, smaller layers and neuron size for decoder for your simple toy problem with small training data set, and also try to make encoder and decoder symmetric in terms of architecture and random weights initially. AE solely focuses on reconstruction error while VAE with KL term limits the freedom of encoder/decoder thus causes posterior collapse.
Comments.UserDisplayName: cinch
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-12-09T21:23:05.213Z
Comments.Text: Also if your 7-d latent space is too limited, your encoder might not have enough capacity to encode meaningful variations, try to enlarge its dimension. Hope you can report back with some good news after more trial and error.
Comments.UserDisplayName: cinch
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-12-13T11:43:09.677Z
Comments.Text: So the encoder and decoder are symmetric. I use ray tune to screen all relevant parameters such as layer sizes, number of layers, learning rate the whole bunch. Not a single one of those has yielded a functional VAE. At this point I am trying to add a normalizing flow as a prior instead of a MVN, but this has so far also not worked.
Comments.UserDisplayName: Patrickens
Tags: <autoencoders><variational-autoencoder><kl-divergence><data-compression>
Title: KL annealing for a VAE does not work, what now?
UserDisplayName: Patrickens
ThreadText: # KL annealing for a VAE does not work, what now?
  
  (Asked by: username_0 on 2024-12-09)
  
  I am trying to train a variational auto-encoder where `x ≈ f_VAE(x) = x_hat`. In my real problem, I have 100-400 dimensional data that I would like to compress to around 30 latent variables. For now I have a toy problem where I have 14 dimensional data that I would like to compress to 7 dimensions. All dimensions are real numbers and have been standardized (zero mean, unit standard deviation). To give some idea of the data here is a scatter plot:
  [![enter image description here](https://i.sstatic.net/VhDuzJth.png)](https://i.sstatic.net/VhDuzJth.png)
  
  I train the VAE by minimizing the following loss `loss = MSE + beta * KL_div`, where `MSE` is the mean-squared-error, `KL_div` is the KL-divergence term and `beta` is a scaling factor. I keep running into the **posterior collapse problem**. In brief, the `MSE ≈ 1.0` and does not drop lower no matter how long you train and what model architecture you use. The `KL_div ≈ 0.0`, thus vanishing and the reconstructed features `x_hat` is a vector of values around 0 (mean of the standardized data) *regardless of input `x` / latent variables `z`*. Posterior collapse has been described in the literature and there are several proposed solutions such as beta-VAE [[2](https://openreview.net/forum?id=Sy2fzU9gl)] and cyclical annealing VAE [[3](https://arxiv.org/abs/1903.10145)]. For my problem, I have tried the following:
  
  * normal VAE (beta = 1)
  * beta VAE (beta < 1)
  * linear annealing with warmup (`beta = 0` for 4 epochs and then rises to 1 over the other 16 epochs and is incremented every mini-batch)
  * cyclical annealing VAE with warmup (`beta = 0` for 4 epochs and then goes through 4 cycles between 0 and 1 with increments every mini-batch)
  
  At the end of the day, none of these strategies work. During warmup, when `beta = 0`, the MSE goes down to around 0.2, actually yielding decent reconstructions `x_hat`. As soon as the KL-regularization kicks in, the network stops caring about the `MSE` reconstruction error and the `KL_div` vanishes. Using ray Tune, I tried many different VAE architectures with differing numbers of hidden layers, different sizes of latent and hidden layers and other learning parameters. One interesting observation is that training a normal auto-encoder (AE, where `beta = 0` for 20 epochs) does seem to work with decent reconstructions. The AE even performs somewhat OK on the validation data.
  
  **Questions** in decreasing order of importance:
  
  * What else can I try to make VAEs work for my problem
  * Why is is possible to train an AE for this data, but as soon as you introduce some KL-div regularization, the whole task becomes impossible
  * What could possible explanations be for the KL regularization leading to posterior collapse
  * What other machine learning approaches exist for latent variable learning that I could look into
  
  **Tags:** autoencoders, variational-autoencoder, kl-divergence, data-compression
  
  > **Comments:**
  > **username_1**: On common cause of posterior collapse is decoder is too powerful to overfit the reconstruction, try use dropout, weight decay, smaller layers and neuron size for decoder for your simple toy problem with small training data set, and also try to make encoder and decoder symmetric in terms of architecture and random weights initially. AE solely focuses on reconstruction error while VAE with KL term limits the freedom of encoder/decoder thus causes posterior collapse.  
  > **username_1**: Also if your 7-d latent space is too limited, your encoder might not have enough capacity to encode meaningful variations, try to enlarge its dimension. Hope you can report back with some good news after more trial and error.  
  > **username_0**: So the encoder and decoder are symmetric. I use ray tune to screen all relevant parameters such as layer sizes, number of layers, learning rate the whole bunch. Not a single one of those has yielded a functional VAE. At this point I am trying to add a normalizing flow as a prior instead of a MVN, but this has so far also not worked.  
  
  ---
  
  ## 0 answers
  
  