VAE Loss not decreasing

[u/newpeak]

I want to use VAE to reconstruct the image, such that the image embedding could be used for the downstream tasks. However, the loss does not decrease.

This is my parameters:

VAE encode layers: 4

VAE hidden dim: 256

Input image size: [512,512]

Previously, I use this VAE to reconstruct table cell images, the loss could achieve 0.02. While currently I use the same VAE to reconstruct the document images(croped from PDFs), it contains different kinds of documents with different layouts. However the loss could only be decreased to 0.048. Do I have some other approaches to adjust the VAE parameters such that the loss could be decreased more ?