Manifold-Aware Perturbations for Constrained Generative Modeling

Katherine Keegan; Lars Ruthotto

cs.LGJan 30, 2026

Manifold-Aware Perturbations for Constrained Generative Modeling

Authors: Katherine Keegan, Lars Ruthotto

Organizations: Department of Mathematics, Emory University, Atlanta, Georgia, USA.

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 & tables

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.

		Method	Train time	Sampling time	COV $\uparrow$	JSD $\downarrow$	TVD $\downarrow$
Plane	DMs	PDM	0.0012	0.4991	0.8678	0.0555	0.2270
		PIDM	0.0019	0.1959	0.8119	0.0873	0.2913
		$\tilde{p}_{\sigma}$	0.0012	0.2908	0.8736	0.0515	0.2015
		DDPM	0.0012	0.1960	0.8632	0.0475	0.1967
		DDPM (proj.)	0.0012	0.1966	0.8852	0.0442	0.1893
		DDPM (proj., iso.)	0.0012	0.1966	0.8096	0.0821	0.2817

Extrinsic metrics.

Panel of Figure 2 — Figure 2 : Metrics across varied $\sigma$ for diffusion model and NF approaches on toy tasks. In all cases, learning $p_{\sigma}$ consistently outperforms or is competitive against learning $p_{0}$ and post-projecting samples, with the expected possible performance decrease as $\sigma\to\textup{reach}(\mathcal{M})$ .

Panel of Figure 3 — Figure 3 : Image examples for plane and sphere tasks. In all cases, $p_{\sigma}$ samples consistently visually outperform or are competitive against the other techniques. The label “(G)” refers to Glow and the label “(R)” refers to RealNVP.

Panel of Figure 4 — Figure 4 : 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 E .

Panel of Figure 5 — 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.

Method	Train time	Sampling time	COV $\uparrow$	JSD $\downarrow$	TVD $\downarrow$
PDM	0.0013	8.4991	0.8626	0.1912	0.4318
PIDM	0.0033	0.2196	0.5952	0.1298	0.3248
$p_{\sigma}$ (ours)	0.0012	0.3663	0.8506	0.1484	0.3471
DDPM	0.0013	0.2246	0.7104	0.1266	0.3090
DDPM (proj.)	0.0013	0.2415	0.8643	0.1480	0.3458

Table 2 : Mesh task metrics at

\sigma=0.0005

. Learning

p_{\sigma}

is consistently competitive with other methods while avoiding outliers and off-manifold samples through projection.

Figure 6 : Metrics across varied

\sigma

on complex tasks. Learning

p_{\sigma}

consistently improves upon the projected DDPM baseline on all advanced problems. We remark that the coverage and FID metrics are computed in the embedding space due to the challenges of distributional evaluation of images, and their competitive performance combined with consistently outperforming class JSD illustrates the utility of our approach for constrained image generation.

Figure 7 : Performance improvement with respect to number of samples and $\sigma$ on MNIST task. Learning $p_{\sigma}$ at the tested $\sigma$ levels consistently improves performance on all metrics, with an expected increase as the number of training samples increases.

Method	Train time	Sampling time	COV $\uparrow$	FID $\downarrow$	Class JSD $\downarrow$
$p_{\sigma}$ (ours)	0.0083	39.0945	0.8765	$1.512\times 10^{-5}$	0.0730
PDM	0.0083	78.1216	0.8457	$2.218\times 10^{-5}$	0.0874
DDPM	0.0083	39.0982	0.8954	$7.434\times 10^{-6}$	0.0868
PIDM	0.0089	39.0990	0.4713	$1.665\times 10^{-4}$	0.2174
DDPM (proj.)	0.0083	39.1013	0.8954	$\mathbf{7.379\times 10^{-6}}$	0.0869

Table 3 : MNIST metrics at

\sigma=0.01

with 10,000 training samples.

Panel of Figure 8 — Figure 8 : Backbone fragment with atom coordinates as generated variables. In our setting, a single data sample consists of $L$ residues ( $L=3$ in this figure) with each residue made up of 3 backbone atoms (N, CA, C), each represented by $(x,y,z)$ -coordinates.

Method	Train time	Sampling time	COV	Pairwise RMSD	MMD
$p_{\sigma}$ (ours)	0.0040	5.5825	0.2200	1.6262	$3.37\times 10^{-4}$
DDPM	0.0041	1.6013	0.2180	0.7414	$\mathbf{3.35\times 10^{-4}}$
PDM	0.0041	270.3082	0.3630	1.3534	$3.80\times 10^{-4}$
DDPM (proj.)	0.0041	6.8878	0.2880	0.8445	$3.36\times 10^{-4}$
PIDM	0.0070	1.6486	0.2240	1.4216	$3.37\times 10^{-4}$

Table 4 : Protein backbone fragment metrics at

\sigma=0.001

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Spaces and Manifolds
Symbol	Description
$\mathbb{R}^{d}$	Ambient space (dimension $d$ )
$\mathcal{M}$	$m$ -dimensional constraint manifold embedded in $\mathbb{R}^{d}$
$m$	Intrinsic dimension of $\mathcal{M}$
$k=d-m$	Codimension of $\mathcal{M}$
$T_{x}\mathcal{M}$	Tangent space to $\mathcal{M}$ at point $x$

Appendix

Table 5 : Summary of notation used throughout this paper.

	Plane	Sphere	Mesh	Images	Protein
Time embedding	Scalar	Scalar	Scalar	MLP	MLP
Time emb. dim.	1	1	1	64	16
# Training samples	100,000	100,000	100,000	10,000	20,000
Batch size	64	64	64	32	32
Epochs	200	200	200	1000	1000
Hidden dimension	64	64	128	1024	1024

Appendix

Table 6: Combined diffusion model training configurations for all tasks.

# Parameters	Plane	Sphere	Mesh	Images	Protein
RealNVP	42,780	42,780	–	–	–
Glow	28,836	28,836	–	–	–
DDPM	4,675	4,675	17,539	1,093,057	112,353

Appendix

Table 7: Total number of trainable parameters for each model across all tasks.

Panel of Figure 9 — Figure 9 : Training and validation losses and accuracy for MNIST classifier.

Panel of Figure 13 — Figure 13 : Examples of effect of $\sigma$ with 100 training data samples.

Panel of Figure 14 — Figure 14 : Examples of effect of $\sigma$ with 1,000 training data samples.

Panel of Figure 15 — Figure 15 : Examples of effect of $\sigma$ with 10,000 training data samples.

Explore similar work

Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling
May 12, 2026Jacob K. Christopher, James E. Warner, Ferdinando FiorettoGenerative Models Contrastive Flow Matching
University of Virginia · NASA Langley Research Center
Integration Matters: Rollout-Based Training for Constrained Diffusion Models
Jul 15, 2026Xiaoxuan Liang, Saeid Naderiparizi, Berend Zwartsenberg +1Diffusion Models Diffusion Alignment
University of British Columbia · Inverted AI · Alberta Machine Intelligence Institute
Efficient Diffusion Models under Nonconvex Equality and Inequality constraints via Landing
Apr 20, 2026Kijung Jeon, Michael Muehlebach, Molei TaoDiffusion Sampling Inequality Constraints
Georgia Institute of Technology · Max Planck Institute for Int. Sys.

May 12, 2026, cs.LGJ/K move · Enter open · S save

Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling

Jacob K. Christopher, James E. Warner, Ferdinando Fioretto

University of Virginia · NASA Langley Research Center

Generative Models Contrastive Flow Matching Training-Inference Mismatch Decisions

Open paper AskPDF arXiv

Deep generative models provide state-of-the-art performance across a wide array of applications, with recent studies showing increasing applicability for science and engineering. Despite a growing corpus of literature focused on the integration of physics-based constraints into the generation process, existing approaches fail to enforce strict constraint satisfaction while maintaining sample quality. In particular, training-free constrained sampling methods, while providing per-sample feasibility guarantees, introduce a fundamental mismatch between the training objective and the constrained sampling procedure, often leading to performance degradation. Identifying this training-sampling misalignment as a central limitation of current constrained generative modeling approaches, this paper proposes Constraint-Aware Flow Matching, a novel end-to-end framework that explicitly incorporates constraint projections into the training objective. By aligning the model's learned dynamics with the constrained sampling process, the proposed method mitigates distributional shift induced by projection-based corrections, enabling high-quality constrained generation. The proposed approach is evaluated on three challenging real-world benchmarks, illustrating the generality and efficacy of the method.

Abstract

Explore similar work

Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling

Integration Matters: Rollout-Based Training for Constrained Diffusion Models

Efficient Diffusion Models under Nonconvex Equality and Inequality constraints via Landing