Normalizing Flows are Capable Generative Models

Normalizing Flows are Capable Generative Models

Speakers: Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista Martin, Navdeep Jaitly, Joshua M Susskind

Conference: ICML 2025

YouTube: https://slideslive.com/39044094

Overview

This talk, presented by Shuangfei Zhai from Apple's machine learning research team, challenges the long-held notion that Normalizing Flows (NFs) are inferior generative models compared to more contemporary approaches like Diffusion Models or Generative Adversarial Networks (GANs). Historically, NFs have struggled to produce high-fidelity, diverse samples, often yielding "barely recognizable" images even on standard benchmarks like ImageNet. The core message of this presentation is to convincingly demonstrate that, with the right architectural innovations and training methodologies, normalizing flows can not only achieve state-of-the-art likelihood performance but also generate samples of competitive visual quality.

The researchers introduce Tarflow, a novel architecture built upon Transformer Autoregressive Flows, which integrates several crucial techniques to unlock the full potential of NFs. These include a powerful transformer-based design, noise augmented training, score-based self-denoising, and guidance mechanisms adapted from other generative modeling paradigms. By meticulously combining these elements, Tarflow achieves unprecedented performance for normalizing flows, marking a significant milestone in the field and suggesting that NFs could become a viable foundation for scalable generative models. This work re-establishes normalizing flows as a powerful and principled approach, pushing their capabilities far beyond previous limitations.

Background

▶ Watch: Normalizing Flows: Background and historical limitations (0:00)

Normalizing flows are a class of likelihood-based generative models that operate on the principle of transforming a complex, target data distribution into a simpler, tractable distribution (typically a standard Gaussian) through a series of invertible and differentiable transformations. The key advantage of this approach is its direct optimization of the exact data likelihood, providing a deterministic and theoretically sound loss function. This loss comprises two main components: an L2-like term that encourages the transformed data to resemble the target simple distribution, and a log-determinant of the Jacobian term that accounts for the change in volume induced by the transformation. After training, the model can generate new samples by simply drawing noise from the simple distribution and passing it through the inverse of the learned transformation.

Despite their theoretical elegance and principled foundation, normalizing flows have historically lagged behind other generative models in terms of sample quality. Early models, even well-known ones like Flow architectures, produced samples that were often blurry, lacked detail, and struggled to capture the complexity of real-world datasets such as ImageNet. This underperformance led to a perception that NFs, while mathematically sound, were not practical for high-fidelity image generation, especially when compared to the rapidly advancing capabilities of GANs and, more recently, diffusion models. The challenge has always been to design an invertible function f that is expressive enough to model complex distributions while maintaining tractability for Jacobian computation and efficient inversion. This talk directly addresses this long-standing problem, aiming to bridge the performance gap and re-establish normalizing flows as a competitive generative modeling paradigm.

Key Findings

▶ Watch: Critical technique 1: Noise augmented training (4:00)

The central finding presented in this talk is that normalizing flows, when equipped with appropriate architectural designs and training strategies, are indeed capable generative models, achieving performance previously thought to be exclusive to other paradigms. The research introduces Tarflow (Transformer Autoregressive Flows), which delivers two primary breakthroughs:

1. State-of-the-Art Likelihood Performance: Tarflow achieves a sub-3 BPD (bits per dimension) score on the unconditional ImageNet 64x64 benchmark for the first time. This significantly surpasses all previous methods, including autoregressive models like PixelCNN and even contemporary diffusion models, demonstrating the architectural strength and efficiency of Tarflow in modeling complex data distributions with high precision.
2. Competitive Generative Quality: Beyond likelihood, Tarflow generates high-fidelity and diverse samples that are visually comparable to those produced by early diffusion models. This marks a crucial first for normalizing flow-based models, which have traditionally struggled with sample quality. The integration of noise augmented training, score-based self-denoising, and guidance mechanisms proved critical in bridging this gap, allowing Tarflow to produce sharp, recognizable images across various resolutions.

These findings collectively indicate that normalizing flows are no longer a niche, theoretically interesting but practically limited, generative modeling approach. Instead, Tarflow demonstrates their potential as a foundation for scalable generative models, capable of both precise likelihood estimation and high-quality sample generation.

Technical Deep Dive

▶ Watch: Critical technique 2: Score-based denoising pipeline (6:00)

The success of Tarflow hinges on a synergistic combination of a powerful architecture and several critical training and sampling techniques.

Tarflow Architecture: Transformer Autoregressive Flows

The core of Tarflow is its novel architecture, which is built on a stacked causal transformer with alternating directions. This design is specifically engineered to be both invertible and possess a trivial Jacobian determinant, simplifying the training loss. For image modeling, the architecture processes inputs by first partitioning them into non-overlapping patches, similar to Vision Transformers. These patches are then fed through T stacked Tarflow blocks.

Each Tarflow block is fundamentally a causal transformer implemented with a standard causal mask. The key innovation within each block is the introduction of a permutation layer at the input. This layer effectively reverses the order of autoregression between successive blocks. For instance, if the first layer processes from left-to-right, the next layer will process from right-to-left. This alternating directional processing enhances the model's ability to capture complex dependencies across the input space. The output of the causal transformer within each block is then transformed into an affine transformation, producing a bias term mu and a scale term alpha, which parameterize the transformation at that stage. Stacking T such blocks forms the complete Tarflow model, which learns to map the input data distribution to a simple Gaussian noise distribution.

Because of the invertible nature and trivial Jacobian of the Tarflow architecture, the training loss function simplifies considerably. It consists of an L2-like term that pushes the output of the function towards zero (mimicking Gaussian noise) and a sum of linear terms derived from the intermediate layers, which arises from the log-determinant of the Jacobian computation.

Noise Augmented Training

A critical technique for achieving robust performance is noise augmented training. The idea is deceptively simple: instead of training the model on the pristine data distribution (P_data), a small but fixed amount of Gaussian noise is added to every input during training. Mathematically, this is equivalent to convolving the true data distribution P_data with a Gaussian kernel.

The purpose of this augmentation is to densify the training distribution. Real-world data often lies on a low-dimensional manifold within a high-dimensional space. Adding a small amount of noise effectively "smears" this manifold, making the training distribution denser and smoother. This regularization significantly improves the model's generalization capabilities, particularly during the sampling phase, where the model needs to synthesize novel data points from the learned distribution. The amount of noise is small but perceptible, ensuring the model learns from slightly perturbed, rather than perfectly clean, inputs.

Score-Based Self-Denoising

A direct consequence of noise augmented training is that the raw samples generated by the model will inherently be noisy, as the model was trained on noisy data. To obtain clean, high-fidelity samples, the researchers employ a score-based self-denoising step. This technique leverages the inherent properties of likelihood-based models. Since normalizing flows learn the exact data density, it is possible to analytically derive the score function (the gradient of the log-likelihood with respect to the input).

Inspired by the diffusion model literature, specifically the Tweedie formula, this derived score can be used to perform a single step of denoising. Crucially, this does not require training a separate denoising model. The full sampling pipeline is thus: first, sample a Gaussian noise vector z; second, pass z through the inverse of the learned Tarflow function (f^-1) to obtain a noisy sample y; and third, apply one step of score-based denoising using the model's own derived score to yield the final clean sample x. This deterministic mapping transforms Gaussian noise into a clean, high-quality output.

Guidance Mechanisms

Finally, guidance is integrated to trade diversity for fidelity, a technique widely recognized for its effectiveness in modern generative models like diffusion models. Tarflow adapts this by extrapolating between conditional and unconditional predictions at every step of the generation process. For conditional generation (e.g., class-conditional image generation), the model computes two versions of its prediction: one conditioned on the class label and one unconditional. A guidance weight w then determines the extent of extrapolation between these two predictions, pushing the generated sample closer to the desired conditional distribution.

A notable bonus is the ability to apply guidance even to unconditional models. This is achieved by manipulating a temperature term injected into all attention layers of the transformer. When the temperature is set to one, the model behaves as originally trained. When it deviates from one, it effectively represents an "inferior" or perturbed version of the model. By extrapolating between the prediction from the original model (temperature=1) and this perturbed version, guidance can be applied to enhance fidelity in unconditional generation settings, providing a flexible and powerful control mechanism.

Experimental Setup & Results

▶ Watch: Critical technique 3: Classifier-free guidance for flows (7:00)

The effectiveness of Tarflow was rigorously evaluated on standard benchmarks, focusing on both likelihood estimation and generative sample quality.

Likelihood Estimation

The primary benchmark for likelihood estimation was the unconditional ImageNet 64x64 dataset. This benchmark provides a historical progression of generative model performance. Earlier models, such as PixelCNN, achieved certain BPD scores, which were later improved by more advanced transformer-based autoregressive models around 2020. Subsequently, diffusion models demonstrated further improvements in likelihood.

Tarflow significantly advanced the state-of-the-art in this domain, achieving a sub-3 BPD score for the first time. This result represents a substantial reduction in negative log-likelihood compared to all previous methods, including leading diffusion models. The authors consider this a pivotal demonstration of the inherent strength and expressiveness of the Tarflow architecture in accurately modeling complex data distributions.

Generative Modeling (FID Scores)

Beyond likelihood, the visual quality and diversity of generated samples were quantitatively assessed using the Fréchet Inception Distance (FID) metric, also on ImageNet 64x64. Qualitatively, Tarflow generates diverse and high-fidelity images that are visually comparable to those from other modern generative models, a significant first for normalizing flow-based approaches.

Quantitatively:

• A baseline Tarflow model achieved an FID of approximately 4 on ImageNet 64x64, which is a respectable starting point.
• Performance was further improved by reducing the patch size, which corresponds to increased computational resources and a more powerful model. This brought the FID down.
• Increasing the overall model size also led to further reductions in FID, showcasing the scalability of the architecture.
• The absolute best FID achieved by Tarflow was "pretty comparable to some of the early results from diffusion models," solidifying its position as a competitive generative model in terms of sample quality.

Ablation Studies

Detailed ablation studies were conducted to understand the contribution of each critical component:

• Noise Augmented Training and Denoising: The impact of injected Gaussian noise on FID was analyzed.
• Raw samples (without denoising): As the amount of injected noise increased, the FID score predictably worsened, indicating a degradation in sample quality.
• Denoised samples (with score-based self-denoising): A distinct U-shaped curve emerged. This indicates a "sweet spot" for the amount of noise to inject during training. Too little noise leads to numerical instability (as noted in the Q&A, related to the manifold hypothesis), while too much noise causes a loss of information. Importantly, the FID scores for denoised samples were significantly lower than raw samples across the board, demonstrating the critical role of the denoising step.
• Guidance: The effect of the guidance weight on FID was also studied for both conditional and unconditional models. In both scenarios, a U-shaped curve was observed, which is typical for guidance mechanisms. This indicates that there's an optimal guidance strength that balances fidelity and diversity, with overly strong guidance potentially leading to mode collapse or reduced diversity.

Qualitative Results

Visual examples further reinforced the quantitative findings. Raw samples generated by Tarflow without denoising or guidance were "reasonable but not great." Applying guidance to these samples (from the same random seed) resulted in "more coherent" images with "recognizable objects," though the inherent noise from training was still visible. Finally, applying the one-step score-based denoising process dramatically improved sample clarity, removing the noise while preserving the underlying signals and details, producing clean and high-fidelity outputs. The presentation also included high-resolution samples, indicating the model's potential for scaling to larger image sizes.

Practical Implications

▶ Watch: Likelihood estimation results on ImageNet 64x64 (8:30)

The advancements demonstrated by Tarflow carry significant practical implications for practitioners, infrastructure teams, and model builders working with generative AI.

Firstly, the achievement of state-of-the-art likelihood on ImageNet 64x64, combined with competitive generative quality, re-establishes normalizing flows as a viable and powerful generative modeling paradigm. This means that teams previously dismissing NFs due to their historical performance issues might now consider them for tasks requiring both high-fidelity generation and precise likelihood estimation, such as anomaly detection, density estimation, or uncertainty quantification, where exact likelihoods are highly valuable.

Secondly, the modular nature of Tarflow's components – particularly the transformer-based architecture, noise augmentation, self-denoising, and guidance – offers a clear blueprint for future research and development in NFs. The speaker explicitly noted in the Q&A that the noise augmentation and denoising tricks are not specific to the Tarflow architecture but are generally applicable to other normalizing flow models, suggesting a broader impact across the NF family. This generality makes these techniques valuable tools for improving existing NF implementations.

For infrastructure teams, the efficiency and deterministic nature of NFs could offer advantages in certain deployment scenarios compared to iterative sampling methods like those used in diffusion models. While the talk did not delve into inference speed comparisons, the single-step denoising process after inverse transformation implies a potentially faster generation pipeline than multi-step diffusion sampling.

However, some limitations and trade-offs remain. The "exact functionality of noise," while empirically effective, is still an area requiring further theoretical understanding. Practitioners will need to carefully tune the amount of noise and guidance weights, as demonstrated by the U-shaped curves in the ablation studies, to find the optimal balance for their specific applications. The computational cost associated with larger model sizes and smaller patch sizes for improved FID also needs to be considered, as these enhancements come with increased resource demands. Despite these considerations, the work "opens the door for scalable normalizing flow-based foundation models," hinting at a future where NFs could underpin large-scale generative applications.

Key Takeaways

• Normalizing Flows (NFs) are now capable generative models: Tarflow challenges historical perceptions by achieving state-of-the-art performance.
• Tarflow leverages a powerful Transformer-based architecture: Its Transformer Autoregressive Flows design, with alternating causal transformers and permutation layers, ensures invertibility and a trivial Jacobian.
• Noise Augmented Training is crucial for generalization: Adding small, fixed Gaussian noise to inputs densifies the training distribution, significantly improving sampling quality.
• Score-Based Self-Denoising yields clean samples: The model's learned density allows for a single, efficient denoising step using the Tweedie formula, eliminating the need for a separate denoiser.
• Guidance enhances sample fidelity: Adapted from other generative models, guidance (including for unconditional models via temperature manipulation) improves sample quality by extrapolating between predictions.
• Tarflow achieves sub-3 BPD on ImageNet 64x64: This unprecedented likelihood score demonstrates the model's exceptional ability to accurately capture data distributions.
• Generative quality is competitive with early diffusion models: Tarflow produces high-fidelity, diverse images, marking a significant milestone for normalizing flows.

About the Speaker(s)

The primary speaker for this talk was Shuangfei Zhai, a distinguished researcher associated with the machine learning research team at Apple. Zhai's work, along with his co-authors, focuses on advancing the capabilities of generative models, particularly in demonstrating the latent potential of methods like normalizing flows. The presentation reflects a deep understanding of generative modeling principles and a commitment to pushing the boundaries of what is achievable with principled likelihood-based approaches.