8 第八章
流匹配和归一化流目的是相同的,都是想把已知的简单分布变成复杂的目标分布:
只不过归一化流是通过一系列可逆函数映射的方式实现的,流匹配则是通过相流让原始噪声数据运动到有意义的数据。
8.1 训练相流
从 图 8.2 可以看出,相流或向量场决定了采样点的移动方向。对RGB彩色图片来说,把所有像素展开成一维向量,就可以把图片看成是从相同维度的高维分布\(Pr_{data}\)中采样出来的一个数据点。
对于从高维正态分布中随机初始化的采样点 \(x\),流匹配的目标在于训练一个向量场(速度场)\(\mu_\theta(x)\)。当采样点沿着该向量场随时间 \(t \in [0, 1]\) 演化时,其运动轨迹 \(\phi(x, t)=\sum_{t=0}^{1/\Delta t} \mu_\theta(x_t)\Delta t+x\) 满足初始条件 \(\phi(x, t=0) = x\),并在 \(t=1\) 时恰好移动到有意义的图片数据点 \(\phi(x, t=1) = z\)。
定义随机采样点\(x\)与目标点\(z\)之间的理想运动轨迹\(x_t\):
\[ x_t=(1-t)x+tz \tag{8.1}\]
式 8.1 本质上就是一个线性插值:
则\(x_t\)关于时间的导数\(u_t\)就是运动速度:
\[ u_t=\frac{\partial x_t}{\partial t}=z-x \tag{8.2}\]
对于最优参数 \(\hat\theta\),想要让向量场\(\mu_{\hat\theta}(x=x_t)\) 尽可能接近真实的 \(u_t\),\(\hat\theta\) 应该满足:
\[ \hat\theta = \arg\min_{\theta} E(\| \mu_\theta(x=x_t) - u_t \|^2) \tag{8.3}\]
考虑多样本 \(z_1, z_2, \cdots, z_N\),随机高斯采样 \(x_1, x_2, \cdots, x_N\) 构成训练样本 \((x_i, z_i)\), 每个训练样本的\(x_i^t\)满足\(x_i^t = (1-t)x_i + t z_i\),速度为 \(v_i^t = z_i - x_i\),则可以定义损失函数: \[ \text{loss} = E({\| \mu_\theta(x_i=x_i^t) - (z_i - x_i) \|^2}) \tag{8.4}\]
从公式来看,它要求训练\(\mu_\theta()\)在不同\(x_i^t\)下的输出需要尽可能接近\(z_i-x_i\),而 \(z_i-x_i\) 运动出来的轨迹是一条直线,那我们是否可以简单认为:
\[ \phi(x_i, t=1)=\sum_{t=0}^{1/\Delta t} \mu_\theta(x_i^t)\Delta t+x_i= x_i + \mu_{\theta}(x=x_i^{t=0}) · 1 \approx z_i \]
为了回答这个问题,我们可以尝试构建一个MLP网络,给定训练集\((x_1,z_1),(x_2,z_2),\cdots,(x_N,z_N)\),为了可视化,假设\(x\)和\(z\)的维度都是二维的。
8.2 单训练样本
先考虑单一训练样本的情况,假设需要从(0,0)移动到(2,3),构建一个具有两个隐藏层,每个隐藏层五个隐藏单元的MLP网络,使用 式 8.3 的损失函数。理论上\(\mu_\theta()\)对于输入(0,0)会学习到速度向量(2,3),并对其他没见过的输入,向量场方向同样指向(2,3)(因为没别的训练数据让向量场有别的方向了)。
实际上确实是这样,训练出来的真实向量场如下:
即使对于其他点的输入,得到速度的方向和大小也是相同的。
但是如果训练样本有两个呢?且出现交叉的情况。
8.3 双训练样本
我们特别地构造一对路径有交叉的样本,分别是从(0.4,0.4)到(0.8,-0.4)和(0.25,-0.2)到(1.0,0.8),它的理想运动轨迹是这样的:
保持相同的MLP架构使用 式 8.4 ,我们先看看训练出来的向量场吧:
大家可以自行尝试勾勒出两个起点的运动轨迹,然后会发现,真实情况和 图 8.6 所展示出来的不一样:
它们的终点被交换了,根本原因在于,向量场里不允许同一个点有两个方向,因此 图 8.6 线条的交叉点在向量场里是不合理的,同一个点的速度方向只能有一个方向。
为了最小化损失函数,网络会选择交换终点来解决这个问题,或者说这是损失函数导致的必然结果(换个角度理解,所有的\(z_i-x_i\)在出现交叉情况时构成了一个超现实的向量场,损失函数在寻找一个最接近它的真实向量场,并能补全训练数据没见过的\(x_j\)的速度向量)。
不过也许其他的损失函数能训练出这种路径:
总之上述内容揭示了两个问题:
多样本训练时,向量场是复杂的,\(x_i + \mu_{\theta}(x=x_i^{t=0}) · 1 \approx z_i\)可能并不是当前输入信号真正得到的输出,需要通过\(\Delta t\)一步步计算。
在多训练样本时,因为初始点随机采样,理想轨迹会出现交叉,交叉会导致网络训练出的终点发生交换,为了解决这个问题,训练网络必须引入时间\(t\)进行升维。
8.4 引入时间
引入时间为什么行,因为二维平面上交叉点在三维空间上不一定相交:
如 图 8.10 所示,引入时间后,原本在二维平面上相交的两条轨迹,在三维时空中就分开了,这样三维空间的理想相流变成了一个“合法”的相流,在不改变损失函数本身的情况下,模型能更好地学习到变换特征。
那会不会存在一种情况,虽然引入了时间,但还是相交了呢?对于原始低维度数据来说是有可能的,这种情况下也会发生结果交换,但是对高维数据而言,两条随机生成的直线在同一时刻 \(t\) 撞在一起,其几何概率在数学上为0。
而加入时间后,也不需要改变损失函数,把时间当成额外的维度就行。
8.5 2D演示
以这样一个图形为例,它像一个夹子,里面有2000个数据点,坐标范围在\(\pm 16\)之间。
然后再从二维高斯分布(0均值、4.5方差)中抽取2000个点,随机配对得到训练集。
如 图 8.12 所示,由于是随机配对,轨迹在空间中自然交叉。正如上文所说,模型将通过引入时间 \(t\) 这一额外维度来尽量减少交叉现象。
之后开始训练,对于训练好的模型,从高斯分布 \(p_0\) 出发可视化\(t\)从0到1的过程中粒子的运动:
进一步地,我们可以来看三维空间下设定的轨迹和训练后的轨迹是什么样的。
以底面 \(t=0\) 为初始高斯分布、顶面 \(t=1\) 为目标分布,绘制出由 \(x_t = (1-t)x_0 + t x_1\) 定义的直连线段。
同样地对训练好的模型,一步步累计出真正的运动轨迹。
图 8.14 中红色的点是发生了位置交换的点。可以注意到左图中,纯随机配对的直线彼此剧烈穿插,虽然它们不一定是直接相交了,但相互靠近、交错得很剧烈。可以想象,模型没办法在一个小区域内平滑地拟合出方向各异的向量场。
因此即使在没有真正相交的情况下,近似相交或剧烈的相互穿插也会导致目标发生交换,模型始终趋向于训练出一个平滑插值的向量场,如右图所示,左侧区域倾向于移动到左边图案去,右侧区域也倾向于移到右边,保持轨迹整体方向一致性。
8.6 整流流
对于流匹配来说,一个优化目标是尽量“一步到位”,也就是不要产生太复杂的向量场,不然得靠一个个时间步\(\Delta t\)去更新位置(如 图 8.14 右图中弯曲的轨迹,人们更希望它是直线的)。
虽然加入时间维度在很大程度上能避免交叉路径,然而回到神经网络或模型本身,即使路径冲突已经很少或没有了,网络训练时对于损失函数的逼近误差和网络本身泛化能力导致的空间插值都会使训练出来的向量场不理想,最终导致相流会发生弯曲,也就是演化路径弯曲。
整流流想做的是让相流进一步变直一些,方法很简单,首先随机配对训练出一个模型\(\mu_{\theta 1}()\)。
然后再随机高斯采样\(x_0^1,x_1^1,\dots,x_N^1\)(其中上标表示训练轮次),经过\(\mu_{\theta 1}\)后计算得到新的\(z_0^1,z_1^1,\dots,z_N^1\),分别配对\(x_i^1\)与\(z_i^1\),用这\(N\)对训练数据训练\(\mu_{\theta 2}()\)。
得到\(\mu_{\theta 2}()\)后再计算一组\((x_i^2,z_i^2)\)用于训练\(\mu_{\theta 3}()\)。这样重复下去好几轮后,相流会变的越来越直。
为什么它会有效呢?因为流训练还是在逼近一个由训练数据定义的相空间并插值,第一次得到的\(\mu_{\theta 1}()\)某种程度上已经比较“规整”了,然后在此基础上得到的训练数据集\((x_i^1,z_i^1)\)天然地也会规整很多,就这样迭代下去相空间向量的方向性会越来越对齐,最终路径轨迹也会越来越直,只需要一两步就能得到结果了。