Mirror Diffusion Models for Constrained and Watermarked Generation
Authors: Guan-Horng Liu
Tianrong Chen
Evangelos Theodorou
Molei Tao
Venue: NeurIPS 2023 poster (2023)
Keywords: diffusion models
constrained generation
constrained manifold
mirror map
watermarked generation
generation privacy
TL;DR: A new class of dual-space diffusion models that generate samples on constrained sets/manifolds while retaining all tractability as in standard Euclidean-space diffusion models

Abstract: Modern successes of diffusion models in learning complex, high-dimensional data distributions are attributed, in part, to their capability to construct diffusion processes with analytic transition kernels and score functions. The tractability results in a simulation-free framework with stable regression losses, from which reversed, generative processes can be learned at scale. However, when data is confined to a constrained set as opposed to a standard Euclidean space, these desirable characteristics appear to be lost based on prior attempts. In this work, we propose Mirror Diffusion Models (MDM), a new class of diffusion models that generate data on convex constrained sets without losing any tractability. This is achieved by learning diffusion processes in a dual space constructed from a mirror map, which, crucially, is a standard Euclidean space. We derive efficient computation of mirror maps for popular constrained sets, such as simplices and $\ell_2$-balls, showing significantly improved performance of MDM over existing methods. For safety and privacy purposes, we also explore constrained sets as a new mechanism to embed invisible but quantitative information (i.e., watermarks) in generated data, for which MDM serves as a compelling approach. Our work brings new algorithmic opportunities for learning tractable diffusion on complex domains.

Mirror Diffusion Models for Constrained and Watermarked Generation
Guan-Horng Liu, Tianrong Chen, Evangelos A. Theodorou', Molei Tao Georgia Institute of Technology, USA
{ghliu, tianrong.cher evangelos. theodorou, mtao) edu
Abstract
Modern successes of diffusion models in learning complex, high-dimensional data distributions are attributed, in part, to their capability to construct diffusion pro- cesses with analytic transition kernels and score functions. The tractability results in a simulation-fis framework with stable regression losses, from which reversed, generative processes can be learned at scale. However, when data is confined a constrained set as opposed to standard Euclidean space, these desirable character- isties appear to be lost based on prior attempts. this work, we propose Mirror Diffusion Models (MDM), new class of diffusion models that generate data on convex constrained sets without losing any tractability. This is achieved by learning diffusion processes in a dual space constructed from a mirror map, which, crucially, is a standard Euclidean space. We derive efficient computation of mirror maps for popular constrained sets, such as simplices and f--balss, showing significantly improved performance MDM over existing methods. For safety and privacy purposes, also explore constrained sets but quantitative information (i.e., watermarks) in generated data, for which MDM serves compelling approach. Our work brings algorithmic opportunities for learning tractable diffusion on complex domains.
mechanism to embed invisible
Introduction
L- 0.1
-- 0.3
AFHOV2
FRIQ
(MManfold
Corresponding
J
MDM: learss tractable diffusions in dual spaces
Figure 1: Mirror Diffusion Models (MDM) a new class of diffusion models for convex constrained manifolds M C Rd. (left) Instead of learning score-approximate diffusions on M, MDM applies a mirror map Vó and learns tractable diffusions in its unconstrained dual-space (Φ(M) = R (right) We also present MDM for watermarked generation, where generated contents (e.g., images) live in a high-dimensional token constrained set M that is certifiable only from the private user.
Diffusion models [1-3] have emerged powerful generative models with their remarkable successes synthesizing high-fidelity data such as images [4, 5], audio [6. 7], and 3D geometry [8, 9]. These models work by progressively diffusing data to noise, and learning the score functions (often
Equal advising.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Table 1: Comparison of different diffusion models for constrained generation on M C R Rather than learning reflected diffusions on .M with approximate scores, our Mirror Diffusion constructs a mirror map Ve and lift the diffusion processes to the urconstrained dual space (ø(M) = Rd inheriting favorable features from standard Fuclidean-space diffusion models. Constraints satisfied by construction via the inverse map Vø*; ; that is, V (y) E
M for all E Vo(M).
Tractable conditional score
Simulation-free training
Regression objective
Diffusion models
Domain
Reflected Diffusion
Fishman et al. [18]
M
x x
x
x
Lou and Ermon [25]
M
A
/
Mirror Diffission
This work (ours)
VΦ(M)
parameterized by neural networks) to reverse the processes [10]: the reversed processes thereby provide transport maps that generate data from noise. Modern diffusion models [11-13] often employ diffusion processes whose transition kernels are analytically available. This characteristic enables simulation-free training. bypassing the necessity to simulate the underlying diffusion processes by directly sampling the diffused data. It also leads to tractable conditional score functions and simple regression objectives, facilitating computational scalability for high-dimensional problems. In standard Euclidean spaces, tractability can be accomplished by designing the diffusion processes as linear stochastic differential equations, but doing for non-Euclidean spaces is nontrivial.
Recent progress of diffusion models has expanded their application to non- non--uulideann spaces, such as Riemannian manifolds [14, 15], where data live in curved geometry. The development of generative models on manifolds has enabled various new applications, such as modeling earth and climate events [16] and learning densities on meshes [17]. In this work, we focus on generative modeling on constrained sets (also called constrained manifolds [18]) that are defined by a set of inequality constraints. Such constrained sets, denoted by M C Rdd, are also ubiquitous in several scientific fields such as computational statistics [19], biology [20]. and robotics [21, 22].
Previous endeavors in the direction of constrained generation have primarily culminated in reflected diffusions [23, 24], which reflect the direction at the boundary aM to ensure that samples remain inside the constraints. Unfortunately, reflected diffusions do not possess closed-form transition kernels [18, 25], thus necessitating the approximation of the conditional score functions that can hinder learning performance. Although simulation-free training is still possible", it comes with computational overhead due to the use of geometric techniques [25]. From a theoretical standpoint, reflected Langevin dynamics are widely regarded for their extended mixing time [26, 27). This raises practical concerns when simulating reflected diffusions, as prior methods often adopted special treatments such thresholding and discretization [18, 25].
An alternative diffusion for constrained sampling is the Mirror Langevin Dynamics (MLD [28]). ML.,, as subclass of Riemannian Langevin Dynamics [29, 30), is tailored specifically to convex constrained sampling. Specifically, MLD constructs a strictly convex function Φ so that its gradient map, often referred to as the mirror map VΦ M Edd defines nonlinear mapping from the initial constrained set to an unconstrained dual space. Despite extensive investigation of MLD [31 131-34,, its potential as foundation for designing diffusion generative models remains be comprehensively examined. Thus, there is a need to explore the potential of MLD for designing new, effective, and tailored diffusion models for constrained generation.
With this in mind, we propose Mirror Diffusion Models (MDM), a new class of diffusion generative models for convex constrained sets. Similar to MLD, MDM utilizes mirror maps Vø that transform data distributions on M to its dual space (a(M). From there, MDM learns a dual-space diffusion between the mirrored distribution and Gaussian prior in the dual space. Note that this is in contrast to MLD, which constructs invariant distributions on the initial constrained set; MDM instead considers dual-space priors. Since the dual space is also unconstrained Euclidean space ( = R&d, MDM
Louu and Ermon [25] proposed semi- simulation free training with the aid of geometric techniques.

---

User-defined tokens
External uscrs
Dataset
V
Dual-Space Diffusion Models
V Watermark
Constrained Manifold
Pretrained (primul-space)
M
(x E o s ax vi
Diffusion Models
Figure 2: MDM for watermarked generation: (left) We first construct constrained set M based set of user-defined tokens private to other users. (right) MDM can be instantiated by either learning the corresponding dual-space diffusions, projecting pretrained, i.e., unwatermarked, diffusion models onto M. In both cases, MDM embeds watermarks that are certifiable only from the user.
preserves many important features of standard diffusion models, such as simulation free training and tractable conditional scores, which yields simple regression objectives; see Table 1. After learning the dual-space distribution, MDM applies the inverse mapping Vø transforming mirrored samples back to the constrained set. We propose efficient construction of mirror maps and demonstrate that MDM outperforms prior reflected diffusion models on common constrained sets, such as f--balls and simplices.
Moreover, we show that MDM also stands a novel approach for watermarked generation, technique that aims to embed undetectable information (i.e., watermarks) into generated contents prevent unethical usage and ensure copyright protection. This newly emerging goal nevertheless attracted significant attention in natural language generation [35-37]. while what we will present is general algorithmic strategy demonstrated for image generations.
To generate watermarked contents that are private to individual users, we begin by constructing a constrained set M based user-defined tokens (see Figure 2). MDM can then be instantiated to learn the dual-space distribution. Alternatively, MDM can distill pretrained diffusion models [38] by projecting the generated data onto the constraint contents that can be certified only by users who know the tokens. Our empirical results suggest that MDM serves as a competitive approach to prior watermarking techniques for diffusion models, such as those proposed by Zhao et al. [39], by achieving lower (hence better) FID values.
In either case, MDM generates watermarked
In summary, present the following contributions:
We introduce Mirror Diffusion Models (MDM), a new class of diffusion models that utilizes mirror maps enable constrained generation with Fuclidean-space diffusion models.
We propose efficient computation of mirror maps common constrained sets, including balls and simplices. Our results demonstrate that MDM consistently outperforms previous reflected diffusion models.
As novel application of constrained generation, we explore MDM's potential for water- marking and copyrighting diffusion models where constraints serve as private tokens. We show that it yields competitive performance compared to prior methods.
2 Preliminary on Euclidean-space Diffusion Models
Notation MC dd denotes convex constrained set. We preserve E .M and y E Rd to better distinguish variables on constrained sets and standard d-dimensional Euclidean space. I E gddd denotes the identity matrix. I; denotes the i-th element of a vector æ. Similarly, [A]é denotes the the element of matrix at i-th row and j-th column.
Diffusion process Given E Rd sampled from some Euclidean sepiideen--pace distribution, conventional diffusion models define the diffusion process as a forward Markov chain with the joint distribution:
T
q(vi.1000 (90) = II ( 1u-1), q(ulb1-1) := N(v /1 1, 1=1
(1)
which progressively injects Gaussian noise to yo such that, for sufficiently large T and a properly chosen noise schedule 3, e (0, 1), yr approaches approximate Gaussian, i.e., q(vr) se I).
3

---

Equation (1) has tractable marginal densities, and two of which that are of particular interest are:
Q(va vo) =N(yi Varyoo (1 -z) ((V lyt, vo)
1: (yt yo), B)
(2)
(yt,, yo) v yo v1 Be ā yts
B. 1 04-1 Be,
II (1 (.)).).
Ot
Okt
0t
The tractable marginal ((0))30) enables direct sampling of y without simulating the forward Markov chain (I), i... Ve Mo can be sampled in simulation-free manner. It also suggests a closed-fom score function V log (Bo)))) The marginal q(3--1 M3. yo) hints the optimal reverse process given Yo-
Generative process The generative process, which aims to reverse the forward diffusion process (I), is defined by backward Markov chain with the joint distribution:
po (Vo:T) (U())) II po po(ve-1lye), po(ue- 199)) Ni 1), 3,.).
Here, the mean μel is typically parameterized by some DNNs, e.g. U-net [40], with 0. Equations (I) and (3) imply an evidence lower bound (ELBO) for maximum likelihood training, indicating that the optimal prediction of (30 t) should match ju, (1t, Uo) in Equation (2).
Parameterization and training There exists many different ways to parameterize (o t), and each of them leads distinct training objective. For instance, Ho al. [2] found that
Br
(o t) = VT-
VI-ai (3(
(4)
achieves better empirical results compared to directly predicting or tı, This parameterization aims predict the "noise" injected to yı ((ya the ELBO reduces simple regression:
(())= = y.,, [A(()]l= (v),
where yt Vāzyo VI- QrE, e ~ .V I), and sampled uniformly from {1, T). The
weighting X(t) is often set to 1. Once eg is properly trained, we can generate samples in Fuclidean space by simulating Equation (3) backward from yr ~ (00 I).
3 Mirror Diffusion Models (MDM)
In this section, we present a novel class of diffusion models, Mirror Diffusion Models, which are designed to model probability distributions Paam(>) (x:) supported on convex constrained set M Rd while retraining the computational advantages of Euclidean-spart diffusion models.
3.1 Dual-Space Diffusion
Mirror map Following the terminology in Li et al. [34), let ø M - R be a twice differentiable function that strictly convex and satisfying that lim || [V)())] + oo and V Rd, We call its gradient map Vá M R2 the mirror map and its image, Vo(M) Rd. as its dual/mirror
space. Let o² Ied - IR be the dual function of defined by
Φ (v) = sup (x, v) ((x)
(6)
A notable property of this dual function d' (v) is that its gradient map reverses the mirror map. In other words, we have VØ' = (Vo)=² implying that
VΦ (Vo(z)) for all E M, Vo(V (v)) for all Vo(M) Rd.
The functions (Vo, Vo*) define nonlinear, yet bijective, mapping between the convex constrained set M and Fuclidean space RRd. As a result, MDM does not require building diffusion models on M, as done in previous approaches, but rather they learn standard Fuclidean-space diffusion model in the corresponding dual space vΦ(M) _dd
We note that standard mirror maps require twice differeetiability to induce a (Riemannian) metric on the mirror space for constructing MLD. For our MDM, however. twice differentiability is needed only for Equation (8), and continuous differentiabilltyy ie.. Cd,, would suffice for training.

---

Ball Jr
Ball 9g(.A!
Singuex
Simplex VALAN
0.2
0.4
0.5
04
0.2
18
0.0
02
0A
0.6
0.8
10
5
Figure 3: Examples of how mirror maps VΦ pushforward constrained distributions to unconstrained < R} and (right) simplex A3 where
ones, for (left) balllll where M {ar RR M == (a E TR22 Ei= T2 I 1, 2 0}. Note that for simplex, 23 J1 - I2 is redundant coordinate; see Section 3.2 for more details.
||
Essentially, the goal of MDM is to model the dual-space distribution:
Po (y) = ([V (p)PP
where $ is the push-forward operation and y (xx(x) is the mirror of data E M..
Training and generation MDM follows the same training procedure of Fuclidean-space diffusion models, except with an additional push- push-forward given by Equation (7). Similarly, its generation pro- cedure includes an additional step that maps the dual-space samples yo, generated from Equation (3), back to the constrained set via Vø (yo) E M..
Tractable variational bound Given a datapoint x E M, its mirror yo = VΦ(x) E R and the dual-space Markov chain (yi, yr the ELBO of MDM can be computed by
(L) = log det 2² ddt (yo) ) TgLBO(U)),
where apply the change of variables theorem [41] w.r.t. the mapping dual-space ELBO LE)BO(y)) shares the same formula with the Euclidean-space ELBO [42, 43),
Vó (36). The
T
LELBO(3O) DKL (g(yT vo) p(gT)) +
DRL DKL(g(vele++1.2) 30) (p log Ps (3olv1).
(9)
It should be noted that (8.9) provide a fractable variational bound for constrained generation. This stands contrast to concurrent methods relying on reflected diffusions, whose ELBO entails either intractable [18] or approximate [25] scores that could be restricted when dealing with more complex constrained sets. While, practice, MDM trained using the regression objective Equation (5), computing tractable ELBO might be valuable for independent purposes such likelihood estimation.
3.2 Efficient Computation of Mirror Maps
What remains to be answered pertains to the construction of mirror maps for common constraint sets such as fz-balls, simplices, polytopes, and their products. In particular, we seek efficient, preferably closed-form, computation for Vá. VØ" and 220"
| < R} denote the ball radius R in Rd. We consider the Ball Let M := (a E Rd
common log-barrier function:
(ean(x) == - log R - ||>11²)
(10)
where 7 E I The mirror map, its inverse, and the Hessian can be computed analytically by
2yæ
Ry
V Thhall())
V Vebaa
|
Rl
(11)
R
1
V Viohan (v)
yy
RV 1lvl3
+ 72 +
R RR
T 7
Rl
7
We refer to Appendix A for detailed derivation. Figure visualizes such mirror map for b2-ball in 2D. It worth noting that while the mirror map of log-barriers generally does not admit an analytical inverse, this is not the case for the ball constraints.

---

Simplex Given a
simplex, Aat1 == fo E Rá+ - = 1,,; 2 0}. we
follow standard practices [18, 25] by constructing the constrained set M := for R Ei= 34
1, 2, 2 O} and define 2d+1 := 1 - Σ i=1 Ty
While conventional barrierssss may remain viable option, preferred alternative that is widely embraced in the MLD literature for enforcing simplices [28, 44] is the entropic function [45, 46]:
(1- Σ
@simplex (x) = Σ zi lo:( +
log 1
Σ
(12)
i=1
i=1
The mirror map, its inverse, and the hessian of the dual function can be computed analytically by
d
el
V siiplix (ximplex(x)) (z)] = log ti - -log 1 1 Σ
[ Cimpleex (y)l:
+
Σi=
en
1=1
(13)
e/
c
V Simplex (2)les
l(i j)
+
21=1 e:
(1+ evi)
)
2
where I1()) is an indicator function. Again, we leave the derivation of Equation (13) to Appendix A. Figure visualizes the entropic mirror map for 3-dimensional simplex Ag.
Polytope Let the constrained set be M := {ar E Rdd G. < a. < bú, Vi E {1,... ,)), where Ca bç E I2 and az E I2d are linearly independent to each other. We consider the standard log-barrier:
ppolytope (x) == Σ (log ((a;, x) c.) + log (b, -(a,,)))) +
Σ j=++1
jay,z)
(14)
2
m
V @polstope (x) = Σ 1=1
ss si((a,,))) )))a +
Σ
(aj,
(15)
j=m+1
where the monotonic function 8; : (ce, bb) R is given by 8, [a:=))-;; [. and { do not impose any con-
Hyperbolic Tangent Eq.18)
Log-barrier
straints and can be chosen arbitrarily as long as they span Red
s((a,xx))
100
2.5
with {a:} (see below for more details); hence poloytto (-M)
0.0
Rd. While its inverse, populytpee does not admit closed-form
1000
2.5
0
-2
in general, it does when all d constraints are orthonormal:
Figure 4: Comparison between 8;
induced by standard log- bagrierriers vs.
m
d
V popoyyoop (1)
Σ s ((a,, 1=1
Σ j=m+1
(a..8)ag. (16)
hyperbolic tangents in Eq. (18).
Here, s,' R + (ca, by) is the inverse of st. Essentially, (15,16) manipulate the coefficients of the bases from which the polytope set is constructed. These manipulations are nonlinear yet bijective, defined uniquely by s; and s, for each orthonormal basis a, or, equivalently, for each constraint.
Naively implementing the mirror map via (15,16) can be problematic due (i) numerical instability of s, boundaries (see Figure 4), (ii) the lack of closed-form solution for its inverse s and (iii) an one-time computation of orthonormal bases, which scales as (dd) for orthonormalization methods such as Gram-Schmidt, Householder or Givens [e.g., 47). Below, we devise a modification for efficient and numerically stable computation, which will be later adopted for watermarked generation.
Improved computation for Polytope (15,16) Given the fact that only the first orthonormal bases {ar.... ... am } matter in the computation of (15,16), as the remaining (d - m) orthonormal bases merely involve orthonormal projections, we can simplify the computational (15,16) to
Vopoly =
(s: (s:((as, x)) (0;, as, Vo (v) = y +
(5 ((ar, v\) (ax, v\) ai
(17)
'The transformation allows to be bounded in constrained set M C RRd rather than hyperplane in gd+1

---

which, intuitively, add and subtract the coefficients of the first m orthonormal bases while leaving the rest intact. This reduces the complexity of computing orthonormal bases to (md2²) = (dd)) and improves the numerical accuracy of inner- inner-product multiplication, which can be essential for high-dimensional applications when d is large.
To address the instability of 5; and intractability of s, for log-barriers, we re-interpret the mirror maps (15,16) as the changes of coefficient bases. As this implies that any nonlinear bijective mapping with tractable inverse would suffice, we propose a rescaling of the hyperbolic tangent:
scl(at, .)) == (lanh -1 rescale;) ((02, .)); 87 ((a,, v)) == (rescale, o tanh) ((ar,)) (18)
where rescale,(z) := 2 [ 1 rescales the range from (ct, bs) to (-1,1). It is clear from Figure that, compared to
Table 2:
Complexity of VΦ and VØ" for each constrained set.
the induced by the log- bog-barrier, the mapping (18) is numerical stable at the boundaries and admits tractable inverse.
B--Ball Simplex Polytope
Complexity Table 2 summarizes the complexity of the mirror
((d)
O(d)
((md)
maps for each constrained set. We note that all computations are parallelizable, inducing nearly no computational overhead.
Remark Equations (17) and (18) can be generalized to non-orthonormal constraints by adopting (ā., x) and (ā., y), where ã; the i-th row of the matrix À == (ATA) A) A and A := [ai.-.. 0mm]
Indeed, when A is orthonormal, we recover ā; a₁. We leave more discussions to Appendix B.
4 Related Work
Constrained sampling Sampling from a probability distribution on constrained set has been long-standing problem not only due to its extensive applications in, e.g., topic modeling [48, 49] and Bayesian inference [50, 51J, but its unique challenges in non-asymptotic analysis and algorithmic design [28, 32-34, 52, 53J. Mirror Langevin is a special instance of endowing Langevin dynamics with a Riemannian metric [30, 54], and it chooses a specific reshaped geometry so that dynamics have to travel infinite distance in order to cross the boundary of the constrained set. It is a sampling generalization of the celebrated Mirror Descent Algorithm [45] for convex constrained optimization, and the convexity of constrained set leads to the existence of a mirror map. However, Mirror Langevin designed draw samples from an unnormalized density supported on a constrained set. This stands in contrast to our MDM, which instead tackles problems of constrained generation, where only samples, not unnormalized density, are available, sharing much similarity to conventional generative modeling [3, 55]. In addition, Mirror Langevin can be understood to go back and forth between primal (constrained) and mirror (unconstrained) spaces, while MDM only does one such round trip.
Diffusion Models in non-Euclidean spaces There has been growing interest in developing diffu- sion generative models that can operate on domains that are not limited to standard Euclidean spaces, e.g., discrete domains [56, 57] and equality constraints [58]. Seminar works by De Bortoli et al. [14] and Huang et al. [15] have introduced diffusion models on Riemannian manifolds, which have shown promising results in, e.g., biochemistry [59, 60] and rotational groups [61, 62]. Regarding diffusion models on constrained manifolds, most concurrent works consider similar convex constrained sets such as simplices and polytopes [18, 25, 63]. Our MDM works on the same classes of constraints and includes additionally falball, which is prevalent in social science domains such as opinion dynamics [64, 65]. It is also worth mentioning that our MDM may be conceptually similar to Simplex Diffusion (SD) [63] which reconstructs samples on simplices similar to the mirror mapping. However, SD designed specifically for simplices and Dirichlet distributions, while MDM is applicable to broader class of convex constrained sets. Additionally, SD adopts Cox- Cox-Ingersool-PR Ross processes [66], yiclding framework that, unlike MDM, is not simulation
key component to our MDM is the construction of mirror map, which allows us to define diffusion models in different space than the original constrained set, thereby alleviating many computational difficulties. This makes MDM a particular type of LDMs [5, 67, 68], which generally consider structurally advantageous spaces for learning diffusion models. While LDMs typically learn such latent spaces either using pretrained model or on the fly using, e.g., variational autoencoders, MDM instead derives analytically latent (dual) space through mirror map tailored specifically to the constrained set. Similar to LDD [69], the diffusion processes of MDM on the initial constrained set are implicitly defined and can be highly nonlinear.
Latent Diffusion Models (LDMs)
7

---

Table 3: Results of f2-ball constrained sets five synthetic datasets of various dimension d. We report Sliced Wasserstein [70] w.r.t. 1000 samples, averaged over three trials, and include constraint violations for unconstrained diffusion models. Note that Reflected Diffusion [18] and MDM satisfy constraints by design. Our MDM clearly achieves similar or better performance to standard diffusion models while fully respecting the constraints. Other metrics, e.g., W1, are reported in Appendix C. so
d= = 2
d=
d d=6 6
d = 8
d=20
Sliced Wasserstein 1
DDPM [2]
0,0704 = 0.0095
0.0236 0.0048
0.0379
0.0015
0.0231
0.0020
0.0200 + 0.0034
Reflected [18]
0.0642 = 0.0181
0.0491
0.0103
0.0609
0.0055
0.0881 = 0.0010
0.0574 + 0.0065
MDM (ours)
0.0544 0.0070
0.0214 0.0025
0.0467
0.0096
0.0292 0.0017
0.0159 0.0044
Constraint violation (%)
DDPM [2]
0.00
0.00
8.67 = 0.87
13.60 = 0.62
19.33
Table 4: Results of simplices constrained sets on five synthetic datasets of various dimension d.
3
d
d =
d =
d-20
Sliced Wasserstein +
DDPM [2]
0.0089 + 0.0002
0.01 10 0.0032
0.0047 0.0004
0.0053 + 0.0003
0.0031 + 0.0003
Reflected [18]
0.0233 0.0019
0.0336
0.0411
0.0897 0.0112
0.0231
MDM (ours)
0.0074 + 0.0008
0.0169
0.0016
0.0005 0.0006
0.0040 0.0003
0.0027
0.0000
Constraint violation (%) 4
DDPM [2]
0.73
14.40 1.39
11.63 0.90
27.53
68.83
Grand
MDM
Reffected[18] Ground
MDM
Roflectad[18]
Table 5: Runtime and memory
complexity w.r.t. DDPM [2].
Runtime Memory
Figure 5: Comparison between MDM and Reflected Diffusion [18] in modeling Dirichlet distribution on a 7-dimensional simplex. We visualize the joint densities between 23.4 and 25:66.
MDM Reflected Diff. [18]
108%
100%
> 1200%
905%
Experiment
5.1 Constrained Generation on Balls and Simplices
Setup We evaluate the performance of Mirror Diffusion Model (MDM) common constrained generation problems, such as f2-ball and simplex constrained sets with dimensions d ranging from 2 to 20. Following Fishman et al. [18], we consider mixtures of Gaussian (Figure 1) and Spiral (Figure 3) for ball constraints and Dirichlet distributions [48] with various concentrations for simplex constraints. We compare MDM against standard unconstrained diffusion models, such as DDPM [2), and their constrained counterparts, such as Reflected Diffusion [18], using the same time-embedded fully-connected network and 1000 sampling time steps. Evaluation metrics include Sliced Wasserstein distance [70] and constraint violation. Other details are left to Appendix C.
MDM surpasses Reflected Diffusion on all tasks. Tables 3 and 4 summarize our quantitative results. For all constrained generation tasks, our MDM surpasses Reflected Diffusion by large mar- gin, with larger performance gaps the dimension d increases. Qualitatively, Figure 5 demonstrates how MDM better captures the underlying distribution than Reflected Diffusion. Additionally, Table reports the relative complexity of both methods compared to simulation-free diffusion models such as DDPM. While Reflected Diffusson requires extensive computation due to the use of implicit score matching [71] and assertion of boundary reflection at every propagation step, our MDM constructs efficient, parallelizable, mirror maps in unconstrained dual spaces (see Section 3), thereby enjoying preferable simulation-fiee complexity and better numerical scalability.
should be highlighted that, while DDPM remains comparable to MDM in terms of distributional metric, the generated samples often violate the designated constrained sets. As shown in Tables and 4, these constraint violations worsen as the dimension increases. In contrast, MDM is by design violation- ffee;; hence marrying the best of both.
MDM matches DDPM without violating constraints.

---

Table 6: The 50k-FID for unconditional watermarked generation. Wa- termark precision denotes the percentage of constraint-satisfied images that are generated by MDMss, controlled by loosing/iightening the con- strained sets. As MDM-dual is trained on constraint-projeet we also report its FIDs* w.r.t. the shifted distributions. The prior wa- termarked diffusion model [39] reported 5.03 on FFHQ and 4.32 on
proj
FFHQ
images,
AFHQv2.
AFHQV2
FFHQ 64 64
AFHQv2 64 66
Precision
59.3%
71.8%
93.3%
56.9%
75.0%
92.7%
MDM-proj 2.54
2.59
3.08
2.10
2.12
2.30
Figure Uncurated water- marked samples by MDMs. More samples can found
2.96
4.57
15.74
2.23
2.86
6.79
MDM-dual 2.65*
2.93*
4.40"
2.21"
2.32
3.05
in Appendix C.
Histogram of (a.. en constriant range [-0.9,0.9] Original
MDM-proj
MDM-dual
2.3
2.3
Precision
1.0
2.2
0.75
2.2
0.7
2.1
2.1
1.0
1.5
0.5
1.0
1.5
0.9
Number of constraints, In log-scale
Folytope constraint range, b,-
Histogram of (a,,a)) for orig- inal and MDM-watermarked images x. Dashed lines indicate the constrained set.
Figure 7:
Figure 8: Ablation studies how (top) number of con-
straints m and (bottom) their ranges, i.e, (ci, b;), affect the FID and precision of MDM-pro) on AFHQv2.
5.2 Watermarked Generation on High-Diienniionll Image Datasets
new approach for watermarked generation, where the watermark is attached to samples within an orthonormal polytope M := {z E R C < a, x < b;. Vi}. These parameters fa,, by e /m serve as the private tokens visible only individual users (see Figure 2). While our approach is application-agnon datasets, including both unconditional and conditional generation. Given a designated watermark precision, defined the percentage of constraint-satiifie images that are generated by MDMs, we randomly construct tokens as m orthonormalized Gaussian random vectors and instantiate MDMs by either projecting unwatermarked diffusion models onto M (MDM-proj), or learning the corresponding dual-space diffusions (MDM-dual), More precisely, MDM-proj projects samples generated by pretrained diffusion models to a constraint set whose parameters (i.e., tokens) are visible only to the private user. In contrast, MDM-dual learns dual-space diffusion model from the constrained- constrained-rrjjeceed projected samples; hence, like other MDMs in Section 5.1, it is constraint-dependderr Other details are left to Appendix C.
Setup We present MDM
we demonstrate MDM mainly on image
Unconditional watermark generation on FFHQ and AFHQv2 We first test out MMD on FFHQ [72] and AFHQv2 [73] on unconditional 64x64 image generation. Following Zhao al. [39], we adopt EDM parameterization [38] and report in Table 6 the performance of MDM in generating watermarked images, quantified by the Frechet Inception Distance (FID) [74]. Compared to prior watermarked approaches, which require additional training of latent spaces [39], MDM stands competitive approach for watermarked generation by directly constructing analytic dual spaces from the tokens (i.e., constraints) themselves. Intriguingly, despite the fact that MDM-dual learns a smoother dual dual-space distribution compared to the truncated one from MDM- (see Figure 7), the latter consistently achieves lower (hence better) FIDs. Since samples generated by MDM-dual still remain close to the watermarked distribution, as indicated in Table 6 by the low FIDs" w.r.t. the shifted training statistics, we conjecture that the differences may be due to the significant distributional shift induced by naively projecting the training data onto the constraint set. This is validated, partly, in the ablation study from Figure 8, where we observe that both the FID and watermark precision improve the number of constraints m decreases, and the constraint ranges (c,, by) loosen. We highlight this fundamental trade off between watermark protection and generation quality. Examples of watermarked images are shown in Figure and Appendix C.
Conditional generation on ImageNet256 Finally, we consider conditional watermarked gen- eration on ImageNet 256x256. Specifically, we focus on image restoration tasks, where the goal

---

Degraded (JPEG QF=5) Input y
Ground Truth Reference (x
MDM-dual Watermarked Output (x M MDM-proj Watermarked Output (x E M)
Figure 9: Conditional watermarked generation on ImageNet 256x 256. Specifically, we consider the JPEG restoration task, where, given degraded, low-quality inputs y (upper-left), we wish to generate their corresponding clean images (upper-right) by learning p(x g). It is clear that both MDM-dual and MDM-proj are capable of solving this conditional generation task, generating clean images that additionally embed invisible watermarks, i.e., E .M.
is to generate clean, watermarked, images conditioned on degraded-restern with JPEG in this case- case-inputs. Similar to unconditional generation, we consider a polytope constraint set whose parameters are chosen such that the watermark yields high precision (> 95%) and low false positive rate (< 0.001 %).. Specifically, we set m 100 and - 1.2. We initialize the networks with pretrained checkpoints from Liu et al. [75].
Figure reports the qualitative results. is clear that both MDM-dual and MDM-proj capable of solving this conditional generation task, generating clean images that additionally embed invisible watermarks, i.e.. x E .M. Note that all non-MDM- --Deerreettedd images, despite being indistinguishable, actually violate the polytope constraint, whereas MDM-generae images always satisfy the con- straint. Overall, our results suggest that both MDM-dual and MDM-pro scale high-dimensional applications and are capable of embedding invisible watermarks in high-resolution images.
Conclusion and Limitation
We developed Mirror Diffusion Models (MDMs), a dual-space diffusion model that enjoys simulation- free training for constrained generation. We showed that MDM outperforms prior methods in standard constrained generation, such as balls and simplices, and offers a competitive alternative in generating watermarked contents. It should be noted that MDM concerns mainly conver constrained sets, which, despite their extensive applications, limits the application of MDM to general constraints. It will be interesting to combine MDM with other dynamic generative models, such as flow-based approaches.
Broader Impact
Mirror Diffusion Models (MDMs) advance the recent development of diffusion models to complex domains subjected to convex constrained sets. This opens up new possibilities for MDMs to serve as preferred models for generating samples that live in e.g. simplices and balls. Additionally, MDMs introduce an innovative application of constrained sets as a watermarking technique. This has the potential to address concerns related to unethical usage and safeguard the copyright of generative models. By incorporating constrained sets into the generating process, MDMs offer means to prevent unauthorized usage and ensure the integrity of generated content.
10

---

Acknowledgements
The authors would like thank Bogdon Vlahov for sharing the computational resources. GHL, TC, ET are supported by ARO Award #W911NF201011 and DoD Basic Research Office Award HQ003421 10002. MT partially supported by NSF DMS-1847802. NSF DCCS-1942522 Cullen- Peck Scholarship, and Gm--r Almmmniy Award.
References
[1] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning (ICML), 2015.
[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances Neural Information Processing Systems (NeurIPS). 2020.
[3] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021.
[4] Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[5] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In IEEE Conference on Computer Vision and Patlern Recognition (CVPR), 2022.
[6] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations (ICLR), 2021.
[7] Jungil Kong, Jachyeon Kim, and Jackyoung Bae. HiFi-GAN: Generative adversarial networks for efficient and high fidelity speech synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[8] Ben Poole, Ajay Jain, Jonathan T' Barron, and Ben Mildenhall. Dreamfusion: Text-to-3D using 2D diffusion. In International Conference on Learning Representations (ICLR), 2023.
[9] Chen-Hsuan Lin, Jun Gao, Luming Tang, Towaki Takikawa, Xiaohui Zeng, Xun Huang, Karsten Kreis, Sanja Fidler, Ming- Yu Liu, and Tsung- Yi Lin. Magic3D: High-resolution text-to-3I) content creation. In JEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2023.
[10] Brian DO Anderson. Reverse-ting time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982.
[11] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference Learning Representutions (ICLR), 2021.
[12] Yang Song and Stefano Ermon. Improved techniques for training score- boore-baed generative models. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[13] Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings, pages 1-10, 2022.
[14] Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thorton,, Yee Whye Tch, and Arnaud Doucet. Riemannian score-based generative modeling. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[15] Chin-Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Riemannian diffusion models. In Advances in Neural Information Processing Systems (NeurIPS), 2022.

---

[16] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[17] Ricky T. Q. Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03666 2023.
[18] Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson. Diffusion models for constrained domains. Transactions on Machine Learning Research (TMLR), 2023.
[19] Ben Morris. Improved bounds for sampling contingency tables. Random Structures & Algorithms, 21(2):135-146, 2002.
[20] Nathan E Lewis, Harish Nagarajan, and Bernhard O Palsson. Constraining the metabolic genotype-phenotype relationship using phylogeny of in silico methods. Nature Reviews Microbiology, 10(4):291-305, 2012.
[21] Li Han and Lee Rudolph. Inverse kinematics for serial chain with joints under distance constraints. In Robotics: Science and Systems (RSS), 2006.
[22] Tsuneo Yoshikawa. Manipulability of robotic mechanisms. The international journal of Robotics Research, 4(2):3-9, 1985.
[23] Ruth J Williams. Reflected brownian motion with skew symmetric data in a polyhedral domain. Probability Theory and Related Fields, 75(4):459-485 1987.
[24] Andrey Pilipenko. An introduction to stochastic differential equations with reflection, volume 1. Universitätsverlag Potsdam, 2014.
[25] Aaron Lou and Stefano Ermon. Reflected diffusion models. In International Conference on Machine Learning (ICML), 2023.
[26] Sebastien Bubeck, Ronen Eldan, and Joseph Lehec. Finite-time analysis of projected langevin monte carlo. In Advances Neural Information Processing Systems (NeurIPS), 2015.
[27] Kanji Sato, Akiko Takeda, Reüchiro Kawai, and Taiji Suzuki. Convergence error analysis of reflected gradient langevin dynamics for globally optimizing non-convex constrained problems. arXiv preprint arXiv:2203. 10215, 2022.
[28] Ya-Ping Hsieh, Ali Kavis, Paul Rolland, and Volkan Cevher. Mirrored langevin dynamics. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
[29] Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.
[30] Mark Girolami and Ben Calderhead. Riemann manifold langevin and hamiltonian monte carlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73(2): 123-214, 2011.
[31] Kelvin Shuangjian Zhang, Gabriel Peyré, Jalal Fadili, and Marcelo Pereyra. Wasserstein control of mirror langevin monte carlo. In Conference on Learning Theory (COLT), 2020.
[32] Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, Philippe Rigollet, and Austin Stromme. Exponential ergodicity of Harror-Langevin diffusions. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[33] Kwangjun Ahn and Sinho Chewi. Efficient constrained sampling via the Lanror-Langevin algorithm. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[34] 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 (ALT), 2022.
12

---

[35] Michael Liebrenz, Roman Schleifer, Anna Buadze. Dinesh Bhugra, and Alexander Smith. Generating scholarly content with chatgpt: ethical challenges for medical publishing. The Lancet Digital Health, 503)=e105-e106, 2023.
[36] John Kirchenbauer, Jonas Geiping, Yuxin Wen, Jonathan Katz, Ian Miers, and Tom Goldstein. A watermark for large language models. In International Conference on Machine Learning (ICML), 2023.
[37] Jürgen Rudolph, Samson Tan, and Shannon Tan. Chatgpt: Bullshit spewer or the end of traditional assessments in higher education? Journal of Applied Learning and Teaching, 6(1), 2023.
[38] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems (NeurlPS), 2022.
[39] Yunqing Zhao, Tianyu Pang, Chao Du, Xiao Yang, Ngai-Man Cheung, and Min Lin. recipe for watermarking diffusion models. arXiv preprint arXiv:2303.10137, 2023.
[40] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional networks for biomedical image segmentation. In International Conference on Medical Image Computing and Computer-assisted Intervention. Springer, 2015.
[41] Peter D Lax. Change of variables in multiple integrals. The American mathematical monthly, 106(6):497-51,, 1999.
[42] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[43] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[44] Sébastien Bubeck and Ronen Eldan. The entropic barrier: a simple and optimal universal self-concordant barrier. In Conference on Leurning Theory (COLT), 2015.
[45] Arkadij Semenovië Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency optimizaaion. Wiley-Interssienn 1983.
[46] Amir Beck and Marc Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-175, 2003.
[47] Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003.
[48] David M Blei, Andrew Y Ng, and Michael I Jordan. Latent dirichlet allocation. Journal of Machine Learning Research (JMLR), 2003.
[49] Arthur Asuncion, Max Welling, Padhraic Smyth, and Yee Whye Teh. On smoothing and inference for topic models. In Conference on Uncertainty in Artificial Intelligence (UAI), 2009.
[50] Lixin Lang, Wen-shiang Chen, Bhavik R Bakshi, Prem K Goel, and Sridhar Ungarala. Bayesian estimation via sequential monte carlo sampling-constieee (9):1615-1622, 2007.
dynamic systems. Automatica, 43
[51] Michael A Gelbart, Jasper Snoek, and Ryan P Adams. Bayesian optimization with unknown constraints. In Conference on Uncertainty in Artificial Intelligence (UAI), 2014.
[52] Yann Brenier. Décomposition polaire réarrangement monotone des champs de vecteurs. CR Acad. Sci. Paris Sér. Math,, 005:805-808, 1987.
[53] Nicolas Brosse, Alain Durmus, éric Moulines, and Marcelo Pereyra. Sampling from a log- concave distribution with compact support with proximal Langevin Monte Carlo. In Conference on Learning Theory (COLT), 2017.
13

---

[54] Sam Patterson and Yee Whye Teh. Stochastic gradient Riemannian Langevin dynamics on the probability simplex. In Advances in Neural Information Processing Systems (NeurIPS). 2013.
[55] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In International Conference on Learning Representations (ICLR), 2023.
[56] Ricky TQ Chen, Brandon Amos, and Maximilian Nickel. Semi-discrete normalizing flows through differentiable tessellation. In Advances in Neural Information Processing Systems (NeurlPS), 2022.
[57] Xingchao Liu, Lemeng Wu, Mao Ye, et al. Learning diffusion bridges on constrained domains. In International Conference on Learning Representations (ICLR), 2022.
[58] Mao Ye, Lemeng Wu, and Qiang Liu. First hitting diffusion models for generating manifold, graph and categorical data. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[59] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[60] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In International Conference on Learning Representations (ICLR), 2023.
[61] Adam Leach, Sebastian M Schmon, Matteo T Degiacomi, and Chris G Willcocks. Denoising diffusion probabilistic models on SO(3) for rotational alignment. In International Conference on Learning Representations (ICLR), Workshop Track, 2022.
[62] Julen Urain, Niklas Funk, Georgia Chalvatzaki, and Jan Peters. (E(3)-DiffusionFields Learn- ing cost functions for joint grasp and motion optimization through diffusion. In IEEE Interna- tional Conference on Robotics and Automation (ICRA), Workshop Track, 2022.
[63] Pierre H Richemond, Sander Dieleman, and Arnud Doucet. Categorical SDEs with simplex diffusion. arXiv preprini arXiv:2210. 14784444 2022.
[64] Marco Caponigro, Anna Chiara Lai, and Benedetto Piccoli. A nonlinear model of opinion formation on the sphere. Discrete & Continuous Dynamical Systems-A, 35(9):4241-4268, 2015.
[65] Jason Gaitonde, Jon Kleinberg, and Ëvv Tardos. Polarization in geometric opinion dynamics. In ACM Conference on Economics and Computation, pages 499-519, 2021.
[66] John C Cox, Jonathan E Ingersoll Jr, and Stephen A Ross. A theory of the term structure of interest rates. In Theory of valuation, pages 129-164. World Scientific, 2005.
[67] Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In Advances in Neural Information Processing Systems (NeurlPS), 2021.
[68] Kushagra Pandey, Avideep Mukherjee, Piyush Rai, and Abhishek Kumar. Diffusevae: Efficient, controllable and high-fidelity generation from low-dimensional latents. Transactions on Machine Learning Research (TMLR), 2022.
[69] Dongjun Kim. Byeonghu Na, Se Jung Kwon, Dongsoo Lee, Wanmo Kang, and II-Chul Moon. Maximum likelihood training implicit nonlinear diffusion models. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[70] Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasser- stein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22-45, 2015.
[71] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Conference on Uncertainty in Artificial Intelligence (UAI), 2020.
14

---

[72] Tero Karras, Samuli Laine, and Timo Aila. A -stylebbased generator architecture for generative adversarial networks. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
[73] Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha. Stargan v2: Diverse image synthe- sis for multiple domains. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
[74] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to local nash equilibrium. In Advances in Neural Information Processing Systems (NeurIPS). 2017.
[75] Guan-Horg Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. SB Image-to-Image Schrödinger bridge. In International Conference on Machine Learning (ICML), 2023.
[76] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf. Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 4024-8035 Curran Associates, Inc., 2019. URL http:: neurrrs...u...p.. cc/paper/ 9015-pytor an- imperative- style--sssle-eeeee high hereyyll-hhhh---ffrmmmaeeee pwwformmnc-- dewpee Learning 1ibrary.. pdf.
[77] lya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019.
[78] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
[79] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3-19, 2018.
15

---

A Derivation of Mirror Mappings
Here, we provide addition derivation of VΦ*. Computation of VΦ(z) and V²0*(y) (y) follow straight- forwardly by differentiating ((x) and V (v) w.r.t. I and respectively.
lg- U--Ball Since the gradient map also reverses the mirror map, we aim to rewrite R- ||111
2y
as
x ((y) = (v). Solving the second-order polynomial,
2 18
2
11
| 1
(19)
100
|
yields
2y
| = R +
7 v 12 +,2
(20)
11y 12
With that, we can rewrite Equation (11) by
R - |
(20) v RRlyll +
7
R
27
1
R
12
+
7
y.
Simplex
Standard calculations convex analysis [28] shows
immplex (1) = log +
Σ e
(21)
Differentiating Equation (21) w.r.t. y yields V implex in Equation (13).
Polytope
Since the gradient map also reverses the mirror map, we aim inverse
d
s;( x) )a, +
Σ j=m+1
()j,
(22)
When all d constraints are orthonormal, taking inner product between and each yiclds
(a;, v) s;((a,,))) (a,.0) = (a,,z). z).
(23)
Therefore, we can reconstruct from y via
d
ais x)a +
Σ
(aj,
i=m+1
m
(23)
((as,)))) +
Σ (aj,y)a,,
j=m+1
which defines
V polyytope (3). For completeness, the Hessian can be presented compactly
22 Ppulytgpe (V) = I AEAT,
(24)
where I is the identity matrix, A la,.... am) is d-by-m matrix whose column vector a, corresponds to each constraint, and Σ E mmxm is a diagonal matrix with leading entries
0x7 (z)
== (a. - (18) b, (1 tanh2 ((a;, v))) 1.
d2
16

---

B Additional Remarks on Polytope
Derivation of Equation (17) other, can rewrite (15)
Since the subspaces spanned by {a,) and {ay) } are orthogonal each
V Opalytope (x)
8i ((( Ja); +
(0s 1c)s
+
(si((a;)2)) - ())) di.
polytope (1) follows similar derivation.
Generalization to non-orthonormal constraints The mirror maps of a polytope, described in Equations (15) to (17), can be seen as operations that manipulate the coefficients associated with the bases defined by the constraints. This understanding allows us extend the computation to non-orthonormal constraints by identifying the corresponding "coefficients" through change of bases, utilizing the reproducing formula:
(ā., 2)a,, where ã; is the i-th row of (A (ATA)- A)-
and A = [a1, am Similarly, we have Ei (ās, y)az Applying similar derivation leads to Vopoly (x) = x (s, 2)) (ā;, ()), ai Vo = + Σ
(st ((ā., v)) (ā;, vi)
2m
C Experiment Details & Additional Results
Table 7: The concentration parameter of each Dirichlet distribution in simplices constrained sets.
d
7
9
20
a 12,48] 51
[1,2, 4, 8,8,8) [1. 0.5, 2, 0.6, 4, 8,
[0.2, 0.4, ... 4, 4,4
Table 8: Hyperparameters of the polytope M == {u E RRd :4 < a. < b.) for each dataset and watermark precision. Note that we fix b = b, = - G in practice.
FFHQ 64 (uncon)
AFHQv2 64 64 (uncon)
Precision
59.3%
71.8%
93.3%
56.9%
75.0%
92.7%
Number of constraints
20
100
20
100
Constraint range 1.05
1.05
1.05
0.9
0.9
0.9
Dataset & constrained
f2-balls construined sets: For d = 2, we consider the Gaussian Mixture Model (with variance 0.05) and the Spiral shown respectively in Figures 1 and 3. For d {6, 8, 20). we place d
isotropic Gaussians, each with variance 0.05, corner of each dimension, and reject samples outside the constrained sets.
Simplices constrained sets:
We consider Dirichlet distributions [48], Dir(a), with various con- centration parameters a detailed in Table 7.
Hypercube constrained sets: For all dimensions d = {2, 3, 6, 20), we place d isotropic Gaussians, each with variance 0.2, the comeer of each dimension, and either reject (d {2,3, 6, 83) or reflect (d 20) samples outside the constrained sets.
Watermarked datasets and polytope constrained sets: We follow the same data preprocessing from EDM' [38] and rescale both FFHQ and AFHQv2 64x64 image resolution. For the polytope constrained sets M := {i E Rd G; < a, < b,, Vi], we construct a; from orthonormalized Gaussian random vectors and detail other hyperparameters in Table 8.
'nttps://githul
www.eo/oam//ediieeeeeo http://ggthuub.cco///aaaa///e/// released under Nvidia Source Code License.
17

---

Implementation All methods are implemented in PyTorch [76]. We adopt ADM and EDM [38] respectively as the MDM's diffusion backbones for constrained and watermarked generation. We implemented Reflected Diffusion [18] by ourselves as their codes have not yet been made available by the submission time (May 2023), and used the official implementation" of Reflected Diffusion [25] in
Table 1.. We also implemented Simplex Diffusion [63], but as observed in previous works [25], it encountered computational instability especially when computing the modified Bessel functions.
Training For constrained generation, all methods are trained with Adam [77] and exponential moving average with the decay rate of 0.99. As standard practices, we decay the learning rate by the decay rate 0.99 every 1000 steps. For watermarked generation, we follow the default hyperparameters from EDM [38]. All experiments are conducted on two TITAN RTXs and one RTX 2080.
Network For constrained generation, all networks take (y. 1) inputs and follow
out out out_mod (norm __yddmd( y t__mod(timestepppomtooe (timesteep_emeedddnngg >)))
where timestep
simestep_embedding. embeddinggn (.) is the standard sinusoidal embedding. t_mod and out_mod consist
of fully-connected layers (Linear) activated by the Sigmoid Linear Unit (SiLU) [78]:
t_mod out_mod Linear -> SiLU - Linear
and y_mod consists of residual blocks,
y y_mod(y) y + res_mod (norm (z)), ) where
res_mod Linear SiLU Linear SiLU Linear SiLU Linear
All Linear's have 128 hidden dimension. We group normalization [79] for all norm. For watermarked generation, we use EDM parameterization [38].
Evaluation We compute the Wasserstein and Sliced Wasserstein distances using the geomloss" and ot packages, respectively. The Maximum Mean Discrepancy (MMD) is based on the popular package https://githe: watermarked generation, we follow the same evaluation pipeline from EDM* [38] by first generating 50,000 watermarked samples and computing the FID w.r.t. the training statisties.
"ZongxianLee
_oss_ Pytorch, which is unlicensed. For
Palse-positive rate Similar to Kirchenbauer et al. [36], we reject the null hypothesis and detect the watermark if the sample produces no violation of the polytope constraint, i.e., x: E M. Hence, the false-pos possitvv samples are those that are actually true null hypothesis (i.e., not generated by MDM) yet accidentally fall into the constraint set, hence being mistakenly detected as watermarked. Specifically, the palse-positive rates of our MDMs are respectively 0.07% and 0.08% for FFHQ and AFHQv2. Lastly, we note that the fact that both MDM-proj and MDM-dual generate samples that always satisfy the constraint readily implies 100% recall and 0% Type II error.
C.1 Additional Results
Tractable variational bound Equation (8) Figure demonstrates how MDM faithfully captures the variational bound to the negative log- likelihoooo (NLL) 2-dimensional GMM.
Groundtruth
MOM NLL
More constrained sets, distributional metrics, & baseline Ta- bles 9 and 10 expand the analysis in Tables 3 and with additional distributional metrics such as Wasserstein-1 (W)) and Maximum Mean Discrepancy (MMD). Additionally, Table 11 reports the results of hypercube [0, 1] constrained set, special instance of polytopes, and includes additional baseline from Lou and Ermon [25], which approximate the
Figure 10: Tractable variational bound by our MDM.
"ttppp:///ttaab..c.oooppp///ggga httppp///gtaaaucco//////arreoeea Diffusion latest commit (65d05c6) at submission, un- licensed. http:s///gtaabbboo//////tttt///p released under MIT License.
released under MIT License.
"nttps://pytthooo..ttuee git.hb..ttu...... https//iooooooomleee/eeeeeeeeeee www.moodue/oo..ooooo sliced. tot. sliced. sliced_
vasserstein_distance released under MIT License.
18

---

intractable scores in Reflected Diffusion using cigenfunctions tailored specifically to hypereubes, rather than implicit score matching as Fishman et al. [18]. Consistently, our findings conclude that the MDM is the only constrained-base based diffusion model that achieves comparable better performance to DDPM. These results afhrm the effectiveness of MDM generating high-quality samples within constrained settings, making it a reliable choice for constrained generative modeling.
More watermarked samples Figures and 12 provide additional qualitative results on watermarked samples generated by MDMs.
Figure 11: FFHQ 64x64 unconditional watermarked samples generated by (left) MDDM-rroj and (right) MDM-dual from the same set of random seeds. Despite the fact that some images, such the one the first row and sixth column, were altered possibly due to the change of dual-space distribution (see Figure 7), they look realistic and remain close to the data distribution.
Figure 12: AFHQv2 64x64 unconditional watermarked samples generated by (left) MDM -poj and (right) MDM-dual from the same set of random seeds. Despite the fact that some images, such the in the fifth and first column, were altered possibly due to the change of dual-space distribution (see Figure 7), they all look realistic and remain close to the data distribution.
19

---

Table 9: Expanded results of (2-ball constrained sets, where we include additional distributional metrics such as Wi and Maximum Mean Discrepancy (MMD), computed with 1000 samples and averaged over three trials. Consistently, our findings conclude that the MDM is the only constrained- based diffusion model that achieves comparable or better performance to DDPM.
d 2
=
d=
=8
d=20
Wi | (unit: 10 10
DDPM [2]
0.66 0.15
0.03
0.52 0.00
0.58 + 0.10
3.45 0.30
Reflected [18]
0.55 0.29
0.46 0.17
3.11
10.13 #.21
19.42
MDM (ours)
0.46 0.07
0.12 a 0.04
0.72 0.39
1.05 0.26
2.63 0.31
MMD | (unit: 10 10-3)
DDPM [2]
0.67 0.23
0.23 0.07
0.37 a 0.19
0.75 a 0.24
0.98 0.42
Reflected [18]
0.58 0.46
5.03 1.17
2.34 0.14
28.82 0.66
14.83 0.62
MDM (ours)
0.52 0.36
0.27 0.19
0.54 0.12
0.35 * 0.23
0.50 0.17
Constraint violation (%) 4
DDPM [2]
0.00 0.00
0.00 0.00
8.67 0.87
13.60 0.62
19.33 1.29
Table 10: Expanded results of simplices constrained sets.
=
d
==7
d=20
Wi (unit: 10 3)
DDPM [2]
0.01 = 0.00
0.02 0.01
0.03 = 0.00
0.05 = 0.00
0.11 0.00
Reflected [18]
0.06 0.01
0.12 0.00
0.62 a 1.08
3.57 a 0.05
0.02
MDM (ours)
0.01 = 0.00
0.01 0.01
0.03 0.00
0,05 0.00
0.13 0.00
MMD (unit: 10 2)
DDPM [2]
0.72 2 0.07
0.72 a 0.30
0.74 = 0.10
0.97 = 0.22
1.12 = 0.07
Reflected [18]
3.91
15.12
16.48
131.44
57.90 2.07
MDM (ours)
0.44 0.16
0.50 0.26
0.42 a 0.08
0.55 # 0.13
0.61 0.03
Constraint violation (%) 4
DDPM
0.73
14.40 1.39
11.63
27.53 0.57
68.83 1.66
Table 11: Results of hypercube [0, 1]d constrained sets.
d= 2
d 3
d = 6
d 8
d=20
Sliced Wasserstein | (unit: 10-3)
DDPM [2]
= 1.22
2.17 0.65
2.05 0.41
2.01 = 0.16
1.54 : 0.01
Reflected [25]
3.75 1.20
6,58 1.18
2.77 0.06
3.50 0.69
3,37 0.46
Reflected [18]
1.51
17.16 = 0.88
11.90 0.43
7.49
2013
4.32 0.23
MDM (ours)
3,00 0.72
1.92 0.81
1.75 0.17
1.85 0.34
3.35 0.64
Wi (unit: 10 1
DDPM [2]
0.07 0.05
0.22 a 0.07
1.65 0.144
3,30 + 0.16
16.74 0.12
Reflected [25]
0.20
1.21 4..99
2.53 0.04
4.82
25.47
Reflected [18]
4,40 + 0.57
6.01 0.977
9,34 + 0.56
9,84 + 0.24
25.27 0.36
MDM (ours)
0.08
0.20 0.07
1.57 + +.088
3.34 0.23
20.59
MMD (unit: 10
DDPM [2]
0.27 0.26
0.32 0.14
0.69 0.21
0.81 0.23
0.73 0.07
Reflected [25]
0.92 0.53
3.56 + 1.31
1.16 0.04
2.09 + 0.70
2,83 : 0.58
Reflected [18]
32.26 3.19
26.64 5.07
29.83
15.84 0.89
7.21 = 0.68
MDM (ours)
0.27 + 0.09
0.29 0.17
0.39 0.14
0.61 + 0.23
0.62 = 0.05
Constraint violation (%) 4
DDPM [2]
9.37 0.12
17.57 1.27
41.70 1.30
59.30 1.39
94.47 0.64
20


Reviews:
[{"review_id": "0SV9k9e0aq", "reviewer": "Reviewer_i5Cc", "summary": "The submission \"Mirror Diffusion Models for Constrained and Watermarked Generation\" describes a new approach to generate constrained data with diffusion models. Using mirror maps, the diffusion process proceeds as usual in an unconstrained space, but the generated data can be converted into the constraint space through the bijective mirror map. \n\nThe submission discusses (toy) applications in constrained generation onto a ball and a simplex, and an application of this approach to watermarking of diffusion model outputs.", "questions": "* Just for my understanding: For MDM-proj, the output of a given, pretrained diffusion model, is project onto the constraint set after generation? And for MDM-dual, a new model is trained/fine-tuned is trained on contraint-projected samples?\n* The polytope constraint for the watermark is hard for me to understand intuitively, what is being constrained here? The pixel space of the generated image is constrained to fulfil a random projection into a random interval? Intuitively, it is not clear to me why the impact of this constraint on the generated images is not larger?\n* Related to question 1, does this mean that only MDM-proj can use multiple tokens for multiple users, and MDM-dual needs to retrain a new model for every new token?\n* Why is the precision of the watermark less than 100%? From my understanding of the preceding sections, a constraint violation should be excluded? Why are images being generated that violate the constraints?\n\nOverall, I hope the authors could clarify these questions and rewrite Section 5.2 to be easier to understand.\n\nA few more questions/comments:\n\n* The functions phi are almost Legendre functions in the sense that they are essentially smooth and essentially strictly convex. Is the twice-differentiability needed in principle (aside from Eq.(8)?\n* The reduction to a bjiective map in Eq.(18) is a bit unsatisfying. I think it would be neater to derive the base function for this tanh shift\n* Slightly related, Legendre functions do not seem to be strictly needed to generate the mirror maps in this work? Technically, a strictly monotone operator (as defined e.g. in \"Convex Analysis and Monotone Operator Theory in Hilbert Spaces\" ) would be sufficient? (This is more of a side question concerning the notation, I don't think that too much is gained by generalizing this way)\n* Is it possible to derive p-values for the watermarked model to quantify the uncertainty about the watermark, as was done in Kirchenbauer 2023? From a practical perspective watermark, precision is not the primary objective for a watermark, but precision with a low false-positive rate of the resulting detection scheme.", "limitations": "The authors sufficiently discuss broader impact in the appendix.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "The proposed approach using tools from convex optimization to encode convex constraints into the generation process is simple and elegant. The submission also spends a good amount of time on exposition of examples for commonly-used mirror maps, which I think will be very beneficial to the wider readership. \nThe application to watermarking is a suprising, but interesting connection that the submission draws, discussing an immediate beneficial application of the proposed approach. Overall I consider this a strong submission.", "weaknesses": "I think I fully understood the submission up to Section 5.2, which is quite compressed (probably due to space reasons), and I found it hard to fully understand. I'll enumerate these questions below in the questions section.\n\n\nOtherwise, I see no major weaknesses. There are some typos, which I'll briefly mention:\n* In standard Euclidean spaces, tractable marginal can be\n* The marginal q(yt−1 |yt , y0 ) hints the optimal reverse\n\nAppendix:\n* Boarder impact\n*  polytop", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1. Clarification on watermark generation in Sec 5.2**\n\n- The reviewer’s understanding of MDM-proj and MDM-dual is correct: MDM-proj projects samples generated by pretrained diffusion models to a constraint set whose parameters (i.e., tokens) are visible only to the private user. In contrast, MDM-dual learns a dual-space diffusion model from the constrained-projected samples. Hence, as conjectured by the reviewer, MDM-proj allows multiple tokens for multiple users, whereas MDM-dual, like other MDMs in Sec 5.1, is constraint-dependent.\n- We view images as samples in the vectorized image space $\\mathbb{R}^d$. To watermark a, e.g., 64x64 image ($d$=3x64x64=12288), we consider a polytope constraint set, $\\mathcal{M}:=${$x\\in \\mathbb{R}^d: c_i < a_i^\\top x < b_i, \\forall i \\in [m]$}, where $a_i \\in \\mathbb{R}^d$ are orthonormal vectors in the **vectorized image space**. Hence, projecting $x$ to $\\mathcal{M}$ constrains the signed distances of $x$ to these linear-independent hyperplanes (defined by $a_i$’s) to be within the intervals $(c_i, b_i)$.\n- The impact of the polytope constraint on generation quality depends on the choice of {$a_i, b_i, c_i$}$_{i=1}^m$, as illustrated by the ablation study in Fig 8. For image applications, we find that watermarks can be invisibly embedded with high-frequency $a_i$’s, larger interval of $c_i$’s and $b_i$’s, and a larger $m$. This is because high-frequency perturbations often preserve the semantic structure of images. While a larger interval improves the generation quality at the cost of loosening the constraint set, the overall precision is tightened up by increasing the number of constraints, $m$.\n- Similar to Kirchenbauer 2023, we reject the null hypothesis and detect the watermark if the sample produces no violation of the polytope constraint, i.e., if $x \\in \\mathcal{M}$. The reviewer is correct that both MDM-proj and MDM-dual generate samples that always satisfy the constraint. This readily implies 100% recall (`TP/(TP+FN)`) and 0% Type II error (`FN`), yet *not* necessarily 100% precision (`TP/(TP+FP)`) due to false-positive (`FP`) samples. Specifically, `FP` samples are those that are actually true null hypothesis (i.e., *not* generated by MDM) yet accidentally fall into the constraint set, hence being mistakenly detected as watermarked.\n(Notation: `TP`,`FP`,`TN`,`FN` respectively denote the numbers of True Positives, False Positives, True Negatives, and False Negatives.) \n- In the table below, we report the precision, false-positive rate (FPR), and accuracy of our MDMs that are used to generate all image figures (Fig 1,6,10,11). We stress that on both datasets, our MDMs achieve high accuracy & precision with low FPR. \n\n    |  | Precision (`TP/(TP+FP)`) | FPR (`FP/(FP+TN)`) | Accuracy (`(TP+TN)/(TP+FP+TN+FN)`) |\n    | --- | --- | --- | --- |\n    | FFHQ | 93.3% | 0.072% | 96.4% |\n    | AFHQv2 | 92.7% | 0.079% | 96.1% |\n\n- We do acknowledge that Sec 5.2 could be stated clearly with additional paragraphs, yet we were limited by the space constraint at the submission time. Admitted that we are still not allowed to revise the submission, we will rewrite Sec 5.2 and include these discussions in the later revision. We thank the reviewer again for raising these comments.\n\n---\n\n**2. Relation of $\\phi$ to Legendre function**\n- We first thank the reviewer for raising this interesting comment. We based our notation (in Sec 3) on the literature of Mirror Langevin Dynamics (MLD), e.g., [1]. There $\\phi$ was indeed required to be of “Legendre type [2]” and twice-differentiable. This was because the twice differentiability was used to induce a (Riemannian) metric on the mirror space, from which MLD is constructed. For our MDM, however, twice differentiability is indeed needed only for Eq 8 as the reviewer pointed out. Continuous differentiability, i.e., $C^1$, would suffice for training. We very much appreciate this comment and will clarify this detail in the revision.\n- In this view, we may indeed be able to generalize our approach to strictly monotone operators. Meanwhile, one handy feature of using a convex $\\phi$ is, the inverse of its Hessian (inverse in the sense of matrix inverse) is the Hessian of its dual. Algorithmically, this might not be necessarily, but how it may affect the performance is not yet clear to us. Therefore, this generalization, despite beyond the scope of our submission, will be an interesting future direction, and we also think other applications such as constrained generation in function spaces would be very interesting. Thank you for asking!\n\n---\n\n**3. Typo & other clarifications**\n- All noted typos will be fixed in the revision. Additionally, we will rewrite Eq 18, the tanh shift, in the form of base functions. We thank the reviewer for the meticulous reading!\n\n---\n\n[1] The Mirror Langevin Algorithm Converges with Vanishing Bias  \n[2] Convex Analysis (Rockafellar 1970)"}, {"review_id": "MMkvHDiJji", "reviewer": "Reviewer_kFs5", "summary": "This paper studies how to learning diffusion model when the data is in a constrained domain. The idea is to map the data into an unconstrained domain using the mirror map and conduct the diffusion on that mirror space. Once the generation finishes, map the data back to the original space.", "questions": "See above.", "limitations": "See above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Using the idea of mirror map seems a novel approach.\nThis approach also allows for likelihood computation, which is helpful for model evaluation.", "weaknesses": "The method seems to be limited to several types of special domains. What's the challenge of generalizing it into more general domain constraints?\n\nThe experiment session can be improved. For example, to evaluate the model's performance on generating the data on simplex, we can consider generating the segmentation maps. Some larger scale experiment would be helpful.\n\nSome very related literature on learning diffusion models on constrained domains are missing [1, 2] and need to be discussed.\n\n[1] Learning Diffusion Bridges on Constrained Domains\n[2] First Hitting Diffusion Models for Generating Manifold, Graph and Categorical Data", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1. Discussion on [1,2]**\n\n- We first thank the reviewer for bringing up these missing references, which are indeed relevant to our Mirror Diffusion Model (MDM). Both [1,2] and our MDM generate samples in constrained domains. While [1,2] can be applied to, e.g., discrete domains [1,2], equality constraints (e.g., sphere boundary) [2], and products of 1D bounded intervals [1], our MDMs instead focus on “**constrained sets”**, i.e., a particular class of constrained domains that are specified by “**inequality constraints**”, $\\mathcal{M}:=${$x\\in\\mathbb{R}^d: f_i(x)<0, \\forall i $}. Hence, MDM stands in parallel to [1,2]. Similar to [1,2], MDM also enjoys simulation-free training and adopts regression objectives (see Table 1), making MDM superior to prior inequality-constrained diffusion models [3,4], as evidenced by Sec 5.1 (Table 3,4,5, Fig 5) and Appendix C.1 (Table 9,10,11). Finally, we note that our MDM is the first to explore inequality-constrained domains as a new mechanism for watermarked generation, which is otherwise absent in [1,2,3,4].\n\n- As we do acknowledge that [1,2] are important references, we will include them, along with the above discussions, in the following revision (kindly note that we are not allowed to revise the submission at the moment).\n\n---\n\n**2. Generalization to other constrained domains**\n\n- Following the discussion in **1.**, we re-emphasize that the goal of our MDM is to generate samples confined to “**inequality constraints**”. In principle, MDM can generate samples confined to *any* convex constraint set given its mirror map. In Sec 3.2, we exemplify three mirror maps, each for a different type of constraint set, mainly to demonstrate how efficient, closed-form mirror maps can be constructed for most inequality constraints considered in prior works [3,4]. This, as also recognized by Reviewer i5Cc, can be beneficial to a broader audience. For general convex constraint sets, MDM still remains applicable by constructing, e.g., log-barriers. While this may induce additional cost at inference, it introduces *no* computational overhead at training, which, crucially, preserves all desired computational advantages from Euclidean-space diffusion models (see Table 1).\n\n- As mirror maps are, by construction, built from convex functions, our MDM is also subjected to their domains (see L283-284 in Sec 6). We note, however, that MDM may still be applicable to general (non-convex), yet compact, constraint sets by adopting the diffeomorphism discussed in [4] (see their Fig 2 (iv)). Constructing simulation-free diffusions (like our MDM) for more general inequality constraints is an interesting future direction worth pursuing, and we thank the reviewer for raising these comments.\n\n---\n\n**3. Large-scale experiments on watermarking ImageNet 256x256**\n\n- We appreciate the reviewer's comment. In the **PDF attached above in our response to all reviewers**, we present additional results of our MDMs on large-scale image datasets, specifically ImageNet 256x256, for both conditional and unconditional watermarked generation. For conditional generation, we focus on image restoration tasks, generating clean, watermarked, images conditioned on degraded inputs, using both MDM-dual and MDM-proj. For unconditional generation, we include mainly MDM-proj due to time constraints during rebuttal. Both MDM-proj and MDM-dual consider a polytope constraint set whose parameters are chosen such that the watermark yields high precision (>95%) and low false positive rate (< 0.001%). Specifically, we set $m$=100, $b$=1.2, $c$=-1.2, and $a_i \\in \\mathbb{R}^{196608}$ orthogonal Gaussian random vectors. Similar to Sec 5.2, we initialize networks with pretrained checkpoints [5,6].\n\n- Our qualitative results suggest that both MDM-dual and MDM-proj scale to high-dimensional applications and are capable of embedding invisible watermarks in high-resolution images. Note that all non-MDM-generated images, despite being indistinguishable, actually violate the polytope constraint, whereas MDM-generated images always satisfy the constraint. These results highlight the scalability of our MDM for large-scale, high-dimensional applications. We will release our codes upon publication.\n\n---\n\n[3] Diffusion Models for Constrained Domains  \n[4] Reflected Diffusion Models  \n[5] https://github.com/openai/guided-diffusion  \n[6] https://github.com/NVlabs/I2SB"}, {"review_id": "dg5Qelu0fr", "reviewer": "Reviewer_Bc59", "summary": "When the data distribution is constrained in some boundary, \nThis paper introduced Mirror Diffusion Models (MDM), where the diffusion process runs not in the distribution of the (constrained) primal space, but in the distribution of the (unconstrained) dual space. For constrained datasets such as simplex, polytopes, and balls. While the primal space is kept constrained, the dual space is unconstrained and the standard Gaussian diffusion process can be used without any restrictions. When adequate dual functions and mirror maps are given, this paper achieved better sampling quality when generating from a constrained set, both in small-dimensional synthetic datasets and real-world datasets. And this offers competitive results in generating watermarked images.", "questions": "* How many sampling steps did you used for generating in each dataset?\n\n=====\n\nCorrection\n* One of the xlabels in Figure 6 should be `MDM-dual`. Now both two are `MDM-proj`.\n* Line 514: polytop --> polytope\n* Figure 8: top and bottom --> left and right", "limitations": "The authors adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper has good clarity in introducing theorems, examples, and applications.\n* To the best of our knowledge, this is the first work that necessitates and uses a primal-dual algorithm in the diffusion model literature. Given well-defined dual functions and mirror maps, the unconstrained diffusion model can be used for generating constrained datasets. The experiments with synthetic datasets (Dirichlet distribution) provides evidence to use MDM in categorical distribution, and competitive results in watermarked distribution provides further direction in privacy issues in diffusion-based generative models.", "weaknesses": "* The performance gap between MDM-proj and MDM-dual is not narrowed. When the U-Net-based network for images are specialized in generating pixel-based images, the dual-space distribution can be less specialized to be sampled using the same network architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1. Performance of dual-space diffusion models**\n\n- We first thank the reviewer for raising the comment. While we do notice a gap between `MDM-proj` and `MDM-dual` (first 2 rows in Table 6), we stress that these FID values are evaluated w.r.t. the original, *constraint-violated*, training set distribution, which differs from the dual-space, *constraint-satisfied*, distribution from which `MDM-dual` was trained. An alternative, arguably more suitable, metric for evaluating `MDM-dual` is the FID w.r.t. the dual-space distribution, which we include in the last row of Table 6 (highlighted in gray). There, `MDM-dual` behaves similarly to `MDM-proj`, exhibiting a FID-Precision trade-off. Qualitatively, Fig 6,10,11 show that `MDM-dual` is able to generate watermarked images with good quality.\n\n- Nevertheless, whether the current parametrization, using U-net, best suits learning dual-space diffusion models is an interesting question, and we believe it will depend strongly on the choice of (polytope) constraint sets, since dual-space samples essentially change the coefficient bases, defined by the constraint set, of primal-space samples (see Eq 17). For image applications, we find that MDM can best embed invisible watermarks using high frequency $a_i$, as it preserves the semantic structure of images. Co-designing parametrization with the constraint set will be an interesting future direction to pursue, and we thank the reviewer for raising the comment.\n\n---\n\n**2. Typo & other clarifications**\n\n- We thank the reviewer for the meticulous reading. There’s indeed a typo in Fig 6. The right 2 columns should be `MDM-dual` rather than `MDM-proj`. All noted typos will be fixed in the revision (kindly note that we are not allowed to revise the submission at the moment).\n\n- Regarding sampling steps, all datasets and diffusion models in Sec 5.1 generate samples with 1000 steps, as mentioned in L239. For image datasets, FFHQ and AFHQv2, in Sec 5.2, we generate watermarked images with 79 steps, following the setup from EDM [1]. We note, however, that we did not perform extensive tuning on this hyper-parameter.\n\n---\n\n[1] Elucidating the design space of diffusion models"}, {"review_id": "QXFtAxJh7K", "reviewer": "Reviewer_VdTP", "summary": "This paper proposes a new class of diffusion models called the Mirror Diffusion Model (MDM), which confines the generation to a constrained convex set. The MDM transforms the generation from a constrained original space to an unconstrained dual space. With this transformation, MDM can be trained and sampled like the unconstrained Euclidean-space diffusion models, such as DDPM, with an additional step mapping back to the original space. In experiments, the authors show that the general quality of MDM on constrained sets outperforms the reflected diffusion models in terms of quality and efficiency. Furthermore, the authors demonstrate an important application of MDM in watermark generation.", "questions": "(Q1) Can the authors provide the proofs that gradient of the strongly convex function maps the constrained convex set to the whole $R^d$? \n\n(Q2) Can the authors provide more experimental results on more realistic data, like CIFAR-100 and ImageNet?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "(S1) This paper proposes a new class of diffusion models that confine the generation to a constrained convex set. It achieves this by transforming the generation from the constrained original space to an unconstrained dual space, enabling efficient training and sampling similar to unconstrained diffusion models.\n\n(S2) The mechanism relies on a strongly convex function defined on the convex set. This function's gradient space spans $R^d$, and the gradient norm approaches infinity near the boundary of the convex set. The author provides clear instructions on how to design such a convex function for various shapes, including the $\\ell_2$-ball, simplex, and general polytopes.\n\n(S3) The experiments demonstrate that the proposed model outperforms reflected diffusion models in terms of generation quality and efficiency. It also achieves generation quality comparable to that of unconstrained DDPM without violating any constraints.", "weaknesses": "(W1) From my understanding, the proposed model requires the design of a strongly convex function whose gradient maps the constrained convex set to the entire $R^d$ space, enabling the application of an unconstrained diffusion model in the dual space. However, it is not immediately apparent to me how the authors ensure that this condition holds for the strongly convex functions they provide.\n\n(W2) It would indeed be beneficial to observe experiments conducted on more realistic datasets, such as CIFAR-100 and ImageNet. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1. Conditions when $\\nabla\\phi(\\mathcal{M})= \\mathbb{R}^d$**\n\n- We first note that the gradient map of a strictly convex function $\\phi$ needs *not* span $\\mathbb{R}^d$, unless additional conditions are satisfied. For mirror maps, we follow the literature (e.g., [1,2]) and require $\\phi$ to be additionally *(i)* of Legendre type [1] (i.e., $\\lim_{x\\rightarrow\\partial\\mathcal{M}}\\|\\nabla\\phi(x)\\|\\rightarrow \\infty$, see L118 Sec 3.1) and *(ii)* continuous differentiable (see L117 Sec 3.1). When $\\phi$ also satisfies these two conditions, its gradient map will be surjective with range $\\mathbb{R}^d$.\n\n- The above surjectivity statement follows from convex analysis (e.g., [1]). Here is how we understood it intuitively. For gaining insight consider a 1D case. Since $\\phi$ is strictly convex, its derivative is strictly increasing (see, e.g., [1,3]). Then, the above two conditions ensure that the derivative not only approach $(-\\infty,\\infty)$ at the boundary $\\partial\\mathcal{M}$, but is also continuous in its domain (i.e., there exists no hole or jump). Hence, the gradient map has range $\\mathbb{R}$.\n\n- It is, however, a valid and in fact critical question, whether such requirements can be satisfied for any convex constraint set (i.e. the mirror map approach works). Unfortunately, to the best of our knowledge, how to explicitly construct a mirror map (with both $\\nabla \\phi$ and $(\\nabla \\phi)^{-1}$ analytically given) for an arbitrary convex set is still an open problem, although existence is less an issue.\n\n---\n\n**2. Large-scale experiments on watermarking ImageNet 256x256**\n- We appreciate the reviewer's comment. In the **PDF attached above in our response to all reviewers**, we present additional results of our MDMs on large-scale image datasets, specifically ImageNet 256x256, for both conditional and unconditional watermarked generation. For conditional generation, we focus on image restoration tasks, generating clean, watermarked, images conditioned on degraded inputs, using both MDM-dual and MDM-proj. For unconditional generation, we include mainly MDM-proj due to time constraints during rebuttal. Both MDM-proj and MDM-dual consider a polytope constraint set whose parameters are chosen such that the watermark yields high precision (>95%) and low false positive rate (< 0.001%). Specifically, we set $m$=100, $b$=1.2, $c$=-1.2, and $a_i \\in \\mathbb{R}^{196608}$ orthogonal Gaussian random vectors. Similar to Sec 5.2, we initialize networks with pretrained checkpoints [4,5].\n\n- Our qualitative results suggest that both MDM-dual and MDM-proj scale to high-dimensional applications and are capable of embedding invisible watermarks in high-resolution images. Note that all non-MDM-generated images, despite being indistinguishable, actually violate the polytope constraint, whereas MDM-generated images always satisfy the constraint. These results highlight the scalability of our MDM for large-scale, high-dimensional applications. We will release our codes upon publication.\n\n---\n\n[1] Convex analysis. (Rockafellar 1970)  \n[2] The Mirror Langevin Algorithm Converges with Vanishing Bias  \n[3] https://math.stackexchange.com/questions/999550/strictly-convex-if-and-only-if-derivative-strictly-increasing  \n[4] https://github.com/openai/guided-diffusion  \n[5] https://github.com/NVlabs/I2SB  \n"}]

Author rebuttal:
### **Author response to all reviewers**

We thank the reviewers for their valuable comments. We are excited that the reviewers identified the novelty of using mirror maps in learning constrained diffusion models ( **Reviewers** ****i5Cc****, ****kFs5, Bc59****, ****VdTP****), acknowledged our superior empirical results over prior diffusion models [1,2] and distinct application to watermarked generation (**Reviewers** ****i5Cc****, ****Bc59****, ****VdTP****), and found the paper well-written (**Reviewers** ****kFs5****, ****Bc59****, ****VdTP****). We believe MDM takes a significant step toward a new class of tractable diffusion models for constrained sets and watermarked generation. 

As *all* reviewers recognized our technical contribution, one of the criticisms stemmed from the insufficient evaluation of our MDM on larger-scale datasets (raised by Reviewers kFs5, VdTP). In the **attached PDF (below)**, we present additional results of our MDMs, on ImageNet 256x256, for both conditional and unconditional watermarked generation. Our qualitative results suggest that MDMs can be scaled to high-dimensional application, and are capable of embedding invisible watermarks in high-resolution images. While all non-MDM-generated images, despite being indistinguishable, violate the constraint/token, our MDM satisfies the constraint by construction. We highlight these new results, gained uniquely in MDMs by constructing efficient mirror maps, that are otherwise absent in prior diffusion methods [1,2].

We tried our best to resolve all raised questions in the individual responses below. If you have any additional questions/comments/concerns, please let us know. We always appreciate the reviewer's precious time in providing their valuable feedback.

[1] Diffusion Models for Constrained Domains  
[2] Reflected Diffusion Models

---

**Please find additional figures in the PDF below**:

Meta-review: In this paper, the authors studied diffusion generative modeling when the data are confined to some constrained set, which goes beyond the standard Euclidean space. A new mechanism called the mirror diffusion model (inspired by mirror descent in the optimization literature) has been proposed, which employs suitable mirror maps to exploit the structure of various constrained sets. Intriguing performance has been shown for the application of watermarking. The paper would benefit from more experiments on larger-scale real-world datasets.
Decision: Accept (poster)