Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling

2024-04-19

Neural Information Processing Systems

AI Summary

Traditional diffusion models rely on fixed forward processes, resulting in complex latent marginal distributions, challenging reverse-generation procedures, and high inference costs. To address these limitations, we propose Neural Flow Diffusion Models (NFDMs), the first framework to enable end-to-end learnable parameterization of the forward process. NFDMs support modeling non-Gaussian, structured, deterministic linear trajectories and bridging arbitrary source-target distributions. Our method integrates neural ordinary differential equations (neural ODEs) for implicit dynamical modeling, flow matching, variational upper-bound optimization, and simulation-free gradient estimation. On standard benchmarks, NFDMs achieve state-of-the-art log-likelihood scores. Moreover, they significantly enhance generation controllability and dynamical system modeling capability, demonstrating strong generalization and effectiveness in distribution transfer and trajectory-controllable generation tasks.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Social Networks and Social Media: Generative AI / large language models and their impact on social systems

Abstract

Conventional diffusion models typically relies on a fixed forward process, which implicitly defines complex marginal distributions over latent variables. This can often complicate the reverse process' task in learning generative trajectories, and results in costly inference for diffusion models. To address these limitations, we introduce Neural Flow Diffusion Models (NFDM), a novel framework that enhances diffusion models by supporting a broader range of forward processes beyond the standard Gaussian. We also propose a novel parameterization technique for learning the forward process. Our framework provides an end-to-end, simulation-free optimization objective, effectively minimizing a variational upper bound on the negative log-likelihood. Experimental results demonstrate NFDM's strong performance, evidenced by state-of-the-art likelihood estimation. Furthermore, we investigate NFDM's capacity for learning generative dynamics with specific characteristics, such as deterministic straight lines trajectories, and demonstrate how the framework may be adopted for learning bridges between two distributions. The results underscores NFDM's versatility and its potential for a wide range of applications.

Problem

• Enhancing diffusion models with learnable forward processes
• Improving generative trajectories via flexible forward processes
• Enabling versatile applications through optimized distribution learning

Innovation

• Learnable forward process enhances diffusion models
• Novel parameterization technique for forward process
• End-to-end simulation-free optimization objective