Neural Approximate Mirror Maps for Constrained Diffusion Models
作者: Berthy T. Feng, Ricardo Baptista, Katherine L. Bouman
分类: cs.LG, cs.CV, eess.IV
发布日期: 2024-06-18 (更新: 2025-04-09)
备注: ICLR 2025
💡 一句话要点
提出神经近似镜像映射，用于约束扩散模型以满足复杂约束条件。
🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)
关键词: 约束扩散模型 镜像扩散模型 神经近似镜像映射 生成模型 逆问题
📋 核心要点
- 现有约束扩散模型方法在处理复杂约束（如非凸约束）时存在局限性，解析镜像映射难以推导。
- 提出神经近似镜像映射（NAMM），通过学习近似镜像映射和逆映射，将数据转换到无约束空间进行扩散建模。
- 实验表明，基于NAMM的镜像扩散模型（MDM）在满足各种约束方面优于无约束扩散模型，并可用于解决约束逆问题。
📝 摘要（中文）
扩散模型在生成视觉上逼真的图像方面表现出色，但常常难以满足训练数据中固有的细微约束。这些约束可能是基于物理的（例如，满足偏微分方程）、几何的（例如，尊重对称性）或语义的（例如，包含特定数量的对象）。当训练数据都满足特定约束时，在扩散模型上强制执行此约束使其更可靠地生成有效的合成数据并解决约束逆问题。然而，现有的约束扩散模型方法在它们可以处理的约束方面受到限制。例如，最近的工作提出了学习镜像扩散模型（MDM），但解析镜像映射仅存在于凸约束中，并且可能难以推导。我们提出了用于一般（可能非凸）约束的神经近似镜像映射（NAMM）。我们的方法只需要一个来自约束集的可微距离函数。我们学习一个近似镜像映射，将数据转换成一个无约束空间，以及一个相应的近似逆映射，将数据映射回约束集。然后，可以在学习的镜像空间中训练生成模型（例如，MDM），并且可以通过逆映射将其样本恢复到约束集。我们在各种约束上验证了我们的方法，表明与无约束扩散模型相比，基于NAMM的MDM显著提高了约束满足度。我们还展示了如何轻松地将现有的基于扩散的逆问题求解器应用于学习的镜像空间，以解决约束逆问题。
🔬 方法详解
问题定义：现有的约束扩散模型方法，特别是基于镜像扩散模型（MDM）的方法，在处理非凸约束或难以推导解析镜像映射的约束时面临挑战。这些方法通常依赖于特定的约束形式，缺乏通用性。因此，需要一种能够处理更广泛约束类型的扩散模型框架。
核心思路：核心思路是学习一个近似的镜像映射，将受约束的数据映射到一个无约束的空间，在这个空间中可以训练标准的扩散模型。同时，学习一个近似的逆映射，将扩散模型生成的样本映射回原始的约束空间。这样，就可以在无约束空间中利用扩散模型的强大生成能力，同时保证生成的样本满足原始的约束条件。
技术框架：整体框架包含两个主要模块：神经近似镜像映射（NAMM）和镜像扩散模型（MDM）。NAMM由一个编码器（近似镜像映射）和一个解码器（近似逆映射）组成，它们都是神经网络。编码器将受约束的数据映射到无约束空间，解码器将无约束空间的数据映射回约束空间。MDM是在NAMM学习到的无约束空间中训练的扩散模型。生成样本时，首先由MDM在无约束空间生成样本，然后通过NAMM的解码器将其映射回约束空间。
关键创新：最重要的创新点在于使用神经网络来近似镜像映射和逆映射，从而摆脱了对解析镜像映射的依赖，使得可以处理更广泛的约束类型，包括非凸约束。此外，该方法只需要一个可微的距离函数来衡量数据点与约束集的距离，这使得该方法更易于应用。
关键设计：NAMM的编码器和解码器可以使用各种神经网络结构，例如MLP或卷积神经网络。损失函数通常包括两部分：一部分是重构损失，用于保证编码器和解码器能够相互逆转；另一部分是约束损失，用于惩罚解码器生成的样本不满足约束的情况。约束损失通常基于数据点到约束集的距离函数。具体的网络结构、损失函数权重等超参数需要根据具体的约束类型进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，基于NAMM的MDM在各种约束条件下，显著提高了约束满足度。例如，在满足特定数量对象的约束下，NAMM-MDM能够更准确地生成包含指定数量对象的图像。与无约束扩散模型相比，NAMM-MDM在约束满足度方面取得了显著提升，验证了该方法的有效性。
🎯 应用场景
该研究成果可广泛应用于需要生成满足特定约束条件的数据的领域，例如：物理模拟（生成满足物理定律的场景）、计算机图形学（生成具有特定几何属性的3D模型）、图像编辑（在保持图像语义一致性的前提下进行修改）以及科学研究（生成符合特定科学假设的数据）。该方法能够提升生成数据的可靠性和有效性，并为解决约束逆问题提供新的思路。
📄 摘要（原文）
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 makes it more reliable for generating valid synthetic data and solving constrained inverse problems. However, existing methods for constrained diffusion models are restricted in the constraints they can handle. For instance, recent work proposed to learn mirror diffusion models (MDMs), but analytical mirror maps only exist for convex constraints and can be challenging to derive. We propose neural approximate mirror maps (NAMMs) for general, possibly non-convex constraints. Our approach only requires a differentiable distance function from the constraint set. We learn an approximate mirror map that transforms 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.