I am currently experiencing a suspected model collapse in a Variational Autoencoder (VAE) model I am working with. Below are details on the project setup and the issue at hand:

Project Goal: Exploring if a the same VAE architecture can independently handle two tasks - reconstruction and regression - using real-world production line data.


  • Dataset: Approximately 80,000 samples

  • Features: Sensor readings such as flow, temperature, force, etc. (both binary and continuous data)

  • Target for regression: The measurement of a product property at the end of the line (can be divided in approximately 50 classes)

  • Preprocessing: Data normalized between 0 and 1 before training

Model Architecture:

  • Encoder: Input layer of 800 and a hidden layer of 600

  • Latent Space: Size of 400 (both mu and logvar being 400)

  • Decoder: Hidden layer of 600 and an output layer of 800 (for reconstruction) or 1 (for regression)

  • Activation Function: Sigmoid in the hidden layers, sigmoid for reconstruction output, and identity function for regression output

  • Training: Each task trained separately, using a learning rate of 1e-3

  • Loss Calculation: MSE for continuous features, BCE with logits loss for binary features. The losses are summed and divided by the batch size.

Training Setup:

  • Early stopping when the validation loss is consistent (patience of 20)

  • Gradient clipping is implemented

  • Training is performed on an Nvidia A100 GPU (64GB) using PyTorch Lightning framework, with logging done via Weights and Biases (wandb)


  • For the reconstruction task, t-SNE visualization of samples shows a normally distributed sphere with no distinct internal clusters for the different classes. The BCE and MSE losses are decreasing, displaying a 'normal' loss curve, but the KLD loss is continually rising.

  • The issue does not exist for the regression task where distinct clusters are observed for each class and KLD loss eventually drops after an initial rise.

  • Anomalies in bias and gradients: Except for the second layer in the encoder and the second layer in the decoder, the bias for all layers is very close to 0. Gradients for all layers are distributed around zeroes, except for these two layers.

Attempted Solutions:

  • Experimented with different activation functions, units in hidden layers, number of hidden layers, sizes of the latent space

  • Tried training on continuous features only

  • Tested different parts of the dataset (odd/even samples), used fewer columns

  • Swapped LSTM layers in place of linear ones

  • Employed techniques like cyclical annealing, adjusting BCE-MSE ratio, and utilizing a warm-up factor for KLD

  • Unfortunately, these attempts had little to no effect on the issue

Any help or advice you could offer would be greatly appreciated. If additional information such as graphs, figures, or parts of the implementation is needed, I would be more than happy to provide them.

I'm really at the end of my rope here so I'm curious about hearing your suggestions.


2 Answers 2


From having some experience with VAEs and tabular data, I can give you some lines of potential improvements:

  • Do you do the reconstruction and regression tasks together ? I would imagine that you have a few (or just 1) linear layers on the output of your encoder for the regression task ?
  • What is your loss, have you tried weighing its different components and optimizing the weights ? I imagine your loss to be $$w_1 * L_{KLD} + w_2 * L_{Reconstruction} + w_3 * L_{Regression}$$
  • The noise introduced in the reparametrization trick is often too strong for the model and it fails to produce a valid embedding, you should try decreasing it (having it as an hyper-parameter is ideal).
  • Finally try to implement your model with a simple dense architecture to make sure that problem comes from the VAE architecture and not somewhere else.

And by the way, it is normal for the KLD to increase to some extent. But it should not keep increasing indefinitely and should stabilize.

  • $\begingroup$ To give answers for each of your suggestions: $\endgroup$
    – Bae Browns
    Commented Jul 11, 2023 at 9:02
  • $\begingroup$ @BaeBrowns your answer seems to have been cut $\endgroup$
    – colin
    Commented Jul 11, 2023 at 9:33
  • $\begingroup$ Thanks. I'll try again. To answer each of your suggestions: - I'm doing the tasks separately. I have either the original input as a target (reconstruction) or a single target value of the product of the line (regression). Training of both these tasks is being done separately. On the output of the encoder, after sampling from the latent representation, I have a single or two layer(s). $\endgroup$
    – Bae Browns
    Commented Jul 11, 2023 at 10:32
  • $\begingroup$ - My loss is KLD + reconstruction or KLD + regression, depending on the task. I have not yet tried the weighting of the KLD and reconstruction loss. I did cyclical annealing and a warmup of KLD (i.e. setting KLD to 0 for the first few epochs). But with no result. I feared that using weights would disrupt the learning balance or lead to overfitting. - By decreasing the noise, do you mean that I should use a factor that constrains the network to output a low variance? $\endgroup$
    – Bae Browns
    Commented Jul 11, 2023 at 10:32
  • $\begingroup$ - I'm already using a simple dense layer network indeed. I have already tried a direct linear/dense layer between the input and the latent space (layers for mu and logvar), which connect directly to the output. I'm also mostly curious about why the architecture would work very well for the one task, but fails very badly at the other. $\endgroup$
    – Bae Browns
    Commented Jul 11, 2023 at 10:32

From your answers to my first answer, I think I have a better idea of your issues and will try to improve my answer:

  • Your loss for the regression task being $L_{KLD} + L_{regression}$ doesn't make sense to me. The KLD isn't too useful in this case, it could be seen as a form of regularization, but that's it.

  • When you say that the t-SNE of the latent space for the reconstruction task shows a single sphere with no clusters, I am not surprised. Have you looked at the reconstructed output ? Is it similar to the input and is the reconstruction loss low ?

I have the intuition that most of your issues are from the reparametrization trick: When you sample $\epsilon \sim \mathcal{N}(0, 1)$ to produce the embedding $z = \mu + \epsilon \times \sigma$, the "noise" from $\epsilon$ is most likely too strong and your model fails to reconstruct and find an appropriate representation. You could try to pick $\epsilon \sim \mathcal{N}(0, eps\_w)$ with $eps\_w \in [0, 1]$. So you could simply try to multiply $\epsilon$ by a scalar in [0, 1] to reduce it. The closer your scalar is to 0, the closer your architecture will be to a simple dense autoencoder, with an additional KLD regularization.

I don't really get why you seek to find if a VAE can handle reconstruction and regression independently, ideally you should do both together, with an architecture like this:


Regarding the anomalies in the gradient, it will be hard to tell without looking at the code.

Note: Finally, you should consider that your code is working just fine, and your data simply doesn't have any natural clusters. So it is normal for you to find a single sphere of data in your latent space. This is plausible as your sensor data has a very large number of features that maybe don't vary too much for clusters to form naturally. To check, simply run a t-SNE on the original data, do you see clusters there ?

  • $\begingroup$ Regarding the question on why I would do the reconstruction and regression independently. I wanted to know if the designed architecture could be used for both of the tasks. What would be the advantage of using them together? Creating a more representative latent space? Additionally, I have used t-SNE on my initial dataset, which does show clusters for the different classes. $\endgroup$
    – Bae Browns
    Commented Jul 11, 2023 at 15:59
  • $\begingroup$ Your architecture is definitely capable of both the regression and reconstruction tasks. You already succeeded at regression, and you should be able to reconstruct your data by reducing the noise of $\sigma$ or increasing the size of your latent space. The advantage of doing both together really depends on your final objective, which I still don't get. If you had clusters in the original data, it means that your training for reconstruction task failed for the reasons stated before. Best of luck. $\endgroup$
    – colin
    Commented Jul 12, 2023 at 8:03
  • $\begingroup$ Many thanks for all of the feedback! I was just mainly wondering why it would fail on one task and not on the other. That's why I was also doing things separately, but again, I will try to combine the tasks and see if it improves. $\endgroup$
    – Bae Browns
    Commented Jul 12, 2023 at 11:01

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.