[论文解读] Mirror Diffusion Models for Constrained and Watermarked Generation

Guan-Horng Liu, Tianrong Chen+2 more|arXiv (Cornell University)|Oct 2, 2023

Markov Chains and Monte Carlo MethodsMathematics被引用 3

一句话总结

本文提出镜像扩散模型（MDM），一种新型扩散模型，通过镜像映射将扩散过程提升至无约束对偶空间，在凸约束集（如单纯形和ℓ₂-球）上生成数据。该方法实现无模拟、可解析训练，具有精确的转移核，性能优于以往的反射扩散方法，同时可通过基于约束的嵌入实现不可见、可量化的水印技术。

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. Our code is available at https://github.com/ghliu/mdm

研究动机与目标

• 解决现有针对单纯形和ℓ₂-球等凸集约束数据的可解析、无模拟扩散模型缺失的问题。
• 克服反射扩散方法的局限性，后者缺乏闭式转移核且需近似得分函数。
• 在不损害训练可扩展性的前提下，实现在复杂几何约束下的鲁棒、高保真生成。
• 探索利用约束集作为机制，在生成数据中嵌入不可见、可量化的水印，以实现版权保护。
• 证明利用镜像映射可将约束生成问题转化为可解析的欧氏空间学习问题。

提出的方法

• 通过严格凸函数ϕ的梯度∇ϕ，将定义在约束集𝒪 ⊆ ℝᵈ上的数据分布映射至无约束对偶空间。
• 在对偶空间ℝᵈ中执行扩散建模，其中转移核与得分函数具有解析可计算性，从而实现无模拟训练。
• 利用逆镜像映射∇ϕ⁻¹将生成样本映射回原始约束集，通过构造保证约束被满足。
• 基于已知的凸共轭对偶性，高效计算常见约束集（如单纯形和ℓ₂-球）的镜像映射。
• 通过在对偶空间中将水印编码为约束，实现水印嵌入，确保内容完整性不可见但可验证。
• 实现两种变体：MDM-dual（在完整对偶空间中学习）与MDM-proj（将训练数据投影至约束集），并在图像生成基准上比较性能。

实验结果

研究问题

• RQ1 当在单纯形和ℓ₂-球等凸约束集上生成数据时，能否使扩散模型具备可解析性与无模拟训练能力？
• RQ2 在对偶空间中使用镜像映射如何在保持几何约束的同时维持生成质量？
• RQ3 能否利用约束集作为机制，在生成数据中嵌入不可见、可量化的水印？
• RQ4 在约束扩散模型中，水印鲁棒性与生成保真度之间的权衡关系如何？
• RQ5 MDM的两种变体（MDM-dual与MDM-proj）在无条件与条件图像生成任务中的性能表现如何？

主要发现

• 与以往的反射扩散方法相比，MDM在无条件图像生成任务（如CelebA-HQ、AFHQv2）中显著提升了FID得分，表明样本质量更优。
• 尽管MDM-dual学习的是更平滑的对偶空间分布，但其FID始终低于MDM-proj，表明投影导致的分布偏移会损害性能。
• 在条件ImageNet 256×256恢复任务中，水印精度超过95%，误报率低于0.001%，证实了水印嵌入的有效性与鲁棒性。
• MDM-dual与MDM-proj生成的图像均满足多面体约束，而非MDM基线模型即使在视觉上难以区分，仍会违反约束。
• 消融实验表明，减少约束数量（m）并放宽约束范围（cᵢ, bᵢ）可同时提升FID与水印精度，揭示了保护强度与生成质量之间的权衡。
• 该方法实现了无模拟训练与精确回归目标，而以往的反射扩散方法需依赖几何近似或模拟。