PixelFlow：基于流的像素空间生成模型 原创

摘要

我们提出了 PixelFlow，这是一系列直接在原始像素空间中运行的图像生成模型，与主流的隐空间模型形成对比。这种方法通过消除对预训练变分自编码器（VAE）的需求，简化了图像生成过程，并使整个模型能够端到端训练。通过高效的级联流建模，PixelFlow 在像素空间中实现了可承受的计算成本。在 256×256 ImageNet 类条件图像生成基准上，它实现了 1.98 的 FID。定性的文本到图像结果表明，PixelFlow 在图像质量、艺术性和语义控制方面表现出色。我们希望这种新范式能为下一代视觉生成模型带来启发并开辟新的机会。代码和模型可在​​https://github.com/ShoufaChen/PixelFlow​​获取。

1. 引言

“如无必要，勿增实体。”

—— 奥卡姆的威廉

在 Stable Diffusion（SD）模型系列 [17,46,47,50] 成功的推动下，隐扩散模型（LDMs）[50] 已成为跨多种模态生成建模的事实上的标准，涵盖图像 [17,35,45]、视频 [7,8,23,66,69]、音频 [18,39] 和 3D [57,67] 等领域。如图 1（a）所示，LDMs 使用预训练的变分自编码器（VAEs）将原始数据压缩到紧凑的隐空间中。这种压缩降低了计算需求，并促进了高效的扩散去噪。尽管取得了广泛的成功，但 LDMs 将 VAE 和扩散组件分离，阻碍了联合优化并使整体诊断复杂化。

另一种方法是在原始像素空间中实现扩散模型。虽然直观，但由于处理每个像素相关性需要大量资源，对于高分辨率图像来说，这在计算上是不可承受的。考虑到这一点，先前的研究 [20,22,44,51,52] 通常采用级联方法：首先生成低分辨率图像，然后使用额外的上采样器生成高质量输出，低分辨率图像作为条件输入，如图 1（b）所示。然而，这些级联方法也为不同阶段引入了独立的网络，仍然限制了端到端设计的优势。

在这项工作中，我们引入 PixelFlow，这是一个简单而有效的端到端框架，用于在原始像素空间中直接生成图像，无需像 VAE 或上采样器这样的独立网络。如图 1（c）所示，PixelFlow 使用统一的参数集通过流匹配 [38,40] 对级联分辨率的多尺度样本进行建模。在去噪的早期阶段，当噪声水平较高时，PixelFlow 在较低分辨率的样本上运行。随着去噪的进行，分辨率逐渐增加，直到在最后阶段达到目标分辨率。这种渐进策略避免了在全分辨率下执行所有去噪步骤，从而显著降低了生成过程的总体计算成本。

在训练期间，不同时间步的跨尺度样本通过以下方式构建：（1）将图像调整为连续的尺度，并向每个缩放后的图像添加高斯噪声；（2）在相邻尺度的噪声图像之间进行插值作为模型输入并进行速度预测。整个模型使用从所有阶段均匀采样的训练示例进行端到端训练。在推理期间，过程从最低分辨率的纯高斯噪声开始。然后模型逐步对图像进行去噪和上采样，直到达到目标分辨率。

图片

图1. 基于隐变量的扩散模型（LDMs）、基于像素的扩散模型（PDMs）与PixelFlow的设计范式对比：(a) LDMs将训练分为两个独立阶段——先独立训练现成的VAE模型，再在预训练VAE提取的token上训练扩散模型；(b) 传统PDMs通常需训练两个独立模型：针对低分辨率图像的扩散模型和用于高分辨率合成的上采样器；(c) 相比之下，PixelFlow提供了端到端的像素级生成方案，兼具高效性与强大生成性能。

我们在类条件和文本到图像生成任务上对 PixelFlow 进行了评估。与已建立的隐空间扩散模型 [42,45,50] 相比，PixelFlow 提供了有竞争力的性能。例如，在 256×256 ImageNet 类条件生成基准上，PixelFlow 实现了 1.98 的 FID。对于文本到图像生成，PixelFlow 在广泛使用的基准上进行了评估，在 GenEval [19] 上得分为 0.64，在 DPG-Bench [26] 上得分为 77.93。此外，图 5 和图 6 的定性结果表明，PixelFlow 具有很强的视觉保真度和文本 - 图像对齐，突出了像素空间生成在未来研究中的潜力。

PixelFlow 的贡献总结为以下三点：

•通过消除对预训练 VAE 的需求，我们直接在原始像素空间中建立了端到端可训练的图像生成模型。

•通过从低分辨率到高分辨率的级联流建模，我们的模型在训练和推理中都实现了可承受的计算成本。

•PixelFlow 在视觉质量方面取得了有竞争力的性能，包括在 256×256 ImageNet 类条件生成基准上的 1.98 FID，以及在文本到图像生成上的吸引人的特性。

2. 相关工作

2.1 隐空间扩散 / 流模型

变分自编码器（VAEs）已成为许多最新生成模型 [16,17,35,47,48,50,59,66] 的核心组件，能够将视觉数据从像素空间映射到低维的感知等效隐空间。这种紧凑的表示便于更高效的训练和推理。然而，VAEs 常常牺牲高频细节 [47]，导致生成输出中不可避免的低级伪影。出于对算法简单性和完全端到端优化的渴望，我们放弃了 VAE，直接在像素空间中操作。

2.2 像素空间扩散 / 流模型

早期的扩散模型 [2,21,56] 主要直接在像素空间中运行，旨在单阶段捕获图像的分布。然而，这种方法对于高分辨率图像生成既具有挑战性又效率低下，导致了级联模型 [20,22,30,52] 的发展，这些模型通过一系列阶段生成图像。这些级联模型通常从生成低分辨率图像开始，然后通过超分辨率模型进行上采样以获得更高的分辨率。然而，基于扩散的超分辨率过程通常需要从纯噪声开始，以低分辨率输出为条件，导致耗时且低效的生成过程。此外，在孤立阶段训练这些模型阻碍了端到端优化，并需要精心设计的策略来确保超分辨率阶段。

此外，像素空间生成的最新进展引入了创新的架构。Simple Diffusion [24,25] 提出了一种用于高分辨率图像合成的简化扩散框架，通过调整模型架构和噪声调度在 ImageNet 上取得了强大性能。FractalGen [37] 通过递归调用原子生成模块构建分形生成模型，形成自相似架构，在逐像素图像生成中表现出强大性能。TarFlow [68] 提出了一种基于 Transformer 的归一化流架构，能够直接建模和生成像素。

3. PixelFlow

3.1 预备知识：流匹配

流匹配算法 [1,38,40] 逐步将样本从先验分布（通常是标准正态分布）转换为目标数据分布。这是通过定义一个前向过程来实现的，该过程由一系列直接连接先验分布样本和目标分布对应样本的线性路径组成。在训练期间，通过首先采样目标样本x_1，从标准正态分布中抽取噪声x_0∼N(0,1)，并选择时间步t∈[0,1]来构建训练示例。然后通过线性插值定义训练示例：

模型被训练来近似由常微分方程（ODE）定义的速度v_t = \frac{dx_t}{dt}，使其能够有效地引导中间样本x_t到真实数据样本x_1的转换。

流匹配的一个显著优点是它能够在两个任意分布之间进行插值，而不仅限于使用标准高斯作为源域。因此，在图像生成任务中，流匹配超越了噪声到图像的场景，可以有效地应用于图像到图像翻译等各种应用。

3.2 像素空间中的多尺度生成

PixelFlow 通过多阶段去噪过程逐步增加图像的分辨率来生成图像。为此，我们通过在每个尺度上以因子 2 递归下采样目标图像x_1来构建其多尺度表示。如图 2 所示，PixelFlow 将图像生成过程分为 S 个阶段。每个阶段s\in\{0,1,\ldots,S-1\}在由起始和结束状态(x_{t_0^s}, x_{t_1^s})定义的时间间隔内运行。在S=1的退化情况下，PixelFlow 简化为用于图像生成的标准单阶段流匹配方法，类似于最近的工作 [17,42]，但关键是在像素空间而非隐空间中操作。

图2. 基于像素空间的级联图像生成框架PixelFlow。我们将完整生成过程划分为多级分辨率阶段：每个分辨率阶段起始时，对前阶段含较多噪声的结果进行上采样，并将其作为当前阶段的起始点。通过这种渐进式优化机制，随着分辨率提升可获得更精细的生成样本。

对于每个阶段 s，我们定义起始和结束状态如下：

其中，\text{Down()}和\text{Up()}分别表示下采样和上采样操作。除非另有说明，我们采用双线性插值进行下采样，最近邻插值进行上采样。

为了训练模型，我们通过在起始和结束状态之间进行线性插值来采样中间表示：

其中，\tau = \frac{t - t_0^s}{t_1^s - t_0^s}是第 s 阶段内的重新缩放时间步 [29,65]。

然后，我们的目标是训练一个模型\mu_\theta(\cdot)来预测速度\mu_\theta(x_{t_\tau^s}, \tau)，目标为v_t = x_{t_1^s} - x_{t_0^s}。我们使用均方误差（MSE）损失，形式上表示为：

\mathbb{E}_{s,t,(x_{t_1^s}, x_{t_1^s})} \left\| \mu_\theta(x_{t_\tau^s}, \tau) - v_t \right\|^2

3.3 模型架构

我们使用基于 Transformer 的架构 [62] 实例化\mu_\theta(\cdot)，选择它是因为其简单性、可扩展性和在生成建模中的有效性。具体来说，我们的实现基于标准扩散 Transformer（DiT）[45]，在所有实验中采用 XL 规模配置。为了更好地与 PixelFlow 框架保持一致，我们引入了以下几个修改：

3.3.1 补丁化（Patchify）

遵循视觉 Transformer（ViT）设计 [15,45]，PixelFlow 的第一层是补丁嵌入层，通过线性投影将输入图像的空间表示转换为 1D 令牌序列。与在 VAE 编码的隐空间上操作的先前隐 Transformer [17,42,45] 不同，PixelFlow 直接对原始像素输入进行令牌化。为了支持批处理中多个分辨率的高效注意力，我们应用序列打包策略 [11]，沿着序列维度连接对应不同分辨率的扁平令牌序列。

3.3.2 旋转位置编码（RoPE）

补丁化后，我们将原始的正弦位置编码 [45] 替换为 RoPE [58]，以更好地处理不同的图像分辨率。RoPE 在支持长度外推方面表现出强大性能，尤其是在大型语言模型中。为了将其适应 2D 图像数据，我们通过对高度和宽度维度独立应用 1D-RoPE 来应用 2D-RoPE，每个维度占据隐藏状态的一半。

3.3.3 分辨率嵌入（Resolution Embedding）

由于 PixelFlow 使用共享的模型参数在多个分辨率上运行，我们引入了额外的分辨率嵌入来区分不同的分辨率。具体来说，我们将补丁嵌入后特征图的绝对分辨率作为条件信号。该信号使用正弦位置嵌入 [62] 进行编码，并在传递到模型之前添加到时间步嵌入中。

3.3.4 文本到图像生成

虽然类条件图像生成通常通过自适应层归一化（adaLN）[45] 整合条件信息，但我们通过在每个 Transformer 块内的每个自注意力层之后引入交叉注意力层 [6,7]，将 PixelFlow 扩展为支持文本到图像生成。这种设计使模型能够在生成过程的每个阶段有效地将视觉特征与文本输入对齐。遵循最近的工作 [8,59]，我们采用 Flan-T5-XL 语言模型 [10] 来提取丰富的文本嵌入，作为整个网络的条件信号。

3.4 训练和推理

为了促进高效训练，我们使用公式（4）中定义的插值方案从所有分辨率阶段均匀采样训练示例。此外，我们采用序列打包技术 [11]，允许在单个小批量中联合训练尺度可变的示例，提高了效率和可扩展性。

在推理期间，生成过程从最低分辨率的纯高斯噪声开始，通过多个阶段逐步过渡到更高的分辨率。在每个分辨率阶段内，我们应用标准的基于流的采样，根据速度和准确性的权衡，使用欧拉离散采样器 [17] 或 Dopri5 求解器。为了确保跨尺度的平滑和连贯过渡，我们采用重新去噪策略 [29,60]，有效缓解了多尺度生成管道中常见的跳跃点问题 [4]。

4. 实验

4.1 实验设置

我们在 ImageNet-1K [12] 数据集上评估 PixelFlow 的类条件图像生成能力。除非另有说明，我们在 256×256 分辨率下训练 PixelFlow。所有模型使用 AdamW 优化器 [32,41] 进行训练，学习率为1×10^{-4}。性能主要通过 Fréchet Inception Distance（FID）使用标准评估工具包 1 进行测量。我们还报告了 Inception Score（IS）[53]、sFID [43] 和 Precision/Recall [33]。

对于文本条件图像生成，我们逐步将 PixelFlow 从 256×256 训练到 1024×1024 分辨率。我们与当前最先进的生成模型进行定性比较，并在流行的基准上进行定量评估，如 T2ICompBench [27]、GenEval [19] 和 DPG-Bench [26]。

4.2 模型设计

4.2.1 起始序列长度

原则上，PixelFlow 可以训练为从非常低的分辨率（例如 1×1）逐步增加到目标分辨率。然而，这种方法在实践中效率低下，因为极低分辨率的令牌传达的有意义信息有限。此外，为非常短的序列分配过多的时间步会未充分利用现代 GPU 的计算能力，导致模型 FLOPS 利用率下降。因此，我们探索图像生成开始时的分辨率（我们称为起始图像分辨率）如何影响整体性能。

对于我们的 Transformer 骨干网络，注意力操作中涉及的令牌数量由原始图像分辨率和补丁大小决定。在这个实验中，我们保持 2×2 的补丁大小 [45]，使起始序列长度直接取决于起始图像分辨率。具体来说，我们评估了三种起始序列长度 ——2×2、8×8 和 32×32—— 同时将目标分辨率固定为 64×64。值得注意的是，32×32 设置代表了没有跨分辨率级联的香草像素基方法。

如表 1 所示，在这些配置中，8×8 起始序列长度与 32×32 基线相比，实现了相当甚至稍好的 FID。这表明，从适当较小的分辨率开始生成并逐步扩展，可以在保持生成质量的同时，通过为最大分辨率阶段分配更少的计算来提高计算效率。相反，将起始序列长度进一步减少到 2×2 会导致性能下降，可能是因为极低分辨率的令牌提供的有用信息有限，无法为后续生成步骤提供足够的指导。考虑到生成质量和计算效率，我们因此采用 8×8 作为默认起始序列长度。

4.2.2 补丁大小

接下来，我们在保持起始序列长度为 2×2 的同时，研究补丁大小对模型性能的影响。首先，我们在目标分辨率为 64×64 的情况下进行实验，比较两种补丁大小 ——2×2 和 4×4—— 结果如表 2 的上部所示。我们观察到 PixelFlow 在这两种设置下表现非常相似，4×4 补丁在五个评估指标中的四个上略优于 2×2 补丁。此外，使用 4×4 补丁大小消除了 2×2 补丁大小配置所需的最高分辨率阶段，从而提高了效率。

当扩展到更大的目标分辨率（即 256×256）时，由于资源需求巨大，使用 2×2 补丁大小在计算上变得不可行，限制我们只能进行 100K 次训练迭代（表 2 的中间部分）。这一限制促使我们采用更大的补丁大小。尽管将补丁大小进一步增加到 8×8 显著提高了计算效率，但它导致性能质量明显下降。此外，即使经过扩展训练（1600K 次迭代），这种性能差距仍然存在，如表 2 的底部所示。考虑到生成质量和计算成本，我们因此选择 4×4 作为默认补丁大小。

4.3 推理调度

在表 3 中，我们详细分析了推理配置空间，包括每个分辨率阶段的推理步骤数、ODE 求解器的选择以及无分类器引导（CFG）的调度。

4.3.1 采样步骤数

在表 3a 中，我们评估了每个分辨率阶段的推理步骤数对生成质量的影响。随着步骤数的增加，我们观察到 FID、sFID 和 IS 的持续改进，在 30 步时达到最佳整体性能。超过这个点后，增益饱和甚至略有下降，表明收益递减。

PixelFlow 的一个显著优点是它在推理期间为每个分辨率阶段分配不同数量采样步骤的灵活性。这种自适应配置允许对采样过程进行细粒度控制，实现性能 - 效率权衡。超越统一设置并探索更细粒度的阶段特定步骤分配有可能带来进一步的性能提升。

4.3.2 ODE 求解器

我们进一步研究 ODE 求解器类型对生成质量的影响。如表 3b 所示，我们将一阶欧拉求解器与自适应高阶 Dormand-Prince（Dopri5）求解器 [14] 进行比较。结果表明，Dopri5 在大多数评估指标上始终优于欧拉求解器，实现了更低的 FID 和 sFID 分数、更高的 Inception Score 以及稍好的精度，同时保持相似的召回率。这表明，更准确和自适应的求解器（如 Dopri5）可以更好地捕捉生成动态，从而生成更高质量的样本 —— 尽管通常伴随着更高的计算成本。

4.3.3 CFG 调度

受最近工作 [5,34,63] 的启发，我们提出了一种阶段式 CFG 调度，其中不同阶段应用不同的 CFG 值，并且从早期阶段到后期阶段，值从 1 增加到CFG_{max}。在 4 个阶段的情况下，我们发现0, 1/6, 2/3和1的(CFG_{max}-1)给出了最佳的 FID 性能。全局常数 CFG 和阶段式 CFG 之间的比较如表 3c 所示，其中我们为每种方法搜索最佳 CFG 值。我们提出的阶段式 CFG 将 FID 性能从 2.43 提升到 1.98。

4.4 在 ImageNet 基准上的比较

在表 4 中，我们将 PixelFlow 与基于隐空间和像素空间的图像生成模型在 ImageNet 256×256 基准上进行了比较。PixelFlow 实现了 1.98 的 FID，代表了相对于最先进的隐空间方法的极具竞争力的性能。例如，它优于 LDM [50]（FID 3.60）、DiT [45]（FID 2.27）和 SiT [42]（FID 2.06），同时实现了可比的 IS 和召回分数。这些结果突出了我们设计的有效性，表明 PixelFlow 可以作为高质量视觉生成系统的强大原型。

与最近的像素基模型相比，PixelFlow 实现了卓越的样本质量。它显著优于 FractalMAR-H [37]，并且与 ADM-U [13]、SiD2 [25] 和 VDM++[31] 等强基线相比，提供了具有竞争力或更好的结果。

我们在图 4 中可视化了 256×256 分辨率下 PixelFlow 的类条件图像生成。我们可以观察到我们的模型能够跨广泛的类别生成高视觉质量的图像。

4.5 文本到图像生成

4.5.1 设置

我们为 PixelFlow 的文本到图像生成采用两阶段训练策略。首先，使用 ImageNet 预训练的检查点在 256×256 分辨率下初始化模型，并在 LAION 数据集 [55] 的子集上以相同分辨率进行训练。在第二阶段，我们在精心挑选的高审美质量图像集上以更高的 512×512 分辨率对模型进行微调。所有报告的 PixelFlow 结果均基于此最终的 512×512 分辨率模型。

4.5.2 定量结果

如表 5 所示，PixelFlow 在所有基准上均实现了有竞争力的性能，展示了在自由形式文本到图像生成中的强大组合理解能力。它在 T2I-CompBench 上表现尤为出色，在颜色和纹理绑定方面得分很高，在 GenEval（0.64）和 DPG-Bench（77.93）上取得了坚实的结果，超越了许多已建立的模型。这些结果强调了 PixelFlow 作为基于自然语言条件的像素空间图像生成的有前途方向 —— 展示了其在开放式文本驱动图像合成中的潜力。

4.5.3 可视化

我们在图 3 中可视化了采样过程中的中间结果，具体显示了每个分辨率阶段的最后一步。随着分辨率的增加，出现了明显的去噪趋势 —— 图像在每个阶段逐渐变得更清晰，噪声更少。图 5（512×512）和图 6（1024×1024）中显示了额外的生成样本及其输入文本提示。PixelFlow 展示了高视觉保真度和强文本 - 图像对齐，有效地捕捉了复杂提示中的关键视觉元素及其关系。值得注意的是，它生成了细粒度的细节 —— 如动物毛发、人类头发和帽子纹理 —— 突出了其在像素空间中对细节的强烈关注。

图3. 级联生成阶段的中间结果可视化。我们从四个生成阶段分别提取中间结果进行直接可视化观测，可以清晰观察到不同分辨率阶段逐步去噪的过程。

5. 结论

我们引入了 PixelFlow，这是一种新颖的图像生成模型，通过直接在原始像素空间中操作，重新思考了基于隐空间模型的主导地位。通过在不同分辨率阶段之间直接转换，我们的模型在简单性和端到端可训练性方面展现出引人注目的优势。在类条件图像生成和文本到图像生成基准上，PixelFlow 已被证明与流行的基于隐空间的方法相比，展示了有竞争力的图像生成能力。我们希望这种新视角将激发未来视觉生成模型的研究。

局限性

尽管具有优势，PixelFlow 仍然面临某些局限性。尽管模型避免了所有阶段的全分辨率计算，但最后阶段需要全分辨率注意力，这约占总推理时间的 80%。此外，我们观察到随着序列长度的增加，训练收敛速度变慢。解决这些挑战为未来提高效率和可扩展性提供了机会。

本文转载自公众号AIRoobt ，作者：Shoufa Chen等

原文链接：​​https://mp.weixin.qq.com/s/ktz8qQ0C-l_W_nuQvxkKkw​​