Title: Mirror Diffusion Models for Constrained and Watermarked Generation
Venue: NeurIPS 2023
Accepted: yes
Keywords: diffusion models

constrained generation

constrained manifold

mirror map

watermarked generation

generation privacy
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

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 glyph[lscript] 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.

1 Introduction

Corresponding nonlinear diffusions in

ℳ

MDM: learns tractable diffusions in dual spaces

Diffusion models [1-3] have emerged as powerful generative models with their remarkable successes in 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

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 so for non-Euclidean spaces is nontrivial.

Recent progress of diffusion models has expanded their application to non-Euclidean spaces, such as Riemannian manifolds [14, 15], where data live in a 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⊆ R d , 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 ∂ M 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 2 , it comes with a 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 as thresholding and discretization [18, 25].

An alternative diffusion for constrained sampling is the Mirror Langevin Dynamics (MLD [28]). MLD, as a 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 ∇ φ : M→ R d , defines a nonlinear mapping from the initial constrained set to an unconstrained dual space . Despite extensive investigation of MLD [31-34], its potential as a foundation for designing diffusion generative models remains to 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 ∇ φ that transform data distributions on M to its dual space ∇ φ ( 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 ∇ φ ( M ) = R d ,MDM

∇𝜙

ℳ

User-defined tokens

ℳ≔ 𝑥 ∈ ℝ

"

∶ 𝑐

#

≤ 𝑎

#

$

#

, ∀𝑖

External users

Pretrained (primal-space)

Diffusion Models

Dual-Space

Diffusion Models

ℳ

Watermark

Dataset

∇𝜙

∗ ℳ

ℝ

"

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 ∇ φ ∗ , transforming mirrored samples back to the constrained set. We propose efficient construction of mirror maps and demonstrate that MDMoutperforms prior reflected diffusion models on common constrained sets, such as glyph[lscript] 2 -balls and simplices.

Moreover, we show that MDM also stands as a novel approach for watermarked generation , a technique that aims to embed undetectable information ( i.e., watermarks) into generated contents to 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 a 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 on 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 set. In either case, MDM generates watermarked contents that can be certified only by users who know the tokens. Our empirical results suggest that MDMserves 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 summary, we present the following contributions:

We introduce Mirror Diffusion Models (MDM) , a new class of diffusion models that utilizes mirror maps to enable constrained generation with Euclidean-space diffusion models.

We propose efficient computation of mirror maps on common constrained sets, including balls and simplices. Our results demonstrate that MDM consistently outperforms previous reflected diffusion models.

As a novel application of constrained generation, we explore MDM's potential for watermarking 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 M⊆ R d denotes a convex constrained set. We preserve x ∈ M and y ∈ R d to better distinguish variables on constrained sets and standard d -dimensional Euclidean space. I ∈ R d × d denotes the identity matrix. x i denotes the i -th element of a vector x . Similarly, [ A ] ij denotes the the element of a matrix A at i -th row and j -th column.

Diffusion process Given y 0 ∈ R d sampled from some Euclidean-space distribution, conventional diffusion models define the diffusion process as a forward Markov chain with the joint distribution:

which progressively injects Gaussian noise to y 0 such that, for a sufficiently large T and a properly chosen noise schedule β t ∈ (0 , 1) , y T approaches an approximate Gaussian, i.e., q ( y T ) ≈ N (0 , I ) .

Equation (1) has tractable marginal densities, and two of which that are of particular interest are:

The tractable marginal q ( y t | y 0 ) enables direct sampling of y t without simulating the forward Markov chain (1), i.e., y t | y 0 can be sampled in a simulation-free manner. It also suggests a closed-form score function ∇ log q ( y t | y 0 ) . The marginal q ( y t -1 | y t , y 0 ) hints the optimal reverse process given y 0 .

Generative process The generative process, which aims to reverse the forward diffusion process (1), is defined by a backward Markov chain with the joint distribution:

Here, the mean µ θ is typically parameterized by some DNNs, e.g., U-net [40], with θ . Equations (1) and (3) imply an evidence lower bound (ELBO) for maximum likelihood training, indicating that the optimal prediction of µ θ ( y t , t ) should match ˜ µ t ( y t , y 0 ) in Equation (2).

Parameterization and training There exists many different ways to parameterize µ θ ( y t , t ) , and each of them leads to a distinct training objective. For instance, Ho et al. [2] found that

achieves better empirical results compared to directly predicting y 0 or ˜ µ t . This parameterization aims to predict the 'noise' injected to y t ∼ q ( y t | y 0 ) , as the ELBO reduces to a simple regression:

where y t = √ ¯ α t y 0 + √ 1 -α t glyph[epsilon1] , glyph[epsilon1] ∼ N (0 , I ) , and t sampled uniformly from { 1 , · · · , T } . The weighting λ ( t ) is often set to 1. Once glyph[epsilon1] θ is properly trained, we can generate samples in Euclidean space by simulating Equation (3) backward from y T ∼ N (0 , 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 p data ( x ) supported on a convex constrained set M⊆ R d while retraining the computational advantages of Euclidean-space diffusion models.

3.1 Dual-Space Diffusion

Mirror map Following the terminology in Li et al. [34], let φ : M→ R be a twice differentiable 3 function that is strictly convex and satisfying that lim x → ∂ M ‖∇ φ ( x ) ‖ → ∞ and ∇ φ ( M ) = R d . We call its gradient map ∇ φ : M→ R d the mirror map and its image, ∇ φ ( M ) = R d , as its dual/mirror space . Let φ ∗ : R d → R be the dual function of φ defined by

A notable property of this dual function φ ∗ ( y ) is that its gradient map reverses the mirror map. In other words, we have ∇ φ ∗ = ( ∇ φ ) -1 , implying that

The functions ( ∇ φ, ∇ φ ∗ ) define a nonlinear, yet bijective, mapping between the convex constrained set M and Euclidean space R d . As a result, MDM does not require building diffusion models on M , as done in previous approaches, but rather they learn a standard Euclidean-space diffusion model in the corresponding dual space ∇ φ ( M ) = R d .

1

0

1

1

0

1

3

0

3

3

0

3

(

)

0.0

0.2

0.4

0.6

0.8

1.0

1.0

0.8

0.6

0.4

0.2

0.0

0.0

0.2

0.4

0.6

0.8

1.0

3

0

3

2

1

4

7

(

)

Essentially, the goal of MDM is to model the dual-space distribution :

where glyph[sharp] is the push-forward operation and y = ∇ φ ( x ) is the mirror of data x ∈ M .

Training and generation MDMfollows the same training procedure of Euclidean-space diffusion models, except with an additional push-forward given by Equation (7). Similarly, its generation procedure includes an additional step that maps the dual-space samples y 0 , generated from Equation (3), back to the constrained set via ∇ φ ∗ ( y 0 ) ∈ M .

Tractable variational bound Given a datapoint x ∈ M , its mirror y 0 = ∇ φ ( x ) ∈ R d , and the dual-space Markov chain { y 1 , · · · , y T } , the ELBO of MDM can be computed by

where we apply the change of variables theorem [41] w.r.t. the mapping x = ∇ φ ∗ ( y 0 ) . The dual-space ELBO ˜ L ELBO ( y 0 ) shares the same formula with the Euclidean-space ELBO [42, 43],

It should be noted that (8,9) provide a tractable variational bound for constrained generation. This stands in 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, in practice, MDM is trained using the regression objective in Equation (5), computing tractable ELBO might be valuable for independent purposes such as 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 glyph[lscript] 2 -balls, simplices, polytopes, and their products. In particular, we seek efficient, preferably closed-form, computation for ∇ φ , ∇ φ ∗ and ∇ 2 φ ∗ .

glyph[lscript] 2 -Ball Let M := { x ∈ R d : ‖ x ‖ 2 2 < R } denote the glyph[lscript] 2 -ball of radius R in R d . We consider the common log-barrier function:

where γ ∈ R + . The mirror map, its inverse, and the Hessian can be computed analytically by

We refer to Appendix A for detailed derivation. Figure 3 visualizes such a mirror map for an glyph[lscript] 2 -ball in 2D. It is worth noting that while the mirror map of log-barriers generally does not admit an analytical inverse, this is not the case for the glyph[lscript] 2 -ball constraints.

Simplex Given a ( d +1) -dimensional simplex, ∆ d +1 := { x ∈ R d +1 : ∑ d +1 i =1 x i = 1 , x i ≥ 0 } , we follow standard practices [18, 25] by constructing the constrained set M := { x ∈ R d : ∑ d i =1 x i ≤ 1 , x i ≥ 0 } and define x d +1 := 1 -∑ d i =1 x i . 4

While conventional log-barriers may remain a viable option, a preferred alternative that is widely embraced in the MLD literature for enforcing simplices [28, 44] is the entropic function [45, 46]:

The mirror map, its inverse, and the hessian of the dual function can be computed analytically by

where I ( · ) is an indicator function. Again, we leave the derivation of Equation (13) to Appendix A. Figure 3 visualizes the entropic mirror map for a 3-dimensional simplex ∆ 3 .

Polytope Let the constrained set be M := { x ∈ R d : c i < a glyph[latticetop] i x < b i , ∀ i ∈ { 1 , · · · , m }} , where c i , b i ∈ R and a i ∈ R d are linearly independent to each other. We consider the standard log-barrier:

where the monotonic function s i : ( c i , b i ) → R is given by s i = -1 〈 a i ,x 〉-c i + -1 〈 a i ,x 〉-b i , and { a j } do not impose any constraints and can be chosen arbitrarily as long as they span R d with { a i } (see below for more details); hence ∇ φ polytope ( M ) = R d . While its inverse, ∇ φ ∗ polytope , does not admit closed-form in general, it does when all d constraints are orthonormal :

2

0

2

100

0

100

si

(

ai

,

x

)

Log-barrier

2

0

2

2.5

0.0

2.5

Hyperbolic

T angent in Eq.(18)

Here, s -1 i : R → ( c i , b i ) is the inverse of s i . 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 i and s -1 i , for each orthonormal basis a i or, equivalently, for each constraint.

Naively implementing the mirror map via (15,16) can be problematic due to (i) numerical instability of s i at boundaries (see Figure 4), (ii) the lack of closed-form solution for its inverse s -1 i , and (iii) an one-time computation of orthonormal bases, which scales as O ( d 3 ) 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 m orthonormal bases { a 1 , · · · , a m } 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

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 O ( md 2 ) ≈ O ( d 2 ) and improves the numerical accuracy of inner-product multiplication, which can be essential for high-dimensional applications when d is large.

To address the instability of s i and intractability of s -1 i 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 a tractable inverse would suffice, we propose a rescaling of the hyperbolic tangent:

where rescale i ( z ) := 2 z -c i ( b i -c i ) -1 rescales the range from ( c i , b i ) to ( -1 , 1) . It is clear from Figure 4 that, compared to the s i induced by the log-barrier, the mapping (18) is numerical stable at the boundaries and admits tractable inverse.

Complexity Table 2 summarizes the complexity of the mirror maps for each constrained set. We note that all computations are parallelizable, inducing nearly no computational overhead.

Table 2: Complexity of ∇ φ and ∇ φ ∗ for each constrained set.

Remark 1. Equations (17) and (18) can be generalized to non-orthonormal constraints by adopting 〈 ˜ a i , x 〉 and 〈 ˜ a i , y 〉 , where ˜ a i is the i -th row of the matrix ˜ A := ( A glyph[latticetop] A ) -1 A glyph[latticetop] and A := [ a 1 , · · · , a m ] . Indeed, when A is orthonormal, we recover ˜ a i = a i . We leave more discussions to Appendix B.

4 Related Work

Constrained sampling Sampling from a probability distribution on a constrained set has been a long-standing problem not only due to its extensive applications in, e.g., topic modeling [48, 49] and Bayesian inference [50, 51], but its unique challenges in non-asymptotic analysis and algorithmic design [28, 32-34, 52, 53]. 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 is designed to 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 diffusion 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 glyph[lscript] 2 -ball, 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 is designed specifically for simplices and Dirichlet distributions, while MDM is applicable to a broader class of convex constrained sets. Additionally, SD adopts Cox-Ingersoll-Ross processes [66], yielding a framework that, unlike MDM, is not simulation-free.

Latent Diffusion Models (LDMs) A key component to our MDM is the construction of mirror map, which allows us to define diffusion models in a 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 a latent (dual) space through a mirror map tailored specifically to the constrained set. Similar to LDM [69], the diffusion processes of MDMon the initial constrained set are implicitly defined and can be highly nonlinear.

x

3

4

Ground Truth

x

3

x

4

MDM

x

3

x

4

Reflected[18]

x

5

x

6

Ground Truth

x

5

x

6

MDM

x

5

x

6

Reflected[18]

x

5 Experiment

5.1 Constrained Generation on Balls and Simplices

Setup We evaluate the performance of Mirror Diffusion Model (MDM) on common constrained generation problems, such as glyph[lscript] 2 -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.

MDMsurpasses Reflected Diffusion on all tasks. Tables 3 and 4 summarize our quantitative results. For all constrained generation tasks, our MDM surpasses Reflected Diffusion by a large margin, with larger performance gaps as the dimension d increases. Qualitatively, Figure 5 demonstrates how MDM better captures the underlying distribution than Reflected Diffusion. Additionally, Table 5 reports the relative complexity of both methods compared to simulation-free diffusion models such as DDPM. While Reflected Diffusion 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-free complexity and better numerical scalability.

MDMmatches DDPM without violating constraints. It 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 3 and 4, these constraint violations worsen as the dimension increases. In contrast, MDM is by design violation-free; hence marrying the best of both.

0.9

0.9

Original

0.9

0.9

MDM-proj

0.9

0.9

MDM-dual

Histogram of

ai

,

x

on constriant range [-0.9,0.9]

FFHQ

MDM-proj

MDM-dual

AFHQv2

10 1

10 2

2.1

2.2

2.3

FID

10 1

10 2

0.50

0.75

Precision

0.5

1.0

1.5

2.1

2.2

2.3

FID

0.5

1.0

1.5

0.7

1.0

Precision

Number of constraints,

m

, in log-scale

Polytope constraint range,

bi

=

5.2 Watermarked Generation on High-Dimensional Image Datasets

Setup We present MDM as a new approach for watermarked generation, where the watermark is attached to samples within an orthonormal polytope M := { x ∈ R d : c i < a glyph[latticetop] i x < b i , ∀ i } . These parameters { a i , b i , c i } i m =1 serve as the private tokens visible only to individual users (see Figure 2). While our approach is application-agnostic, we demonstrate MDM mainly on image datasets, including both unconditional and conditional generation. Given a designated watermark precision, defined as the percentage of constraint-satisfied 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 a dual-space diffusion model from the constrained-projected samples; hence, like other MDMs in Section 5.1, it is constraint-dependent. Other details are left to Appendix C.

Unconditional watermark generation on FFHQ and AFHQv2 We first test out MMD on FFHQ [72] and AFHQv2 [73] on unconditional 64 × 64 image generation. Following Zhao et 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], MDMstands as a 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-space distribution compared to the truncated one from MDM-proj (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 as the number of constraints m decreases, and as the constraint ranges ( c i , b i ) loosen. We highlight this fundamental trade off between watermark protection and generation quality. Examples of watermarked images are shown in Figure 6 and Appendix C.

Conditional generation on ImageNet256 Finally, we consider conditional watermarked generation on ImageNet 256 × 256. 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

)

x

)

is to generate clean, watermarked, images conditioned on degraded-restored with JPEG in this 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 b = -c = 1 . 2 . We initialize the networks with pretrained checkpoints from Liu et al. [75].

Figure 9 reports the qualitative results. 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., x ∈ M . Note that all non-MDM-generated images, despite being indistinguishable, actually violate the polytope constraint, whereas MDM-generated images always satisfy the constraint. Overall, our 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.

6 Conclusion and Limitation

We developed Mirror Diffusion Models (MDMs), a dual-space diffusion model that enjoys simulationfree 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 convex 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 a means to prevent unauthorized usage and ensure the integrity of generated content.

References

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning (ICML) , 2015.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

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.

Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2022.

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.

Jungil Kong, Jaehyeon Kim, and Jaekyoung Bae. HiFi-GAN: Generative adversarial networks for efficient and high fidelity speech synthesis. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

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.

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-3D content creation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2023.

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations (ICLR) , 2021.

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACMSIGGRAPH 2022 Conference Proceedings , pages 1-10, 2022.

Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modeling. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

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.

Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Ricky T. Q. Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660 , 2023.

Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson. Diffusion models for constrained domains. Transactions on Machine Learning Research (TMLR) , 2023.

Ben J Morris. Improved bounds for sampling contingency tables. Random Structures & Algorithms , 21(2):135-146, 2002.

Nathan E Lewis, Harish Nagarajan, and Bernhard O Palsson. Constraining the metabolic genotype-phenotype relationship using a phylogeny of in silico methods. Nature Reviews Microbiology , 10(4):291-305, 2012.

Li Han and Lee Rudolph. Inverse kinematics for a serial chain with joints under distance constraints. In Robotics: Science and Systems (RSS) , 2006.

Tsuneo Yoshikawa. Manipulability of robotic mechanisms. The international journal of Robotics Research , 4(2):3-9, 1985.

Ruth J Williams. Reflected brownian motion with skew symmetric data in a polyhedral domain. Probability Theory and Related Fields , 75(4):459-485, 1987.

Andrey Pilipenko. An introduction to stochastic differential equations with reflection , volume 1. Universitätsverlag Potsdam, 2014.

Aaron Lou and Stefano Ermon. Reflected diffusion models. In International Conference on Machine Learning (ICML) , 2023.

Sebastien Bubeck, Ronen Eldan, and Joseph Lehec. Finite-time analysis of projected langevin monte carlo. In Advances in Neural Information Processing Systems (NeurIPS) , 2015.

Kanji Sato, Akiko Takeda, Reiichiro 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.

Ya-Ping Hsieh, Ali Kavis, Paul Rolland, and Volkan Cevher. Mirrored langevin dynamics. In Advances in Neural Information Processing Systems (NeurIPS) , 2018.

Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators , volume 103. Springer, 2014.

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.

Kelvin Shuangjian Zhang, Gabriel Peyré, Jalal Fadili, and Marcelo Pereyra. Wasserstein control of mirror langevin monte carlo. In Conference on Learning Theory (COLT) , 2020.

Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, Philippe Rigollet, and Austin Stromme. Exponential ergodicity of mirror-Langevin diffusions. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Kwangjun Ahn and Sinho Chewi. Efficient constrained sampling via the mirror-Langevin algorithm. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

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.

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 , 5(3):e105-e106, 2023.

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.

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.

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 (NeurIPS) , 2022.

Yunqing Zhao, Tianyu Pang, Chao Du, Xiao Yang, Ngai-Man Cheung, and Min Lin. A recipe for watermarking diffusion models. arXiv preprint arXiv:2303.10137 , 2023.

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.

Peter D Lax. Change of variables in multiple integrals. The American mathematical monthly , 106(6):497-501, 1999.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Sébastien Bubeck and Ronen Eldan. The entropic barrier: a simple and optimal universal self-concordant barrier. In Conference on Learning Theory (COLT) , 2015.

Arkadij Semenoviˇ c Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization . Wiley-Interscience, 1983.

Amir Beck and Marc Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters , 31(3):167-175, 2003.

Yousef Saad. Iterative methods for sparse linear systems . SIAM, 2003.

David M Blei, Andrew Y Ng, and Michael I Jordan. Latent dirichlet allocation. Journal of Machine Learning Research (JMLR) , 2003.

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.

Lixin Lang, Wen-shiang Chen, Bhavik R Bakshi, Prem K Goel, and Sridhar Ungarala. Bayesian estimation via sequential monte carlo sampling-constrained dynamic systems. Automatica , 43 (9):1615-1622, 2007.

Michael A Gelbart, Jasper Snoek, and Ryan P Adams. Bayesian optimization with unknown constraints. In Conference on Uncertainty in Artificial Intelligence (UAI) , 2014.

Yann Brenier. Décomposition polaire et réarrangement monotone des champs de vecteurs. CR Acad. Sci. Paris Sér. I Math. , 305:805-808, 1987.

Nicolas Brosse, Alain Durmus, Éric Moulines, and Marcelo Pereyra. Sampling from a logconcave distribution with compact support with proximal Langevin Monte Carlo. In Conference on Learning Theory (COLT) , 2017.

Sam Patterson and Yee Whye Teh. Stochastic gradient Riemannian Langevin dynamics on the probability simplex. In Advances in Neural Information Processing Systems (NeurIPS) , 2013.

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.

Ricky TQ Chen, Brandon Amos, and Maximilian Nickel. Semi-discrete normalizing flows through differentiable tessellation. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Xingchao Liu, Lemeng Wu, Mao Ye, et al. Learning diffusion bridges on constrained domains. In International Conference on Learning Representations (ICLR) , 2022.

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.

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.

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.

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.

Julen Urain, Niklas Funk, Georgia Chalvatzaki, and Jan Peters. SE (3)-DiffusionFields: Learning cost functions for joint grasp and motion optimization through diffusion. In IEEE International Conference on Robotics and Automation (ICRA), Workshop Track , 2022.

Pierre H Richemond, Sander Dieleman, and Arnaud Doucet. Categorical SDEs with simplex diffusion. arXiv preprint arXiv:2210.14784 , 2022.

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.

Jason Gaitonde, Jon Kleinberg, and Éva Tardos. Polarization in geometric opinion dynamics. In ACM Conference on Economics and Computation , pages 499-519, 2021.

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.

Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

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.

Dongjun Kim, Byeonghu Na, Se Jung Kwon, Dongsoo Lee, Wanmo Kang, and Il-Chul Moon. Maximum likelihood training of implicit nonlinear diffusion models. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision , 51:22-45, 2015.

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.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2019.

Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2020.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. I 2 SB: Image-to-Image Schrödinger bridge. In International Conference on Machine Learning (ICML) , 2023.

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, highperformance deep learning library. In Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019. URL http://papers.neurips.cc/paper/ 9015-pytorch-an-imperative-style-high-performance-deep-learning-library. pdf .

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR) , 2019.

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.

Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pages 3-19, 2018.

A Derivation of Mirror Mappings

Here, we provide addition derivation of ∇ φ ∗ . Computation of ∇ φ ( x ) and ∇ 2 φ ∗ ( y ) follow straightforwardly by differentiating φ ( x ) and ∇ φ ∗ ( y ) w.r.t. x and y , respectively.

glyph[lscript] 2 -Ball Since the gradient map also reverses the mirror map, we aim to rewrite y = 2 γ R -‖ x ‖ 2 2 x as x = f ( y ) = ∇ φ ∗ ball ( y ) . Solving the second-order polynomial,

With that, we can rewrite Equation (11) by

Simplex Standard calculations in convex analysis [28] shows

Differentiating Equation (21) w.r.t. y yields ∇ φ ∗ simplex in Equation (13).

Polytope Since the gradient map also reverses the mirror map, we aim to inverse

When all d constraints are orthonormal, taking inner product between y and each a yields

Therefore, we can reconstruct x from y via

which defines x = ∇ φ ∗ polytope ( y ) . For completeness, the Hessian can be presented compactly as

where I is the identity matrix, A := [ a 1 , · · · , a m ] is a d -bym matrix whose column vector a i corresponds to each constraint, and Σ ∈ R m × m is a diagonal matrix with leading entries

yields

B Additional Remarks on Polytope

Derivation of Equation (17) Since the subspaces spanned by { a i } and { a j } are orthogonal to each other, we can rewrite (15) as

∇ φ ∗ polytope ( y ) follows similar derivation.

Generalization to non-orthonormal constraints The mirror maps of a polytope, as 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 to extend the computation to non-orthonormal constraints by identifying the corresponding 'coefficients' through a change of bases, utilizing the reproducing formula:

and A := [ a 1 , · · · , a m ] . Similarly, we have y = ∑ d i =1 〈 ˜ a i , y 〉 a i . Applying similar derivation leads to

C Experiment Details & Additional Results

Dataset & constrained sets

glyph[lscript] 2 -balls constrained 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 , at the corner of each dimension, and reject samples outside the constrained sets.

Simplices constrained sets : We consider Dirichlet distributions [48], Dir( α ) , with various concentration parameters α detailed in Table 7.

Hypercube constrained sets : For all dimensions d = { 2 , 3 , 6 , 8 , 20 } , we place d isotropic Gaussians, each with variance 0 . 2 , at the corner of each dimension, and either reject ( d = { 2 , 3 , 6 , 8 } ) or reflect ( d = 20 ) samples outside the constrained sets.

Watermarked datasets and polytope constrained sets : Wefollow the same data preprocessing from EDM 5 [38] and rescale both FFHQ and AFHQv2 to 64 × 64 image resolution. For the polytope constrained sets M := { x ∈ R d : c i < a glyph[latticetop] i x < b i , ∀ i } , we construct a i from orthonormalized Gaussian random vectors and detail other hyperparameters in Table 8.

Implementation All methods are implemented in PyTorch [76]. We adopt ADM 6 and EDM 5 [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 7 of Reflected Diffusion [25] in Table 11. 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 AdamW [77] and an 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 5 [38]. All experiments are conducted on two TITAN RTXs and one RTX 2080.

Network For constrained generation, all networks take ( y, t ) as inputs and follow

where timestep_embedding( · ) is the standard sinusoidal embedding. t_mod and out_mod consist of 2 fully-connected layers ( Linear ) activated by the Sigmoid Linear Unit ( SiLU ) [78]:

and y_mod consists of 3 residual blocks, i.e., y_mod( y ) = y + res_mod(norm( y )) , where

All Linear 's have 128 hidden dimension. We use group normalization [79] for all norm . For watermarked generation, we use EDM parameterization 5 [38].

Evaluation We compute the Wasserstein and Sliced Wasserstein distances using the geomloss 8 and ot 9 packages, respectively. The Maximum Mean Discrepancy (MMD) is based on the popular package https://github.com/ZongxianLee/MMD_Loss.Pytorch , which is unlicensed. For watermarked generation, we follow the same evaluation pipeline from EDM 5 [38] by first generating 50,000 watermarked samples and computing the FID w.r.t. the training statistics.

False-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., if x ∈ M . Hence, the false-positive 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 false-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 in Equation (8) Figure 10 demonstrates how MDM faithfully captures the variational bound to the negative log-likelihood (NLL) of 2-dimensional GMM.

More constrained sets, distributional metrics, & baseline Tables 9 and 10 expand the analysis in Tables 3 and 4 with additional distributional metrics such as Wasserstein-1 ( W 1 ) and Maximum Mean Discrepancy (MMD). Additionally, Table 11 reports the results of hypercube [0 , 1] d constrained set, a special instance

MDM NLL

of polytopes, and includes additional baseline from Lou and Ermon [25], which approximate the

intractable scores in Reflected Diffusion using eigenfunctions tailored specifically to hypercubes, rather than implicit score matching as in Fishman et al. [18]. Consistently, our findings conclude that the MDM is the only constrained-based diffusion model that achieves comparable or better performance to DDPM. These results affirm the effectiveness of MDM in generating high-quality samples within constrained settings, making it a reliable choice for constrained generative modeling.

More watermarked samples Figures 11 and 12 provide additional qualitative results on the watermarked samples generated by MDMs.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n(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.\n\nWEAKNESSES:\n(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.\n\nQUESTIONS:\n(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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nWhen 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: 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.

STRENGTHS:
Using the idea of mirror map seems a novel approach.
This 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?

The 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.

Some very related literature on learning diffusion models on constrained domains are missing [1, 2] and need to be discussed.

[1] Learning Diffusion Bridges on Constrained Domains
[2] First Hitting Diffusion Models for Generating Manifold, Graph and Categorical Data

QUESTIONS:
See above.

LIMITATIONS:
See above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThe authors sufficiently discuss broader impact in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}