Neural Network-Based Score Estimation in Diffusion Models: Optimization and Generalization
Yinbin Han, Meisam Razaviyayn, Renyuan Xu
Abstract
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers17
- Evaluating the design space of diffusion-based generative modelsYuqing Wang, Ye He, Molei TaoNeurIPS 2024 · 30 citations
- Constrained Diffusion Models via Dual TrainingShervin Khalafi, Dongsheng Ding, Alejandro RibeiroNeurIPS 2024 · 24 citations
- Unraveling the Smoothness Properties of Diffusion Models: A Gaussian Mixture PerspectiveYingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song et al.ICCV 2025 · 23 citations
- A solvable model of learning generative diffusion: theory and insightsHugo Cui, Cengiz Pehlevan, Yue M. LuNeurIPS 2025 · 11 citations
- Algorithm- and Data-Dependent Generalization Bounds for Diffusion ModelsBenjamin Dupuis, Dario Shariatian, Maxime Haddouche, Alain Durmus et al.NeurIPS 2025 · 5 citations
Builds on29
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Efficiently Modeling Long Sequences with Structured State SpacesAlbert Gu, Karan Goel, Christopher RéICLR 2022 · 3,482 citations
- Improved Techniques for Training Score-Based Generative ModelsYang Song, Stefano ErmonNeurIPS 2020 · 1,527 citations
- Plug and Play Language Models: A Simple Approach to Controlled Text GenerationSumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung et al.ICLR 2020 · 1,166 citations
- Score-based Generative Modeling in Latent SpaceArash Vahdat, Karsten Kreis, Jan KautzNeurIPS 2021 · 903 citations
Related papers
- Implicit Regularisation in Diffusion Models: An Algorithm-Dependent Generalisation AnalysisTyler Farghly, Patrick Rebeschini, George Deligiannidis, Arnaud DoucetICLR 2026 · 9 citations
- Convergence Dynamics of Over-Parameterized Score Matching for a Single GaussianYiran Zhang, Weihang Xu, Mo Zhou, Maryam Fazel et al.ICLR 2026 · 2 citations
- Diffusion Based Representation LearningSarthak Mittal, Korbinian Abstreiter, Stefan Bauer, Bernhard Schölkopf et al.ICML 2023 · 71 citations
- Understanding Generalizability of Diffusion Models Requires Rethinking the Hidden Gaussian StructureXiang Li, Yixiang Dai, Qing QuNeurIPS 2024 · 45 citations
- Maximum Likelihood Training for Score-based Diffusion ODEs by High Order Denoising Score MatchingCheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen et al.ICML 2022 · 109 citations
