Mirror Diffusion Models for Constrained and Watermarked Generation

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.

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. (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. (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. (W2) It would indeed be beneficial to observe experiments conducted on more realistic datasets, such as CIFAR-100 and ImageNet.

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$? (Q2) Can the authors provide more experimental results on more realistic data, like CIFAR-100 and ImageNet?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

When the data distribution is constrained in some boundary, This 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.

Strengths

* The paper has good clarity in introducing theorems, examples, and applications. * 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.

Questions

* How many sampling steps did you used for generating in each dataset? ===== Correction * One of the xlabels in Figure 6 should be `MDM-dual`. Now both two are `MDM-proj`. * Line 514: polytop --> polytope * Figure 8: top and bottom --> left and right

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors adequately addressed the limitations.

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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

See above.

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

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. The 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. Otherwise, I see no major weaknesses. There are some typos, which I'll briefly mention: * In standard Euclidean spaces, tractable marginal can be * The marginal q(yt−1 |yt , y0 ) hints the optimal reverse Appendix: * Boarder impact * polytop

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? * 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? * 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? * 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? Overall, I hope the authors could clarify these questions and rewrite Section 5.2 to be easier to understand. A few more questions/comments: * 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)? * 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 * 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) * 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

2 fair

Contribution

3 good

Limitations

The authors sufficiently discuss broader impact in the appendix.

Thanks

Thank you for providing this detailed response, it would be great if these clarifications could be included in the update. I have no further questions!

Response to the official review

Thank you to the authors for the thoughtful responses. I keep the current score.

Thanks

The rebuttal partially addresses my concerns and I thus increase my score.