Optimal estimation of a factorizable density using diffusion models with ReLU neural networks

Abstract

This paper investigates the score-based diffusion models for density estimation when the target density admits a factorizable low-dimensional nonparametric structure. To be specific, we show that when the log density admits a $d^*$-way interaction model with $\beta$-smooth components, the vanilla diffusion model, which uses a fully connected ReLU neural network for score matching, can attain optimal $n^{-\beta/(2\beta+d^*)}$ statistical rate of convergence in total variation distance. This is, to the best of our knowledge, the first in the literature showing that diffusion models with standard configurations can adapt to the low-dimensional factorizable structures. The main challenge is that the low-dimensional factorizable structure no longer holds for most of the diffused timesteps, and it is very challenging to show that these diffused score functions can be well approximated without a significant increase in the number of network parameters. Our key insight is to demonstrate that the diffused score functions can be decomposed into a composition of either super-smooth or low-dimensional components, leading to a new approximation error analysis of ReLU neural networks with respect to the diffused score function. The rate of convergence under the 1-Wasserstein distance is also derived with a slight modification of the method.