SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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_i^{-1}$ , and  $(iii)$  an one-time computation of orthonormal bases, which scales as  $\mathcal{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,\dots ,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
$$
\nabla \phi_ {\mathrm {p o l y}} (x) = x + \sum_ {i = 1} ^ {m} \left(s _ {i} \left(\left\langle a _ {i}, x \right\rangle\right) - \left\langle a _ {i}, x \right\rangle\right) a _ {i}, \nabla \phi_ {\mathrm {p o l y}} ^ {*} (y) = y + \sum_ {i = 1} ^ {m} \left(s _ {i} ^ {- 1} \left(\left\langle a _ {i}, y \right\rangle\right) - \left\langle a _ {i}, y \right\rangle\right) a _ {i}, \tag {17}
$$
4The transformation allows  $x$  to be bounded in a constrained set  $\mathcal{M}\subseteq \mathbb{R}^d$  rather than a hyperplane in  $\mathbb{R}^{d + 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  $\mathcal{O}(md^2) \approx \mathcal{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_i^{-1}$  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:
$$
s _ {i} \left(\left\langle a _ {i}, x \right\rangle\right) := \left(\tanh ^ {- 1} \circ \operatorname {r e s c a l e} _ {i}\right) \left(\left\langle a _ {i}, x \right\rangle\right), \quad s _ {i} ^ {- 1} \left(\left\langle a _ {i}, y \right\rangle\right) := \left(\operatorname {r e s c a l e} _ {i} ^ {- 1} \circ \tanh  \right. \left(\left\langle a _ {i}, y \right\rangle\right), \tag {18}
$$
where  $\operatorname{rescale}_i(z) := 2\frac{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  $\nabla \phi$  and  $\nabla \phi^{*}$  for each constrained set.
l2-Ball	Simplex	Polytope
O(d)	O(d)	O Md)
Remark 1. Equations (17) and (18) can be generalized to non-orthonormal constraints by adopting  $\langle \tilde{a}_i, x \rangle$  and  $\langle \tilde{a}_i, y \rangle$ , where  $\tilde{a}_i$  is the  $i$ -th row of the matrix  $\tilde{\mathbf{A}} := (\mathbf{A}^\top \mathbf{A})^{-1} \mathbf{A}^\top$  and  $\mathbf{A} := [a_1, \dots, a_m]$ . Indeed, when  $\mathbf{A}$  is orthonormal, we recover  $\tilde{a}_i = a_i$ . We leave more discussions to Appendix B.
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  $\ell_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 MDM on the initial constrained set are implicitly defined and can be highly nonlinear.
Table 3: Results of  $\ell_2$ -ball constrained sets on 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.,  $\mathcal{W}_1$ , are reported in Appendix C. sc
d=2	d=2	d=6	d=8	d=20
Sliced Wasserstein ↓
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	0.00 ± 0.00	8.67 ± 0.87	13.60 ± 0.62	19.33 ± 1.29
Table 4: Results of simplices constrained sets on five synthetic datasets of various dimension  $d$ .
d=3	d=3	d=7	d=9	d=20
Sliced Wasserstein ↓
DDPM [2]	0.0089 ± 0.0002	0.0110 ± 0.0032	0.0047 ± 0.0004	0.0053 ± 0.0003	0.0031 ± 0.0003
Reflected [18]	0.0233 ± 0.0019	0.0336 ± 0.0009	0.0411 ± 0.0039	0.0897 ± 0.0112	0.0231 ± 0.0011
MDM (ours)	0.0074 ± 0.0008	0.0169 ± 0.0016	0.0051 ± 0.0006	0.0040 ± 0.0003	0.0027 ± 0.0000
Constraint violation (%) ↓
DDPM [2]	0.73 ± 0.12	14.40 ± 1.39	11.63 ± 0.90	27.53 ± 0.57	68.83 ± 1.66
Figure 5: Comparison between MDM and Reflected Diffusion [18] in modeling a Dirichlet distribution on a 7-dimensional simplex. We visualize the joint densities between  $x_{3:4}$  and  $x_{5:6}$ .
Table 5: Runtime and memory complexity w.r.t. DDPM [2].
Runtime	Memory
MDM	108%	100%
Reflected
Diff. [18]	>1200%	905%
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  $\ell_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.
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 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.
MDM matches 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.
Table 6: The 50k-FID for unconditional watermarked generation. Watermark precision denotes the percentage of constraint-satisfied images that are generated by MDMs, controlled by loosing/tightening the constrained sets. As MDM-dual is trained on constraint-projected images, we also report its  $FIDs^{*}$  w.r.t. the shifted distributions. The prior watermarked diffusion model [39] reported 5.03 on FFHQ and 4.32 on AFHQv2.
Precision	FFHQ 64×64	AFHQv2 64×64
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
MDM-dual	2.96	4.57	15.74	2.23	2.86	6.79
2.65*	2.93*	4.40*	2.21*	2.32*	3.05*
Figure 7: Histogram of  $\langle a_i, x \rangle$  for original and MDM-watermarked images  $x$ . Dashed lines indicate the constrained set.
Figure 8: Ablation studies on how (top) number of constraints  $m$  and (bottom) their ranges, i.e.,  $(c_i, b_i)$ , affect the FID and precision of MDM-proj on AFHQv2.
Figure 6: Uncurated watermarked samples by MDMs. More samples can be found in Appendix C.
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  $\mathcal{M} := \{x \in \mathbb{R}^d : c_i < a_i^\top x < b_i, \forall i\}$ . These parameters  $\{a_i, b_i, c_i\}_{i=1}^m$  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  $\mathcal{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 \times 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], MDM stands 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 \times 256$ . Specifically, we focus on image restoration tasks, where the goal
Degraded (JPEG QF=5) Input y
Ground Truth Reference  $(x\notin M)$
MDM-dual Watermarked Output  $(x\in \mathcal{M})$
MDM-proj Watermarked Output  $(x\in \mathcal{M})$
Figure 9: Conditional watermarked generation on ImageNet  $256 \times 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  $x$  (upper-right) by learning  $p(x|y)$ . 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 \in \mathcal{M}$ .
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 \in \mathcal{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.
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 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.
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.
The authors would like to thank Bogdon Vlahov for sharing the computational resources. GHL, TC, ET are supported by ARO Award #W911NF2010151 and DoD Basic Research Office Award HQ00342110002. MT is partially supported by NSF DMS-1847802, NSF ECCS-1942523, Cullen-Peck Scholarship, and GT-Emory AI.Humanity Award.
[1] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised 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 in 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 Pattern 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, 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.
[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-3D content creation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2023.
[10] Brian DO Anderson. Reverse-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 on Learning Representations (ICLR), 2021.
[12] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[13] Chitwan Sahara, 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 Thornton, Yee Whye Teh, 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.03660, 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 J 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 a phylogeny of in silico methods. Nature Reviews Microbiology, 10(4):291-305, 2012.
[21] Li Han and Lee Rudolph. Inverse kinematics for a 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. Universitatsverlag 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 in Neural Information Processing Systems (NeurIPS), 2015.
[27] 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.
[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 langlevin 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 Peyre, 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 mirror-Langevin diffusions. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[33] Kwangjun Ahn and Sinho Chewi. Efficient constrained sampling via the mirror-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.