扩散模型和最优传输之间到底存在怎样的联系?对很多人来说还是一个未解之谜。
但有一点很清楚的是:在相似的数据集上训练的不同扩散模型倾向于恢复出相似的映射关系。
这就提出一个问题:如果这些映射关系不是最优传输(OT,Optimal Transport )映射,那么它们到底在什么意义上是最优呢?
2022 年,博科尼大学助理教授 Hugo Lavenant 与里昂第一大学教授 Filippo Santambrogio 合作,在论文《 THE FLOW MAP OF THE FOKKER-PLANCK EQUATION DOES NOT PROVIDE OPTIMAL TRANSPORT 》中探讨了流模型在最优传输框架中的应用,并提供了一个反例,表明在某些情况下,流模型并不能实现最优传输。
论文地址:https://cvgmt.sns.it/media/doc/paper/5469/counterexample_flow_v3.pdf
在这篇文章的摘要部分,作者表示,Khrulkov 和 Oseledets 在先前研究中(论文:Understanding DDPM Latent Codes Through Optimal Transport)提出了一个猜想,该猜想认为通过积分 Fokker-Planck 方程的 Wasserstein 速度得到的 ODE 流,可以获得一个最优传输映射。
然而,在 Kim 和 Milman 的论文中《A generalization of Caffarelli’s contraction theorem via (reverse) heat flow》,这一结果被认为是错误的,但没有提供证明。Hugo Lavenant、Filippo Santambrogio 的这篇论文正好展示了 Khrulkov 和 Oseledets 所声称的结果不能成立。
但这篇文章过于晦涩难懂,全篇论文看下来几乎都是推导公式且篇幅又长。
为此,法国数学家 Gabriel Peyré 在论文《 Diffusion models and Optimal Transport 》中给出了一个很好的概括,文章重述了 Hugo Lavenant 和 Filippo Santambrogio 关于简洁证明的主要内容,即一般情况下,扩散模型不能定义最优传输映射。
地址:https://github.com/mathematical-tours/mathematical-tours.github.io/blob/971ddb3aab5803c7a4abef122f878292f6a6c25d/book-sources/diffusion-models/note-diffusion-ot.pdf
接下来,我们看看这篇文章讲了什么内容。
生成模型旨在在参考分布 α(通常是各向同性高斯分布)和数据分布 β 之间构建传输映射 T。用 T♯α 表示 α 被 T 向前推进(如果 α 是由 Dirac 质量在 x_i 处构成的,那么 T♯α 是由 Dirac 质量在 T(x_i)处构成的)。
因此,目标是找到 T,使得 T♯α = β 。很明显,对于任何 β,这样的映射总是存在的,但找到 T 的明确构造方法却出奇地困难。
这里有两种标准方法,分别是最优传输和集成扩散过程的逆向积分伴随的平流场。
最优传输
最优传输通过求解 Monge 问题求出 T:
1991 年,Brenier 著名定理表明这个映射是存在的,且是唯一的,并且可以写成一个凸函数 T =∇φ 的梯度。根据质量守恒定律,即 T♯α = β ,等价于说 φ 解决了 Monge-Ampère 方程:
逆向 Flow Map
扩散模型需要考虑 β_0 = β 和 β_∞= α =N (0,Id) 之间的差值 β_t ,求解过程定义如下:
请注意,使用 y 来表示空间变量,因为演化是逆向进行的,即从数据 β 到后一个变量 α,它收敛于 β_∞= α。
将方程写成离散形式:
这表明,如果已经计算出了 β_t,那么这种演化可以根据向量场 v 演化粒子来获得。
映射 S_t 就是 flow map :
逆向 Flow Map 不是最优传输
人们很自然地想知道逆向 Flow Map 是否是 (1) 的解。在一维情况下,S_t 定义微分同胚( diffeomorphism),因此是单调的,也是单调的。因此,它是凸函数的梯度,根据 Brenier 定理使其最优。如果 β 是高斯分布,直到空间旋转以使协方差对角化,则扩散映射由沿每个轴的单调映射定义,并且也是最优传输。
Lavenant 和 Santambrogio 通过矛盾证明,一般来说,逆向 flow map 并不是最优传输。他们构造了一个接近各向同性高斯 α 的 β,但他们没有证明 β 的猜想是错误的,而是证明存在一些 t ≥ 0,使得从 α 到 β_t 的逆向 Flow Map T_t 不是最优传输。他们实际上表明,对于某些 t_0 > 0,T_t 并不是所有 t ∈ (0, t_0] 的最优传输。
用 S_t 表示从 β_0 = β 到 β_t 的 Flow Map。如果猜想成立,则从 α 到 β_t 的逆向 Flow Map T_t 是所有 t 的最优传输。根据 Flow Map 的构成规则,该 Map 为:
并且目标是证明:如果 β 选择得当(具体来说,非常接近 α,且特定的二阶和四阶对数密度导数为 0),那么 T_t 是所有 t 的最优传输会导致矛盾。根据 Brenier 定理,T_t 是最优传输意味着它是凸函数的梯度,这相当于:
结合:
对 (7) 对 t 求微分
对流 ODE (4) 对 x 求微分
然后在 t = 0 时评估所获得的方程,Hugo 和 Filippo 通过显式计算表明,这会导致:
利用基本性质:A、B 对称且 AB 对称则 (8) 意味着:
为了达到矛盾,假设 G (y) 和 H (y) 对于所有 y 都是可交换的。由于并且 T 和 S 是逆最优传输映射,因此将表示为 ψ 凸。Monge-Amp`ere 方程 (2) 意味着:
为了使 β 接近 α,请考虑:对于较小的 ε,
在泰勒级数中展开,经过一些计算:
着眼于 y = 0,目标是通过设计 h 来达到矛盾,使得和不能交换。在 0 附近,h 必须至少是 4 次多项式。二维情况下的一个示例是:
产生: