Train autoencoder to reconstruct 2D Gaussian data

Question

I am trying to train an autoencoder to reconstruct 2D Gaussian data. These datapoints are simply sampled from Gaussians with means and covariances chosen randomly. An illustration of what I am trying to do can be visualized below for datapoints from 3 different Gaussians

The blue crosses are the training data and the orange points are the reconstruction of the autoencoder. The goal is to have the orange points match as closely as possible to the blue points. I am using MSE loss for now, although KL divergence between the input and output distributions seem to make more sense, but is trickier to implement for backprop.

For each training set I sample 100 2D coordinates from a different Gaussian, say (x,y) and I stack them into a 200-long feature vector [x1, y1, x2, y2, ... , x100, y100] and pass in mini-batches of these into the autoencoder. I sort these points by the L2 norm of the tuples (x,y), which I've found to help a bit. I've tried basic fully connected autoencoders (200 input/output neurons) and also convolutional autoencoders.

Any tips on how I could make it work better? Thanks!

Answer 1

Training an autoencoder with the MSE loss is particularly tricky because it is highly likely that reconstructions will collapse to simple statistics, such as the sample mean.

Tips on training the autoencoder

• I would suggest you turn off any biases in your linear and convolutional layers, as this will avoid collapsing your reconstructions to the mean value of the input images. This is discussed in footnote 1 in section 3.3 (page 5) in the paper titled Deep One-Class Classification by Lukas Ruff et al.
• Also, consider keeping your batch size small at the beginning of training, to around 8 images per batch. I've observed that a smaller batch size greatly helps training.