Autoencoder doesn't learn 'sparse' input images

Question

I am trying to train an autoencoder with PyTorch on 2D images containing 2D Gaussian densities such as this:

The images are of size 100x100 (I feed them into the autoencoder as 1x10000 tensors). The training data consists of densities with random locations inside the grid and varying standard deviations.

I get good results with my current architecture for such images (nearly identical outputs). But when I try out densities with very small standard deviation, the autoencoder has problems with the reconstruction. Here is an example input image:

Using the same architecture as above and training only on those sparse images, I get results like this:

The location is reconstructed well but not the shape (e.g. no peak in the middle).

And here is the evolution of the loss during training:

According to this thread: Autoencoder for sparse data it shouldn't be a problem that now the input is very sparse (most elements/pixels are zero). I already tried out different learning rates, batch_sizes and architectures but it didn't help.

My current architecture is a fully-connected autoencoder with hidden layer sizes as follows:

10.000 -> 1.024 -> 512 -> 256 -> 64 -> 256 -> 512 -> 1.024 -> 10.000

and ReLu activations in between. I am using MSE as the loss function.

Any ideas what could go wrong here?

Answer 1

I guess your model is overfitting (validation performance is worse than training) and it's probably because of one/all of the following:

• The huge number of parameters (dense architecture) with -I assume- no regularization;
• The distribution/diversity of your data (your gaussian distributions are all centered?);
• The loss metric MSE is averaging over the number of pixels (10.000) even though most pixels are zeroed;

I would suggest starting from the bottom and using an asymmetric loss where the cost of predicting zero when it should be non-zero is different from the cost of predicting non-zero when it should be zero or using a weighted loss where you give more weight to the error on non-zero pixels.