Reflected diffusion models adapt to low-dimensional data

2026-03-25

AI Summary

This work addresses generative modeling on a $d$-dimensional linear subspace embedded within the hypercube $[0,1]^D$, proposing a diffusion model with reflecting boundaries to overcome the limitations of conventional approaches that rely on unbounded Euclidean Gaussian kernels. The score function is derived analytically via an infinite series expansion of the transition density of the reflected diffusion process and approximated using sparse ReLU networks. Theoretical analysis demonstrates that the reflecting boundaries do not compromise statistical efficiency: under Sobolev smoothness assumptions, the model achieves a convergence rate of $n^{-(\alpha+1-\delta)/(2\alpha+d)}$ in the 1-Wasserstein distance for any arbitrarily small $\delta>0$, matching the near-optimal rate of the unconstrained setting and, for the first time, fully adapting to the intrinsic dimension $d$.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

While the mathematical foundations of score-based generative models are increasingly well understood for unconstrained Euclidean spaces, many practical applications involve data restricted to bounded domains. This paper provides a statistical analysis of reflected diffusion models on the hypercube $[0,1]^D$ for target distributions supported on $d$-dimensional linear subspaces. A primary challenge in this setting is the absence of Gaussian transition kernels, which play a central role in standard theory in $\mathbb{R}^D$. By employing an easily implementable infinite series expansion of the transition densities, we develop analytic tools to bound the score function and its approximation by sparse ReLU networks. For target densities with Sobolev smoothness $α$, we establish a convergence rate in the $1$-Wasserstein distance of order $n^{-\frac{α+1-δ}{2α+d}}$ for arbitrarily small $δ> 0$, demonstrating that the generative algorithm fully adapts to the intrinsic dimension $d$. These results confirm that the presence of reflecting boundaries does not degrade the fundamental statistical efficiency of the diffusion paradigm, matching the almost optimal rates known for unconstrained settings.

Problem

• reflected diffusion models
• bounded domains
• low-dimensional data
• score-based generative models
• intrinsic dimension

Innovation

• reflected diffusion models
• intrinsic dimension adaptation
• score-based generative models
• sparse ReLU networks
• Wasserstein convergence rate