[论文解读] Diffusion Models for Constrained Domains

Nic Fishman, Leo Klarner+3 more|arXiv (Cornell University)|Apr 11, 2023

Reservoir Engineering and Simulation Methods被引用 5

一句话总结

本文提出了两种针对由不等式约束定义的约束流形的新颖扩散模型，传统黎曼扩散方法因测地线未定义而失效。该文提出基于对数障碍的加噪过程与反射布朗运动方法，实现了在凸多面体和对称正定矩阵上的基于得分的生成建模，在蛋白质主链和机械臂构象建模中取得实证成功。

ABSTRACT

Denoising diffusion models are a novel class of generative algorithms that achieve state-of-the-art performance across a range of domains, including image generation and text-to-image tasks. Building on this success, diffusion models have recently been extended to the Riemannian manifold setting, broadening their applicability to a range of problems from the natural and engineering sciences. However, these Riemannian diffusion models are built on the assumption that their forward and backward processes are well-defined for all times, preventing them from being applied to an important set of tasks that consider manifolds defined via a set of inequality constraints. In this work, we introduce a principled framework to bridge this gap. We present two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints. As existing diffusion model techniques cannot be applied in this setting, we derive new tools to define such models in our framework. We then demonstrate the practical utility of our methods on a number of synthetic and real-world tasks, including applications from robotics and protein design.

研究动机与目标

• 解决现有黎曼扩散模型的局限性，其要求全局定义的测地线，因此在由不等式约束定义的流形上失效。
• 实现在凸多面体、对称正定矩阵以及机器人与蛋白质设计中流形的生成建模。
• 开发与扩散模型的时间反演和得分匹配框架兼容的合理加噪过程——对数障碍与反射布朗运动。
• 将扩散模型的理论与实践框架扩展至非测地线完备的流形，拓宽其在真实科学领域中的适用性。
• 通过合成与真实世界任务，在高维约束流形上展示所提模型的可扩展性与灵活性。

提出的方法

• 基于对数障碍度量设计前向加噪过程，该度量诱导出尊重不等式约束的黎曼结构，并确保扩散路径的明确定义。
• 提出基于反射布朗运动的加噪过程，使轨迹在约束域边界处发生反射，从而保持流形的几何结构。
• 通过将得分匹配与斯蒂尔得分估计扩展至非欧几里得、约束设置，推导出两种加噪过程的时间反演SDE。
• 通过在约束流形上使用去噪目标训练神经网络以近似后向漂移（斯蒂尔得分），构建基于得分的生成模型。
• 采用低温采样方法，通过温度参数控制生成分布中样本保真度与多样性的权衡。
• 使用最大均值差异（MMD）的U-统计量近似来评估真实样本与生成样本之间的分布相似性。

Figure 1 : The behaviour of different types of noising processes considered in this work defined on the unit interval. $\mathbf{B}_{t}$ : Euclidean (unconstrained) Brownian motion. $\hat{\mathbf{B}}_{t}$ : Log barrier forward noising process. $\bar{\mathbf{B}}_{t}$ : Reflected Brownian motion. All s

实验结果

研究问题

• RQ1 扩散模型能否被扩展至测地线无法在所有时间定义的不等式约束流形？
• RQ2 如何为约束流形构建合理的加噪过程，以保持几何结构并支持时间反演？
• RQ3 对数障碍度量与反射布朗运动能否作为标准测地线基加噪在扩散建模中的可行替代方案？
• RQ4 所提模型在高维约束域（如Birkhoff多面体与蛋白质构象空间）中的表现如何？
• RQ5 模型在机器人与计算生物学等真实世界应用中的泛化能力如何？

主要发现

• 所提出的对数障碍与反射布朗运动加噪过程成功实现了在标准黎曼扩散失效的约束流形上的扩散建模。
• 在超立方体、单纯形与Birkhoff多面体上的实证结果表明，模型实现了低MMD值，表明真实数据与生成数据间具有强分布对齐。
• 在大小为 $ n=4 $ 的Birkhoff多面体上，模型以高保真度生成了有效的双随机矩阵，经成对与边缘样本可视化验证。
• 该方法可推广至高维设置，包括 $ n=10 $ 的蛋白质主链构象与机械臂构型，展示了可扩展性。
• 采用 $ heta_0 \in \{1,2,3,5,10\} $ 的低温采样可控制样本多样性与保真度，较高 $ heta_0 $ 产生更锐利、更集中的样本。
• 在单纯形上的可视化显示，得分网络预测的漂移随时间合理演化，证实了训练的稳定性与对底层得分函数的有意义学习。

Figure 2 : A convex polytope defined by six constraints $\{f_{i}\}_{i\in\mathcal{I}}$ , along with the log barrier potential, and ‘straight trajectories’ under the log-barrier metric and under the Euclidean metric with and without reflection at the boundary.