SimFlow: Simplified and End-to-End Training of Latent Normalizing Flows
This work addresses inefficiencies in normalizing flow training for image generation, offering a simpler and more effective approach that improves state-of-the-art results in a specific domain.
The paper tackles the limitations of prior normalizing flow methods, such as complex pipelines and suboptimal quality from frozen VAE encoders, by fixing the variance to a constant, which simplifies training and improves performance. On ImageNet 256x256 generation, SimFlow achieves a gFID of 2.15, outperforming STARFlow (2.40), and with REPA-E integration, it sets a new SOTA with a gFID of 1.91.
Normalizing Flows (NFs) learn invertible mappings between the data and a Gaussian distribution. Prior works usually suffer from two limitations. First, they add random noise to training samples or VAE latents as data augmentation, introducing complex pipelines including extra noising and denoising steps. Second, they use a pretrained and frozen VAE encoder, resulting in suboptimal reconstruction and generation quality. In this paper, we find that the two issues can be solved in a very simple way: just fixing the variance (which would otherwise be predicted by the VAE encoder) to a constant (e.g., 0.5). On the one hand, this method allows the encoder to output a broader distribution of tokens and the decoder to learn to reconstruct clean images from the augmented token distribution, avoiding additional noise or denoising design. On the other hand, fixed variance simplifies the VAE evidence lower bound, making it stable to train an NF with a VAE jointly. On the ImageNet $256 \times 256$ generation task, our model SimFlow obtains a gFID score of 2.15, outperforming the state-of-the-art method STARFlow (gFID 2.40). Moreover, SimFlow can be seamlessly integrated with the end-to-end representation alignment (REPA-E) method and achieves an improved gFID of 1.91, setting a new state of the art among NFs.