Why is mode() used instead of sample() in VAE during training?

[hwangdonghyun]

Lumina-Image-2.0/finetune.py

Line 602 in be7067b

x[i] = (vae.encode(img[None].bfloat16()).latent_dist.mode()[0] - vae_shift) * vae_scale

I noticed in the code that the VAE uses mode() instead of sample() during training. Could you please explain the reason behind this choice?

From my understanding, most model training setups typically use sample() in VAEs. I’m curious if using mode() provides any particular advantages in this case.

Thank you!

[hwangdonghyun]

Additionally, since this script is for fine-tuning, I would also like to know whether the model was originally trained from scratch in the same way (i.e., using mode() instead of sample()).