【论文笔记】| 定制化生成PuLID

PuLID: Pure and Lightning ID Customization via Contrastive Alignment

ByteDance, arXiv:2404.16022v1
Theme: Customized generation
原文链接：https://arxiv.org/pdf/2404.16022

Main Work

提出了 Pure 和 Lightning ID 定制 (PuLID)，这是一种用于文本到图像生成的一种新颖的无调整 ID 定制方法。通过将Lightning T2I分支与标准扩散分支相结合，PuLID引入了对比对齐损失和准确的ID损失，最大限度地减少对原始模型的破坏，并确保较高的ID保真度。

ID插入前后图像元素(如背景、照明、组成和样式)尽可能地保持一致。

BackGround

作为定制文本到图像 (T2I) 生成的特殊类别，identity(ID)定制允许用户调整预训练的 T2I 扩散模型以与其个性化 ID 对齐。

• Tuning-based Text-to-image ID Customization. 用用户提供的相同ID在几个图像上微调某些参数，从而将ID嵌入到生成模型中（计算成本高）
• Tuning-free Text-to-image ID Customization. 在扩展肖像数据集上预训练ID适配器，使用编码器来提取 ID 特征。然后将提取的特征以特定方式集成到基本扩散模型中（例如，嵌入到交叉注意力层中）。但这种方法面临两种挑战（其实就是如何去平衡加入的ID强度）：
  • Insertion of ID disrupts the original model’s behavior.
    1. 首先，理想的ID插入应该只改变与ID相关的方面（迁移粒度），如人脸、发型和肤色，而图像元素不直接与特定身份相关联，如背景、光照、组成和风格，应该与原始模型的行为一致。
    2. 其次，在 ID 插入之后，它仍然应该保留原始 T2I 模型遵循提示的能力（编辑能力）。在 ID 定制的背景下，这通常意味着通过提示改变 ID 属性（例如年龄、性别、表情和头发）、方向和配件（例如眼镜）的能力。
  • Lack of ID fidelity. Diffusion Model引入ID Loss的难度，当前仅在噪声较小的时间步计算 ID 损失或使用额外的推理步骤预测 $x_0$，但仍有改进的空间。

作者在标准的扩散去噪训练分支旁边引入了一个Lightning T2I分支。利用最近的快速采样方法，Lightning T2I分支可以从具有有限和可管理步数的纯噪声生成高质量的图像。通过这个额外的分支，可以同时解决上述两个挑战。在这个分支中，作者结合了对比对齐损失和 ID 损失，以最小化原始模型上 ID 信息的污染，同时确保保真度。

Method

框架的上半部分展示了传统的扩散训练过程。从同一图像中提取的面部被用作ID条件信息输入。

框架的下半部分展示了PuLID（Text-to-Image）训练分支。该分支利用最近的快速采样方法，通过几次迭代（本文中为4步）从纯噪声逐步去噪到生成高质量的图像。在这个分支中，构建了有和没有ID-Embedding嵌入的对比路径。同时引入了对齐损失，用于控制模型如何在不干扰原始模型行为的情况下插入ID条件。

总的损失函数为：

$$\mathcal{L}={\mathcal{L}}_{\mathrm{diff}}+{\mathcal{L}}_{\mathrm{align}}+{\lambda }_{\mathrm{id}}{\mathcal{L}}_{\mathrm{id}}.$$

1. ID embedding的引入

在扩散过程中，根据预定义的噪声调度对噪声ε进行采样并添加到数据样本$x_0$中。这个过程在时间步$t$产生一个噪声样本$x_t$。相反，在去噪过程中，去噪模型${\epsilon }_{\theta }$ 将$x_t$、$t$ 和可选的附加条件$C$作为输入来预测添加的噪声，优化过程可以表示为：

$${\mathcal{L}}_{\mathrm{diff}}={\mathrm{E}}_{x_0,ϵ,t}(\Vert ϵ-{ϵ}_{\theta }(x_t,t,C)\Vert ).$$

${\epsilon }_{\theta }$主要由残差网络、自我注意层和交叉注意层组成的UNET。提示作为一种条件，嵌入到遵循注意力机制的交叉注意力层中。ID条件$C_{id}$通常是从目标图像x0中裁剪出来的。在这种情况下，ID条件与提示和UNET特征完全对齐，意味着在训练过程中，ID条件并不构成对T2I扩散模型的污染，本质上还是一个重建训练任务。

$$\{\begin{array}{l}\text{Attention}(Q,K,V)=\text{Softmax}(\frac{Q{K}^T}{\sqrt{d}})V\\ K={\text{W}}_K{\tau }_{txt}(C_{txt});V={\text{W}}_V{\tau }_{txt}(C_{txt}),\end{array}$$

$\mathrm{Q}$ 从 UNET 图像特征中投影，${\tau }_{txt}$ 表示一个预训练的语言模型，将提示 $C_{txt}$ 转换为文本特征，${\mathbf{W}}_K$ 和 ${\mathbf{W}}_V$ 是可学习的线性层。

在测试过程中，当我们提供与ID条件发生冲突或未对齐的提示语，例如改变ID属性或更换风格时，这些方法往往会失败。

2. 通过对比对齐实现无污染的ID引入

文章中提出了一种直观的解决方案，即在常规的扩散去噪训练分支之外引入一个 Lightning T2I 训练分支。与测试设置类似，Lightning T2I分支从纯噪声开始，并通过完整的迭代去噪步骤直到达到$x_0$，可以观察到嵌入ID后图像的元素是否发生变化，以及它是否仍然具有提示跟随能力。并且需要的迭代次数是有限且可控的。

作者使用具有 4 个去噪步骤的 SDXL-Lightning，并准备了一个具有挑战性的提示列表，可以很容易地揭示污染

作者构建了从同一提示和初始潜在开始的对比路径。

• 一条路径仅由提示为条件，UNET 特征表示为 $Q_t$
• 另一条路径同时使用 ID 和提示作为条件，UNET 特征表示为$Q_{tid}$

通过在语义上对齐这两个路径上的UNET特征，该模型将学习如何在不影响原始模型行为的情况下嵌入ID。实际上，对齐是在所有层和时间步上进行的。

核心通过构建对齐损失来解决：对齐损失由两个组成部分构成：语义对齐损失（Lalign-sem）和布局对齐损失（Lalign-layout）。我们使用文本特征K来查询UNET特征Q。对于K中的每个token，它将计算与Q的相关性，并进一步基于相关性矩阵聚合Q。这里的注意力机制可以表示为 $Attention(K,Q,Q)$，可以解释为UNET 特征对提示的响应。

语义对齐损失背后的洞察力很简单：如果 ID 的嵌入不影响原始模型的行为，那么 UNET 特征对提示的响应应该在两条路径中是相似的。

$${\mathcal{L}}_{\text{align-sem}}={\Vert \text{Softmax}(\frac{K{Q}_{tid}^T}{\sqrt{d}})Q_{tid}-\text{Softmax}(\frac{K{Q}_t^T}{\sqrt{d}})Q_t\Vert }_2.$$

Lalign-sem的引入显著减轻了ID信息污染模型行为的问题。然而，它不能保证布局的一致性，因此增加了一个布局对齐损失Lalign-layout，其定义如下：

$${\mathcal{L}}_{\text{align-layout}}={\Vert Q_{tid}-Q_t\Vert }_2.$$

最终完整损失函数的构建为：

$${\mathcal{L}}_{\mathrm{align}}={\lambda }_{\mathrm{align}-\mathrm{sem}}{\mathcal{L}}_{\mathrm{align}-\mathrm{sem}}+{\lambda }_{\mathrm{align}-\mathrm{layout}}{\mathcal{L}}_{\mathrm{align}-\mathrm{layout}},$$

3. 更准确的ID保真损失函数

在训练期间引入 ID 损失，但仅使用单步直接预测扩散训练过程中第 t 个时间步的 $x_0$，这将产生嘈杂且有缺陷的预测 $x_0$，从而导致 ID 损失的计算不准确。因此一般仅在噪声较小的阶段应用 ID 损失。但如果只计算部分时间步长，会影响全部有效性。

本实验在引入的 Lightning T2I 分支上计算ID损失${\mathcal{L}}_{\mathrm{id}}$，定义为：

$${\mathcal{L}}_{\mathrm{id}}=CosSim\left(\varphi (C_{id}),\varphi (\mathrm{L}-T2I(x_T,C_{id},C_{txt}))\right),$$

其中$x_T$表示纯噪声，L-T2I 表示Lighting T2I 分支，$\varphi$表示人脸识别主干。为了生成照片逼真的人脸，我们将提示$C_{txt}$固定为“portrait, color, cinematic”。

Experiments

三阶段训练：

1. 使用传统的扩散损失${\mathcal{L}}_{\mathrm{diff}}$来训练模型。
2. 加入ID loss ${\mathcal{L}}_{\mathrm{id}}$ (arcface-50)，在不考虑对原始模型的污染的情况下努力实现最大 ID 保真度。
3. 添加对齐损失${\mathcal{L}}_{\mathrm{align}}$来微调模型。

定量比较。跨不同基础模型和数据集的 ID 余弦相似度与 SOTA 方法的比较。

表 2 说明了与基线相比，使用朴素 ID 损失（直接从当前时间步预测 x0）和本文引入的更准确的 ID 损失 Lid 来提高 ID 保真度。

|

参考文献

Guo Z, Wu Y, Chen Z, et al. PuLID: Pure and Lightning ID Customization via Contrastive Alignment[J]. arXiv preprint arXiv:2404.16022, 2024.