了解 InstaFlow/Rectified Flow

大家好！我通常在 Medium 上发布文章 这里，但我想试试 Hugging Face 博客，所以这将是我的第一篇博客文章！我之所以对 Instaflows/Rectified Flows 感兴趣，是因为

1. 我们在 Eleuther 扩散阅读小组讨论了这个问题，听起来很有趣
2. 我想在 Diffusers 中为此提交一个 PR（在这个 issue）。PR 完成后我会添加代码！

什么是 InstaFlow/Rectified Flow？

Rectified flow 是一种微调扩散模型的方法，只需 1 步即可生成图像，而传统上需要大约 12 步。Instaflow 只是将其应用于稳定扩散。如果您想尝试，请查看 这里 的演示！

那么，让我们首先了解一下 rectified flows 的原理

矫正流

矫正流在论文 "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow" 中提出。

这篇论文探讨的不仅仅是单步生成图像，而是如何让两个图像分布（例如，${\pi }_0\backslash pi\_0$ 可以是马，${\pi }_1\backslash pi\_1$ 可以是斑马）之间的流尽可能地笔直。将其扩展到单步生成意味着，如果我们使噪声分布 ${\pi }_0\backslash pi\_0$ 和真实图像分布 ${\pi }_1\backslash pi\_1$ 之间的流尽可能短和笔直，那么即时生成图像就成为可能！

这被称为图像到图像的转换/传输映射问题。

GAN 中的传输映射问题

传输映射问题在几乎所有涉及图像的领域都备受关注。例如，CycleGAN

及其随后的改进 StarGAN 和 StarGAN2 使用 GAN 来学习一个领域到另一个领域图像之间的映射！实现这一目标的主要方法称为 循环一致性损失，即给定一张斑马的图像，将其转换为马，然后再转换回斑马，并确保它仍然是斑马。欲了解更多详情，请查看链接！

然而，GAN 本身存在严重的训练不稳定问题，并且生成质量不如扩散模型，因此引入了这种方法！

回到矫正流

现在，让我们从分布 ${\pi }_0\backslash pi\_0$ 中抽取 $X_0\sim {\pi }_{0}X\_0\; \backslash sim\; \backslash pi\_0$，从分布 ${\pi }_1\backslash pi\_1$ 中抽取 $X_1\sim {\pi }_{1}X\_1\; \backslash sim\; \backslash pi\_1$。回到我们的类比，这意味着我们将从马的图像分布 ${\pi }_0\backslash pi\_0$ 中获得一张特定的马的图像 $X_{0}X\_0$，并从斑马图像分布 ${\pi }_1\backslash pi\_1$ 中获得一张特定的斑马图像 $X_{1}X\_1$。

现在，给定图像对 $(X_0,X_1)(X\_0,\; X\_1)$。在这里，我们将定义一个参数 t，它介于（包括）0 和 1 之间，表示我们在 $X_{0}X\_0$ 和 $X_{1}X\_1$ 之间有多远。所以 $X_{0.5}X\_\{0.5\}$ 将正好是中点。

现在，虽然在传统扩散模型中我们可以做到这一点，但路径并不完全是直的。事实上，它可能非常曲折，如下所示

所以现在，让我们从路径 X 移动到路径 Z。我们仍然需要相同的端点，即 $X_0=Z_{0}X\_0=Z\_0$ 和 $X_1=Z_{1}X\_1=Z\_1$，但我们希望路径尽可能笔直。这表示为

$$d{Z}_t=v(Z_t,t)dtdZ\_t\; =\; v(Z\_t,\; t)dt$$

首先，d 表示一个非常小的步长。这里的 v 表示在时间步 t 处 $Z_{t}Z\_t$ 点的速度。这意味着如果我们不断地将 $Z_{t}Z\_t$ 推向 $v(Z_t,t)v(Z\_t,\; t)$ 的方向，每次都重新计算 v，那么我们将到达 $Z_{1}Z\_1$，并成功获得一只斑马。

现在，正如我们之前多次提到的，我们想要一条直线。这对于 v 的上下文意味着什么？首先，v 应该是常数，因为无论 t 是什么，直线都应该以相同的量和相同的方向推动 $Z_{t}Z\_t$。其次，它应该是最近的路径。这意味着当我们进行积分（即只将 Z 中整个路径的长度相加）时，它应该等于 $X_1-X_{0}X\_1-X\_0$。

现在，用正式的术语来说，这将意味着

$$\underset{v}{\min }{\int }_0^1\mathbb{E}[\mid \mid (X_1-X_0)-v(X_t,t)\mid {\mid }^2]dt\backslash min\_v\; \backslash int\_0^1\; \backslash mathbb\{E\}[\backslash mid\backslash mid\; (X\_1-X\_0)-v(X\_t,\; t)\; \backslash mid\; \backslash mid^2]dt$$

，其中

$$X_t=t{X}_1+(1-t)X_{0}X\_t\; =\; tX\_1+(1-t)X\_0$$

这也可以写成

$$d{X}_t=(X_1-X_0)dtdX\_t\; =\; (X\_1-X\_0)dt$$

也是！

这里快速提一下，对于 DDIM，这将是

$$X_t={\alpha }_t{X}_0+{\mathrm{B}}_t{X}_{1}X\_t\; =\; \backslash alpha\_t\; X\_0+\backslash Beta\_t\; X\_1$$

如果有人感兴趣，我可以在这里链接一些理论背景！

现在，论文深入探讨了一些非常有趣的数学部分，我将在此博客中跳过，但如果您喜欢数学和微分方程以及为什么上述内容不会陷入一些陷阱，我建议您查看。

现在，在实践中，您可能已经猜到，v 将是我们的稳定扩散模型。$X_{0}X\_0$ 将是初始噪声，$X_{1}X\_1$ 将是输出图像。因此，我理解的一种策略是，我们可以记录大量初始噪声和稳定扩散输出图像的数据集。然后，我们可以微调一个稳定扩散模型，使预测的 epsilon/v 始终在给定 $X_t=t{X}_1+(1-t)X_{0}X\_t=tX\_1+(1-t)X\_0$ 之间形成一条直线。所以整体算法是

关于这一点，一个很棒的事情是，从算法中可以看出，我们将 t 定义在 0 到 1 之间，所以我们只需将 v 乘以 1 加到 $X_{0}X\_0$ 即可得到 $X_{1}X\_1$！

重流

现在，一个问题是，这种拉直是否是那么微不足道的解决方案？是否没有相关的错误？答案是肯定的！解决方案是，一旦您获得了从 $Z_{0}Z\_0$ 到 $Z_{1}Z\_1$ 的最佳路径，您只需反复在该路径上应用矫正流，直到它最终变得笔直，如下所示

算法是

然而，论文提到，虽然重流会使线条更直更短，但代价是获得正确的 $X_{1}X\_1$，因为它偏离得太远。

蒸馏

现在，鉴于我们有一个可以预测速度的重流模型，我们可以对其进行蒸馏。为此，InstaFlow 给出了最佳方程，所以

本质上，这样做的作用是，我们不是试图预测添加到 $X_{0}X\_0$ 之后会变成 $X_{1}X\_1$ 的速度，而是直接预测 $X_{1}X\_1$，这很有趣。

现在，论文中还有一些数学内容涉及非线性问题下的 DDIMs，但由于我们主要关注完成 PR，所以让我们继续讨论 InstaFlow！

InstaFlow

InstaFlow 基本上是将稳定扩散应用于矫正流。从一些统计数据来看，他们使用了 199 个 A100 天（即 4776 个 A100 GPU 小时）进行训练，这对于机构来说成本约为 5000 美元，对于没有协议的机构来说则为 10000 美元。考虑到稳定扩散 2.1 训练使用了 200000 个 GPU 小时（相当于 20 万美元），这相当便宜。它可以在 A100 上以 0.12 秒的速度生成图像，这很合理，因为它是一个 1 步模型。

训练算法

如上所示，算法几乎完全相同，只是我们对文本进行了条件化，而原始的矫正流是无条件的。然后，还有一个额外的蒸馏步骤。作者观察到重流对于高质量非常重要。

Instaflow 训练设置

他们使用了来自 laion2B-en 的部分提示。生成的图像通过 DPM 求解器以 25 步完成，指导比例为 6.0。对于蒸馏，他们使用了 LPIPS 损失，使用了一个我假设是 VGG 的网络来获取图像的高级相似性（面部、物体等）。最后，他们使用了 32 的批大小和 8 个 A100 GPU 进行训练，优化器为 AdamW。

待办事项列表

总的来说，就是这样！因此，PR 的待办事项列表如下：

• 弄清楚如何将 epsilon 映射到速度 我的理解是，我们在矫正流期间忽略 DDPM/epsilon 目标，只让 unet 直接输出 v
• 编写一个脚本来生成潜在噪声、图像和文本以保存到数据集中
• 编写矫正流/重流脚本
• 编写蒸馏脚本