一、引言
在这份技术报告中，作者提出了PIXART-δ，它将LCM和ControlNet合并到PIXART-α中。值得注意的是，PIXART-α 是一种先进的高质量 1024px 扩散转换器文本到图像合成模型，以其通过极其高效的训练过程实现的卓越图像生成质量而闻名。
我们将 LCM 合并到 PIXART-δ 中以加速推理。 LCM 通过将反向扩散过程视为求解增强概率流 ODE (PF-ODE)，在预训练的 LDM 上仅需 2∼4 个步骤即可实现高质量和快速的推理，这使得 PIXART- δ 在（∼4）步内生成样本，同时保持高质量的生成。因此，PIXART-δ 在 A100 GPU 上每张 1024 × 1024 图像只需 0.5 秒，与 PIXART-α 相比，推理速度提高了 7 倍。我们还支持 LCM-LoRA，以获得更好的用户体验和便利
此外，我们将类似 ControlNet 的模块合并到 PIXART-δ 中。 ControlNet展示了在各种条件下对文本到图像扩散模型输出的卓越控制。然而，值得注意的是，ControlNet 的模型架构是为基于 UNet 的扩散模型而复杂设计的，我们观察到将其直接复制到 Transformer 模型中效果较差。因此，我们提出了一种为 Transformer 模型定制的新型 ControlNet-Transformer 架构。我们的 ControlNet-Transformer 实现了明确的可控性并获得了高质量的图像生成。
二、LCM IN PIXART-δ
在本节中，我们采用潜在一致性蒸馏（LCD）在 120K 内部图像文本对上训练 PIXART-δ。在2.1，我们首先提供详细的训练算法和具体修改的消融研究。在2.2，我们说明了PIXART-δ的训练效率和LCM的加速比。最后，在2.3，我们展示了PIXART-δ的训练细节。
2.1 算法和修改
LCD 算法：源自初始的一致性蒸馏（CD）和LCD 算法，我们在算法1中提出了具有无分类器指导（CFG）的PIXART-δ的伪代码。具体来说，如图1所示的训练流程所示，三个模型（教师模型、学生模型和EMA模型）充当ODE求解器的去噪器。分别为\(\textcolor{blue}{\Psi(\cdot, \cdot, \cdot, \cdot), f_\theta,}\) 和 \(\textcolor{blue}{f_{\theta^-}}\)。在训练过程中，我们首先在时间步 \(\textcolor{blue}{t_{n+k}}\) 处采样噪声，其中教师模型用于去噪以获得\(\textcolor{blue}{\hat{z}_{T_{t_0}}}\)。然后，利用ODE求解器\(\textcolor{blue}{\Psi(\cdot, \cdot, \cdot, \cdot)}\) 从\(\textcolor{blue}{z_{t_{n+k}}}\)和\(\textcolor{blue}{\hat{z}_{T_{t_0}}}\)计算\(\textcolor{blue}{\hat{z}^{\Psi,\omega}_{t_n}}\)。然后利用EMA模型进一步去噪，得到\(\textcolor{blue}{\hat{z}_{E_{t_0}}}\)。同时，学生模型在\(\textcolor{blue}{t_{n+k}}\)处对样本\(\textcolor{blue}{z_{t_{n+k}}}\)进行加噪，以导出\(\textcolor{blue}{\hat{z}_{S_{t_0}}}\)。最后一步涉及最小化\(\textcolor{blue}{\hat{z}_{S_{t_0}}}\)和\(\textcolor{blue}{\hat{z}_{E_{t_0}}}\)之间的距离，也称为优化一致性蒸馏目标。
图1：PIXART-δ 的训练流程。该图的上半部分提供了训练过程的高级概述，描述了沿着特定 ODE 轨迹的噪声采样和去噪的连续阶段。序列号标记在映射线上，以清楚地指示这些步骤的顺序。下半部分深入研究了预训练（教师）模型和学生模型的复杂作用，揭示了它们在上半部分训练过程中各自的功能，并标记了相应的序列号，以便于交叉引用。
与原始的 LCM 从指定范围 \(\textcolor{blue}{[ω_{min}, ω_{max}]}\) 中选择可变制导尺度 \(\textcolor{blue}{ω}\) 不同，在我们的实现中，我们将制导尺度设置为常数 \(\textcolor{blue}{ω_{fix}}\)，去除了 LCM 中的制导尺度嵌入操作。
超参数的影响。我们的研究补充了 LCM 训练过程的两个关键方面：CFG 规模和批量大小。使用 FID 和 CLIP 分数作为性能基准来评估这些因素。图2中的术语“\(\textcolor{blue}{bs}\)”、“\(\textcolor{blue}{ω_{fix}}\)”和“\(\textcolor{blue}{ω_{Embed}}\)”分别表示训练批量大小、固定指导尺度和嵌入指导尺度。
- CFG 量表分析：参考图 2，我们检查了三种不同的 CFG 量表：(1) 3.5，在我们的消融研究中使用； (2) 4.5，在 PIXART-α 中产生最佳结果； (3) 不同范围的 CFG 尺度嵌入 (\(ω_{ Embed}\))，这是 LCM 中的标准方法。我们的研究表明，采用恒定的指导尺度而不是更复杂的 CFG 嵌入可以提高 PIXART-δ 的性能并简化实施。
- 批量大小检查：使用 2 个 V100 GPU 和 32 个 V100 GPU 两种配置来评估批量大小对模型性能的影响；每个 GPU 加载 12 个图像。如图 2 所示，我们的结果表明，较大的批量会对 FID 和 CLIP 分数产生积极影响。然而，如图 8 所示，PIXART-δ 也可以快速收敛，并以较小的批量大小获得可比的图像质量。
- 收敛：最后，我们观察到训练过程在大约 5,000 次迭代后趋于收敛。超过这个阶段，进一步的改进是微乎其微的。
图2：FID 和 CLIP Score 对无分类器指导量表 (ω) 的各种策略的消融研究及其对训练期间蒸馏收敛的影响。
Noise Schedule 调整。噪声表是扩散过程中最重要的部分之一。接下来，我们调整了 LCM 中的噪声调度函数以与 PIXART-α 噪声调度保持一致，该噪声调度在蒸馏训练期间具有更高的 logSNR（信噪比）。图 3 可视化了 PIXART-δ 或 LCM 不同选择下的噪声表函数及其各自的 logSNR。值得注意的是，PIXART-δ 可以对更广泛的噪声分布进行参数化，这一功能已被进一步证明可以增强图像生成。
图3：\(β_t\)、噪声调度函数以及 PIXART-δ 和 LCM 之间相应的 logSNR 的实例化。 \(β_t\) 是扩散过程中的系数 \(z_t = \sqrt{\bar{α}_t}z_0 + \sqrt{1 − \bar{α}_t}ε\), \(α_t = 1 − β_t\)
2.2 训练效率和推理时间
对于训练，如表所示。如图 1 所示，我们在 32GB GPU 内存限制内成功地进行了蒸馏过程，同时保留了相同的批量大小并通过 SDXL-LCM 支持高达 1024 × 1024 的图像分辨率。这种训练效率显着使 PIXART-δ 能够在各种消费级 GPU 规范上进行训练。根据第 3.1 节中关于较大批量大小的有益影响的讨论，我们的方法显着使得即使在内存容量有限的 GPU 上也可以利用较大批量大小。
对于推论，如表所示。如图 2 和图 7 所示，我们对我们的模型 PIXART-δ 与其他方法（如 SDXL LCM-LoRA、PIXARTα 和跨不同硬件平台的 SDXL 标准）实现的生成速度进行了比较分析。始终如一，PIXART-δ 在 A100 上的 0.5 秒内实现了 1024x1024 高分辨率图像的生成，并且在 T4 上仅用 3.3 秒、在 V100 上仅用 0.8 秒就完成了该过程，所有这些都在批量大小为 1 的情况下完成。这是一个重要的结果。比其他方法有所改进，例如，SDXL 标准在 T4 上需要 26.5 秒，在 T4 上需要 3.8 秒在 A100 上。 PIXART-δ 的效率是显而易见的，因为与 PIXART-α 和 SDXL 标准分别需要 14 和 25 个步骤相比，它仅用 4 个步骤就在生成速度方面保持了一致的领先优势。值得注意的是，随着8位推理技术的实现，PIXART-δ需要不到8GB的GPU VRAM。这种卓越的效率使 PIXART-δ 能够在各种 GPU 卡上运行，甚至开启了在 CPU 上运行的可能性。
图7：生成的输出示例。上半部分是 PIXART-δ 和 SDXL-LCM 之间的比较，有 4 个采样步骤。下半部分，比较涉及 PIXART-δ 和 PIXART-α（教师模型，使用 DPM-Solver，14 个步骤）。
2.3 训练描述
正如第 2 节中所讨论的。 3.1 中，我们使用具有 120K 图像的高质量内部数据集，在两种分辨率设置（512×512 和 1024×1024）下进行实验。我们利用 PIXART-α 的多尺度图像生成功能（支持 512px 和 1024px 分辨率）顺利地在两种分辨率下训练模型。对于这两种分辨率，PIXART-δ 在达到 5K 次迭代之前都产生了令人印象深刻的结果，此后只观察到了最小的改进。训练在 2 个 V100 GPU 上执行，总批量大小为 24，学习率为 2e-5，EMA 速率 μ = 0.95，并使用 AdamW 优化器。我们采用 DDIMSolver和跳过步骤 k = 20来提高效率。如第 2 节所述。如图 3.1 和图 3 所示，对原始 LCM 调度程序进行了修改，以适应预训练的 PIXART-α 和稳定扩散模型之间的差异。按照 PIXART-α 方法，我们将扩散过程中的 \(\textcolor{blue}{β_t}\) 从缩放线性更改为线性曲线，同时将 βt0 从 0.00085 调整为 0.0001，将 \(\textcolor{blue}{β_{t_T}}\) 从 0.012 调整为 0.02。引导尺度 \(\textcolor{blue}{ω_{fix}}\) 设置为 4.5，在 PIXART-α 中被确定为最佳。虽然在训练过程中省略了 LCM 中 ω 的傅立叶嵌入，但 PIXART-α 和 PIXART-δ 都保持相同的结构和可训练参数。这使我们能够使用与教师扩散模型 (PIXART-α) 相同的参数来初始化一致性函数 \(\textcolor{blue}{f_θ( \hat{z}, ω_{fix}, c, t_n)}\)，而不会影响性能。基于LCM-LoRA的成功，PIXART-δ可以进一步轻松集成LCM-LoRA，增强其对更多样化应用的适应性。
三、ControlNet in PIXART-δ
3.1 架构
ControlNet 最初是为 UNet 架构定制的，采用跳跃连接来增强控制信号的集成。将 ControlNet 无缝融入基于 Transformer 的模型（以 PIXART-δ 为例）带来了独特的挑战。与 UNet 不同，Transformers 缺乏明确的“编码器”和“解码器”块，使得这些组件之间的传统连接不合适。
为了应对这一挑战，我们提出了一种创新方法ControlNet-Transformer，以确保ControlNet与Transformer的有效集成，保持ControlNet在管理控制信息和高质量生成PIXART-δ方面的有效性。
PIXART-δ 包含 28 个 Transformer 块。我们将ControlNet中原来的零卷积替换为零线性层，即权值和偏置都初始化为零的线性层。我们探索以下网络架构：
ControlNet-UNet。为了遵循原始的 ControlNet 设计，我们将前 14 个块视为 PIXART-δ 的“编码器”级别，将后 14 个块视为 PIXART-δ 的“解码器”级别。我们使用 ControlNet 创建 14 个编码块的可训练副本。随后，这些块的输出通过加法集成到 14 个跳跃连接中，这些连接链接到最后 14 个解码器块。网络设计如图4（b）所示。
值得注意的是，由于原始 Transformer 设计中缺乏明确的“编码器”和“解码器”阶段以及跳过连接，这种称为 ControlNet-UNet 的改编遇到了挑战。这种适应背离了 Transformer 的传统架构，这会妨碍有效性并导致次优结果。
ControlNet-Transformer。为了应对这些挑战，我们提出了一种新颖且专门定制的 Transformer 设计，如图 4 (c) 所示。这种创新方法旨在将 ControlNet 结构与 Transformer 架构的固有特征无缝集成。为了实现这种集成，我们有选择地将 ControlNet 结构应用于初始 N 个基本块。在这种情况下，我们生成前 N 个基本块的 N 个可训练副本。第 i 个可训练块的输出复杂地连接到零线性层，然后将所得输出添加到相应的第 i 个冻结块的输出。随后，该组合输出用作后续第 (i + 1) 个冻结块的输入。该设计遵循了 PixArt 的原始数据流，我们的观察强调了 ControlNet-Transformer 在可控性和性能方面的显着增强。这种方法代表了在此类应用中充分发挥基于 Transformer 的模型潜力的关键一步。 我们使用N=13作为最终模型。
图4：PIXART-δ 与 ControlNet 集成。 (b)：ControlNet-UNet。基本块分为“编码器”和“解码器”阶段。控制网络结构应用于 PIXART-δ 的每个编码器级，输出通过跳跃连接连接到解码器级。 (c)：ControlNet-Transformer。 ControlNet 应用于前几个块。每个块的输出被添加到相应冻结块的输出，作为下一个冻结块的输入。
3.2 实验设置
我们使用 PIXART-δ 中的 HED 边缘图作为条件，对 512px 生成进行消融研究，重点关注网络架构变化。具体来说，我们在 ControlNet-UNet 和 ControlNet-Transformer 上进行消融。其他条件，例如canny，将是未来的工作。对于ControlNet-Transformer，我们消融了复制块的数量，包括1、4、7、13和27。我们在内部数据上提取HED，并将梯度累积步骤设置为4的建议是较大的梯度积累会带来更好的结果。优化器和学习率设置与PIXARTδ相同。所有实验均在 16 个 32GB V100 GPU 上进行。实验 ControlNet-Transformer (N = 27) 的每个 GPU 的批量大小设置为 2。对于所有其他实验，批量大小设置为 12。我们的训练集由 3M HED 和图像对组成。
3.3 消融实验
如图 5 所示，ControlNet-Transformer 总体表现优异，表现出更快的收敛速度和更高的整体性能。这种优越性可归因于 ControlNet-Transformer 的设计与 Transformer 架构的固有数据流无缝结合。相反，ControlNet-UNet 在不存在的“编码器”和“解码器”阶段之间引入了概念信息流，偏离了 Transformer 的自然数据处理模式。
