Neural Approximate Mirror Maps for Constrained Diffusion Models
Authors: Berthy Feng
Ricardo Baptista
Katherine Bouman
Venue: ICLR 2025 Poster (2025)
Keywords: generative models
diffusion models
mirror maps
constrained generation
inverse problems
Primary area: generative models

Abstract: Diffusion models excel at creating visually-convincing images, but they often struggle to meet subtle constraints inherent in the training data. Such constraints could be physics-based (e.g., satisfying a PDE), geometric (e.g., respecting symmetry), or semantic (e.g., including a particular number of objects). When the training data all satisfy a certain constraint, enforcing this constraint on a diffusion model makes it more reliable for generating valid synthetic data and solving constrained inverse problems. However, existing methods for constrained diffusion models are restricted in the constraints they can handle. For instance, recent work proposed to learn mirror diffusion models (MDMs), but analytical mirror maps only exist for convex constraints and can be challenging to derive. We propose *neural approximate mirror maps* (NAMMs) for general, possibly non-convex constraints. Our approach only requires a differentiable distance function from the constraint set. We learn an approximate mirror map that transforms data into an unconstrained space and a corresponding approximate inverse that maps data back to the constraint set. A generative model, such as an MDM, can then be trained in the learned mirror space and its samples restored to the constraint set by the inverse map. We validate our approach on a variety of constraints, showing that compared to an unconstrained diffusion model, a NAMM-based MDM substantially improves constraint satisfaction. We also demonstrate how existing diffusion-based inverse-problem solvers can be easily applied in the learned mirror space to solve constrained inverse problems.

Published as a conference paper at ICLR 2025
NEURAL APPROXIMATE MIRROR MAPS FOR CONSTRAINED DIFFUSION MODELS
Berthy T. Feng
Ricardo Baptista
California Institute of Technology bfengêcal wewcccateccc.ee..
California Institute of Technology rsbêca a ttech edu
Katherine L. Bouman
California Institute of Technology klbouman& caltech. edu
ABSTRACT
Diffusion models excel at creating visually-convincing images, but they often struggle to meet subtle constraints inherent in the training data. Such constraints could be physics-bare based (e.g.. satisfying a PDE), geometric (e.g., respecting sym- metry), or semantic (c.g., including a particular number of objects). When the training data all satisfy a certain constraint, enforcing this constraint on a diffusion model makes it more reliable for generating valid synthetic data and solving con- strained inverse problems. However, existing methods for constrained diffusion models are restricted in the constraints they can handle. For instance, recent work proposed to learn mirror diffusion models (MDMs), but analytical mirror maps only exist for convex constraints and can be challenging derive. We propose neural approximate mirror maps (NAMMs) for general, possibly non-convex con- straints. Our approach only requires a differentiable distance function from the constraint set. We learn approximate mirror map that transforms data into an unconstrained space and a corresponding approximate inverse that maps data back to the constraint set. A generative model, such MDM, then be trained in the learned mirror space and its samples restored to the constraint set by the inverse map. We validate our approach pared to an unconstrained diffusion model, a NAMM-based MDM substantially improves constraint satisfaction. We also demonstrate how existing diffussio--assed inverse-problem solvers can be easily applied in the learned mirror space to solve constrained inverse problems.
variety of constraints, showing that com-
INTRODUCTION
Many data distributions follow rule that is not visually obvious. For example, videos of fluid flow obey a partial differential equation (PDE), but human may find it difficult to discern whether a video agrees with the prescribed PDE. We can characterize such distributions as constrained distributions. Theoretically, a generative model trained on constrained image distribution should satisfy the constraint, but in practice due to learning and sampling errors (Daras al. 2024). it may generate visually-convincing convincing images that break the rules. Ensuring constraint satisfaction in spite of such errors would make generative models more reliable for applications such as solving inverse problems.
Diffusion models are popular generative models, but existing approaches for incorporating constraints either restrict the type of constraint or not scale well. Equivariant (Niu et al., 2020), Riemannian
(De Bortoli et al.) 2022 Huang al. (022). reflected (Lou & Ermon 2023) Fishman et al.) 2023a).
log-barrier (Fishman et al. 2023a), and mirror (Liu et al. 2024) diffusion models are all restricted to certain types of constraints, such symmetry groups NNuu al. 2020), Riemannian manifolds
(De Bortoli et al. 2022 Huang et al. 2022), or convex sets (Fishman et al., 2023a: Liu et al.
2024). Generally speaking, these restrictions are in the service of guaranteeing hard constraint, whereas a soft constraint offers flexibility and may be sufficient in many scenarios. One could, for example, introduce guidance term encourage constraint satisfaction during sampling (Graikos

---

Published as a conference paper at ICLR 2025
M
M
data sample
data sumple
DM sample (valid)
f
MDM sample
DM sample (invalid)
inverted MDM sample through fy)
Rd
Rd
RR
(a) Regular diffusion model
(b) Mirror diffusion model with NAMM
Figure 1: Conceptual illustration. (a) Despite being trained on data distribution constrained M,, regular diffusion model (DM) may generate samples that violate the constraint. (b) We propose to learn neural approximate mirror map (NAMM) that entails forward map g and inverse map fø- The forward map transforms the constrained space into an unconstrained ("mirrorr" space. Once and fy are learned, mirror diffusion model (MDM) can be trained on the pushforward of the data distribution through g and its samples mapped back to the constrained space through fg.
ett al. 2022
Bansal al. 2023 Zhang et al., 2023), but far there does not exist principled framework to do so. Previous work suggested imposing a soft constraint during training (Daras) al. 2024| Upadhyay et al. 2023) by estimating the clean image every noisy diffusion step and evaluating its constraint satisfaction. However, approximation (Eiron) 201T) of the clean image from the intermediate noisy image crude at high noise levels and thus unsuitable for constraints that are sensitive to approximation error and noise, such as PDE constraint. Of course it is possible to penalize invalid generated samples during or after training (Huang et al. 2024), but this approach relies on computationally-xxeesssv simulation steps. Instead, we aim for flexible approach to impose general constraints by construction.
Our goal is to find invertible function that maps constrained images into an unconstrained space so that a regular generative model can trained in the unconstrained space and automatically satisfy the constraint through the inverse function. We propose neural approximate mirror maps (NAMMs which bring the flexibility of soft constraints into the principled framework of mirror diffusion models (Liu et al. 2024). A mirror diffusion model (MDM) allows for training completely unconstrained diffusion model in a "mirror" space defined by mirror map. Unconstrained samples from the diffusion model
mapped back the constrained space via inverse mirror map. However, invertible mirror maps are challenging or impossible to derive in closed form for general constraints. We address this by jointly optimizing two networks approximate mirror map and its inverse.
A NAMM encompasses (forward) mirror map go and its approximate inverse map fgz They are trained so that f. = g and f. maps unconstrained points to the constrained space (see Figure ] for a conceptual illustration). Our method works for any constraint that has a differentiable function quantify the distance from an image the constraint set. We parameterize go as the gradient of a strongly input-conver neural network (ICNN) (Amos et al. (017)) to satisfy invertibility. We train the NAMM with conssscons loss (Zhu al. 2017 to ensure f. go and train the inverse map with constraint loss to ensure f. (x) is close to the constraint set for all & that we are interested (we define this formally in Section 3.T). An MDM can trained on the pushforward of the data distribution through Ba. and its generated samples can be mapped to the constraint set via fg- Although not inherently restricted to diffusion models, our approach maintains the many advantages of diffusion models, including expressive generation, simulation- free training (Song et al., 202lb). and tractable computation of probability densities (Liu et al.) 2024). One can also adapt existing diffusion- based inverse solvers for the mirror space and enforce constraints with the inverse map.
Our experiments show improved constraint satisfaction for various physics- based, geometric, and semantic constraints. We also discuss ablation studies and adapt popular diffusion-based inverse solver to solve constrained inverse problems, in particular data assimilation with PDE constraints.
'Code he found http:///ghthbb...oocm/rrttrt////
2

---

Published as a conference paper at ICLR 2025
2
BACKGROUND
2.1
CONSTRAINED GENERATIVE MODELS
Explicitly incorporating a known constraint into generative model poses benefits such as data
efficiency (Ganchev et al. 2010: Batzner et al.
2022), generalization capabilities (Köhler et al.
2020), and feasibility of samples GGannone et al. 2023). Some methods leverage equivariant neural
networks (Satorras et al. 2021| Geiger & Smidt 2022; Thomas et al. (018) for symmetry (Allingham
ct al.
2022, et a 2020 Klein et al. 20224] Song et al. 2024 Boyda et al.] 2021 Rezende et al.
2019
Garcia Satorras et al. 2021 Köhler et al.
2020 Midgley et al. [024]] Dey et al. |020 Xu et al.
2022
Hoogeboom et al. [022] Yim et al. (2023)
Yim
but do not generalize to other types of constraints or
generative models (Klein et al. |024|] Song et al. [2024, Boyda et al. 2021; Rezende et al. 2019
Garcia Satorras et al.] [022] Köhler et al. 2020; Midgley al. 2024; Niu et al.
2020
Xu al.
2022; Hoogeboon et al. 2022 Yim et al. 2023] Dey et al. 2020; Allingham el al.
2022
Previous
methods for constrained diftusion models (De Bortoli al., 2022] Huang et al.
2022 Fishman
et al. 2023a| Lou & Ermon 2023 Liu et al. (024)) make strong assumptions about the constraint,
such as being characterized as a Riemannian manifold (De Bortoli et al. 2022 Huang et al. 2022).
having well-defined reflection operator (Lou & Ermon (023) projection operator (Christopher
ctt al. 2024). or corresponding to a convex constraint set (Fishman et al. 2023a Liu et al. (022)).
Fishman et al. (2023b) proposed diffusion model that incorporates Metropolis- Hetropolis-Hastinnss steps to work with general constraints, but impractically high rejection rates may occur with constraints that challenging to satisfy, such as a PDE constraint.
An alternative approach is to introduce a soft constraint penalty when training the generative model (Ganchev et al., 2010 Mann & McCallum 2007; Chang et al. 2007 Giannone et al. 2023 Daras
et [ Upadhyay et al. 2023). However, evaluating the constraint loss of generated samples during training may be prohibitively expensive. Instead, one could add constraint-violatin training examples (Giannone et al. (2023), but it is often difficult to procure useful negative examples. In
contrast, our approach does not alter the training objective of the generative model.
2.2 MIRROR MAPS
For any convex constraint set C C Rdd, one can define a mirror map that maps from C to Idd This is cwt.ruuuuusyy-dfffenttabblleel and strongly-convex (Bubeck et al. 2015 Tan et al., 2023 .. The mirror map is the gradient V : + Rd (Liu et al. 2024). Every mirror map has inverse (Vo)-1 Rdd - C, which, unlike the forward mirror map, is not necessarily the gradient of a strongly-convex function (Zhoul 2018 Tan et al. 2023). Mirror maps have been used for constrained optimization (Beck & Teboulle 2003) and sampling (Hsieh) et al., 2018: Li et al., 2022 Liu et al. 2024).. Although true mirror maps exist only for convex constraints, we seek generalize the concept learn approximate mirror maps to handle non-convex constraints. Recent work suggested learned mirror maps for convex optimization (Tan et al. 2023) and reinforcement learning (Alfano et al. 2024) but did not tackle constrained generative modeling. Our work proposes novel training objective for learning mirror maps with the goal of constraining generative models to satisfy arbitrary constraints.
done by defining mirror potential ō C + R that is
2.3 DIFFUSION MODELS
A diffusion model learns to generate new data samples through gradual denoising process (Sohl-
Dickstein et al. 2015 Ho et al., 2020 Song & Ermon 2019 Song al. 2021b Kingma al. 2021). The diffusion, or noising, process can be modeled as a stochastic differential equation (SDE) (Song et al. 2021b) that induces a time-dependent distribution Pt, where po Pdata (the target data distribution) and Pr se (O, The diffusion model learns to reverse the denoising process by modeling the score function of Pt. defined as Vx lo Pe (x). In the context of imaging problems, the score function is often parameterized using a convolutional neural network (CNN) with parameters objective (Hyvärinen & Dayan 2005 Vincent 2011 Song al. 202 that allows for simulation-free training (i.e., simulating the forward noising process is not necessary during training). The score function appears in reverse-time SDE (Song al.] 2021b Anderson 1982) that can be used sample from the clean image distribution po by first sampling a noise image from A(O,, I) and then gradually denoising it. Our work addresses the problem that diffusion models are often not sensitive to visually-
0, which we denote by Sg. The score model is trained with a denoising
3

---

Published as a conference paper at ICLR 2025
8
noise
M
M
Figure 2: NAMM training illustration. Given data that lie on constraint manifold M (e.g., the hyperplane of images with the same total brightness), jointly train approximate mirror map go and its approximate inverse fp. After mapping data x Pdata to the mirror space as (x),, we perturb them with additive Gaussian noise whose standard deviation can be anywhere between 0 and max The inverse map fo is trained to map these perturbed samples back onto M.
constraints on the data. Our proposed NAMM allows for diffusion model to be trained in an unconstrained space yet satisfy the desired constraint by construction.
3 METHOD
We now describe neural approximate mirror maps (NAMMs). We focus diffusion models, but any generative model can be trained in the learned mirror space (see Appendix A for results with VAE). We denote the constrained image distribution Pdam and the (not necessarily convex) constraint by M C IRd Images the constrained and mirror spaces are denoted by x and &, respectively. The pushforward of the data distribution Pdata through mirror map go denoted by (go),Pdat.
3.1 LEARNING THE FORWARD AND INVERSE MIRROR MAPS
Let g& and f. be the neural networks modeling the forward and inverse mirror maps, respectively, where ó and vi are their parameters. We formulate the following learning problem:
6t 36 E arg min (Cyyue(Bp (8p. fg.) + Cecomtr-Commte (go: f&) + Aveg R R(ga)
(I)
d.v
where Layele encourages go and fo be inverses of each other; Ceonstr encourages f. map unconstrained points back to the constraint set; and R a regularization term to ensure there is unique solution for the maps. Here Aconsur Aveg E E are scalar hyperparameters..
true inverse mirror map satisfies cycle consistency and constraint satisfaction on all of Rdd so ideally (x) Bg (x) and (x) E M for all & E Rd But since it would be computationally infeasible to optimize f,y over all possible points in IRd instead optimize it over distributions that we would expect the inverse map to face in practice in the context of generative models. That is, we only need fo to valid for samples from an MDM trained on (go)yPdata which we refer as the mirror distribution. To make f. robust to learning/sampling error the MDM, we consider sequence of noisy distributions in the mirror space, each corresponding to adding Gaussian noise to samples from (go) Pdata. which, for a maximum perturbation level T we denote by
(go) (ga)) (N,0, OT))
oE(0,
We train fa to be valid inverse mirror map only for points from these noisy mirror distributions. Since we do not know a priori how much noise the MDM samples will contain, we consider all possible noise levels up to 0 Fmx for robustness (see Appendix fo2 ablation study of this choice).
We define a cycle-consistency loss (Zhu et al. 2017) that covers the forward and inverse directions and evaluates the inverse direction for the entire sequence of distributions defined in Equation 2
Layole (Bo fg)) = Exxpue
|x fo (x)(
EznN [x]B((x) + -o (g.(()) + +oz)))l o2)))11]
(3)
4

---

Published as a conference paper at ICLR 2025
Let Ccoastr IR + R20 be differentiable constraint distance that measures the distance from an input image to the constraint set. We define the following constraint loss to encourage fo to map points from the noisy mirror distributions (Equation
2 to the constraint set:
conssu (gor (gor fp)
E
Ex-pun piun
( [ecasstr (Bs(()) + o2))
(4)
To ensure unique solution, we regularize g to be close to the identity function:
R(gs) == Ex- [||x go (x)]l]
(5)
We use Monte-Cari to approximate the expectations in the objective over the noisy mirror distribu- tions with ~ (10, mux)) and approximately solve Equation with stochastic gradient descent.
Architecture We parameterize go as the gradient of an input-convex neural network (ICNN) following the implementation of Tan al. (2023). For convex constraints, this satisfies the theoretical requirement that go be the gradient of a strongly-conver function. Even for non-convex constraints, this choice brings practical benefits, we discuss in Section 4.3] We note that g is not true mirror map since M is not assumed to be convex, and g& is defined on all of IR instead of just on M. We parameterize fg as a besseet CNN similar to the one used in CycleGAN (Zhu et al. 2017).
3.2 LEARNING THE MIRROR DIFFUSION MODEL
Similarly to Liu et al. (2024), we train an MDM on the mirror distribution (g+), Pdata and map its samples to the constraind space through fg. In particular, we train a score model with the following denoising score matching objective in the learned mirror space (defined the range of go):
[|so (x()). VR(t) log por (())) x(0)) } (6)
0" E arg min E, (())) (2())
where x0) P is obtained as x0) = go gx (x)) for x(0) ~ Pdata- Here Por denotes the transition kernel from x0)) to x() under the diffusion SDE, and X(1) E R is a time-dependent weight. To sample new images, we sample x(T) ,O, run reverse diffusion in the mirror space, and map the resulting x0)) to the constrained space via fup.
3.3 FINETUNING THE INVERSE MIRROR MAP
The inverse map f. is trained with samples from the noisy mirror distributions in Equation 2 but ultimately wish evaluate f. with samples from the MDM. To reduce the distribution shift, it may be helpful finetune f.. with MDM samples. We generate a training dataset of samples from the MDM and then finetune the inverse map to deal with such samples specifically. In the following finetuning objective, we replace x ~ (Pa)a with x where is the distribution of MDM samples in the mirror space:
arg min
Ex-pum |x (g)())))
(7) Finetuning essentially tailors f. to the MDM. The original objective assumes that the MDM will Gaussian-perturbed
Eg [1xx [||x +oz (f. az))ll + Accoarkoome (fe (x oz))
sample Gaussian perturbed images from the mirror distribution, but in reality it samples from slightly different distribution. As the ablation study in 3.3tion|43]/howe finetuning is not an essential component of the method; suggest it as optional step for when it is critical to optimize the constraint distance metric.
RESULTS
We present experiments with constraints ranging from physics-based to semantic. For the considered examples, our method achieves from 38% to as much as 96% improvement constraint satisfaction upon vanilla DM trained on the same data (see Table [T]. Appendices D.I and [F] providde imple- mentation and constraint details, respectively. The following paragraphs introduce the demonstrated constraints f. For each consider image dataset for which the constraint is physically meaningful.
5

---

Published as a conference paper at ICLR 2025
Total Brightness
1D Burgers'
Divergence-free
Periodic
Count
0.069
0.100
0.1
0.007
Ours
Baseline
0.450
0.833
0.934
638
8:00
0:35
2.50
60
0.5
10
0.0
01.5
60
0b5
Reweline
Free.
0.0
1.0 0.5 1.01
1.0
100
10.5
1.0
0.0
10.5
Norm.
Constr. Dist. (4)
Figure 3: Improved constraint satisfaction. Samples from approach are nearly indistinguishable from baseline samples, yet there is a significant difference in their distances from the constraint set. The baseline DM trained on the original constrained dataset. Our approach is to train NAMM and then an MDM in the mirror space induced by ge. Samples are obtained by sampling from the MDM and then passing samples through f.p. The histograms show normalized constraint distances É of 128 samples (normalized so that each constraint has maximum of across the samples from both methods). Our results are from the finetuned NAMM. For each constraint, we made sure that the DM was trained for least as long as the NAMM, MDM, and finetuned NAMM combined.
Total brightness In astronomical imaging, even if source's 's structure is unknown priori, its total brightness, or total flux, is often well constrained (EHTC) 2019). We define f ) as the absolute difference between Ei=1 X; and the true total brightness. We demonstrate with dataset of 64 x 64 images of black-hole simulations (Wong et al. 2022) whose pixel values sum 120. While this constraint is simple warmup example, generic diffusion models perform surprisingly poorly on it.
ID Burgers' We consider Burgers' equation (Bateman 1915 Burgers 1948) for a ID viscous fluid, representing the discretized solution as an fg X fy image X, where nl and vi, are the numbers of grid points in space and time, respectively. The distance burggrss (x) compares each 1D state in the image to the PDE solver's output given the previous state (based on Crank-Nicolson
dineeiiscrttizaiion (Crank) 1947 Kidger 2022)). ). The dataset consists of 64 64 images of Crank-Nicolson
Nicolson
solutions with Gaussian random fields as initial conditions.
Divergence-free A time-dependent 2D velocity field u u= u(z, y, is called divergence-free or incompressiblee if V - = O. We define the constraint distance laiv as the -n-ormm of the divergence and demonstrate this constraint with 2D Kolmogorov flows (Chandler & Kerswell 2013 Boffetta & Hcke 2012) Rozet & Louppe 2024). We represent the trajectory of the 21D velocity, discretized space- time, as a two-channel (for both velocity components) image x with the states appended sequentially. We used jax - cfd (Kochkov et al., 2021 to generate trajectories of eight 64 x 64 states and appended them
4 pattern to create 128 x 256 images.
Periodic We consider images x that are periodic tilings of a unit cell. This type of symmetry appears in materials science, such as when constructing metamaterials out of unit cells (Ogren et al.) 2024p. We use a distance function Eperiodic that compares all pairs of tiles in the image and computes the average fl-norm of their differences. We created a dataset of 64 x 64 images (composed of 32 32 unit cells tiled in 2 fashion) using code from Ogren et al. (202)).
Count Generative models can sometimes generate an incorrect number of objects (Paiss et al. 2023 We formulate differentiable count constraint by relying on CNN to estimate the count of particular object in an image x. Note that using a neural network leads to a non-analytical and highly non-convex constraint. Letting fcNN R IR be the trained counting CNN, use the distance function & (x) := [x - č for a target count č. The dataset consists of 128 x 128 simulated images of exactly eight (8) radio galaxies with background noise (Connor al. 2022

---

Published as a conference paper at ICLR 2025
Total Brightness
ID Burgers'
Divergence-fee free
Periodic
Count
Contoaat Dist (0
Oues
Ches
Rooalls
Figure 4: Training efficiency. For each method, we clocked the total compute time during training (ignoring validation and I/O operations) and here plot the mean + std. dev. of the constraint distances l of 128 generated samples at each checkpoint. The MDM training curve ("Ours w/o FT is offset by the time took train the NAMM. The finetuning curve ("Ours")) is offset by the time took to train the NAMM and MDM and generate finetuning data. For most constraints, the DM has consistently higher constraint distance without any sign of converging to the same performance as that of the MDM. For the count constraint, the MDM performs on par with the DM, but finetuning noticeably accelerates constraint satisfaction. Each run was done on the same hardware (4x A100 GPUs).
4.1 IMPROVED CONSTRAINT SATISFACTION AND TRAINING EFFICIENCY
First and foremost, we verify that our approach leads to better constraint satisfaction than a vanilla diffusion model (DM). We evaluate constraint satisfaction by computing the average constraint distance of generated samples. Since the constraint distance is non-negative, an average constraint distance of 0 implies that the constraint is satisfied almost surely.
For each constraint, we trained NAMM on the corresponding dataset and then trained an MDM on the pushforward of the dataset through the learned gg. We show results from a finetuned NAMM, but as shown in Section 4.3] finetuning is often not necessary. The baseline DM was trained on the original dataset. Figure 3 hhighlighss that MDM samples inverted through f, are much closer to the constraint set than DM samples despite being visually indistinguishable. For the total brightness, ID Burgers', divergence-free free, and periodic constraints, there is a significant gap between our distribution of constraint distances and the baseline's. ss The gap is smaller for the count constraint, which may be due to difficulties in identifying and learning the mirror map for a highly non-convex constraint.
Furthermore, our approach achieves better constraint satisfaction in less training time. In Figure 4 plot constraint satisfaction as a function of compute time, comparing our approach (with and without finetuning) to the DM. Accounting for the time it takes to train the NAMM, our MDM achieves much lower constraint distances than the DM for the three physics- based constraints and the periodic constraint, often reaching a level that the DM struggles achieve. For the count constraint, we find that finetuning is essential for improving constraint satisfaction, and it is more time-efficient to finetune the inverse map than to continue training the MDM.
4.2 SOLVING CONSTRAINED INVERSE PROBLEMS WITH MIRROR DPS
Many methods have been proposed use pretrained diffusion model to sample images from the posterior distribution p(x y) < p(y
x x))(x) (Choi et al.. 2021 Graikos et al., 2022 Song et al. 2022 Jalal et al., 2 Kawar et al. 2022, Song al. (023), given measurements y E Th and a diffusion-model prior (x)). One of the most popular methods is diffusion posterior sampling (DPS) (Chung al. 2022). To adapt DPS for the mirror space, we simply evaluate the measurement
likelihood on inverted mirror images (x) instead of on images x in the original space.
We demonstrate mirror DPS on data assimilation, an inverse problem that aims to recover the hidden state of a dynamical system given imperfect observations of the state. In Figure [5] we show results for data assimilation of a 1D Burgers' system and a divergence- state observations, which can be essentially formulated as manoiise-nddinpaiin problem. For each case, we used mirror DPS with the corresponding NAMM-based MDM. We include two baselines: (1) vanilla DPS with the DM and (2) constraint-guider DPS (CG-DPS) with the DM. The latter incorporates the constraint distance as an additional likelihood term. As FigureBsshows our approach leads to notably less constraint violation (i.e., less deviation from the PDE or less divergence) than both baselines. Appendix [D.3 shows that our method consistently outperforms the baselines for different measuremeet-likelinoee and constraint-guidanee weights used in DPS and CG-DPS.
Kolmogorov flow given few noisy
7

---

Published as a conference paper at ICLR 2025
Observed
Mirror DPS (Ours)
Sample
0.001
Sample
0.025
B
Resiuual
Divergence
PDE
Sample
0.01
Sample
DPS Baseline
Divergence
PDE Residual
CG DPS Baseline
Sample
0:08
3
Sample
0,142
PDE Residual
Divergence
1D Burgers'
(b) Kolmogorov flow (divergence-free)
Figure 5: Data assimilation. We used the same finetuned NAMM, MDM, and DM checkpoints as in
Fig. 3 (a) Given noisy observations of the first eight states, we sampled possible full trajectories of ID Burgers' system. Our solutions have smaller deviation from the PDE than samples obtained with DPS, even those of constraint-guider DPS (CG-DPS). (b) The task to infer the full Kolmogorov flow from noisy observations of the first and last states. Our solution has significantly less divergence.
4.3 ABLATION STUDIES
Finetuning Table shows the improvement in constraint satisfaction after finetuning f. while also verifying that the generated distribution stays close to the true data distribution. We use maximum mean discrepancy (MMD) (Gretton et al. 2012) and Kernel Inception Distance (KID) (Binkowski et al. 2018) measures of distance between distributions. MMD evaluates distance in feature space defined by a Gaussian kernel, and KID uses Inception v3 (Szegedy et al., (015) features (see Appendix E for details). Finetuning does not notably change the the MDM and some cases improves while improving constraint satisfaction. Compared to a vanilla DM, our approach before and after finetuning does not lead to significantly different MMD and even gives better KID while significantly improving constraint distance.
accuracy of
Total Brightness
Burgers"
Divergence-Ire free
Periodic
Count
Ours
w/o FT
MMD (L)
CD (L)
0.0957 0.0130
0.57 1: 0.48
0.09 1 0.01
0.0096
0.0664 = 0.0027
2.51 1 0.21
0.0760
0.04 1 0.05 0.0058
0.1716 0.0062
0.71 1 0.53
KID (L)
O. 0.0026 + 0.0007
0,0010
0.0004
0,0035
0.0004
0.0018 = 0.0005
0.0367 = 0.0010
(U)
0.04
1.57
0.27
Onts
MMD (L)
0.1023 1 0.0131
0.1291 1 0.0096
0.0781
1 0.0023
0.0758 1 0.0058
0.1978 1 0.0062
(U)
0.0027
0.0008
0.0022 + 0.0004
0.0058
4 0.0006
0.0014
0.0005
0.0570
0.0014
Baseline
CD (L)
2.20 1.61
0.45 1: 0.03
6.96 1 0.27
1.04 0.08
0.56 0.44
MMD (L)
0.0956 0.0099
0.0621 + 0.0091
0..5995-
0.0533 0.00-13
0.1276 = 0.0066
KID (L)
0.0462 1 0.0016
0.2308 1 0.0028
0.0053 + 0.0007
0.0064 + 0.0004
0.1084 1 0.0015
Table I: Effect of finetuning. Constr. dist. (CD) = 100 The improvements in mean CD are (left to right, comparing "Ours" to "Baseline"): 78%. 91%, 77%, 96%, 38% for five problems. For all metrics, mean + std. dev. is estimated with 10000 samples. In terms of MMD/KID, finetuning does not significantly impact distribution-matching accuracy but improves constraint distance. DM baseline results are shown for comparison. According to MMD, the baseline gives better distribution- matching accuracy; according to KID, our approach captures the true data distribution better.
8

---

Published as a conference paper at ICLR 2025
Constraint loss There are two hyperparameters for the constraint loss in Equation [ max deter- mines how much noise to add to samples from the mirror distribution, and Acoust is the weight of the constraint loss. Intuitively, higher max means that the inverse map f must map larger regions of IR back to the constraint set, making its learning objective more challenging. We would expect ge to cooperate by maintaining a reasonable SNR in the noisy mirror distributions. Figure shows how increasing max causes (x) for X Pdatn to have larger magnitudes so that the added noise will not hide the signal. However, setting max too high can worsen constraint satisfaction, perhaps due the challenge of mapping larger region of IRd back to the constraint set. On the flip side, setting o max too low can worsen constraint satisfaction because of poor robustness of fg.. Meanwhile, increasing Aconstr for the same max leads to lower constraint distance, although there is tradeoff between constraint distance and cycle-consistency inherent in the NAMM objective (Equation
Tmax 0.001
Finex 0.1
Finax 0.5
Otaex 0.001
Free
Oraex 0.1
Oraex 0.5
Aconstr
0.01
Aconstr 0.1
Aconstr
1.0
11.
9.6
0L5
Aconste 0.01
Free.
Avonstr 0.1
Avonste
0.50 Norm. Constrant Distanee
1.20
1.00
Figure 6: Effect of max and Acunsur demonstrated with ID Burgers'. First row: max increases (keeping Acunstr = 1.0), the mirror image go (x) for x ~ Pama increases in magnitude to maintain similar SNR. Histograms show constraint distances of 128 inverted MDM samples, normalized to have a maximum of 1 across samples from all three settings. Decreasing Omax from 0.5 to 0.1 improves the constraint distances, but further lowering max 0.001 causes them to go back up. This indicates a tradeoff between robustness and performance of fg. Second row: Aconstr increases (keeping max 0.1). go does change as much as when increasing max but the constraint distances decrease (with tradeoff in cycle consistency). For all three settings, the same number of NAMM and MDM epochs was used as in Fig. but without finetuning.
Mirror map parameterization Figure compares parameterizing go as the gradient of ICNN versus as ResNet-based CNN. We demonstrate how the mirror space changes when parameterizing the forward map as ResNet-based CNN. The mirror space becomes less regularized, leading to worse constraint satisfaction of the MDM, perhaps because the MDM struggles to learn less-regularized mirror space. Thus, even when the constraint is non-convex and there are no theoretical reasons to use an ICNN, it may still be practically favorable.
Gradient of ICNN
ResNet
1D Burgers'
(X)
Frem.
30
10
0.150
0L775
(1,000
0.95
0.50
0.25
100
Divergences free
(Dun r/lo FT)
Redile (scjo FT)
Freq.
DM Baddiee
Nor.
Constraint Distance (7)
Norm Constroint Distance (?)
Figure 7: Architecture of g gradient of ICNN vs. ResNet-based CNN. Both approaches preserve visual structure in the mirror space, but the ResNet causes irregularities, such as the patch circled in yellow. The histograms show the normalized constraint distances i of 128 inverted MDM samples (without finetuning). DM histograms from Fig- Frre shown for comparison. An ICNN leads to better constraint satisfaction with fewer outliers. We trained the NAMM and MDM for the same number of epochs in Fig. B without finetuning. We found that even with finetuning, ResNet-based forward map leads to worse constraint satisfaction or noticeably worse visual quality of generated samples.

---

Published as a conference paper at ICLR 2025
5 CONCLUSION
We have proposed method for constrained diffusion models that minimally restricts the generative model and constraint. A NAMM consists of a mirror map g and its approximate inverse fy, which is robust noise added samples in the mirror space induced by gø. One can train mirror diffusion model in this mirror space to generate samples that are constrained by construction via the inverse map. We have validated that our method provides significantly better constraint satisfaction than vanilla diffusion model on physics-based geometric, and semantic constraints and have shown its utility for solving constrained inverse problems. Our work establishes that NAMMs effectively rein in generative models according to visually-subtle yet peysically-meannnghhh constraints.
ACKNOWLEDGMENTS
The authors would like to thank Rob Webber, Jamie Smith, Dmitrii Kochkov, Alex Ogren, and Florian Schäfer for their helpful discussions. BTF and KLB acknowledge support from the NSF GRFP, NSF Award 2048237, NSF Award 2034306, the Pritzker Award, and the Amazon Al4Science Partnership Discovery Grant. RB acknowledges support from the US Department of Energy AEOLUS center (award DE-SCC0113333, the Air Force Office of Scientific Research MURI on "Machine Learning and Physics-Base -aaied Modeling and Simulation" (award FA9550-20-1-0358), and Department of Defense (DoD) Vannevar Bush Faculty Fellowship (award N00014-22-1-2990;
REFERENCES
8bitmp3. hww.aaawwuuiiiananccccsaaaiaiiia http://aww.wooiiiiiooeoeeeeeee/a
com/ 8bi tmp3/
Tutoriaat http://otoooiaoaaoooooooaaaooaoa Lith- 2023.
Carlo Alfano, Sebastian Towers, Silvia Sapora, Chris Lu, and Patrick Rebeschini. Meta-leaming the mirror map in policy mirror descent. arXiv preprint arXiv:2402.051183 2024.
James Urquhart Allingham, Javier Antoran, Shreyas Padhy, Eric Nalisnick, and José Miguel Hernández-Lobato. Learning generative models with invariance symmetries. In NeurIPS 2022 Workshop on Symmetry and Geometry in Neural Representations, 2022.
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155, PMLR, 2017.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982.
Arpit Bansal, Hong- Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 843-852, 2023.
Harry Bateman. Some recent researches on the motion of fluids. Monthly Weather Review, 43(4): 163-170, 1915.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-175, 2003.
Mikotaj Binkowski, Danica J Sutherland, Michael Arbel, and Arthur Gretton. Demystifying mmd gans. arXiv preprint arXiv: 0101, 2018.
Guido Boffetta and Robert Ecke. Two-ddmensioonll turbulence. Annual review of fluid mechanics, 44:427-451, 2012.
Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez Rezende, Michael Albergo, Kyle Cranmer, Daniel C Hackett, and Phiala Shanahan. Sampling using su (n) gauge equivariant flows. Physical Review D, 103(7):074504, 2021.
10

---

Published as a conference paper at ICLR 2025
Francois-Xate Maaviirs-aavier Briol, Alessandro Barp, Andrew B Duncan, and Mark Girolami. Statistical inference for generative models with maximum mean discrepancy. arXiv preprint arXiv:1906.0553 909444000044 2019.
Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231-357, 2015.
Johannes Martinus Burgers. A mathematical model illustrating the theory of turbulence. Advances in applied mechanics, 1:171-199, 1948.
Gary Chandler and Rich R Kerswell. Invariant recurrent solutions embedded in turbulent two- dimensional kolmogorov flow. Journal of Fluid Mechanics, 722:554-595, 2013.
Ming- Wei Chang, Lev Ratinov, and Dan Roth. Guiding semi-supervision with constraint- driven learning. In Proceedings of the 45th annual meeting of the association of computational linguistics, pp. 280-287, 2007.
Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Conditioning method for denoising diffusion probabilistic models. In ICCV. IEEE, 2021.
Jacob K Christopher, Stephen Baek, and Ferdinando Fioretto. Projected generative diffusion models for constraint satisfaction. arXiv preprint arXiv:2402.03553 2024.
Hyungjin Chung. Jeongsol Kim, Michael T Mccann, Mare Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 2022.
Liam Connor, Katherine L Bouman, Vikram Ravi, and Gregg Hallinan. Deep rudio-nnerferrometic imaging with polish: Dsa-2000 and weak lensing. Monthly Notices of the Royal Astronomical Society, 514(2):2614-2666 2022.
John Crank and Phyllis Nicolson. practical method for numerical evaluation solutions of partial differential equations of the heat-conduction type. In Mathematical proceedings of the Cambridge philosophical society, volume 43, pp. 50 50-67. Cambridge University Press, 1947.
Giannis Daras, Yuval Dagan, Alex Dimakis, and Constantinos Daskalakis. Consistent diffusion models: Mitigating sampling drift by learning be consistent. Advances in Neural Information Processing Systems, 36, 2024.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. Advances in Neural Information Processing Systems, 35:2406-2422, 2022.
Neel Dey, Antong Chen, and Soheil Ghafurian. Group equivariant generative adversarial networks. arXiv preprint arXiv:2005.01 0668, 2020.
Bradley Efron. Tweedie's s formula and selection bias. Journal of the American Statistical Association, 106(496):1002-1114 2011.
ЕНTС. First m87 event horizon telescope results. iv. imaging the central supermassive black hole.
The Astrophysical Journal Letters, 875(1),L4, 2019. doi: 10.384772041182213aabee8. URL
org /tt.:///..000000211112222221
Abdul Fatir. Metrics for Evaluating GANs (Pytorch). lhttps:// http:hhu...omaaaauulllccaca
gan wwt.mmerics-pytoree March 2021.
Nic Fishman, Leo Klamer, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson. Diffusion models for constrained domains. arXiv preprint arXiv:2304.05364 2023a.
Nic Fishman, Leo Klarner, Emile Mathieu, Michael Hutchinson, and Valentin De Bortoli. Metropolis sampling for constrained diffusion models. Advances Neural Information Processing Systems, 36, 2023b.
Kuzman Ganchev, Joao Graça, Jennifer Gillenwater, and Ben Taskar. Posterior regularization for structured latent variable models. The Journal of Machine Learning Research, 1 1:2001122499 2010.
11

---

Published as a conference paper at ICLR 2025
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181-4192, 2021.
Damien Garreau, Wittawat Jitkrittumm, and Motonobu Kanagawa. Large sample analysis the median heuristic. arXiv preprint arXiv: 2017.
Mario Geiger and Tess Smidt. e3nm: Euclidean neural networks. arXiv preprint arXiv:2207.09453, 2022.
Giorgio Giannone, Lyle Regenwetter, Akash Srivastava, Dan Gutfreund, and Faez Ahmed. Learning from invalid data: On constraint satisfaction in generative models, arXiv preprint urXiv:2306.15100 2023.
Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. Advances Neural Information Processing Systems, 35:14715-14725 2022.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773) 2012.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. Veta Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017. URL
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-688 6851, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867-8887. PMLR, 2022.
Ya-Ping Hsieh, Ali Kavis, Paul Rolland, and Volkan Cevher. Mirrored langevin dynamics. Advances in Neural Information Processing Systems, 31, 2018.
Chin-Wes Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Riemannian diffusion models. Advances Neural Information Processing Systems, 35:2750-2761, 2022.
Yujia Huang, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli S Sastry, Sid- dharth Gururani, Sageev Oore, and Yisong Yue. Symbolic music generation with non -on-differentiable rule guided diffusion. arXiv preprint arXiv:2402. 14285, 2024.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros Dimakis, and Jonathan Tamir. Robust compressed sensing mri with deep generative priors. NewrIPS, 2021.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration
models. In Advances in Neural Information Processing Systems, 2022.
Patrick Kidger. On neural differential equations. arXiv preprint arXiv:2202.0243: 2022.
Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. arXiv preprint arXiv:2107.00630, 2021.
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36, 2024.
12

---

Published as a conference paper at ICLR 2025
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning-acceleraten accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 11821))c210178411 18, 2021.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pp. 5361-5370. PMLR, 2020.
Ruilin Li, Molei Tao, Santosh S Vempala, and Andre Wibisono. The mirror langevin algorithm converges with vanishing bias. In International Conference on Algorithmic Learning Theory, Pp. 718-742. PMLR, 2022.
Phillip Lippe. UvA Deep Learning Tutorials. http:://waaale.nottaoo......aaaa
Io/en/latestr 2024.
Guan-Horng Liu, Tianrong Chen, Evangelos Theodorou, and Molei Tao. Mirror diffusion models for constrained and watermarked generation. Advances in Neural Information Processing Systems, 36, 2024.
Aaron Lou and Stefano Ermon. Reflected diffusion models. ICML, 2023.
Gideon S Mann and Andrew McCallum. Simple, robust, scalable semi-supervised learning via expectation regularization. In Proceedings of the 24th international conference on Machine learning, pp. 593-600, 2007.
Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hemández-Lobato. Se (3) equivariant augmented coupling flows. Advances in Neural Information Processing Systems, 36, 2024.
Chenhao Niu, Yang Song, Jiaming Song, Shengjia Zhao, Aditya Grover, and Stefano Ermon. Permu- tation invariant graph generation via score-based generative modeling. In International Conference on Artificial Intelligence and Statistics, pp. 4474 4484. PMLR, 2020.
Alexander C Ogren, Berthy T Feng. Katherine I. Bouman, and Chiara Daraio. Gaussian process regression as a surrogate model for the computation of dispersion relations. Computer Methods in Applied Mechanics and Engineering, 420: 16661, 2024.
Roni Paiss, Ariel Ephrat, Omer Tov, Shiran Zada. Inbar Mosseri, Michal Irani, and Tali Dekel. Teaching clip to count to ten. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3170-3180, 2023.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks; deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant hamiltonian flows. arXiv preprint arXiv: /909, 13739, 2019.
François Rozet and Gilles Louppe. Score- data assimilation. Advances in Neural Information Processing Systems, 36, 2024.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pp. 9323-9332. PMLR, 2021.
Maximilian Seitzer. pytorch-fid: -irrch-fidd:a FID Score for PyTorch. Inttps://githut http://ttttuub.coommmmttteeee www.nseiizer// pytorch- fid
August 2020. Version 0.3.0.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In Int. Conf. Machine Learning, pp. 2256 2265.
PMLR, 2015.
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse- diffusion models for inverse problems. In International Conference on Learning Representations, 2023.
URL http:ppeee.peeeeeeeeeeee//e///es 3_MMd===__MMMMKKK
13

---

Published as a conference paper at ICLR 2025
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In NeurlPS, pp. 11895- 11907, 2019.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415-1422
2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021b. URL /ttphtttss//opeeeeeieeee
http://oorummmxxxx1111i/3320 1d-PXTIG12RRHS
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In ICLR, 2022. URL https::/ net
Forum?id=vaRCHv
Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and YiigYing Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36, 2024.
Danica J Sutherland. opt-mmd, https: / //tggggtttb.t///////////hh.hhh https//ww..../ccaaaaaaaaaaaaaaaa master/two_sampll mmd. py] 2018.
Danica J Sutherland, Hsiao- Yu Tung, Heiko Strathmann, Soumyajit De, Aaditya Ramdas, Alex Smola, and Arthur Gretton. Generative models and model criticism via optimized maximum mean 2016.
discrepancy, arXiv preprint
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1-9, 2015.
Hong Ye Tan, Subhadip Mukherjee, Junqi Tang, and Carola-Bibiang Schönlieb. Data-driven mirror descent with input-convex neural networks. SIAM Journal on Mathematics of Data Science, 5(2):
2023.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation -qquivirnnr neural networks for 3d point clouds. arXiv preprint arXiv: wixiv:1802008219 2018.
Rishi Upadhyay, Howard Zhang, Yunhao Ba, Ethan Yang, Blake Gella, Sicheng Jiang, Alex Wong, and Achuta Kadambi. Enhancing diffusion models with 3d perspective geometry constraints. ACM Transactions on Graphics (TOG), 12(5):1-15, 2023.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674 2011.
George N Wong, Ben Prather, Vedant Dhruv, Benjamin R Ryan, Monika Mošcibrodzka, Chi- kwan Chan, Abhishek V Joshi, Ricardo Yarza, Angelo Ricarte, Hotaka Shiokawa, et al. Patoka: Simulating electromagnetic observables of black hole accretion. The Astrophysical Journal Supplement Series, 259(2):64, 2022.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffuuson model for molecular conformation generation. arXiv preprint M.XXi:220330022233 2022.
Jason Yim, Brian I. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint ar arXiv:2302.0223
2023.
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3836-3847, 2023.
Xingyu Zhou. On the fenchel duality between strong convexity and lipschitz continuous gradient. arXiv preprint arXiv: 1803.000333 2018
14

---

Published as a conference paper at ICLR 2025
Jun- Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to -o-gage translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference computer vision, pp. 2223-223222 2017.
15

---

Published as a conference paper at ICLR 2025
Ours w/o FT (MVAE) Samples
Baseline (VAE) Samples
Vorticity
Divergence
Vorticity
Divergence
Free.
20
0.25
0.50
0.75
1.00
Norm.
Constr.
Dist. (4)
Ours w/o FT (MVAE)
11): 5.33 8.53
(1): 33.20
Baseline (VAE)
AND [1): 0.0692 + 0.002
NMD (J): 0.0696 0.0816
KID
(8): 0.1334 = 0.8022
KID (1): 0.1575 0.0025
Figure 8: Mirror VAE (MVAE) vs. VAE, demonstrated on the divergence- free constraint. Baseline samples are obtained from VAE trained in the original space, while MVAE samples are obtained by sampling from VAE trained the mirror space and then mapping those samples back to the original space via the learned inverse mirror map. Three samples divergence on the right) for each method. (Recall that each image consists of eight state snapshots; here have labeled the number of each snapshot.) The vorticity fields show that the visual statistics of both generated distributions are extremely similar, but the corresponding divergence fields are drastically different. The MVAE samples are much closer to satisfying O-divergence everywhere. As further evidence, the histograms show that normalized constraint distances of MVAE samples significantly lower. We also report the mean + std. dev. constraint distance (CD), computed as com as well as the MMD and KID. All three metrics were estimated with 10000 generated and true samples. The MVAE leads to improved constraint satisfaction and distribution matthiing accuracy compared to vanilla VAF. This experiment demonstrates how a NAMM can be used to constrain generative models besides diffusion models.
shown (vorticity
the left and
A NAMMS FOR CONSTRAINED VAES
Our approach is compatible with any generative model, not just diffusion models. Once a NAMM is trained, any generative model can be trained in the learned mirror space and its samples mapped back to the constrained space via the learned inverse mirror map. this appendix, we apply our approach to training variational autoencoder ( (VAE) that satisfies the divergence- free constraint, comparing VAE trained in the learned mirror space ("MVAE") to a VAE trained in the original data space.
A.1 IMPROVED CONSTRAINT SATISFACTION WITH MIRROR VAE
For this experiment, we trained a VAE in the mirror space induced by the learned mirror map that was trained for the divergence-free constraint (without finetuning). We call this the "mirror VAE," or MVAE, approach. As a baseline, we trained the same VAE architecture the original divergence-fiee data without transformation. We note that the same data are used to train both the MVAE and VAE; the only difference is that the MVAE is trained in the mirror space, while the VAE is trained in the original space. The training procedure was otherwise the same for both the MVAE and VAE.
Figure g shows samples from the MVAE and VAE. We ensured that the total training time of the MVAE did not exceed that of the VAE (on the same hardware, 4x A100 GPUs). The VAE was trained for 3500 epochs, and the MVAE was trained for 600 epochs (following 100 epochs of NAMM training). Both approaches produce visually similar samples, yet the images of the divergence field and histograms for this constraint distance show that the MVAE leads to overall better constraint satisfaction. Furthermore, in terms of MMD and KID, the MVAE distribution is even closer to the true data distribution.
A.2 VAE IMPLEMENTATION DETAILS
We used a convolutional autoencoder architecture consisting of five convolutional layers with GELU activation functions in the encoder and decoder. Our implementation is borrowed from the autoencoder
16

---

Published as a conference paper at ICLR 2025
Total Brightness
1D Burgers'
Divergence free
Periodic
Count
Real
Ours
0.49
9-15
[1l:
9-31
CB
0.27
KID (a0: 0.0027 e.eemt
KED (41: 0.0022
0.0064
KID 0.0855 0.0006
280014 0.0005
0.0570 00.0034
Baseline
2.20
0.27
IDO
00-0091
0.0026
9.9065
KID 0331 0.8462 0.0016
KID 5411 0.2368 0.0828
KID 2012 0.0053 0.0007
K00 (331 8.0964 0.0004
KID (331 0.1884 0.0015
Figure 9: Comparisons of generated distributions. Nine (six for the divergence-free constraint) samples are shown for each distribution. The constraint distance (CD), MMD, and KID values repeated here from Tab. forr ease of comparison. Qualitatively, our approach gives samples that are visually very similar real samples and baseline samples.
tutorial of Lippe (2024). We set the number of latent dimensions to 128 and the number of features in the first layer of the encoder to 64. For training, we followed the 3 V-VAE training objective (Higgins likelihood of training data under the VAE probabilistic model and one to minimize the KI. divergence from the latent distribution to
al. 2017), which consists of two terms: one to increase
Gaussian prior. The latter is weighted by a scalar 3 > 0. For our purposes, the maximum maximum-licelihodd term corresponds to rrean-suuaedderrro (MSE) reconstruction loss, and we set B 0.001. We
used the Adam optimizer with a learning rate of 0.0002.
B SAMPLE COMPARISONS
Figure shows random samples from approach and the baseline DM approach, corresponding the results shown in Figure 3 We show more samples in this appendix to give a sense of the visual similarity between the samples generated with our approach and those generated with the baseline approach. We emphasize again that our samples are much closer to the constraint set despite being visually indistinguishable from baseline samples.
C CONSTRAINT HYPERPARAMETERS
Here we provide an ablation study of the constraint hyperparameters, Aconstr and Tmax, for all the demonstrated constraints. We performed hyperparameter sweep across all combinations of Aconstr E [0.01, 0.1, 1.0] and mmax E [0.001, 0.1, 0.5). For each setting, trained a NAMM and then an MDM. and then we evaluated the constraint satisfaction and distribution matching accuracy of the MDM samples. Table 2 reports the constraint distance, MMD, and KID metrics for all hyperparameter settings and constraints.
These results illustrate how performance changes with respect to Aconstr and maax We observe that increasing Aconstr generally leads to lower constraint distances for any constraint. For some constraints, there is also an improvement MMD/KID Acoear increases (e.g., ID Burgers'), whereas in other
17

---

Published as a conference paper at ICLR 2025
cases, there seems to be a tradeoff between constraint satisfaction and distribution-mate accuracy (e.g., Count). We observe that the metrics change non-monotonically as a function of max. This may be attributed to the fact that the training objective has nonlinear dependence max When applying our method to a new constraint, one can run a similar hyperparameter sweep to choose values that lead to the best combination of constraint satisfaction and distribution-matching accuracy.
0.001
Tuax 0.1
0.5
CD (
0.66 + 0.38
0.55 0.45
0.66 + 0.57
Aconti 0.01
MMD (L)
0.0517 + 0.0060
0.05.19 + 0. 0.0067
0.0757 + 0.0055
Total Brightness
KTD (
0.0386 + 0.0025
0.0078
0.0014
0.0062 + 0.0008
CD (
0.03 1 0.02
0.02 1: 0.01
0.22 1 0.03
Aumuu 0.1
MMD (
1.1723 4: 0.0006
1.1723 0.0006
1.1727 +: 0.0006
KID (
0.4372 = 0.0043
0.5275 = 0.0045
0.6396 = 0.0047
CD (
0.07 + 0.02
0.03 + 0.01
0.11 + 0.01
Aamar 1.0
MMD (
1.1721 = 0.000G
1.1722 0.0006
1.1727 = 0.0006
KID (
0.5052 1 0.0044
0.5191
0.0045
0.6151 + 0.0046
(
0.20 0.03
0.27
0,26 +: 0.03
Aconti 0.01
MMD (
0.0488 1 0.0090
0.0631 + 0.0095
0.0945 + 0.0078
ID Burgers'
KID (
0.0483 + 0.0012
0.1111 0.0018
0.1142 + 0.0020
CD
0.20 + 0.03
0.20 + 0.02
0.18 + 0.02
Aumir
MMD (1)
0.0784 0.0075
0.0682
0.0071
0.1069 + 0.0062
KID (
0.0434 1 0.0014
0.0495 1 0.0016
0.0499 1 0.0015
CD (
0.20 + 0.03
0.07 0.01
0.17 + 0.02
Aconir 1.0
MMD (
0.1159 1 0.0078
0. 1092
0.0090
0.2169 1 0.0084
KID (4)
0.0494 + 0.0012
0.0026 0. 0.0001
0.0144 + 0.0008
CD (
4.17 1: 0.29
5.17 0.33
6.01 1: 0.35
Aonsu 0.01
MMD (
0.0493 + 0.0029
0.0484
0.0025
0.0522 + 0.0018
KID (L)
O. 0.0007 + 0.0002
0.0029 = 0.0004
0.0062 + 0.0008
Divergence-foe
CD
3.37 + 0.30
3.16 + 0.33
3.56 + 0.31
Aumar 0.1
MMD (
0.0393 + 0.0026
0.0569 0.0031
0.0588 = 0.0030
KID (L)
0.0029 + 0.0006
0,0009
0.0002
0.0020 + 0.0003
CD (
2.42 # 0.36
2.28 0.18
2.61 + 0.21
Aumar 1.0
MMD (
0.0554 1 0.0037
0.0548 1: 0.0031
0.0593 1 0.0019
KID (
0.0031 + 0.0005
0.0031 0.0004
0.0029 + 0.0004
CD
0.25 + 0.40
0.28 = 0.12
0.30 + 0.15
Aumar 0.01
MMD (
0.2133 4 0.0057
0.0484 0.0041
0.0507 + 0.0039
KID (L)
0.0215 + 0.0018
0.0006 = 0.0002
0.0019 = 0.0004
Periedic
CD (
0.35 + 0.29
0.15 0.09
0.35 + 0.20
Aumitr 0.1
MMD (L)
0. 0.0795 + 0.0059
0.0621 = 0. 00-15
0.1013 + 0.0060
KID (1)
0.0054 0.0008
0,0015 0.0005
0.0011 = 0.0008
CD (
0.30 + 0.11
0.05 0.06
0.00 + 0.00
Aunar 1.0
MMD (
0. 1 0.0056
0.0856 0.0057
0.7822 1 0.0037
KID (
0.0049 + 0.0007
0.0044 = 0.0009
0.8408 + 0.0020
CD
0.73 + 0.55
0.65 0.49
0.61 + 0.17
Aunar 0.01
MMD (
0. 0.0715 + 0.0057
0.0449 + 0.0030
0.2119 + 0.0067
KID (L)
0.0084 + 0.0006
0.0514 = 0. 0.0011
0.0266 + 0.0011
Count
(L)
0.13 + 0.
0.59 0.45
0.52 + 0.41
Aumiar 0.1
MMD (L)
0.1773 1 0.0023
0.0884 0.0058
0.0520 1 0.0032
KID (
0.5006 + 0.0040
0.0594
+ 0.0012
0.0086 + 0.0004
CD
0.01 1 0.01
0.02 1: 0.01
0.03 1 0.01
Avauu 1.0
MMD (L)
0.6346 4 0.0022
0.6774 0.0024
0.7988 + 0.0029
KTD (
0.4995 + 0.0018
0.4957 # 0.0020
0.4786 + 0.0016
Table 2: Ablation study of constraint loss hyperparameters. Constr. dist. (CD) 100 MMD and KID are metrics for distribution-matciin accuracy. The mean + std. dev. of each metric was estimated with 10000 samples. For each constraint, the NAMM and MDM were trained for many epochs (before finetuning) as reported in Table The hyperparameter setting used for the main paper results reported in Table is underlined for each constraint (the results here are slightly different due to randomness in the training runs). These results, which were attained without finetuning, give a sense of how performance changes with respect to Accsnus and maa for each constraint.
18

---

Published as a conference paper at ICLR 2025
D EXPERIMENT DETAILS
D.1 IMPLEMENTATION
MDM score model For training the score model sæ in the learned mirror space, we followed the implementation of Song et al. (021a). We used the NCSN+ architecture with 64 filters in the first and Fmax 20. Training was done using the Adam optimizer with learning rate of 0.0002 and gradient norm clipping with a threshold of 1.
layer and the VP SDE with Smin
NAMM For go. we followed the implementation of the gradient a strongly-conver ICNN of Tan ct al. (2023), contiguring the ICNN to be 0.9- 0.9-strongly convex. Following the settings of CycleGAN (Zhu et al. 2017). f.. was implemented as a ResNet-based generator with 6 residual blocks and 32 filters in the last convolutional layer. For all constraints except the divergence free constraint,
had the ResNet-bbsed generator output the residual image (i.e., fy (&) ResNet( (X) + x). We found inverse map (i.e., fe(()) ResNet(x))
that for the divergence-free constraint, non-residuall
achieves better constraint loss. The NAMM was trained using Adam optimizer with learning rate of 0.001 for the divergence-Fig constraint and learning rate of 0.0002 for all other constraints.
Table 3 shows the hyperparameter choices for each constraint. The regularization weight Xreg in the NAMM objective (Equation was fixed 0.001. We used ICNN layers for images 64 64 or smaller and 2 ICNN layers for images 128 x 128 or larger for the sake of efficiency. These hyperparameter values do not need be heavily tuned, as we chose these settings through coarse parameter search (e.g., trying Aconstr 0.01 X = 1 to see which would lead to reasonable loss curves).
Num. ICNN layers
Omax
Aconstr
Total Brightness
3
0.1
0.01
ID Burgers'
0.1
1
Divergence-free
2
0.5
1
Periodic
0.1
Count
2
0.1
0.01
Table NAMM hyperparameter values for each constraint in our experiments.
The main results shown in Figure 3 were taken from the finetuned NAMM, ensuring that the total training time of the NAMM, MDM, and finetuning did not exceed the total training time of the baseline DM. While kept track of the validation loss, this was not used determine stopping time. We found that the NAMM training and MDM training were not prone to overfitting, so we chose the total number of epochs based on observing a reasonable level of convergence of the loss curves. We found that some overfitting is possible during finetuning but did not perform early stopping. All results were obtained from unseen test data because we fed random samples from the MDM into the inverse map and made sure not to use the same random seed as the one used to generate finetuning data. For all constraints, we generated 12800 training examples from the MDM for finetuning.
Table details the exact number of training epochs for each constraint. Figure in the main results compares the constraint distances of our method without finetuning, our method with finetuning, and the baseline DM as function of compute time.
Mirror map parameterization ablation study For the comparison of parameterizing the mirror map as the gradient of an ICNN versus as ResNet-basee generator (Figure 7 we used ResNet- based generator that outputs the residual image. This means that the inverse mirror map was parameterized as residual- based network (f.p (x) ResNet(x;; v) + x), and so was the ResNet- based forward mirror map (go (x) = ResNet(x; Φ) x).
D.2 A BLATTON OF SEQUENCE OF NOISY MIRROR DISTRIBUTIONS
Recall that the NAMM training objective involves optimizing f. over the sequence of noisy mirror distributions defined in Equation 2 Thus instead of considering single noisy mirror distribution
19

---

Published as a conference paper at ICLR 2025
NAMM epochs
(before FT)
MDM epochs
FT epochs
DM epochs
Total Brightness
30
200
1000
450
ID Burgers'
100
300
700
1500
Divergence-fres
100
700
500
2000
Periodic
50
300
1000
1000
Count
50
300
700
1500
Table 4: Number of training epochs of the NAMM, MDM, finetuning, and DM used for the results
in Fig. 3 These were chosen so that our method (including the NAMM training, MDM training, finetuning data generation, and finetuning) did not take longer train than the DM.
Before Finetuning
After Finetuning
Fixed
30
Integral (Ours)
Count
0.00015
0,0011
3,0015
B1,005
0.0010
BU0055
Constraint
Figure 10: Fixed a vs. integrating over [0, m- The NAMM objective involves optimizing fg over the sequence of noisy mirror distributions defined in Eq. We compare this approach of integrating over a [0, max] to setting a fixed noise standard deviation o = max in the context of the ID Burgers' constraint (here max 0.1). Both before and after finetuning, the constraint distances fourgns of inverted MDM samples are smaller the NAMM was trained with varying a. The histograms show the constraint distances of 128 samples from each method.
(Pa)y (O, I² with fixed noise level we perturb samples from (Bo), with varying levels of Gaussian noise with variances ranging from to o² max. Figure 10 we compare this choice, which involves integrating over [0, maxx], to the use of a fixed 0mxx
D.3 DATA ASSIMILATION WITH MIRROR DPS
Following the original DPS CChung et al. 2022), we use a hyperparameter Ç E Ryo to re-weight the time-dependent measurement likelihood. At diffusion time i, the measurement weight is given by
<(t) == </ lly A(xo)llr:
(8)
where Xo= = fa (E]x(0) {())). Here we assume that the measurement process has the form
y = A(x") +n, n ~~J I"
(9)
for some unknown source image x" E Rid where A: IR + R. is a known forward operator, and m X m is the known noise covariance matrix T. Higher values of Ç impose greater measurement consistency, but setting ( too high can cause instabilities and artifacts. The data assimilation results in Figure 5 Jueed ( = 0.1 and constraint-gui guiddacce weight equal to 200. Figure 11 shows results for the same tasks but different values of in DPS and the constraint constraint-guidance weight in CG-DPS.
E MEASURES OF DISTANCE BETWEEN DISTRIBUTIONS
E
MMD
The maximum mean discrepancy (MMD) between two distributions is computed by embedding both distributions into a reproducing kernel Hilbert space (RKHS) and using samples to estimate the resulting distance. We use the popular Gaussian radial basis function (RBF) kernel to construct the RKHS, setting the length scale as
median (1xx(0-x(0111) /2,
(10)
20

---

Published as a conference paper at ICLR 2025
Mirror DPS (Ours)
Sample
0006
Sample
(= 0.05
Resiual
PDE
Diva
Mirrr DPS (Ours)
Sample
0 1000
Sample
{ == 0.1
Resiival
Divergerca
FDE
Mirror DPS (Ouns)
Sample
0 0006
Sample
0.5
PDE Residual
Dhanganz
0000
Sample
Sample
{ = 0.05
DPS
FDE Resioval
Divergenoe
Sample
Sample
< 0.1
DPS
MDE Residial
Dianginca
Sampie
Sample
=
DPS
POE Resicual
Divergence
Sample
0274
Sample
CG CG-DPS.5T 5.50
= 0.1
PDE Resiual
Divagencc
Sumple
Sample
CG- DE-DPS. 100 X 0.1
PDE Resicaal
Divergance
Sample
Divergnnce Sample
CG.DPS-20 200
{ = 0.1
Resiial
PDE
(a) 1D Burgers
(b) Kolmagorav flow (divergence free)
Figure 11: Data assimilation results for different values Ç and constraint constraint-guidance strengths. E.g., "CG-DPS-50" refers to CG-DPS with constraint -uidannce weight of 50. These highlight that even for different values of Ç. the constraint errors (i.e., PDE residual or absolute divergence) of our solutions much smaller than those of the baseline solutions. Furthermore, changing the constraint-guidance weight of CG-DPS does not meaningfully change the level of constraint satisfaction.
21

---

Published as a conference paper at ICLR 2025
Le., the square root of half the median of the pairwise squared Euclidean distances in the dataset
{xx+))" 1==. This is popular choice in previous work (Briol al. 2019) and has been theoretically and empirically justified Garreau et al. (2017).
Our MMD implementation is based on the code provided for the work of Sutherland et al. (2016) (Sutherland 2018). We estimated mean and standard deviation empirically with 50 random subsets of 1000 samples from each dataset. The length scale was estimated (Equation T0 for each subset using the samples in the true subset. In total, the generated and true datasets contained 10000 samples each. Held-out test images were used true samples for all constraints, except for total brightness (due to a lack of test images in the dataset, training images were used for this constraint only).
E.2 KID
The Kernel Inception Distance (KID) between two distributions is based the Inception v3 fea- tures evaluated for samples from both distributions. Following standard practice, we use the 2048- dimensional final average pooling features. The KID is computed as the squared MMD (using polynomial kernel) between the two embedded distributions. It has several advantages over the Fréchet Inception Distance (FID) (Heusel al., 2017). including being unbiased and more sample- efficient (Binkowski et al. 2018). We note, however, that the Inception network was trained on natural images, so both KID and FID are not perfect metrics for the types of data we consider this work, such as physics-based based simulation outputs.
KID evaluation is based on the gan-met cs -pytorch repository (Fatir 202T). using the same Inception v3 weights those used in the official TensorFlow implementation of FID (Heusel et al., 2017). We evaluated KID with the same samples that were used for MMD. Since the Inception network takes RGB images input, represented the samples as grayscale images converted to RGB. For all the constraint datasets except for the divergence- free Kolmogorov flows, we clipped the image pixel values to [0, 1] before converting them to RGB. For the divergence-free data, clipped the values in the vorticity images to [-20, 20) and then rescaled this range to [0, 1].
E.3 FID
Although FID is biased finite-sample estimator and heavily depends the number of samples, is a popular metric for evaluating generative models. Table 5 reports FID values in addition to the MMD and KID values in Table We find that the rankings of method before finetuning, our method after finetuning, and the baseline method are consistent with the rankings when using KID in Table ! Again we that FID and KID based on Inception features that were tuned to natural images, so they are not the most reliable measures of distance between distributions for these particular image datasets.
Total Brighteess
ID Burgers'
Divergence-free
Periodic
Count
Ours w/o FT
5.02
3.36
2.21
2.26
24.18
Ours
4.58
2.04
3.97
1.99
37.65
Baseline
42.25
140..70
3.41
6.42
69.76
Table 5: FID values (4) associated with the comparisons in Tab. According to FID, our approach consistently outperforms the baseline DM approach in matching the true distribution. Finetuning even sometimes improves FID while improving constraint satisfaction (see Tab. [ for constraint distances).
Our FID implementation is borrowed from the pytorch-fi fid codebase (Seitzer) 2020). Per standard practice, we estimated FID with 50000 generated samples and 50000 true samples. As was the case with MMD and KID, held-out test images were used as true samples, except that training images were used for the total brightness data. We pre-processed the samples for the Inception network in the same way we did for KID.
22

---

Published as a conference paper at ICLR 2025
F CONSTRAINT DETAILS
Here we provide details about each demonstrated constraint and its corresponding dataset.
Total brightness The total brightness, total flux, of a discrete image E IR is simply the sum of its pixel values: Vxx) = Σ X4. We use the constraint distance function
L == [V() V
where E R20 is the target total brightness. The dataset used for this constraint contains images syggeeto-hydoddyamiii (GRMHD) simulations (Wong al.] (022) of Sgr
from general relativistic A with a fixed field of view. The images (originally 400 x 400) were resized to 64 x 64 pixels
and rescaled to have a total flux of 120. The dataset consists of 100000 training images and 100 validation images.
ID Burgers' Burgers' equation (Bateman 1915; Burgers 1948 is nonlinear PDE that useful model for fluid mechanics. We consider the equation for viscous fluid u = u(t, x) in one-dimensional space:
Ou
du
pu Ox.
01 + u
(11)
8x
where u(0,xx) is some initial condition uz(x), and v Rzo the viscosity coefficient. We use Crank-Nicolson (Crank & Nicolson 1947) to discretize and approximately solve Equation bT representing the solution on an n2 X n, grid, where n, the spatial discretization, and n, the number of snapshots in time. Given an nt n4 image, we wish to verify that it could be solution to Equation 11 with the Crank-Nicolson discretization. Letting x E x denote the 2D image, we solver:
formulate the following distance function for evaluating agreement with the
The 2
e burgerss
x) The
Σ t=
|xx[ 1
fc.N (x(x(x(
where SC.N : TR - nt outputs the snapshot at the next time using Crank-Nicolson, and Pythonic notation is used for simplicity. Note that -iiteedifffeeenes loss as proposed for physics-informed neural networks (PINNs) Raissi et al. (2019) would also work, but then our data would have non- negligible constraint distances since Crank-Nicolson solutions do not strictly follow a low-order finite-dite difee-diffeeenes approximation.
Using Crank-Nicolson solver (Crank & Nicolson 1947
implemented with Diffrax (Kidger 2022
we numerically solved the ID Burgers' equation (Equation II with viscosity coefficient // 0.5. The initial conditions were sampled from a Gaussian process based on Matérn kernel with smoothness parameter 1.5 and length scale equal to 1.0. We discretized the spatiotemporal domain into a 64 x 64 grid covering the spatial extent x E [0, 10] and time interval E [0, 8]. We ran Crank-Nicolson with time step of At 0.025 and saved every fifth step for total of 64 snapshots. We followed this process to create our ID Burgers' dataset of 10000 training images and 1000 validation images.
Divergence-free
The study of fluid dynamics often involves incompressible, divergence- fluids. Letting u = u(r, 1) be the time-dependent trajectory of a 2D velocity field, the divergence- free constraint says that V trajectories two-channel (for the two velocity components) images showing each n. ny snapshot for a total of n, snapshots. Such an image x has a corresponding image of the divergence field div(x). which has the same size as and represents the divergence of the trajectory. We formulate the following distance function that penalizes non-zero divergence:
duz + du, 0. We assume an nt. x nl spatial grid and represent by
Laiv (x) = |div(x)ll
We created dataset of Kolmogorov flows, which satisfy a Navier-Stokes PDE, to demonstrate the divergence-free constraint. The Navier-Stokes PDEs are ubiquitous in fields including fluid dynamics, mathematics, and climate modeling and have the following form:
Ou
--v
V Re
Vp+fV == 0,
0t
23

---

Published as a conference paper at ICLR 2025
where ((2 t. is the 2D velocity field at spatial location (z, y) and time t. Re is the Reynolds number, p is the density, p is the pressure field, and is the external forcing. Following Kochkov et al. (2021) and Rozet & Louppe )2024). set Re 10a p 1, and corresponding to Kolmogorov forcing (Chandler & Kerswell 2013, Boffetta & Ecke 2012) with linear damping. We consider the spatial domain [0, 2 with periodie boundary conditions and discretize it into 64 64 uniform grid. We used jax- (Kochkov et al., (021) to randomly sample divergence-free spectrally filtered initial conditions and then solve the Navier-Stokes equations with the forward Euler integration method with At = 0.01 time units. We saved snapshot every 20 time units for total of 8 snapshots in the time interval [3. 4.6). We represent the solution two-channel 128 x 256 image showing the snapshots in left-to-right order. In total, the dataset consists of 10000 training images and 1000 validation images.
Periodic Assuming the constraint that every image is periodic tiling of Ttiles unit cells, formulate the following constraint distance for a given image x:
Dilles
Thcles
periodic (x)
Thiles
which compares each pair of tiles, where (x) denotes the i-th tile in the image. For our experiments, we consider x 32 unit cells that tiled in a 2 pattern to create 64 64 images. Using the unit-cell generation code of Og et al. (2024), we created a dataset of 30000 training images and 300 validation images.
Count For the count constraint, we rely CNN to estimate the count of particular object. Letting fcNN Rd + IR the trained counting CNN, we turn to the following constraint distance function for target count č:
fcomm == [| a|
We demonstrate this constraint with astronomical images that contain a certain number galaxies. In particular, simulated 128 128 images of radio galaxies with background noise Connor et al. (2022), each which has exactly eight (c = 8) galaxies with an SNR 2 15 dB. The dataset consists of 10000 training images and 1000 validation images.
To train the counting CNN, we created a mixed dataset with images of 6, 8, 9, or 10 galaxies that includes 10000 training images and 1000 validation images for each of the five labels. The CNN architecture was adapted from simple MNIST classifier (8bitmp3 2023) with two convolutional layers followed by two dense layers with ReLU activations. The CNN was trained minimize the mean squared error between the vellev estimated count and the ground-trus count.
24


Reviews:
[{"review_id": "5ggjwEPeoL", "reviewer": "Reviewer_y9Gb", "summary": "The paper proposes neural approximate mirror maps for constraint data generation with diffusion models. Compared to typical mirror diffusion models, the mirror map is parameterized by the gradient of ICNN and learned via penalizing the differentiable constraint distance, thereby being applicable to general non-convex constraints. The forward and inverse mirror maps are learned by a combination of cycle consistency loss, constraint loss and regularization loss. Experiments in several settings ranging from physics-based to semantic demonstrate the effectiveness on constraint satisfaction, training efficiency and constrained inverse problem solving.", "questions": "- The regularization loss is introduced to ensure a unique solution according to the paper. What is the meaning of unique? As the mirror map is parameterized as the gradient of ICNN, the reversibility is already ensured.\n- Is the method scalable? For example, can it be applied to the image settings in Reflected Diffusion Models?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method is applicable to more general constraints than previous works.\n- The cycle-consistency loss tailored for mirror maps and diffusion models is sound.\n- The experiments are conducted on diverse settings, including constrained DPS.\n- The ablation studies are comprehensive.", "weaknesses": "- The experiments are primarily toy. It is not clear whether the proposed method can scale to high dimensions and apply to domains such as images.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nshFaZ245r", "reviewer": "Reviewer_iSD2", "summary": "This paper presents neural approximate mirror maps (NAMMs) to enforce soft constraints for diffusion models (DMs). NAMMs employ two neural networks to learn a mirror map that maps constrained points into the mirror space, and its inverse that transforms data back to the constraint set. A mirror diffusion model (MDM) can be trained in the learned mirror space, and its generated samples can be mapped to the constraint set via the inverse map. This method is tested on five benchmark problems, ranging from physics-based, geometric to semantic constraints, and the results show the proposed method improve constraint satisfaction compared to a vanilla unconstrained DM. And, this paper also demonstrates NAMMs leads to less constraint violation when solving constrained inverse problems.", "questions": "1.\tIn this paper, the robustness of the model is enhanced by introducing noise into the mirror space. What does robustness mean here, and are there any experimental results that support the robustness of the method?\n2.\tIn order to fully demonstrate the performance of the method, it is necessary to supplement its experiments on five benchmark problems and an additional baseline model.\n3.\tIf fine-tuning is considered to be one of the important components of the method and one of the contributions of this paper, more experimental support is needed.\n4.\tOn line 1097, there is a clerical error, “a la”.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. NAMMs generalize the concept of true mirror maps to learn approximate mirror maps to handle non-convex constraints.\n2. NAMMs can handle physics-based, geometric and semantic constraints, while existing methods are restricted in the types of constraints they can handle.\n3. NAMMs not only help diffusion models, but also help VAEs to improve constraint satisfaction, showing the potential to be compatible for other generative models. And NAMMs are also helpful to diffusion-based inverse-problem solvers for solving constrained inverse problems.", "weaknesses": "1. Theoretically, NAMMs lack the guarantee of the existence and uniqueness of the mirror maps when applied to non-convex problems.\n2. The proposed method is validated on five benchmark problems in the main experiments to show the superiority of applying NAMMs to diffusion models in generating constrained data. However, in the experiments to solve inverse problems, ablation experiments, and the experiments applied to VAE, this method is only carried out on partial problems and does not fully demonstrate its performance on the three types of constraints mentioned.\n3. Finetuning is an important part of the proposed method introduced in section 3. But, it is mentioned in subsection 4.1 that “We show results from a finetuned NAMM, but as shown in Section 4.3, finetuning is often not necessary”. Moreover, in the ablation studies of constraint loss and mirror map parameterization, and experiments about the VAE, fine tuning is not used.\n4. The basic unconstrained model is used for comparison, but the comparison with another existing methods dealing with constraints is lacking.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uy0lmVvApG", "reviewer": "Reviewer_FTXm", "summary": "This paper proposes a technique to learn a mirror map for general constrained diffusion generation. The resulting neural approximate mirror map transforms the constrained problem domain to an unconstrained space, where the diffusion model is trained. The training loss encourages that the inverse of the mirror map lies within the constraint set. Numerical experiments on various problems demonstrate the efficacy of the proposed technique in enforcing the constraints.", "questions": "- It appears in Fig. 6 that constraint distance is fairly sensitive to $\\sigma_{max}$ and $\\lambda_{constr}$. Can authors propose systematic ways to tune the hyperparameters? How does the sensitivity depend on the dimensionality of the problem or the complexity of the constraint?\n- Minor comment: the citations in Section 2.1 take up a large portion of the paragraph somewhat reducing readability (top of page 3).\n- Minor comment 2: Fig. 2 is never directly referenced in the text.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method is well motivated. Proper constraint satisfaction is challenging in diffusion generation, limiting many important applications in engineering, physics and computer vision. \n\n- The proposed approach is sensible, and to the best of my knowledge novel. It improves the flexibility of mirror diffusion models by obviating the need for analytical mirror maps.\n\n- The experimental results are promising on a wide array of problems. \n\n- The paper overall is well-written and easy to follow.", "weaknesses": "- It is unclear how complex of a constraint the proposed method can handle, as the gap between NAMM and the baseline is less significant for the semantic problem. \n- It is unclear if there is a systematic way to tune the introduced hyperparameters, and how sensitive the performance is in higher dimensions to $\\sigma_{max}$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)