一文通解：VAE、Diffusion 与 Flow 模型原理与联系

原创

languageX

修改于 2025-08-25 23:26:05

9970

文章被收录于专栏：多模态模型系列多模态模型系列

最近在学习何凯明大佬的MeanFlow。借此机会，我对从 VAE、扩散模型、Flow Matching 到 MeanFlow 的技术思路进行了梳理，整理成技术文档。如果文章中有错误或疏漏，欢迎评论区指出讨论~

VAE

VAE论文：Auto-Encoding Variational Bayes

AE VS VAE

AE（AutoEncoder）自编码器是将输入数据映射到一个低维的隐含表示，然后从这个低维表示中重建出原始输入数据。

VAE（Variational AutoEncoder）变分自编码器，不仅关注于重构输入数据，还试图将学到的隐含表示遵循某种概率分布，通常是一个标准正态分布。这使得VAE不仅能够进行数据压缩和特征提取，还可以用于生成新的类似样本。

为什么不用AE作为生成模型呢？因为AE的隐变量层的值是带有确定性的，它没有任何概率解释，就是在重建过程中获取的特征表示层，没办法基于它获取一个概率分布来进行采样生成。而VAE是让隐变量服从一个概率分布，对它进行了约束，基于分布进行采样来生成新的数据。

VAE

抛开公式推导用白话进行阐述，VAE的思路就是基于样本通过网络学到符合正态分布的均值和方差。得到隐变量的正态分布后，从分布采样变量提供生产器生成样本。

但是上图可以发现Z1和X1是没有相关性的，我们在Loss对比中只考虑最终样本对齐，对分布的约束并没有参与训练。所以在训练的过程中，应该将基于样本X得到的隐空间Z这个后验假设是正态分布的假设加入训练。

也就是论文中公式：

我们假设存在一个专属于X_k 的分布p(Z|X_k|) ，并进一步假设这个后验分布是（独立的、多元的）正态分布。

\log q_{\phi}(z | \mathbf{x}^{(i)}) = \log \mathcal{N}(z; \mu^{(i)}, \sigma^{2(i)} \mathbf{I})

直接计算样本的均值和方差不显示，但是我们可以构建两个神经网络来计算:

\mu_k = f_1(X_k), \quad \log \sigma_k^2 = f_2(X_k)

但是这样训练会出现的问题是，现在Z是通过分布采样得到的，为了让X更好的重构，Z的方差网络在学习的过程中肯定为了偷懒，方差往0靠近，只学均值，这就退化为AE了...

所以，我们还得增加一个约束，让噪声不为0保证模型有生产能力。办法就是让Z往标准正态分布学习。

除了重构loss，怎么保证Z往标准正态分布看齐呢？加loss

\mathcal{L}_{\mu} = \| f_1(X_k) \|^2 \quad \text{和} \quad \mathcal{L}_{\sigma^2} = \| f_2(X_k) \|^2

分别代表了均值\mu_k 和方差的对数\log \sigma_k^2 ，达到标准正态分布就是就是希望二者尽量接近于0了。但是这两个loss的权重不好取，论文直接用一般（各分量独立的）正态分布与标准正态分布的KL散度作为Loss

KL\left(\mathcal{N}(\mu, \sigma^2) \middle\| \mathcal{N}(0, I)\right)

\begin{align*} \mathcal{L}_{\mu, \sigma^2} &= \mathcal{L}_{\mu} + \mathcal{L}_{\sigma^2} \\ \mathcal{L}_{\mu} &= \frac{1}{2} \sum_{i=1}^{d} \mu_{(i)}^2 = \frac{1}{2} \| f_1(X) \|^2 \\ \mathcal{L}_{\sigma^2} &= \frac{1}{2} \sum_{i=1}^{d} \left( \sigma_{(i)}^2 - \log \sigma_{(i)}^2 - 1 \right) \end{align*}

计算结果：

\mathcal{L}_{\mu, \sigma^2} = \frac{1}{2} \sum_{i=1}^{d} \left( \mu_{(i)}^2 + \sigma_{(i)}^2 - \log \sigma_{(i)}^2 - 1 \right)

推理过程：

\begin{align*} KL\left(N(\mu, \sigma^2) \middle\| N(0, 1)\right) &= \int \frac{1}{\sqrt{2\pi\sigma^2}} e^{-(x-\mu)^2/2\sigma^2} \left( \log \frac{e^{-(x-\mu)^2/2\sigma^2}/\sqrt{2\pi\sigma^2}}{e^{-x^2/2}/\sqrt{2\pi}} \right) dx \\ &= \int \frac{1}{\sqrt{2\pi\sigma^2}} e^{-(x-\mu)^2/2\sigma^2} \log \left\{ \frac{1}{\sqrt{\sigma^2}} \exp \left\{ \frac{1}{2} \left[ x^2 - (x-\mu)^2/\sigma^2 \right] \right\} \right\} dx \\ &= \frac{1}{2} \int \frac{1}{\sqrt{2\pi\sigma^2}} e^{-(x-\mu)^2/2\sigma^2} \left[ -\log \sigma^2 + x^2 - \frac{(x-\mu)^2}{\sigma^2} \right] dx \end{align*}

第一项就是-{ \log \sigma^2} 乘以概率密度的积分=1；第二项是正态分布的二阶矩=\mu^2 + \sigma^2 ；第三项就是负方差/方差=-1~

另外在训练过程中还用到一个重参数技巧。从\mathcal{N}(0,I) 中采样\varepsilon ,然后通过Z = \mu + \varepsilon \times \sigma 进行采样。

最终将整个流程串起来就是下图所示

简单来就是训练过程中重构loss希望不要有噪声，而KL loss希望有噪声，所以约束隐变量往标准正态分布学习。但是学到的均值方差也不可能真的是完全符合标准的正态分布，不然分布和X样本就没关系了。两个loss博弈保证最后生成的图像效果有一定X的信息而且生成效果也不差。

Diffusion Model

论文：Denoising Diffusion Probabilistic Models

扩散模型不像VAE是一步编码+解码，而是就是多步加噪+去噪的过程。在VAE生成数据中，由于对数据有压缩以及数据分布学习中的平滑性问题，图像生成效果不是很理想。

扩散模型是对图片一步步进行加噪，得到X_T 是噪声图像；反向对特定分布的噪声图像，一步步去噪生成图片。

训练和生成过程

扩散模型的目的是基于一张噪声图，一步一步的去掉噪声，生成样本图片，那我们需要直到的就是每个步骤的噪声分布是什么，所以模型训练的目的其实就是在学习噪声。

公式推理在后面简单描述，先看结论如何进行训练和推理。

训练过程：

step1: 从样本集采样一张图片x

step2: 从时间序列采样一个时间t

step3: 从标准正态分布采样一个noise图片

step4: 基于t和t时刻生成的图像输入网络生成的噪声图和目标噪声图（真实噪声图生产的t时刻的图像）计算loss，网络是在预测噪声图

step5: 迭代step1-step4过程直到收敛。

训练过程很简单，要学噪声图像，其中最重要的就是step4的理解，后面再展开讲。

推理过程:

step1: 生成服从标准正态分布的图像X_T

step2: 时间采样从T到1，迭代step3-step5

step3: 如果t>1, 生成标准正态分布噪声z

step4: 如上图逻辑，基于上一个迭代生成的x_t 和时间t来预测噪声，再经过一堆操作得到x_{t-1}

step5: 跌倒直到生成x_0

推理过程很简单，去掉噪声，其中最重要的就是step4的理解，后面再展开讲。

加噪和去噪过程

加噪过程

我们基于下图了解加噪过程，重点理解上面训练步骤step4的公式。

生成的x_t 图像由于，x_{t-1} （图像）和{\varepsilon}_{t-1} （噪声）加权插值得到，公式如下：

\mathbf{x}_t = \sqrt{\alpha_t} \, \mathbf{x}_{t-1} + \sqrt{1 - \alpha_t} {\varepsilon}_{t-1}

也就是x_{t-1} 生成x_t 生成的概率分布符合均值为 \sqrt{\alpha_t} \, \mathbf{x}_{t-1} ，方差为1-\alpha_t 的正态分布。

其中\alpha_t 是设定的值，在加噪过程中x_t 的权重越来越小，噪声\varepsilon 的权重越来越大。

但是每次加噪都迭代计算了太大，其实可以通过一步计算直接通过x_0 跳跃式加噪。

以上就是加噪的过程，回头看训练步骤，基于x_0 一次性增加噪声生成x_t 的概率分布和网络预测的噪声进行KL散度，解释了训练中的step4的逻辑,下面我们看怎么去噪。

去噪过程

从x_t 加噪生成x_{t-1} 好理解，从x_{t-1}去噪为x_0 就无法求解了。直观理解，从一张噪声图片，没有任何先验或者条件情况下，怎么知道去噪方向？所以论文提出，我们可以告知x_0 ，问题转为求q(x_{t-1}| x_t,x_0)概率分布。

上图展示很清晰，基于x_0，x_t 生成x_{t-1} 的概率分布其实可以转换为三个高斯分布。所以如果知道x_0，p(x_{t-1}| x_t,x_0)是可以计算得到的。

如果知道x_0 ,生成x_{t-1} 的概率分布符合的高斯分布q({x_{t-1}|x_t,x_0}) 的方差是固定值，那我们训练的目标就是神经网络训练的均值和分布计算的均值的KL散度最小。

但是我们是基于噪声生成图片，其实是不知道x_0 的,得想办法去掉：

x_t = \sqrt{\overline{\alpha}_t} \, x_0 + \sqrt{1 - \overline{\alpha}_t} \, \varepsilon_t \quad \rightarrow \quad x_0 = \frac{x_t - \sqrt{1 - \overline{\alpha}_t} \, \varepsilon_t}{\sqrt{\overline{\alpha}_t}}

x_0 可以用x_t和\varepsilon 表示，带入分布均值\mu_q(x_t,t)公式

回过头来看推理过程，上面公式也就是生成过程中的step4的推理过程。

MeanFlow

Flow Matching

论文：FLOW MATCHING FOR GENERATIVE MODELING

上面介绍的去噪扩散概率模型DDPM，从标准高斯分布中采样，然后再逐步去噪生成一个目标分布的样本。其中需要建模思路就是从数据分布和先验分布中构建一个逐步加噪的过程，然后对噪声进行预测来学习逆过程（去噪）。

所以生成式模型的本质其实都是从一个已知的分布中进行采样，然后通过某种可学习的变换将其映射到目标数据的分布。

核心问题：从一个已知分布P0出发，通过神经网络的引导，逐步转向目标分布P1。

FlowMatching的原理过于复杂，要理清来龙去脉可以详读论文和B站视频，本章只做简单思路整理。

我们假设x_0 是噪声，x_1是目标图像，从x_0 到x_1 最简单的路径就是直线。

通过线性插值，x_t 随着时间从x_0 到x_1 , 那么x_t = t x_1 + (1 - t) x_0

对x_t 求导，就是x在t时刻的速度矢量，是每个时刻应该前进的方向，也就是我们需要知道的分布的变换路径。

\frac{\mathrm{d}}{\mathrm{d}t} x_t = \frac{\mathrm{d}}{\mathrm{d}t} \left( t x_1 + (1 - t) x_0 \right) = \frac{\mathrm{d}}{\mathrm{d}t} \left( t x_1 + x_0 - t x_0 \right) = x_1 - x_0

我们定义一个表示从噪声变为数据的函数x_t = \psi_t(x) ，x_t 表示在时间点t 的数据。

再定义u_t(x_t) 表示对这个流函数的导数

\frac{d}{{d}t} \psi_t(x) = u_t(\psi_t(x))

\frac{d}{dt} x_t = u_t(x_t)

这个导数就可以理解x_t 的变化速度，就是速度向量场u_t(x_t) 。

所以我们要做的就是学一个神经网络学习v_t去逼近真实的速度场u_t。

\mathbb{E}_{x_t \sim p_t(x_t)} \left[ \| v_t(x_t) - u_t(x_t) \|^2 \right]

而真实采样路径是定义了一个条件高斯路径

p_t(x_t \mid x_1) = \mathcal{N}(x_t; \mu_t, \sigma_t^2 I)

整个流程如下：

配合代码查看更好理解，参考代码

# 训练
x_1, y = x_1.to(device), y.to(device)
            # Compute the probability path samples
            x_0 = torch.randn_like(x_1)
            t = torch.rand(x_1.size(0), device=device, dtype=x_1.dtype)
            x_t, dx_t = path_sampler.sample(x_0, x_1, t)
 
            flow.zero_grad(set_to_none=True)
 
            # Compute the conditional flow matching loss with class conditioning
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
            # v_t通过网络获取，dx_t通过采样得到
                vf_t = flow(t=t, x=x_t, y=y)
                loss = F.mse_loss(vf_t, dx_t)
 
            # Gradient scaling and backprop
            scaler.scale(loss).backward()
            torch.nn.utils.clip_grad_norm_(flow.parameters(), max_norm=1.0)  # clip gradients
            scaler.step(optimizer)
            scaler.update()
            
# 路径采样
from torch import Tensor
 
from flow_matching.utils import expand_t_like_x
class PathSampler:
    def __init__(self, sigma_min: float = 0.0):
        # Minimal Gaussian distribution std for target distribution
        self.sigma_min = sigma_min
 
    def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> Tensor:
        """Samples points x_t from the distribution p_t(x_t | x_1)
 
        According to eq. (20) in the paper (https://arxiv.org/abs/2210.02747), we have:
 
            mu_t = t * x_1, and sigma_t = 1 - (1 - sigma_min) * t
 
        and the conditonal flow phi_t(x_0) is given by (see eq. 22):
 
            phi_t(x_0) = mu_t + sigma_t * x_0
 
        We treat x_t as the flow of phi_t(x_0), then we have:
 
            x_t = t * x_1 + (1 - (1 - sigma_min) * t) * x_0
 
        And, the target conditional vector field u_t(x_1|x_0) equals to the derivative of x_t w.r.t t:
 
            u_t(x_1|x_0) = dx_t = x_1 - (1 - sigma_min) * x_0
 
        We can use this to construct the conditional flow matching loss:
 
            loss = loss_fn(vf_t(x_t), dx_t)
 
        This function returns the samples x_t ~ p_t(x_t | x_1) and dx_t to compute the above loss.
 
        Args:
            x_0: Samples from base distribution, (batch_size, ...)
            x_1: Samples from target distribution, (batch_size, ...)
            t: Time steps, (batch_size,)
 
        Returns:
            Tensor: Samples from the distribution p_t(x_t | x_1), (batch_size, ...)
        """
        t = expand_t_like_x(t, x_1)
        mu_t = t * x_1
        sigma_t = 1 - (1 - self.sigma_min) * t
 
        # Draw a sample from the probability path: x_t ~ p_t(x|x_1)
        x_t = mu_t + sigma_t * x_0
 
        # Compute the target conditional vector field
        dx_t = x_1 - (1 - self.sigma_min) * x_0
 
        return x_t, dx_t
        
#推理生成
@torch.no_grad()
def generate_at_multiple_times(num_samples=512, t_points=None):
    """
    在多个时间点 t 生成样本
    t_points: 如 [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
    """
    model.eval()
    z0 = torch.randn(num_samples, 2).to(device)  # 初始噪声
 
    if t_points is None:
        t_points = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
    t_tensor = torch.tensor(t_points).to(device)
 
    # 使用 odeint 获取每个 t_point 的输出
    # 注意：这里odeint内部用更高级的数值方法（如 Dopri5、RK4），所以每一步都以上一步的状态为输入的，代码上看似“独立”而已
    traj = odeint(
        func=model,
        y0=z0,
        t=t_tensor,
        method='dopri5',
        rtol=1e-5,
        atol=1e-7
    )  # 返回形状: [len(t_points), num_samples, 2]
 
    # 转为字典：{t: samples}
    samples_at_t = {}
    for i, t in enumerate(t_points):
        samples_at_t[t] = traj[i].cpu().numpy()
 
    return samples_at_t

MeanFlow

论文：Mean Flows for One-step Generative Modeling

终于到本文真的想要介绍的meanflow了，何凯明的论文, 利用平均速度场一步到位的生成模型。

Flow Matching 建模的是瞬时速度（instantaneous velocity）。MeanFlow建模的是平均速度：在一段时间间隔内的总位移除以时间间隔。

论文从平均速度的定义出发，推导出MeanFlow恒等式（MeanFlow Identity），定义平均速度与瞬时速度及其时间导数之间的关系。训练目标就是预测平均速度场。

再回顾前面介绍的：

扩散模型：定义一个从数据到噪声的「加噪」过程，然后学习一个「去噪」过程来反转它。去噪过程通常被建模为一个随机微分方程（SDE）或常微分方程（ODE）。

流匹配（Flow Matching）：学习一个从噪声点流动到数据点的速度场，这个过程也被建模为一个 ODE。

但是数值求解 ODE 就是把整个过程分解成很多小段，每小段都根据当前的瞬时速度走一小步，这就导致要获得精确的求解结果，需要非常小的步长，意味着你需要走很多步。每一步都需要调用一次神经网络来预测当前状态的速度（或者扩散模型中的噪声/得分），这被称为一次函数评估（Function Evaluation, FE）。总共需要多少次函数评估，就是 NFE (Number of Function Evaluations)。高质量的扩散模型或 Flow Matching 模型通常需要几百甚至上千的 NFE，才能生成清晰锐利的图像。

所以如何减少 NFE，实现 1-NFE 生成就是MeanFlow解决的问题。

定义平均速度的公式(t - r) u(zt, r, t) = \int{r}^{t} v(z_\tau, \tau) \, d\tau

其中u(z_t,r,t) 是从r时刻到t时刻这段时间区间的平均速度；v(z_τ,τ) 代表的是t时刻的瞬时速度。公式两边求导：

\frac{d}{dt}(t - r) u(zt, r, t) = \frac{d}{dt} \int{r}^{t} v(z_\tau, \tau) \, d\tau

公式左侧，根据乘积法则进行求导

\frac{d}{dt}((t - r)u(z_t, r, t)) = \frac{d(t - r)}{dt}u + (t - r)\frac{du}{dt} = 1 \cdot u + (t - r)\frac{du}{dt}

公式右侧，根据微积分基本定理，对积分上限t 求导就是被积函数在t的值，所以右侧在t 时刻就是v(z_t,t) 。

所以公式整合为

u(z_t, r, t) + (t - r) \frac{d}{dt} u(z_t, r, t) = v(z_t, t)

整合公式，它提供了一个不含积分的、可以用来构建训练目标的形式

\underbrace{u(z_t, r, t)}_{\text{平均速度}} = \underbrace{v(z_t, t)}_{\text{瞬时速度}} - \underbrace{(t - r)\frac{d}{dt}u(z_t, r, t)}_{\text{对时间求导加速度}}

训练过程：

根据以上的分析，我们可以知道神经网络要学习的目标场就是平均速度场

u_{\mathrm{tgt}} = v(z_t, t) - (t - r) \frac{d}{dt} u\theta(z_t, r, t)

第一项瞬时速度：根据 x_t = (1-t)x1 +t(x0) 两边求导可以容易得到 v_t = x0- x1（其中x_0是噪声，x_1是样本）

第二项求全导（注意不是偏导，是全导数因为u不仅对依赖t,还依赖t时刻的状态z_t，而z_t也是随着t变化的）

代入 \frac{dz_t}{dt} = v(z_t, t), \frac{dr}{dt} = 0, \frac{dt}{dt} = 1，得到:

\frac{d}{dt} u(z_t, r, t) = \frac{\partial u}{\partial z_t} \frac{dz_t}{dt} + \frac{\partial u}{\partial r} \frac{dr}{dt} + \frac{\partial u}{\partial t} \frac{dt}{dt} = v(z_t, t) \frac{\partial u}{\partial z_t} + \frac{\partial u}{\partial t}

\frac{d}{dt} u_\theta(z_t, r, t) \approx v_t \partial_z u_\theta + \partial_t u_\theta

其中v_t \partial_z u_\theta 是神经网络输出的 u_\theta 对输入的z_t的雅可比矩阵\partial_z u_\theta ，再和向量v_t相乘。这就是雅可比-向量积（jacobian-Vector Product，JVP）可以使用jvp库计算。

Loss计算

训练的目标训练目标是最小化网络预测u_{\theta}与构建的 u_{tgt}之间的均方误差：

\mathcal{L}(\theta) = \mathbb{E}\left\| u_\theta(z_t, r, t) - \text{sg}(u_{\mathrm{tgt}}) \right\|_2^2

（这里有个sg表示Stop-gradient，指在计算梯度的时候将u_{tgt}部分视为常数。因为本来u_{tgt}里就包含了求导，如果计算loss再求一次导会变为二阶求导增加计算量而且训练不稳定，Stop-gradient 避免了这个问题，将训练过程简化为一阶优化。）

训练过程中，通过采样不同时间点对（r,t），网络被强制学习如何在不同时间间隔内预测平均速度，以及这些平均速度与瞬时速度之间的关系。当 r=t时， t-r项为零，目标退化为u_{tgt} = v_t ，网络学习匹配瞬时速度。当 r \neq t

时，导数项发挥作用，指导网络学习跨时间间隔的平均行为。消融实验证明，包含非零比例的r \neq t 样本对于训练出有效的一步生成模型至关重要。

以下是论文给出的meanFlow的伪代码，通过以上推导应该能更好理解如何进行训练和一步生成了~