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