介绍

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

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

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

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

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）。

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

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

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