Numerical instability in VAE #331
Description
Description
The current VAE implementation is potentially unstable, causing NaN losses for some combinations of datasets/hyperparameters. This is due to exploding standard deviations in the reparameterization trick.
Unbounded log variances in the output of the VAE encoder explode when exponentiated in the reparameterization trick. This causes numerical instability when training the VAE (for some datasets / hyperparameters).
Expected Behavior
Clamp the logvar (output of VAE encoder) to a reasonable value. We can take inspiration from, e.g., the Hugging Face library, which clamp logvar to [-30,20].
See: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/vae.py