SimFlow: Simplified and End-to-End Training of Latent Normalizing Flows
Abstract
Normalizing Flows (NFs) learn invertible mappings betweenthe data and a Gaussian distribution. Prior works usually suffer fromtwo limitations. First, they add random noise to training samples orVAE latents as data augmentation, introducing complex pipelines in-cluding extra noising and denoising steps. Second, they use a pretrainedand frozen VAE encoder, resulting in suboptimal reconstruction and gen-eration quality. In this paper, we find that the two issues can be solvedin a very simple way: just fixing the variance (which would otherwise bepredicted 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 tokensand the decoder to learn to reconstruct clean images from the augmentedtoken distribution, avoiding additional noise or denoising design. On theother hand, fixed variance simplifies the VAE evidence lower bound, mak-ing it stable to train an NF with a VAE jointly. On the ImageNet 256×256generation task, our model SimFlow obtains a gFID score of 2.15, outper-forming the state-of-the-art method STARFlow (gFID 2.40). Moreover,SimFlow can be seamlessly integrated with the end-to-end representa-tion alignment (REPA-E) method and achieves an improved gFID of1.91, setting a new state of the art among NFs.