Manifold-Aware Perturbations for Constrained Generative Modeling

TL;DR

This work addresses equality-constrained generative modeling by recognizing that constrained data lie on a manifold $\\mathcal{M}$ and are degenerate in the ambient space.The authors introduce a manifold-aware perturbation $p_{\\sigma}$ that adds anisotropic Gaussian noise in the normal directions to $\\mathcal{M}$, trains unconstrained generative models on this lifted distribution, and then projects samples back to the manifold with a nearest-point projection $\\Pi$.Key theoretical results show perfect recovery for linear constraints and a total-variation bound $\\text{TV}(\\tilde{p}_{\\sigma}, p_0) \le C_1 e^{-C_2 r^2/\\sigma^2}$ for nonlinear manifolds, while empirical evaluations demonstrate improved sampling stability and distributional fidelity across plane, sphere, mesh, MNIST flux, and protein backbone tasks.The approach offers a computationally light, flexible, and provably-constrained learning paradigm that leverages standard diffusion and normalizing-flow architectures while guaranteeing constraint adherence via projection.These results suggest practical applicability to science-focused generative modeling, with potential extensions to spatially varying noise, curvature-aware perturbations, and diffusion-SDE designs tailored to specific constraints.

Abstract

Generative models have enjoyed widespread success in a variety of applications. However, they encounter inherent mathematical limitations in modeling distributions where samples are constrained by equalities, as is frequently the setting in scientific domains. In this work, we develop a computationally cheap, mathematically justified, and highly flexible distributional modification for combating known pitfalls in equality-constrained generative models. We propose perturbing the data distribution in a constraint-aware way such that the new distribution has support matching the ambient space dimension while still implicitly incorporating underlying manifold geometry. Through theoretical analyses and empirical evidence on several representative tasks, we illustrate that our approach consistently enables data distribution recovery and stable sampling with both diffusion models and normalizing flows.

Figures (15)

• Figure 1: An illustration of our proposed manifold-aware perturbation and comparison to isotropic perturbations showing exact recovery for linear manifolds and lower error for the spiral dataset.
• Figure 2: Sampling stability for plane and sphere tasks. We observe the expected reduction in score and Jacobian log-determinant magnitude across both tasks and generative modeling paradigms. We present similar results for the complex tasks in Appendix \ref{['app:sampling_stability']}.
• Figure 3: Metrics across varied $\sigma$ for diffusion model and NF approaches on toy tasks. In all cases, learning $p_{\sigma}$ consistently outperforms learning $p_{0}$ and post-projecting samples, with the expected possible performance decrease as $\sigma \to \textup{reach}(\mathcal{M})$.
• Figure 5: Mesh task samples. The projected $p_{\sigma}$ samples are close to $p_{0}$ and do not risk off-manifold samples as others do.
• Figure 6: Metrics across varied $\sigma$ on complex tasks. Learning $p_{\sigma}$ consistently improves upon the projected DDPM baseline on all advanced problems.

Theorems & Definitions (25)

• Definition 3.1: $p_{\sigma}$
• Definition 4.1: Medial Axis
• Definition 4.2: Reach
• Theorem 4.3: Geometric Non-degeneracy of $p_{\sigma}$
• proof
• Definition 4.4: Distribution of the Projected Samples
• Theorem 4.5: Perfect Recovery under Linear Constraints
• proof
• Theorem 4.6: Total Variation Bound
• proof