浅谈autoencoder（2）——细水长flow

前言

上文我们概括了VAE的形成及其弊端。诚然，它的学习目标仍然是autoencoder比较传统的路子：通过降维升维来学习数据个性化的表示形式。但这样的学习方式太过固定，且并未更多的和实际结合起来。我们同样思考变工况退化这个案例，我们本质需要学习的是系统最基本的退化情况，外在工况条件改变会改变它的退化情况，而这个“退化情况”可以从物理中“场”这个抽象概念思考：工况改变场会变化，进而退化的概率路径会变化。一个场会有许多不同的概率路径，而学习的目标就是找到最适合当前数据目标的概率路径，对应到案例中即退化分布。

很自然的，flow-matching和diffusion model被引出，他们就是从“场”这个概念建立的学习目标。下图清晰的展示了我上述的抽象概念：

其中：

• 青色箭头表示向量场
• 浅绿色曲线表示概率路径

无论是flow matching还是diffusion model，学习的目标是拟合一个向量场$u_t$，使得能够得到对分布进行变换的概率路径$p_t$。有了这一概念，接下来就详细讲讲他们的实现形式

模型概念

flow models

向量场$u_t$，在数学概念上就是一个常微分方程（ODE），它的解就是对分布进行变换的概率路径$p_t$。具体而言，在每个时间 t 和位置 x 上，都会有向量$u_t(x)$指定空间中的速度。ODE对轨迹施加了一个条件：我们想要一个“沿着”向量场$u_t$的轨迹$X$，从点$x_0$开始。我们可以将这样的轨迹正式化为方程的解：

\[\begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} X_t &= u_t(X_t) \\ X_0 &= x_0 \end{aligned}\]

我们可以想象，当以$x_0$，$t=0$开始时，轨迹运动到$t$点时的位置$X_t$如何表示？这就需要引入flow 模型处理： \(\psi : \mathbb{R}^d \times [0,1] \to \mathbb{R}^d, \quad (x_0, t) \mapsto \psi_t(x_0)\)

\[\frac{\mathrm{d}}{\mathrm{d}t} \psi_t(x_0) = u_t(\psi_t(x_0))\] \[\psi_0(x_0) = x_0\]

这里ODE的轨迹自然形成为$X_t=\psi_t(X_0)$，至此，我们可以定义：向量场定义了ODE，其解是flow。接下来自然会思考到，它的解是否存在且唯一? 在Theory of ordinary differential equations中已然得到证明。

但现实世界中，向量场通常不会如此简单。幸运的是数值分析中存在许多解复杂向量场得到flow的方法，这里介绍个最经典的Euler方法。

我们首先初始化$X_0 = x_0$，通过下式更新：

\[X_{t+h} = X_t + h\,u_t(X_t) \quad (t = 0, h, 2h, 3h, \ldots, 1-h)\]

其中$h=1/n$表示步长，$n$表示要模拟的步数，他在数值分析中已经是简单且好用的方法了。

接下来需要把上述的数学框架纳入到神经网络中，本质在于从一个简单的分布（比如标准高斯）出发，沿着一条由神经网络定义的flow，把样本”推”到目标数据分布。

\[\begin{aligned} X_0 &\sim p_{\text{init}} &&\blacktriangleright \text{ random initialization} \\ \frac{\mathrm{d}}{\mathrm{d}t} X_t &= u_t^\theta(X_t) &&\blacktriangleright \text{ ODE} \end{aligned}\]

其中神经网络用$u_t^\theta$表示，$p_\text{init}$通常取标准高斯分布

而目标就是让目标分布尽可能拟合为数据真实分布：

\[X_1 \sim p_{\text{data}} \quad \Leftrightarrow \quad \psi_1^\theta(X_0) \sim p_{\text{data}}\]

这里的ODE的解，也就是flow为$\psi_t^\theta$，它需要通过Euler一步一步积分得到：$X_{t+h} = X_t + h \cdot u_t^\theta(X_t)$

diffusion models

它和flow model的区别就是从ODE扩展到SDE，将确定性的轨迹扩展为随机过程，直观上说它更加通用，但相应的计算过程也会增加。

这里简单说下如何从ODE扩展到SDE。首先，我们需要重写下ODE的轨迹表达式：

\[\frac{\mathrm{d}}{\mathrm{d}t} X_t = u_t(X_t) \quad \blacktriangleright \text{ expression via derivatives}\] \[\stackrel{(i)}{\Leftrightarrow} \quad \frac{1}{h}(X_{t+h} - X_t) = u_t(X_t) + R_t(h)\] \[\Leftrightarrow \quad X_{t+h} = X_t + h\,u_t(X_t) + h\,R_t(h) \quad \blacktriangleright \text{ expression via infinitesimal updates}\]

其中$R_t(h)$是无穷小量。上述推导更清晰的展示了ODE的轨迹$(X_t){0≤t≤1}$在每个时间步骤中，都会朝着$u_t（X_t）$方向迈出一小步。我们现在可以修改最后一个部分，使其成为随机方程：SDE的轨迹$（X_t）{0≤t≤1}$，在每个时间步骤中，都朝着$u_t（X_t）$方向走一小步，加上布朗运动的一些贡献：

\[X_{t+h} = X_t + \underbrace{h\,u_t(X_t)}_{\text{deterministic}} + \sigma_t\underbrace{(W_{t+h} - W_t)}_{\text{stochastic}} + \underbrace{h\,R_t(h)}_{\text{error term}}\]

其中$\sigma_t \geq 0$ 是扩散系数，$R_t(h)$ 是随机误差项。我们用更常见的符号再次表示它们：

\[\begin{aligned} \mathrm{d}X_t &= u_t(X_t)\mathrm{d}t + \sigma_t\mathrm{d}W_t &&\blacktriangleright \text{ SDE} \\ X_0 &= x_0 &&\blacktriangleright \text{ initial condition} \end{aligned}\]

然而，这样的推导会让我们发现，它不再有$\phi_t$了。这是由于 $X_t$ 不再由初始分布 $X_0 \sim p_{\text{init}}$ 表示，因为它是随机过程。只是这个缺憾并不影响我们将它转化为神经网络的算法。这里给出算法架构：

Algorithm 2 Sampling from a Diffusion Model (Euler-Maruyama method)

Require: Neural network $u_t^\theta$, number of steps $n$, diffusion coefficient $\sigma_t$

1. Set $t = 0$
2. Set step size $h = \frac{1}{n}$
3. Draw a sample $X_0 \sim p_{\text{init}}$
4. for $i = 1, \ldots, n$ do
5. $\quad$ Draw a sample $\epsilon \sim \mathcal{N}(0, I_d)$
6. $\quad X_{t+h} = X_t + h\,u_t^\theta(X_t) + \sigma_t\sqrt{h}\,\epsilon$
7. $\quad$ Update $t \leftarrow t + h$
8. end for
9. return $X_1$