🤖 machine learning

Stochastic Transition-Map Distillation for Fast Probabilistic Inference

本文介绍了随机转移图蒸馏（STMD），这是一种无需教师模型的框架，通过将采样随机微分方程的完整转移图蒸馏为条件平均流模型，从而加速扩散模型的推理，使其能够在无需预训练教师或复杂优化的情况下实现高效的一步或少步随机采样。

原作者： George Rapakoulias, Peter Garud, Lingjiong Zhu, Panagiotis Tsiotras

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

原作者： George Rapakoulias, Peter Garud, Lingjiong Zhu, Panagiotis Tsiotras

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一个机器人如何画出一只猫的图画。

旧方法：缓慢的、按部就班的雕刻家
传统的 AI 模型（称为“扩散模型”）的工作原理就像一位雕刻家：他从一块大理石（纯噪声）开始，一点点地凿去微小的碎屑，从而显露出猫的形象。为了获得良好的结果，雕刻家必须进行成千上万次微小的凿刻。这需要很长时间和大量的计算能力。

一些较新的方法试图通过教机器人进行更大、更快的凿刻来加速这一过程。然而，大多数这些快速方法就像是一个只知道一种特定方式去雕刻猫的机器人。如果你让它再次雕刻这只猫，它会生成一只像素级完全相同的猫。这是确定性的。但在现实世界中，猫是各不相同的！有的有条纹，有的有斑点，有的眼睛颜色不同。我们希望机器人能够创造出许多不同的、独特的猫，而不仅仅是一个复制品。

当前“快速”方法的问题
其他研究人员试图让机器人既快速又随机（随机性），但他们通常需要一个已经完美的“教师”机器人来指导他们如何做。这就像你需要一位主厨站在你肩头，在你学会快速烹饪之前先纠正你的烹饪过程。这既昂贵又复杂。

新解决方案：STMD（“地图”方法）
本文作者提出了一种名为**随机转换映射蒸馏（Stochastic Transition-Map Distillation, STMD）**的新方法。其工作原理如下，使用一个简单的类比：

旅程与目的地：
想象将噪声转化为猫的过程，是从一片迷雾森林（噪声）到一片阳光草地（猫）的旅程。
- 旧的快速方法试图直接学习目的地。它们会说：“如果你从这里开始，你就会到达那里。”
- STMD 学习的是旅程的地图。它学习雾气如何消散、路径如何在每一步发生变化的规则。它理解的是河流的流向，而不仅仅是起点和终点。
“无教师”技巧：
通常，为了快速学习这张地图，你需要一位大师级教师。STMD 的特殊之处在于它是无需教师的。它通过观察旅程的规则（噪声背后的数学原理）来自学，并弄清楚如何从一个步骤跨越到下一个步骤，实现单次巨大飞跃，而不是迈着小步前行。
“平均流”指南针：
该论文使用了一个名为“平均流”（Mean Flow）的概念。想象你在人群中行走。你不需要确切知道每个人要去哪里；你只需要知道人群移动的平均方向。STMD 学习这种“平均方向”（即流向），以便即使跳过中间许多步骤，也能预测图像下一步应该走向何方。

为什么这很重要？

速度快： 它只需一步或几步就能生成高质量图像，而不是成千上万步。
具有随机性（这是好事）： 因为它学习的是旅程的流向，所以它可以走不同的路径到达同一个目的地。这意味着它可以生成许多不同的、独特的猫，而不仅仅是一个复制品。
无需教师： 它不需要一个预先训练好的“完美”模型作为学习对象。它是从零开始构建自己的地图。

他们在什么上进行了测试？
研究人员在三个著名的图像数据集上测试了这种方法：

MNIST： 简单的手写数字（如 0 和 1）。
CIFAR-10： 日常物体（汽车、鸟类、猫）的小型彩色图片。
CelebA： 人脸照片。

在所有这些测试中，他们的方法生成的图像质量与缓慢的传统方法一样好，但速度快得多。他们还展示了该方法可用于“图像修复”（填补图片中缺失的部分，就像修复一张破损的照片）。

核心结论
这篇论文提出了一种方法，使 AI 图像生成器变得超级快速，同时不丧失创造多样化、独特图像的能力。它通过教导 AI 理解转换过程的“流向”来实现这一点，使其能够跳过枯燥缓慢的步骤，直接跃升至结果，而无需一位大师级教师来指导它如何操作。

技术摘要：用于快速概率推断的随机转移映射蒸馏

问题陈述

与生成对抗网络（GANs）和变分自编码器（VAEs）等前身模型相比，扩散模型在生成质量、多样性和分布覆盖方面已确立了最先进的性能。然而，这种性能往往伴随着高昂的推理成本，这是因为需要在许多步骤上积分连续时间的常微分方程（ODE）或随机微分方程（SDE）。

尽管现有的加速技术（如蒸馏（例如一致性模型）和基于最优传输的方法）已经存在，但它们通常在生成质量/多样性与推理复杂度之间面临权衡。当前蒸馏方法的一个关键局限性在于它们依赖于确定性推理（例如通过概率流 ODE）。给定一个固定的噪声样本，输出是唯一的。这限制了它们在本质上需要随机推理的下游任务中的实用性，例如用于逆问题的扩散后验采样、基于能量的微调，或改善与用户查询的对齐。此外，许多现有的蒸馏框架需要预训练的教师模型、双层优化，或轨迹的模拟与缓存，这增加了训练的复杂性和扩展难度。

方法论：随机转移映射蒸馏（STMD）

作者提出了随机转移映射蒸馏（STMD），这是一个无需教师模型的框架，旨在加速扩散模型推理，同时保留底层过程的随机转移结构。

核心概念

与通常建模后验分布均值的基于分数的扩散模型不同，STMD 蒸馏与采样 SDE 相关的完整转移映射。该方法学习一个条件生成器，用于在单步或少数几步内预测 $t' < t$ 时的转移概率 $p(x_{t'} | x_t)$ ，从而避免昂贵的反向 SDE 数值积分。

技术表述

条件平均流（MF）： 该框架使用条件平均流模型对 SDE 转移进行参数化。它不是学习确定性流映射，而是学习扩散过程轨迹上的平均速度。
x0 参数化： 为了建模转移 $p(x_{t'} | x_t)$ ，作者利用 $x_0$ 参数化。他们将转移表示为后验 $p(x_0 | x_t)$ 与已知桥接测度 $p(x_{t'} | x_0, x_t)$ 的积分。由于方差保持（VP）SDE 的桥接测度具有闭式解，该问题简化为学习条件模型 $p_\theta(x_0 | x_t)$ 。
训练目标： 模型 $u_\theta$ 被训练以逼近条件平均速度。作者定义了一个**条件平均流（CMF）**损失函数：
$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$
其中 $z_s$ 是数据 $x_0$ 与先验样本 $z_1$ 之间的直线插值，以扩散状态 $x_t$ 和时间 $t$ 为条件。 $\text{sg}$ 算子表示停止梯度。
无需教师的训练： 训练过程不需要预训练的教师模型。它依赖于联合分布 $p(x_0, x_t)$ 的样本，这些样本可以通过从数据分布中采样 $x_0$ 并应用前向 SDE 核（已知具有闭式解）来生成。

推理

推理通过向后积分学习到的条件平均流来完成。该过程涉及：

采样初始噪声 $x_1$ （或 $z_1$ ）。
迭代应用学习到的速度场 $u_\theta$ ，从时间 $s$ 步进到 $s-\Delta s$ 。
在每个扩散时间步 $t$ ，从桥接分布 $p(x_{t-\Delta t} | x_0, x_t)$ 中采样以更新状态，从而有效地在大步长中模拟随机转移。

主要贡献

新颖框架： 提出了 STMD，这是一个基于条件平均流的蒸馏框架，它能够通过回归目标，利用单步或少数几步生成器学习高维扩散 SDE 的转移映射，且无需任何预训练教师。
理论收敛性： 推导了平均流算法在2-Wasserstein 距离下的收敛界。据作者所知，这是该度量下对平均流算法的首次收敛分析。他们进一步扩展证明了其学习到的条件转移映射收敛于底层反向 SDE 的转移核。
实证验证： 在 MNIST、CIFAR-10 和 CelebA 数据集上的图像生成和图像修复任务中演示了该方法，表明其在保留随机性的同时，性能与基线模型具有竞争力。

实验结果

作者在三个数据集上评估了 STMD：

MNIST： 使用 U-Net 骨干网络，将该模型与原始平均流算法和去噪扩散概率模型（DDPM）进行了比较。作者利用了分类器潜在空间中的弗雷歇距离（FD）（因为 FID 对 MNIST 不可靠）。结果表明，STMD 在支持随机采样的同时，表现与基线模型相当。
CIFAR-10： 使用 U-Net 进行了类似的实验。在固定数量的训练迭代下，该模型取得了与 DDPM 和平均流基线模型具有竞争力的弗雷歇初始距离（FID）分数，同时保持了扩散采样器的完全随机性。
CelebA： 实验利用了潜在 DiT 架构。该模型生成了无条件样本，FID 为 8.28（使用 4 个推理步骤和 2 个平均流步骤），并成功执行了图像修复任务。

在所有情况下，模型均从头开始训练，使用与基线模型相同的网络骨干和迭代次数，以确保公平比较。

意义与主张

该论文声称，STMD 解决了快速确定性蒸馏与下游任务对随机推理需求之间的差距。通过蒸馏转移映射而不仅仅是流映射，STMD 保留了扩散过程的随机结构。

作者强调，他们的方法是无需教师的，消除了对双层优化或轨迹缓存的需求，这使得训练过程更加高效且可扩展。建立 2-Wasserstein 收敛界的理论贡献为平均流方法提供了坚实的基础，弥合了理论保证与实际应用之间的差距。

该工作将自己定位为需要概率推断的任务（如逆问题和后验采样）的解决方案，在这些任务中，确定性蒸馏方法往往表现不足。作者对当前性能保持谦逊，指出有限的训练迭代次数阻碍了获得最先进结果，但比较研究验证了该方法相对于现有基线的有效性。