Manifold-Aware Perturbations for Constrained Generative Modeling

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.

💡 Research Summary

The paper addresses a fundamental obstacle that arises when generative models are applied to data constrained by equalities h(x)=0, which forces the data to lie on a lower‑dimensional manifold M within a high‑dimensional ambient space. Standard diffusion models (DDPM, SBDM) and normalizing flows (NF) assume that the target distribution p₀ and the latent Gaussian p_T share the same full dimension, because they rely on the change‑of‑variables formula and on a well‑defined score field. When p₀ is singular with respect to the d‑dimensional Lebesgue measure (i.e., supported on an m‑dimensional manifold with m<d), the Jacobian determinant becomes zero and the score explodes, leading to numerical instability and poor sample quality.

To overcome this, the authors propose a “Manifold‑Aware Perturbation” that lifts p₀ to a full‑dimensional distribution p_σ by adding Gaussian noise only in the normal directions of M. Formally, for each data point z∈M, they compute an orthonormal basis of the normal space N_zM (k = d‑m dimensions) and draw N∼𝒩(0,σ²I_k). The perturbed point is X = z + N, and the law of X is denoted p_σ. Because the noise is confined to the normal bundle, the resulting distribution has support that fills a tubular neighbourhood of M, guaranteeing a non‑degenerate density in ℝ^d.

The training pipeline then proceeds as follows: (1) generate a dataset {X_i} from p_σ; (2) train any off‑the‑shelf unconstrained generative model (diffusion, NF, etc.) on this dataset; (3) at generation time, draw a sample \hat X from the trained model; (4) project \hat X back onto the manifold using the nearest‑point projection Π(·)=arg min_{z∈M}‖·‑z‖. The projection can be performed analytically when M has a simple form or via a few gradient descent steps for more complex constraints.

Theoretical contributions are threefold. First, Theorem 4.3 shows that p_σ has full‑dimensional support (the “tubular distribution”), eliminating Jacobian singularities. Second, Theorem 4.5 proves that when the constraint is linear (i.e., M is an affine subspace), the projection of p_σ exactly recovers the original distribution p₀, guaranteeing perfect reconstruction. Third, for general nonlinear manifolds, Theorem 4.6 provides a bound on the total variation distance between the projected distribution \tilde p_σ = Π#p_σ and p₀:
 TV( \tilde p_σ, p₀ ) ≤ C₁ exp(‑C₂ r²/σ²),
where r is the reach of M (the minimal distance to the medial axis) and C₁, C₂ depend only on the codimension k. This bound implies exponential decay of the distortion as the noise scale σ gets smaller relative to the manifold’s curvature. A localized version (Theorem 4.7) refines the bound using pointwise reach, showing that even in regions of higher curvature the error remains controlled if σ is chosen below the local reach.

Empirical validation spans five tasks of increasing complexity. (i) 3‑D point clouds on a 2‑D plane, a sphere, and a mesh surface illustrate the basic geometry. (ii) 784‑dimensional MNIST images demonstrate that the method works for high‑dimensional data whose intrinsic dimension is far lower. (iii) 90‑dimensional protein backbone coordinates showcase a realistic scientific application where bond‑length and angle constraints define M. For diffusion models, the authors compare three baselines: (a) training directly on p₀ (standard DDPM), (b) training on isotropically noised p₀ with the same σ (to isolate the effect of anisotropy), and (c) training on the proposed p_σ followed by projection. They also evaluate two recent constrained diffusion approaches—Projected Diffusion Models (PDM) and Physics‑Informed Diffusion Models (PIDM). For normalizing flows, RealNVP and Glow are trained on p₀, isotropic p₀, and p_σ.

Results consistently favor the manifold‑aware perturbation. On the plane (linear constraint) the projected DDPM trained on p_σ achieves JSD and TV virtually identical to the ground truth, while isotropic noise leads to substantial distortion. On the sphere and mesh, despite curvature, the TV after projection stays below 0.05 when σ is set below the reach, confirming the theoretical bound. Normalizing‑flow experiments show that the log‑determinant of the Jacobian remains bounded (≈0) for p_σ training, whereas isotropic noise causes exploding determinants and unstable training. In the MNIST experiment, the FID improves from ~18 (direct p₀ training) to ~12 with p_σ, indicating better sample fidelity. The protein backbone task demonstrates that generated structures respect bond constraints and retain realistic dihedral angle distributions, while sampling speed remains comparable to unconstrained baselines.

Overall, the paper delivers a simple, computationally cheap, and mathematically justified strategy to reconcile the dimensional mismatch between latent and target distributions in equality‑constrained settings. By perturbing only along the normal bundle, the method preserves manifold geometry, guarantees exact recovery for linear constraints, and provides exponential TV guarantees for nonlinear manifolds. Its compatibility with any off‑the‑shelf generative architecture makes it immediately applicable to a wide range of scientific domains where physical laws impose strict equality constraints.