GeomFlow：一种基于赫斯矩阵（Hessian）信息的、具备几何感知能力的自适应扩散模型

《Neurocomputing》：GeomFlow: Geometry-aware adaptive diffusion model via hessian information

提出几何感知自适应扩散模型GeomFlow，通过全局噪声调度与局部几何复杂度估算模块解决标准扩散模型在异构数据流形上的效率问题，理论证明其等价于黎曼流形预条件拉格朗日动力学，实验在CIFAR-10和CelebA-HQ上验证其语义一致性及细节恢复优势。

　　

郝卓超|王晓瑞|于立阳|丁王才

青海大学，西宁，810016，青海，中国

摘要

基于分数的生成模型通过使用随机微分方程（SDEs）逐步将数据转换为噪声，从而实现了显著的逼真度。然而，标准方法存在一个根本的几何不匹配问题：它们在具有高度异质曲率的数据流形上应用空间均匀的扩散动态。这种“一刀切”的策略导致采样效率低下——在平坦区域过度计算，而在复杂的高频细节上则处理不足。为了解决这个问题，我们提出了GeomFlow，这是一种新颖的几何感知自适应扩散模型。GeomFlow结合了两种关键机制：一个全局可学习的噪声调度器，用于优化宏观噪声的进展；以及一个几何复杂性估计器，它利用对Hessian迹的稳健随机近似来动态调节局部扩散强度。理论上，我们证明了我们的几何感知逆过程等同于黎曼预处理的朗之万动力学（Riemannian preconditioned Langevin dynamics），从而实现了加速收敛并更好地逃离鞍点。广泛的实验表明，该模型在CIFAR-10（FID 2.14）和CelebA-HQ任务上表现出极具竞争力的性能，在条件生成和图像修复任务中展现了卓越的结构理解能力，并在保持语义一致性和恢复缺失的纹理细节方面取得了显著改进。

引言

随着基于扩散的方法的出现，生成建模取得了前所未有的突破，从根本上改变了人工智能在各个领域的应用。从合成逼真图像[19]、[21]到预测复杂的蛋白质结构，扩散模型已经建立了新的黄金标准。这些进展的理论基础在于将分数匹配（score matching）与朗之万动力学（SMLD）[30]和去噪扩散概率模型（DDPM）[7]在随机微分方程（SDEs）[28]的框架下统一起来。通过学习到的分数函数逆转逐步的噪声注入过程，这些连续时间模型允许灵活采样和精确的似然计算[29]、[32]。最近的理论分析进一步表明，这些模型的成功源于它们能够迭代地将数据分布收缩到低维流形上[2]、[13]、[20]、[31]、[34]、[37]。尽管它们的理论优雅性是一致的，但标准的扩散SDEs存在一个关键限制：扩散过程与数据结构之间的几何不匹配。现有框架通常采用仅依赖于时间的空间均匀漂移和扩散系数，隐含地假设数据流形是均匀的。然而，现实世界的数据分布（如自然图像）具有高度异质的复杂性。高频区域具有高局部曲率和内在维度，而低频区域在几何上则是平坦的。一种“一刀切”的各向同性扩散过程无法解释这种异质性，从而导致了困境：模型要么在简单区域过度计算，浪费资源，要么在复杂区域处理不足，导致细节模糊和模式崩溃。这种低效率不仅仅是一个计算问题，而是与流形假设的根本性不一致。为了弥合这一差距，我们提出了GeomFlow，这是一种新颖的几何感知自适应扩散模型。GeomFlow的核心直觉是引入对扩散过程的几何感知抵抗力。类似于车辆在急转弯时必须减速，但在直路上可以加速一样，我们的模型根据局部曲率动态调节扩散强度。具体来说，在几何复杂性高的区域（Hessian迹较大），GeomFlow细化采样步骤以捕捉复杂细节；相反，在平坦区域，它加速流动以确保效率。这种状态依赖的适应将标准的各向同性噪声去除转变为与流形对齐的传输过程。在这项工作中，我们通过严格的理论和实证框架实现了这一范式。我们的具体贡献如下：

• •
  几何复杂性估计器（GCE）： 我们引入了一个轻量级模块，使用Hessian矩阵的迹来估计数据流形的局部曲率。关键的是，我们采用Hutchinson的随机估计器来近似这个迹，确保我们的方法在计算上保持高效（$$的开销），并且能够扩展到高维数据，同时证明其足够稳健，可以指导采样过程。
• •
  黎曼解释和理论严谨性： 我们提供了理论分析，证明我们的自适应逆时间SDE在数学上等同于黎曼预处理的朗之万动力学（PLD）。通过将自适应方差解释为预处理器，我们展示了GeomFlow自然地提高了收敛速率，并帮助采样过程比标准SGD类采样更有效地逃离鞍点。
• •
  在逆问题中的广泛应用： 除了无条件生成之外，我们还证明了GeomFlow对几何的理解在具有挑战性的下游任务中表现出更优的性能。在类别条件生成和图像修复上的新实验表明，GeomFlow在保持语义一致性和恢复缺失的结构细节方面显著优于基线方法，验证了其作为通用生成引擎的潜力。

相关工作

相关工作

基于扩散的生成模型的发展迅速推进，从基本的概率框架发展到高度优化的、具有几何感知的系统。我们的综述围绕三个关键主题展开：统一的SDE框架和最近的加速技术、自适应采样策略以及生成建模中的几何结构。

方法

我们提出了GeomFlow，这是一个旨在解决标准扩散过程中几何不匹配问题的原则性框架。GeomFlow通过引入两种协同的适应机制来推广统一的SDE公式：一个全局时间调度器，用于优化噪声的进展速率；以及一个局部空间调节器，用于使扩散动态与数据流形的曲率对齐。考虑标准的正向SDE，它将数据$$随时间间隔$$转换为噪声