Normalizing Flows with Iterative Denoising

1. Problem Statement

The paper addresses the limitations of current generative modeling paradigms, specifically focusing on Normalizing Flows (NFs) and their comparison to Diffusion Models and Autoregressive (AR) models.

• The Noise Dilemma in NFs: While recent NF models like TARFlow have shown promise, they suffer from a fundamental trade-off known as the "noise dilemma."
  • Low Training Noise: If the noise injected during training is too small, the model generates images with rich local textures but poor global structure. Additionally, training becomes unstable due to the unbounded nature of log-likelihood near the data manifold.
  • High Training Noise: If the noise is too large, the model captures better global structures but produces overly smooth, blurry images with visible artifacts upon self-denoising (a phenomenon predicted by Tweedie's Lemma).
• Limitations of Existing Solutions: Previous attempts to mitigate this involved using moderate noise regimes with auxiliary VAEs or curriculum training, but these often failed to fully resolve the trade-off between structural coherence and fine-grained detail.
• Comparison to Diffusion: Diffusion models achieve high quality but require hundreds of iterative steps at inference, leading to high computational overhead. NFs offer exact likelihoods and faster sampling but have historically lagged in visual fidelity.

2. Methodology: Iterative TARFlow (iTARFlow)

The authors propose iTARFlow, a framework that combines the autoregressive structure of TARFlow with an iterative denoising procedure inspired by diffusion models, while maintaining a fully end-to-end likelihood-based objective.

Core Architecture

• Backbone: Uses a Transformer-based Autoregressive Flow (TARFlow) composed of stacked causal Transformer blocks. It maps data to a simple prior (Gaussian) via an invertible transformation.
• Multi-Noise Training: Unlike standard NFs trained on a single noise level, iTARFlow is trained across a continuum of noise levels $t \in [t_{min}, t_{max}]$.
  • The model learns to maximize the log-likelihood $\log p_\theta(x_t, t)$ for varying $t$.
  • The input is defined as $x_t = x + t\epsilon$, where $\epsilon \sim \mathcal{N}(0, I)$.
  • A noise-dependent reweighting term $\gamma_t$ (set to $t$) is applied to the loss.

Inference Pipeline

The sampling process is divided into two distinct phases:

1. Autoregressive Generation (Coarse):
   • The model generates a noisy image $x_{t_{max}}$ patch-by-patch in an autoregressive manner, similar to standard TARFlow.
   • This phase leverages the global coherence learned from high noise levels.
   • Classifier-Free Guidance (CFG) is applied only during this phase to control generation.
2. Iterative Denoising (Refinement):
   • Once the noisy sample is generated, the model performs parallel, iterative denoising to refine the image from $t_{max}$ down to $t_{min}$.
   • The denoising follows a probabilistic ODE: $dx_t = -t \nabla_x \log p_\theta(x_t, t) dt$.
   • The score function $\nabla_x \log p_\theta(x_t, t)$ is computed directly via automatic differentiation of the learned likelihood parameterization.
   • This step is fully parallelizable (unlike the autoregressive generation), making it computationally efficient.

3. Key Contributions

• Identification of the Noise Dilemma: The paper systematically characterizes the trade-off between noise levels in NF training, demonstrating that high noise aids global structure while low noise aids texture, but neither alone is sufficient.
• Iterative TARFlow (iTARFlow): A novel framework that trains NFs across multiple noise scales and employs an iterative, likelihood-induced score estimation for denoising. This allows the model to generate globally coherent structures first and refine details later.
• Bridging the Performance Gap: The method significantly narrows the performance gap between Normalizing Flows and state-of-the-art Diffusion/AR models.
• Hybrid Compatibility: Unlike previous NF variants that struggled to integrate with diffusion backbones, iTARFlow's score functions are shown to be compatible with DiT (Diffusion Transformer) denoisers, suggesting a path toward unified hybrid frameworks.

4. Experimental Results

The authors evaluated iTARFlow on ImageNet at resolutions of 64, 128, and 256 pixels, using Fréchet Inception Distance (FID) as the primary metric.

• ImageNet-64:
  • iTARFlow-B (770M params) achieved an FID of 1.68, outperforming TARFlow (2.66) and approaching the performance of top diffusion models (e.g., EDM-SDE at 1.55) despite using a fundamentally different paradigm.
  • It achieved this with significantly fewer parameters than many diffusion baselines.
• ImageNet-128:
  • iTARFlow-L achieved an FID of 2.44, surpassing TARFlow (5.03) and competing with diffusion models like Simple Diff (1.94) and RIN (2.75).
• ImageNet-256:
  • iTARFlow-XL achieved an FID of 3.32, improving upon TARFlow (5.56) and STARFlow (4.69).
  • It outperformed recent autoregressive NF models (e.g., FARMER-Patch8 at 3.60).
• Efficiency:
  • The iterative denoising process requires very few steps (5 steps for high visual quality, 10 for optimal FID), making inference much faster than diffusion models requiring 50–100 steps.
  • Training converges in fewer epochs (600) compared to DiT (1400).

5. Significance and Limitations

Significance:

• Revitalizing Normalizing Flows: The work proves that NFs can achieve competitive visual quality with diffusion models while retaining the benefits of exact likelihoods and autoregressive scalability.
• Theoretical Insight: It provides a new perspective on how to handle noise in likelihood-based models, moving away from single-point noise injection to a continuum approach.
• Scalability: By leveraging Transformer architectures and autoregressive structures, it offers a promising alternative to diffusion models that may scale more predictably.

Limitations:

• Artifacts: The paper identifies two specific failure modes: "collapsed backgrounds" (entirely black backgrounds) and "low-frequency blur" (particularly in the first autoregressive token). These are attributed to CFG scale issues and the lack of conditioning context for the first token.
• Computational Cost: While inference is fast, the iterative denoising step requires automatic differentiation through the entire network at each step, leading to high GPU memory consumption.
• Hyperparameter Tuning: There is currently no theoretical guidance for selecting the optimal maximum noise level ($t_{max}$), requiring empirical ablation.

In conclusion, iTARFlow represents a significant advancement in Normalizing Flow research, successfully integrating iterative denoising to overcome historical limitations in visual fidelity and establishing NFs as a viable, high-performance alternative to diffusion models.