Localized Diffusion Models

TL;DR

This paper introduces localized diffusion models that leverage low-dimensional locality structures to reduce sample complexity and improve training efficiency in high-dimensional generative tasks.

Contribution

It proposes a localized score matching approach that exploits locality structures, enabling diffusion models to bypass the curse of dimensionality with theoretical and numerical support.

Findings

Localized diffusion models reduce sample complexity.

Moderate localization radius balances errors effectively.

Parallel training enhances scalability for large datasets.

Abstract

Diffusion models are state-of-the-art tools for various generative tasks. Yet training these models involves estimating high-dimensional score functions, which in principle suffers from the curse of dimensionality. It is therefore important to understand how low-dimensional structure in the target distribution can be exploited in these models. Here we consider locality structure, which describes certain sparse conditional dependencies among the target random variables. Given some locality structure, the score function is effectively low-dimensional, so that it can be estimated by a localized neural network with significantly reduced sample complexity. This observation motivates the localized diffusion model, where a localized score matching loss is used to train the score function within a localized hypothesis space. We prove that such localization enables diffusion models to circumvent the curse of dimensionality, at the price of additional localization error. Under realistic sample size scaling, we then show both theoretically and numerically that a moderate localization radius can balance the statistical and localization errors, yielding better overall performance. Localized structure also facilitates parallel training, making localized diffusion models potentially more efficient for large-scale applications.