this is not overfitting but something else, right?

Question

I trained a VAE on a dataset containing 1k images. The VAE itself is convolutional, downsamples 256x256 rgb images four times before reconstruction and uses both relu and BatchNorm Layers as well as ResNet-like Skip Connection prior and after the bottleneck. Input is normalized to unit scale and no data augmentation takes place. The modal consists in total roughly of about 5e6 parameters and the visual inspection of the reconstructed images look decent enough. This is represented for both training and validation runs by the orange curve. No overfitting occurs as far as I can tell. However, if I use the entire dataset, i.e. ~ 35k training images and 5k validation images, the same model performs like shown by the green curve. Taking into account only the lower validation plot, I had said, this is overfitting, although it would have struck me as odd, that the same model that performed decently on the 1k images dataset now overfits on a much larger dataset. But also the training loss shows a clear inflection point around epoch 15.

Can somebody tell me a likely source of error here? Is the models capacity not large enough to sufficiently capture the complexity of the data? I dont think so, since I am using a sota convolutional VAE used by OpenAI for a dataset comprising millions of images.

The context is, that I am trying to learn a discrete vocabulary of latent codes, i.e. have a discrete learnable embedding to sort of quantize the otherwise continuous latent outputs of the encoder that are then used to reconstruct the input image via the decoder cf. this code snippet. So the idea is not to generate random sampled from noise but to learn an efficient notebook, i.e. bottleneck that captures the essentials of the data set. The decoder then outputs a prob distribution for every pixel over the 255 possible values 8 bit images can take o. The KL (assuming a uniform prior to encourage uniform use of all possible vocabulary entries) is currently weighted with 1.

Any help would be appreciated. Thank you in advance

Answer 1

The purpose of VAE is generating images and not reconstructing the same image. i.e. at the end of the process of VAE you should be able to generate new images out of random sampled data. In your task, if i understand correctly the validation part, you are interested in reconstructing the same image, using both the encoder and the decoder, and not generating a new one (i.e. decoder only). So, it makes sense to me that you will overfit, as the loss of the VAE is not 100% suitable for this task.

One thing you can try is to give more focus to the reconstruction loss. The VAE loss have two parts:

• Reconstruction loss
• KL divergence between the prior and the encoder's output - This is the regularization term that will make it possible to generate new images out of any sample from the random latent space.

The final loss is the sum of these two losses. In your case if you want a good performance on the validation set i.e. generate the same image, you should put lower weight on the last term of the loss.

i.e. instead of $$loss = reconstructionLoss + KL(p(z), q_x(z))$$

use: $$loss = reconstructionLoss + \beta * KL(p(z), q_x(z))$$ and try different $\beta$ values that are less than 1.0 - i'd start with 0.1 and see if it improves.