How to reduce checkerboard artifacts while preserving detail in VAE finetuning?

#3
by Cutecat123 - opened

Hi spacexl,

I’d like to ask how you reduce checkerboard artifacts while still preserving fine details.

Are you adjusting the weight of the LPIPS loss when tuning only the decoder?

Also, if you fine-tune both the encoder and decoder simultaneously, do you apply any constraints to prevent the VAE from drifting too far from the original distribution, so that it remains compatible with the DiT?

Thank you very much!

Did you try adding Discriminator to the flow? It's the common solution for those kind of artifacts

A discriminator can counteract LPIPS artifacts, but it's better to not introduce them in the first place. The checkerboard artifacts originate from the strided convolutions with uneven overlap in VGG, so replacing VGG with a ConvNext or ViT model can help. In this case I used DINOv2-B, on a couple of early hidden states. I also tried several varieties of edge and FFT losses, but they didn't make as much of a difference.

I didn't train the encoder and decoder together, only one at a time. First decoder only, then encoder only. That prevents either from drifting away from the original latent space, although some finetuning of the diffusion model would still be necessary to take advantage of any encoder improvements.

Sign up or log in to comment