Terminally constrained flow-based generative models from an optimal control perspective
作者: Weiguo Gao, Ming Li, Qianxiao Li
分类: cs.LG, math.OC
发布日期: 2026-01-14
备注: 59 pages, 9 figures
💡 一句话要点
提出TOCFlow，通过最优控制解决Flow模型的终端约束采样问题
🎯 匹配领域: 支柱七：动作重定向 (Motion Retargeting)
关键词: 生成模型 流模型 最优控制 终端约束 Hamilton-Jacobi-Bellman方程 黎曼几何 约束优化
📋 核心要点
- 现有基于流的生成模型在处理具有终端约束的采样问题时面临挑战，难以保证生成样本满足特定约束条件。
- 论文提出TOCFlow，利用最优控制理论，在采样过程中引入几何感知的引导，确保生成样本满足终端约束。
- 实验表明，TOCFlow在Darcy流、约束轨迹规划和湍流快照生成等任务中，显著提升了约束满足度，同时保持了生成质量。
📝 摘要（中文）
本文通过最优控制公式，解决了预训练的基于流的生成模型中终端约束分布的采样问题。理论上，我们通过Hamilton-Jacobi-Bellman方程表征了值函数，并将最优反馈控制推导为相关哈密顿量的最小化器。我们证明，随着控制惩罚的增加，受控过程恢复参考分布，而随着惩罚的消失，终端定律收敛到约束流形上的广义Wasserstein投影。在算法上，我们引入了带有基于流的模型的终端最优控制（TOCFlow），这是一种用于预训练流的几何感知采样时间引导方法。在跟踪参考轨迹的终端同动框架中解决控制问题，可以得到沿黎曼梯度的闭式标量阻尼因子，从而在没有矩阵求逆的情况下捕获二阶曲率效应。因此，TOCFlow以标准梯度引导的计算成本匹配了高斯-牛顿更新的几何一致性。我们在三个高维科学任务上评估了TOCFlow，这些任务涵盖了等式、不等式和全局统计约束，即Darcy流、约束轨迹规划和具有Kolmogorov谱缩放的湍流快照生成。在所有设置中，TOCFlow在保持参考模型生成质量的同时，改进了欧几里德引导和投影基线上的约束满足。
🔬 方法详解
问题定义：论文旨在解决预训练的基于流的生成模型在生成满足特定终端约束的样本时遇到的困难。传统的基于流的模型难以直接控制生成样本的最终状态，导致无法满足例如物理约束、轨迹规划约束等实际应用需求。现有方法，如欧几里德引导和投影，在满足约束的同时，可能会牺牲生成样本的质量，偏离原始数据分布。
核心思路：论文的核心思路是将终端约束采样问题转化为一个最优控制问题。通过设计一个控制策略，引导生成过程在满足终端约束的同时，尽可能地保持与原始数据分布的一致性。该方法利用Hamilton-Jacobi-Bellman方程来描述值函数，并通过最小化哈密顿量来获得最优反馈控制。
技术框架：TOCFlow的整体框架包含以下几个主要步骤：1) 使用预训练的基于流的生成模型作为参考模型。2) 将终端约束采样问题建模为最优控制问题。3) 利用Hamilton-Jacobi-Bellman方程求解值函数和最优控制策略。4) 在采样过程中，使用最优控制策略引导生成过程，使其满足终端约束。5) 在终端同动框架下，计算黎曼梯度上的标量阻尼因子，以捕捉二阶曲率效应。
关键创新：TOCFlow的关键创新在于引入了几何感知的采样时间引导方法，该方法利用终端同动框架，在黎曼梯度上计算标量阻尼因子，从而在没有矩阵求逆的情况下，匹配了高斯-牛顿更新的几何一致性。这种方法能够在满足终端约束的同时，更好地保持生成样本的质量，避免了传统方法中常见的生成质量下降问题。
关键设计：TOCFlow的关键设计包括：1) 使用预训练的基于流的模型，例如RealNVP或Glow，作为参考模型。2) 定义合适的控制惩罚项，以平衡约束满足和生成质量。3) 在终端同动框架下，计算黎曼梯度上的标量阻尼因子，该因子可以根据局部曲率自适应地调整引导强度。4) 使用合适的优化算法，例如梯度下降法，求解最优控制策略。
🖼️ 关键图片
📊 实验亮点
实验结果表明，TOCFlow在Darcy流、约束轨迹规划和湍流快照生成等任务中，显著优于欧几里德引导和投影等基线方法。例如，在Darcy流问题中，TOCFlow能够更精确地满足边界条件，同时保持流场的物理合理性。在湍流快照生成任务中，TOCFlow能够更好地满足Kolmogorov谱缩放，生成更逼真的湍流快照。
🎯 应用场景
TOCFlow在多个领域具有广泛的应用前景，例如：在流体动力学中，可以生成满足特定边界条件的流场快照；在机器人学中，可以生成满足终端位置和姿态约束的轨迹；在分子动力学中，可以生成满足特定能量和动量约束的分子构象。该方法能够有效地解决具有复杂约束的生成问题，为科学研究和工程应用提供强大的工具。
📄 摘要（原文）
We address the problem of sampling from terminally constrained distributions with pre-trained flow-based generative models through an optimal control formulation. Theoretically, we characterize the value function by a Hamilton-Jacobi-Bellman equation and derive the optimal feedback control as the minimizer of the associated Hamiltonian. We show that as the control penalty increases, the controlled process recovers the reference distribution, while as the penalty vanishes, the terminal law converges to a generalized Wasserstein projection onto the constraint manifold. Algorithmically, we introduce Terminal Optimal Control with Flow-based models (TOCFlow), a geometry-aware sampling-time guidance method for pre-trained flows. Solving the control problem in a terminal co-moving frame that tracks reference trajectories yields a closed-form scalar damping factor along the Riemannian gradient, capturing second-order curvature effects without matrix inversions. TOCFlow therefore matches the geometric consistency of Gauss-Newton updates at the computational cost of standard gradient guidance. We evaluate TOCFlow on three high-dimensional scientific tasks spanning equality, inequality, and global statistical constraints, namely Darcy flow, constrained trajectory planning, and turbulence snapshot generation with Kolmogorov spectral scaling. Across all settings, TOCFlow improves constraint satisfaction over Euclidean guidance and projection baselines while preserving the reference model's generative quality.