ICLR2024
Neural Network-Based Score Estimation in Diffusion Models: Optimization and Generalization
Yinbin Han, Meisam Razaviyayn, Renyuan Xu
被引用 33 次
摘要
Diffusion models have emerged as a dominant paradigm in generative AI, rivaling GANs in producing high-fidelity and robust samples. A core component of these models is learning the score function of perturbed data distribution via denoising score matching. While recent theoretical works have established strong statistical guarantees for diffusion models, they predominantly rely on algorithm-agnostic assumptions, presuming access to a theoretical oracle that perfectly minimizes the empirical risk. In practice, however, score functions are parameterized by highly non-convex neural networks and trained via gradient descent (GD). It remains a major open question whether practical gradient-based algorithms can navigate the optimization landscape of score matching to achieve provable accuracy. As a first step toward answering this question, this paper establishes a mathematical framework for analyzing score estimation using neural networks trained by GD. Our analysis covers both the optimization and the generalization aspects of the learning procedure. In particular, we propose a novel parametric formulation that reduces denoising score matching to a regression problem with inherently noisy labels. Unlike standard supervised learning, the score-matching problem introduces unique theoretical challenges, including unbounded input, vector-valued output, and an additional time variable, preventing existing techniques from being applied directly. We address these challenges by showing that, with proper designs, the evolution of GD-trained neural networks can be accurately approximated by a sequence of localized kernel regression problems. Our analysis is grounded in a novel parametric form of the neural network and an innovative connection between score matching and regression analysis, which facilitate the application of advanced statistical and optimization techniques. Furthermore, since prolonged training on noisy labels causes catastrophic overfitting, we derive a novel extension of early-stopping rules for unbounded domains. This, in turn, allows us to establish the first minimax-optimal generalization error (sample complexity) bounds for GD-trained neural networks in diffusion models. Finally, we validate our theory-inspired optimization framework on a real-world Credit Default dataset, demonstrating that our principled approach achieves performance comparable to heavily tuned heuristic training schemes in generating high-fidelity financial tabular data.