Diffusion models excel at creating visually-convincing images, but they often struggle to meet subtle constraints inherent in the training data. Such constraints could be physics-based (e.g., satisfying a PDE), geometric (e.g., respecting symmetry), or semantic (e.g., including a particular number of objects). When the training data all satisfy a certain constraint, enforcing this constraint on a diffusion model not only improves its distribution-matching accuracy but also makes it more reliable for generating valid synthetic data and solving constrained inverse problems. However, existing methods for constrained diffusion models are inflexible with different types of constraints. Recent work proposed to learn mirror diffusion models (MDMs) in an unconstrained space defined by a mirror map and to impose the constraint with an inverse mirror map, but analytical mirror maps are challenging to derive for complex constraints. We propose neural approximate mirror maps (NAMMs) for general constraints. Our approach only requires a differentiable distance function from the constraint set. We learn an approximate mirror map that pushes data into an unconstrained space and a corresponding approximate inverse that maps data back to the constraint set. A generative model, such as an MDM, can then be trained in the learned mirror space and its samples restored to the constraint set by the inverse map. We validate our approach on a variety of constraints, showing that compared to an unconstrained diffusion model, a NAMM-based MDM substantially improves constraint satisfaction. We also demonstrate how existing diffusion-based inverse-problem solvers can be easily applied in the learned mirror space to solve constrained inverse problems.

翻译：扩散模型在生成视觉上逼真的图像方面表现出色，但在满足训练数据中固有的细微约束方面常常存在困难。这些约束可以是基于物理的（例如满足偏微分方程）、几何的（例如保持对称性）或语义的（例如包含特定数量的对象）。当训练数据均满足特定约束时，在扩散模型中强制执行该约束不仅能提高其分布匹配的准确性，还能使其在生成有效合成数据及解决约束反问题方面更加可靠。然而，现有约束扩散模型的方法对不同类型约束的适应性不足。近期研究提出在由镜像映射定义的无约束空间中学习镜像扩散模型，并通过逆镜像映射施加约束，但解析镜像映射对于复杂约束难以推导。本文针对一般约束提出神经近似镜像映射方法。我们的方法仅需约束集的可微距离函数。我们学习将数据推入无约束空间的近似镜像映射，以及将数据映射回约束集的对应近似逆映射。随后可在学习到的镜像空间中训练生成模型（如镜像扩散模型），并通过逆映射将样本恢复至约束集。我们在多种约束条件下验证了该方法，结果表明：与无约束扩散模型相比，基于神经近似镜像映射的镜像扩散模型显著提升了约束满足度。我们还展示了现有基于扩散的反问题求解器如何在学习到的镜像空间中轻松应用于解决约束反问题。