把Rectified flow应用到了实践, 声称表现比diffusion更好. 文本和图片使用分离的参数, 允许双向流的信息流动交流, 可以获得更好的文本理解. 声称这种结构具有可预测的拓展趋势和较低的验证损失, 改善了text to image生成能力, 达到SOTA水平.

介绍

先介绍了一下diffusion, 正向路径从真实数据趋于随机噪声, 训练的目的是反转这个正向路径. 结合神经网络的近似和泛化特性, 可以生成训练数据中没有的数据, 但是符合训练数据的分布. 这种方法很高效, 是生成高质量图片视频的主流方法, 但是因为迭代的天性, 推理阶段超长的采样时间, 我们需要更快生成的方法.

前向采样的路径选择很重要, 这个过程如果没有完全移除噪声, 可能造成训练分布和测试分布的差异, 出现伪影; 同时这个过程会影响反向采样过程的效率. 当弯曲的路径造成很多步的采样, 一条直线路径是最快的, 累积误差最小.

Rectified Flow, 在一条直线上连接真实数据和噪声. 尽管这种方法在理论上更好, 还没有在实践中证明(只有一些中小规模的实验), 该论文和diffusion比较, 声称效果更好.

他么觉得在text-image合成中, 一个固定的text表示直接注入模型(例如通过cross attention)并不好, 他们同时为图像和文本注入可以学习的流, 这种方法和Rectified Flow相结合, 评估发现下评估更好.

总结

前面没写清楚, 但是大致意思明白了. 就diffusion他的采样路径是弯曲的嘛, 他们觉得这样在推理阶段反向采样就需要更多的步骤, 于是他们引入Rectified Flow, data和noise之间连一条直线, 这样反向采样就简单多了, 这是第一个点; 第二个点, 他们观察到在diffusion方法中, text和image表示在刚开始就通过cross attention等方法融合了, 他们为text和image分别设计了可学习的流, 觉得这样有更好的模态理解, 更好的拓展性.

注意

下面两节的公式推导为AI翻译.

Simulation-Free Training of Flows

我们考虑生成模型，这些模型通过常微分方程（ODE）定义从噪声分布 ${p}_{1}$ 中的样本 ${x}_{1}$ 到数据分布 ${p}_{0}$ 中的样本 ${x}_{0}$ 之间的映射，

d{y}_{t} = {v}_{\Theta }\left( {{y}_{t},t}\right) {dt}, \tag{1}

其中速度 $v$ 由神经网络的权重 $\Theta$ 参数化。Chen等（2018）的先前工作建议通过可微分ODE求解器直接求解方程(1)。然而，这一过程计算成本高昂，特别是对于参数化 ${v}_{\Theta }\left( {{y}_{t},t}\right)$ 的大规模网络架构而言。一种更高效的替代方案是直接回归一个向量场 ${u}_{t}$ ，该向量场生成 ${p}_{0}$ 和 ${p}_{1}$ 之间的概率路径。为了构建这样的 ${u}_{t}$ ，我们定义了一个前向过程，对应于 ${p}_{0}$ 和 ${p}_{1} = \mathcal{N}\left( {0,1}\right)$ 之间的概率路径 ${p}_{t}$ ，如下所示：

{z}_{t} = {a}_{t}{x}_{0} + {b}_{t}\epsilon \;\text{ where }\epsilon \sim \mathcal{N}\left( {0,I}\right) . \tag{2}

对于 ${a}_{0} = 1,{b}_{0} = 0,{a}_{1} = 0$ 和 ${b}_{1} = 1$ ，边缘分布，

{p}_{t}\left( {z}_{t}\right) = {\mathbb{E}}_{\epsilon \sim \mathcal{N}\left( {0,I}\right) }{p}_{t}\left( {{z}_{t} \mid \epsilon }\right) , \tag{3}

与数据和噪声分布一致。

为了表达 ${z}_{t},{x}_{0}$ 和 $\epsilon$ 之间的关系，我们引入 ${\psi }_{t}$ 和 ${u}_{t}$ 如下：

{\psi }_{t}\left( {\cdot \mid \epsilon }\right) : {x}_{0} \mapsto {a}_{t}{x}_{0} + {b}_{t}\epsilon \tag{4}

{u}_{t}\left( {z \mid \epsilon }\right) \mathrel{\text{:=}} {\psi }_{t}^{\prime }\left( {{\psi }_{t}^{-1}\left( {z \mid \epsilon }\right) \mid \epsilon }\right) \tag{5}

由于 ${z}_{t}$ 可以写成常微分方程 ${z}_{t}^{\prime } = {u}_{t}\left( {{z}_{t} \mid \epsilon }\right)$ 的解，初始值 ${z}_{0} = {x}_{0},{u}_{t}\left( {\cdot \mid \epsilon }\right)$ 生成 ${p}_{t}\left( {\cdot \mid \epsilon }\right)$ 。值得注意的是，可以构建一个边缘向量场 ${u}_{t}$ ，它生成边缘概率路径 ${p}_{t}$ （Lipman等，2023）（见B.1），使用条件向量场 ${u}_{t}\left( {\cdot \mid \epsilon }\right)$ ：

{u}_{t}\left( z\right) = {\mathbb{E}}_{\epsilon \sim \mathcal{N}\left( {0,I}\right) }{u}_{t}\left( {z \mid \epsilon }\right) \frac{{p}_{t}\left( {z \mid \epsilon }\right) }{{p}_{t}\left( z\right) } \tag{6}

在使用流匹配目标回归 ${u}_{t}$ 时

{\mathcal{L}}_{FM} = {\mathbb{E}}_{t,{p}_{t}\left( z\right) }{\begin{Vmatrix}{v}_{\Theta }\left( z,t\right) - {u}_{t}\left( z\right) \end{Vmatrix}}_{2}^{2}. \tag{7}

由于公式6中的边缘化，直接处理是不可行的，条件流匹配，

{\mathcal{L}}_{CFM} = {\mathbb{E}}_{t,{p}_{t}\left( {z \mid \epsilon }\right) ,p\left( \epsilon \right) }{\begin{Vmatrix}{v}_{\Theta }\left( z,t\right) - {u}_{t}\left( z \mid \epsilon \right) \end{Vmatrix}}_{2}^{2}, \tag{8}

使用条件向量场 ${u}_{t}\left( {z \mid \epsilon }\right)$ 提供了一个等价且可处理的目标。

为了将损失转换为显式形式，我们将 ${\psi }_{t}^{\prime }\left( {{x}_{0} \mid \epsilon }\right) = {a}_{t}^{\prime }{x}_{0} + {b}_{t}^{\prime }\epsilon$ 和 ${\psi }_{t}^{-1}\left( {z \mid \epsilon }\right) = \frac{z - {b}_{t}\epsilon }{{a}_{t}}$ 代入（5）

{z}_{t}^{\prime } = {u}_{t}\left( {{z}_{t} \mid \epsilon }\right) = \frac{{a}_{t}^{\prime }}{{a}_{t}}{z}_{t} - \epsilon {b}_{t}\left( {\frac{{a}_{t}^{\prime }}{{a}_{t}} - \frac{{b}_{t}^{\prime }}{{b}_{t}}}\right) . \tag{9}

现在，考虑信噪比 ${\lambda }_{t} \mathrel{\text{:=}} \log \frac{{a}_{t}^{2}}{{b}_{t}^{2}}$ 。利用 ${\lambda }_{t}^{\prime } = 2\left( {\frac{{a}_{t}^{\prime }}{{a}_{t}} - \frac{{b}_{t}^{\prime }}{{b}_{t}}}\right)$ ，我们可以将公式（9）重写为

{u}_{t}\left( {{z}_{t} \mid \epsilon }\right) = \frac{{a}_{t}^{\prime }}{{a}_{t}}{z}_{t} - \frac{{b}_{t}}{2}{\lambda }_{t}^{\prime }\epsilon \tag{10}

接下来，我们使用公式（10）将公式（8）重新参数化为噪声预测目标：

{\mathcal{L}}_{CFM} = {\mathbb{E}}_{t,{p}_{t}\left( {z \mid \epsilon }\right) ,p\left( \epsilon \right) }{\begin{Vmatrix}{v}_{\Theta }\left( z,t\right) - \frac{{a}_{t}^{\prime }}{{a}_{t}}z + \frac{{b}_{t}}{2}{\lambda }_{t}^{\prime }\epsilon \end{Vmatrix}}_{2}^{2} \tag{11}

= {\mathbb{E}}_{t,{p}_{t}\left( {z \mid \epsilon }\right) ,p\left( \epsilon \right) }{\left( -\frac{{b}_{t}}{2}{\lambda }_{t}^{\prime }\right) }^{2}{\begin{Vmatrix}{\epsilon }_{\Theta }\left( z,t\right) - \epsilon \end{Vmatrix}}_{2}^{2} \tag{12}

其中我们定义了 ${\epsilon }_{\Theta } \mathrel{\text{:=}} \frac{-2}{{\lambda }_{t}^{\prime }{b}_{t}}\left( {{v}_{\Theta } - \frac{{a}_{t}^{\prime }}{{a}_{t}}z}\right)$ 。

注意，当引入时间相关权重时，上述目标的最优值不会改变。因此，可以推导出各种加权损失函数，它们提供了通向期望解的信号，但可能会影响优化轨迹。为了统一分析不同方法，包括经典扩散公式，我们可以将目标写成以下形式（遵循Kingma & Gao（2023））：

{\mathcal{L}}_{w}\left( {x}_{0}\right) = - \frac{1}{2}{\mathbb{E}}_{t \sim \mathcal{U}\left( t\right) ,\epsilon \sim \mathcal{N}\left( {0,I}\right) }\left\lbrack {{w}_{t}{\lambda }_{t}^{\prime }{\begin{Vmatrix}{\epsilon }_{\Theta }\left( {z}_{t},t\right) - \epsilon \end{Vmatrix}}^{2}}\right\rbrack ,

其中 ${w}_{t} = - \frac{1}{2}{\lambda }_{t}^{\prime }{b}_{t}^{2}$ 对应于 ${\mathcal{L}}_{CFM}$ 。

Flow Trajectories

在这项工作中，我们考虑了上述形式体系的不同变体，并在下文中简要描述。

矫正流矫正流（RFs）（Liu等，2022；Albergo & Vanden-Eijnden，2022；Lipman等，2023）将前向过程定义为数据分布与标准正态分布之间的直线路径，即

{z}_{t} = \left( {1 - t}\right) {x}_{0} + {t\epsilon }, \tag{13}

并使用 ${\mathcal{L}}_{CFM}$ ，然后对应于 ${w}_{t}^{\mathrm{{RF}}} = \frac{t}{1 - t}$ 。网络输出直接参数化速度 ${v}_{\Theta }$ 。

EDM EDM（Karras等，2022）使用的前向过程形式为

{z}_{t} = {x}_{0} + {b}_{t}\epsilon \tag{14}

其中（Kingma & Gao，2023） ${b}_{t} = \exp {F}_{\mathcal{N}}^{-1}\left( {t \mid {P}_{m},{P}_{s}^{2}}\right)$ ， ${F}_{\mathcal{N}}^{-1}$ 是均值为 ${P}_{m}$ 、方差为 ${P}_{s}^{2}$ 的正态分布的分位数函数。注意，此选择导致

{\lambda }_{t} \sim \mathcal{N}\left( {-2{P}_{m},{\left( 2{P}_{s}\right) }^{2}}\right) \;\text{ for }t \sim \mathcal{U}\left( {0,1}\right) \tag{15}

网络通过 $\mathbf{F}$ 预测（Kingma & Gao，2023；Karras等，2022）进行参数化，损失可以表示为 ${\mathcal{L}}_{{w}_{t}^{\mathrm{{EDM}}}}$ ，其中

{w}_{t}^{\mathrm{{EDM}}} = \mathcal{N}\left( {{\lambda }_{t} \mid - 2{P}_{m},{\left( 2{P}_{s}\right) }^{2}}\right) \left( {{e}^{-{\lambda }_{t}} + {0.5}^{2}}\right) \tag{16}

Cosine（Nichol & Dhariwal，2021）提出了一种形式的前向过程

{z}_{t} = \cos \left( {\frac{\pi }{2}t}\right) {x}_{0} + \sin \left( {\frac{\pi }{2}t}\right) \epsilon . \tag{17}

结合 $\epsilon$ 参数化和损失，这对应于权重 ${w}_{t} = \operatorname{sech}\left( {{\lambda }_{t}/2}\right)$ 。当与v预测损失（Kingma & Gao，2023）结合时，权重由 ${w}_{t} = {e}^{-{\lambda }_{t}/2}$ 给出。

(LDM-)Linear LDM（Rombach等，2022）使用了DDPM调度（Ho等，2020）的修改版本。两者均为方差保持调度，即 ${b}_{t} = \sqrt{1 - {a}_{t}^{2}}$ ，并根据扩散系数 ${\beta }_{t}$ 定义 ${a}_{t}$ 为离散时间步 $t = 0,\ldots ,T - 1$ 的 ${a}_{t} = {\left( \mathop{\prod }\limits_{{s = 0}}^{t}\left( 1 - {\beta }_{s}\right) \right) }^{\frac{1}{2}}$ 。对于给定的边界值 ${\beta }_{0}$ 和 ${\beta }_{T - 1}$ ，DDPM使用 ${\beta }_{t} = {\beta }_{0} + \frac{t}{T - 1}\left( {{\beta }_{T - 1} - {\beta }_{0}}\right)$ ，而LDM使用 ${\beta }_{t} ={\left( \sqrt{{\beta }_{0}} + \frac{t}{T - 1}\left( \sqrt{{\beta }_{T - 1}} - \sqrt{{\beta }_{0}}\right) \right) }^{2}$ 。

针对RF模型的定制SNR采样器

RF损失在 $\left\lbrack {0,1}\right\rbrack$ 中的所有时间步上均匀地训练速度 ${v}_{\Theta }$ 。然而，直观地说，生成的速度预测目标 $\epsilon - {x}_{0}$ 在 $\left\lbrack {0,1}\right\rbrack$ 中间的时间步 $t$ 上更困难，因为对于 $t = 0$ ，最优预测是 ${p}_{1}$ 的均值，而对于 $t = 1$ ，最优预测是 ${p}_{0}$ 的均值。通常，将 $t$ 上的分布从常用的均匀分布 $\mathcal{U}\left( t\right)$ 更改为密度为 $\pi \left( t\right)$ 的分布，等价于加权损失 ${\mathcal{L}}_{{w}_{t}^{\pi }}$ ，其中

{w}_{t}^{\pi } = \frac{t}{1 - t}\pi \left( t\right) \tag{18}

因此，我们旨在通过对中间时间步进行更频繁的采样来赋予它们更多权重。接下来，我们将描述用于训练模型的时间步密度 $\pi \left( t\right)$ 。

Logit-Normal采样一种对中间步骤赋予更多权重的分布选项是logit-normal分布（Atchison & Shen，1980）。其密度为

{\pi }_{\ln }\left( {t;m,s}\right) = \frac{1}{s\sqrt{2\pi }}\frac{1}{t\left( {1 - t}\right) }\exp \left( {-\frac{{\left( \operatorname{logit}\left( t\right) - m\right) }^{2}}{2{s}^{2}}}\right), \tag{19}

其中 $\operatorname{logit}\left( t\right) = \log \frac{t}{1 - t}$ ，有一个位置参数 $m$ 和一个尺度参数 $s$ 。位置参数使我们能够将训练时间步偏向数据 ${p}_{0}$ （负 $m$ ）或噪声 ${p}_{1}$ （正 $m$ ）。如图11所示，尺度参数控制分布的宽度。

在实践中，我们从正态分布 $u \sim \mathcal{N}\left( {u;m,s}\right)$ 中采样随机变量 $u$ ，并通过标准逻辑函数对其进行映射。

重尾模式采样对数正态密度总是在端点0和1处消失。为了研究这是否对性能有不利影响，我们还使用了在 $\left\lbrack {0,1}\right\rbrack$ 上具有严格正密度的时间步长采样分布。对于尺度参数 $s$ ，我们定义

{f}_{\text{mode }}\left( {u;s}\right) = 1 - u - s \cdot \left( {{\cos }^{2}\left( {\frac{\pi }{2}u}\right) - 1 + u}\right) . \tag{20}

对于 $- 1 \leq s \leq \frac{2}{\pi - 2}$ ，此函数是单调的，我们可以用它从隐含密度 ${\pi }_{\text{mode }}\left( {t;s}\right) =$ $\left| {\frac{d}{dt}{f}_{\text{mode }}^{-1}\left( t\right) }\right|$ 中采样。如图11所示，尺度参数控制采样过程中是更倾向于中点（正 $s$ ）还是端点（负 $s$ ）。此公式还包括一个均匀加权 ${\pi }_{\text{mode }}\left( {t;s = 0}\right) = \mathcal{U}\left( t\right)$ 用于 $s = 0$ ，这在之前关于修正流（Rectified Flows）的研究中已被广泛使用（Liu等，2022；Ma等，2024）。

CosMap 最后，我们还在RF设置中考虑了第3节中的余弦调度（Nichol&Dhariwal，2021）。具体来说，我们在寻找一个映射 $f : u \mapsto f\left( u\right) =$ $t,u \in \left\lbrack {0,1}\right\rbrack$ ，使得log-snr与余弦调度相匹配： $2\log \frac{\cos \left( {\frac{\pi }{2}u}\right) }{\sin \left( {\frac{\pi }{2}u}\right) } = 2\log \frac{1 - f\left( u\right) }{f\left( u\right) }$ 。求解 $f$ ，我们得到 $u \sim \mathcal{U}\left( u\right)$

t = f\left( u\right) = 1 - \frac{1}{\tan \left( {\frac{\pi }{2}u}\right) + 1}, \tag{21}

由此我们得到密度

{\pi }_{\text{CosMap }}\left( t\right) = \left| {\frac{d}{dt}{f}^{-1}\left( t\right) }\right| = \frac{2}{\pi - {2\pi t} + {2\pi }{t}^{2}}. \tag{22}

文本到图像架构

对于文本条件下的图像采样，我们的模型必须同时考虑文本和图像两种模态。我们使用预训练模型来获得合适的表示，然后描述扩散主干的架构。图2展示了其概览。

我们的总体设置遵循LDM（Rombach等，2022），在预训练自动编码器的潜在空间中训练文本到图像模型。与将图像编码为潜在表示类似，我们也遵循先前的方法（Saharia等，2022b；Balaji等，2022），并使用预训练、冻结的文本模型对文本条件 $c$ 进行编码。详细信息见附录B.2。

多模态扩散主干我们的架构基于DiT（Peebles&Xie，2023）架构。DiT仅考虑类条件图像生成，并使用调制机制使网络同时以扩散过程的时间步和类标签为条件。类似地，我们使用时间步 $t$ 和 ${c}_{\text{vec }}$ 的嵌入作为调制机制的输入。然而，由于池化后的文本表示仅保留了关于文本输入的粗粒度信息（Podell等，2023），网络还需要来自序列表示 ${c}_{\text{ctxt }}$ 的信息。

我们构建了一个由文本和图像输入嵌入组成的序列。具体来说，我们添加位置编码并将 $2 \times 2$ 个 $x \in {\mathbb{R}}^{h \times w \times c}$ 潜在像素表示的补丁展平为长度为 $\frac{1}{2} \cdot h \cdot \frac{1}{2} \cdot w$ 的补丁编码序列。在将此补丁编码和文本编码 ${c}_{\text{ctxt }}$ 嵌入到共同维度后，我们将两个序列连接起来。然后我们遵循DiT并应用一系列调制注意力和多层感知机（MLP）。

rectifiedFlow

图2. 我们的模型架构。连接用 $\odot$ 表示，逐元素乘法用 $*$ 表示。 $Q$ 和 $K$ 的RMS归一化可以添加以稳定训练过程。建议放大查看。

由于文本和图像嵌入在概念上差异很大，我们为这两种模态使用了两组独立的权重。如图2b所示，这相当于为每种模态拥有两个独立的Transformer，但在注意力操作中将两种模态的序列结合在一起，使得两种表示可以在各自的空间中工作，同时考虑另一种表示。

在我们的扩展实验中，我们通过设置隐藏层大小为 ${64} \cdot d$ （在MLP块中扩展为 $4 \cdot {64} \cdot d$ 通道），以及注意头数量等于 $d$ ，以模型深度 $d$ （即注意块的数量）来参数化模型的大小。

感受

遇到困难我就摆烂, 这篇文章绝对有研究的价值, 但是现在不想静下心来读, 先占个位置.

Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

介绍

Simulation-Free Training of Flows

Flow Trajectories

针对RF模型的定制SNR采样器

文本到图像架构

On this page

介绍

Simulation-Free Training of Flows

Flow Trajectories

针对RF模型的定制SNR采样器

文本到图像架构