[论文解读] Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks

David Alvarez-Melis, Yair Schiff+1 more|Digital Access to Scholarship at Harvard (DASH) (Harvard University)|Jun 1, 2021

Topological and Geometric Data Analysis参考文献 49被引用 9

一句话总结

该论文提出 JKO-ICNN，一种利用输入-凸神经网络（ICNNs）参数化 JKO 方案中凸函数的可扩展方法，用于在概率测度空间上优化泛函。该方法实现了稳定、高维的泛函优化，并具备收敛性保证，已在低维 PDE 和高维分子生成任务中得到验证，显著提升了药物相似性（QED）得分。

ABSTRACT

Gradient flows are a powerful tool for optimizing functionals in general metric spaces, including the space of probabilities endowed with the Wasserstein metric. A typical approach to solving this optimization problem relies on its connection to the dynamic formulation of optimal transport and the celebrated Jordan-Kinderlehrer-Otto (JKO) scheme. However, this formulation involves optimization over convex functions, which is challenging, especially in high dimensions. In this work, we propose an approach that relies on the recently introduced input-convex neural networks (ICNN) to parametrize the space of convex functions in order to approximate the JKO scheme, as well as in designing functionals over measures that enjoy convergence guarantees. We derive a computationally efficient implementation of this JKO-ICNN framework and experimentally demonstrate its feasibility and validity in approximating solutions of low-dimensional partial differential equations with known solutions. We also demonstrate its viability in high-dimensional applications through an experiment in controlled generation for molecular discovery.

研究动机与目标

• 解决在高维空间中优化概率测度上泛函的计算挑战。
• 克服在高维空间中 JKO 方案下凸函数优化的困难。
• 开发一种可扩展的、可微的框架，用于在 Wasserstein 空间上实现梯度流，结合深度学习方法。
• 实现对复杂概率泛函（如分子生成中的泛函）的实际优化，并具备理论收敛性保证。
• 在具有已知解的低维 PDE 与高维实际应用中，验证该方法的可行性。

提出的方法

• 使用输入-凸神经网络（ICNNs）参数化 JKO 方案中的凸势函数，确保其关于输入变量的凸性。
• 通过 Brenier 定理将 JKO 变分问题重新表述为对凸函数的优化，从而实现可微优化。
• 通过有限样本和自适应梯度下降法实现 JKO-ICNN，结合自动微分完成端到端训练。
• 采用凸代理模型（残差 ICNN）从分子嵌入中预测 QED 得分，确保势泛函的凸性。
• 将分布距离度量（Sinkhorn 或 MMD）与可学习权重结合，引导流向目标分布。
• 通过学习可泛化的传输映射来分摊计算成本，实现在未见样本上的高效推理。

实验结果

研究问题

• RQ1 ICNN 是否能有效参数化 JKO 方案中高维概率优化所需的凸函数空间？
• RQ2 JKO-ICNN 框架在实现高维可扩展、可微优化的同时，是否仍保持收敛性保证？
• RQ3 JKO-ICNN 在近似具有解析解的低维 PDE 已知解方面表现如何？
• RQ4 JKO-ICNN 是否能成功优化高维分子生成中的复杂泛函（如药物相似性 QED）？
• RQ5 不同的分布距离度量（Sinkhorn 与 MMD）及加权策略对模式崩溃与优化稳定性有何影响？

主要发现

• JKO-ICNN 成功近似了具有已知解析解的低维 PDE 解，表现出良好的准确性和收敛性。
• 在 QM9 数据集上，使用 MMD 且 λ₂ = 1,000 时，QED 中位数从 0.315 提升至 0.452，相对提升 44.8%。
• 使用 Sinkhorn 距离且 λ₂ = 10,000 时，QED 中位数达到 0.419，相比初始分布提升 33.0%。
• 在分子生成中，方法实现了 92.7% 的有效性与 81.9% 的唯一性，表明模式覆盖良好且结构质量高。
• 实验表明，高 λ₂ 值（10,000）或基于 MMD 的距离会导致模式崩溃，凸显对超参数选择的敏感性。
• 所学习的传输映射可泛化至未见样本，实现无需微调的高效推理，验证了计算分摊的有效性。