In this project we researched the effect of regularizations and augmentations on the Soft-Introspective VAE presented in "Soft-IntroVAE: Analyzing and Improving the Introspective Variational Autoencoder." by Daniel, Tal, and Aviv Tamar, in order to improve its generation ability. Using differentiable augmentations on both real and fake data which also act as a regularization, we managed to improve the FID score reported in the paper of 4.30 on the CIFAR-10 dataset, to a score of 2.96.

This repository contains the implementation of the different methods we used integrated with the code from the official repository.

Project created by: Yuval Baruch & Dan Haramati

General Description: The recently introduced Soft Introspective Variational Auto-Encoder (Soft-IntroVAE) is an explicit deep generative model that enjoys the good traits of variational auto-encoders (VAEs) and generative adversarial networks (GANs) by proposing a variational-based approach to adversarial training, and it exhibits outstanding performance in various tasks such as density estimation, image generation and more. However, in adversarial training, it is quite common that the discriminatory module (discriminator in GANs, encoder in Soft-IntroVAE) may overfit at some point to the training data, preventing further improvements for the generative module (generator in GANs, decoder in Soft-IntroVAE). In this project, we researched and implemented different data augmentations and regularization methods with the aim of improving the Soft-Intro VAE model in terms of generation quality on the CIFAR-10 dataset. We achieved our best results using differentiable augmentations on both real and fake data which also act as a regularization and managed to improve the FID score from 4.30 to 2.96 on the CIFAR-10 dataset, further narrowing the gap between VAE based models and GANs. We include our method’s theoretical and empirical justification as well as an analysis of the exact source of our significant improvement. In addition, we present results on downstream tasks our model can be useful for as a consequence of being trained with augmented data.

augmentation-enhanced-Soft-Intro-VAE

1. Datasets
2. Training
3. Method Overview
4. Results
   • Generation
   • Reconstruction
   • Interpolation in the Latent Space
5. Files and Directories in the Repository
6. Credits

We used the CIFAR-10 dataset to optimize and evaluate our model on.
For the evaluation and analysis we also used MNIST which is more simple and structured than CIAFR-10.

main.py --help

You should use the main.py file with the following arguments:

Command with recommended hyperparameters:

CIFAR-10 - python main.py --dataset cifar10 --device 0 --lr 2e-4 --num_epochs 1000 --beta_kl 1.0 --beta_rec 1.0 --beta_neg 256 --z_dim 128 --batch_size 32 --p_augment 0.5

MNIST - python main.py --dataset mnist --device 0 --lr 2e-4 --num_epochs 500 --beta_kl 1.0 --beta_rec 1.0 --beta_neg 256 --z_dim 32 --batch_size 128 --p_augment 0.5

Notes:

• Our best experiment was with --p_augment 0.5, but you can experiment with the other methods as well.
• The optimizer uses a step scheduler for the learning rate, the milestone is set to epoch 350.
• During training:
  • figures of samples and reconstructions are saved locally.
  • statistics are printed (reconstruction error, KLD, expELBO).
  • at the end of each epoch, a summary of statistics will be printed.
  • FID can be calculated only on 3 channels images (rgb)
• Tips:
  • KL of fake/rec samples should be >= KL of real data.
  • It is usually better to choose beta_kl >= beta_rec.
  • FID calculation is not so fast, so turn it off if you don't care about it.
• The baseline model code supports additional datasets that we did not optimize our method for.

The code in this repository includes implementations of several adjustable regularization methods:

• Dropout in different layers of the encoder and decoder
• Gaussian noise added to the original data or to intermediate layers of the encoder and decoder
• Label switching (Real <-> Fake)
• Differentiable data augmentations

The method that has achieved the most significant improvement in the FID score was data augmentation.

Differentiable Data Augmentation
We used kornia for the implementation of differentiable augmentations to the data by incorporating them as initial layers in the encoder module. We applied them on both real and fake data in all training stages, as proposed in "Differentiable augmentation for data-efficient gan training.", by Zhao, Shengyu, et al. and portrayed in the following figure (applied in (i), (ii) and (iii)):

The augmentations that worked best for us and were applied with a probability of 0.5 were:

• random cutout - block of up to half the size of the image
• random translation - of up to 1/8 of the width/height of the image

We based our choice on our own trials and on "Training generative adversarial networks with limited data.", by Karras, Tero, et al.. An important difference in our setting is that the model has access to the original images through the reconstruction loss calculation, making it robust to augmentations leaking to the generated data.

We believe there were two main sources for this method's success in this setting:

• The augmentations enrich our dataset (CIFAR-10 is a low data benchmark based on "Training generative adversarial networks with limited data.").
• As both augmentations we used effectively delete blocks of pixels from the image, they also act as a regularization, creating a better balance between the two competing modules.

All of the results can be conveniently produced using the provided S_IntroVAE_Analysis.ipynb notebook by loading trained models from checkpoints and running the relevant cells.

We improved the Fréchet inception distance (FID) score on CIFAR-10 from 4.30 to 2.96.

These are a few samples generated by our model:

We were able to maintain good reconstruction although generation significantly improved.

In-Painting
A downstream task our model performs well due to how it was trained is reconstructing images with missing blocks by encoding and then decoding them.

Our model is much more robust to this kind of noise.

Interpolation between two images in the latent space by decoding for several values of .

• Daniel, Tal, and Aviv Tamar. "Soft-IntroVAE: Analyzing and Improving the Introspective Variational Autoencoder." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.
• Zhao, Shengyu, et al. "Differentiable augmentation for data-efficient gan training." Advances in Neural Information Processing Systems 33 (2020): 7559-7570.
• Karras, Tero, et al. "Training generative adversarial networks with limited data." Advances in Neural Information Processing Systems 33 (2020): 12104-12114.