SimFlow: Simplified and End-to-End Training of Latent Normalizing Flows

TL;DR

SimFlow introduces a simple approach fixing VAE variance to enable end-to-end training of latent normalizing flows, resulting in improved image generation quality and state-of-the-art performance on ImageNet 256x256.

Contribution

The paper proposes fixing the VAE variance to simplify training and improve the quality of latent normalizing flows, enabling end-to-end training without complex pipelines.

Findings

Achieves a gFID of 2.15 on ImageNet 256x256, outperforming STARFlow.

Integrating with REPA-E further improves gFID to 1.91.

Simplifies training by fixing variance, avoiding additional noise or denoising steps.

Abstract

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.