How to Resolve Variational Autoencoder (VAE) Model Collapse in Reconstruction Task Using Sensor Data?

Question

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.

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)

Issue:

• 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.

Answer 1

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.

Answer 2

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 ?