DreamBooth Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

TL; DR：使用 3-5 张目标概念的图片微调整个 t2i 模型，实现定制化生成。使用 t2i 模型自己生成一些目标类别的图像样本，来丰富类别样本的多样性，避免语言漂移和多样性降低的问题。

---

导语

DreamBooth 不仅做到了 subject-driven generation，即根据给定的几张图片，定制化地生成目标概念，而且还能够在生成目标概念的新图时，使得生成背景、风格等丰富多样。

先来看一下结果，可以看到，DreamBooth 将目标 subject 小狗生成了出来，保真度很高。而且，背景也很丰富多样，根据不同的 prompt，将小狗生成到了不同的场景中。

方法

本文主要研究的问题有两个，一是如何选择合适的 identifier 来作为特定 subject 的触发词，二是如何避免微调过程过拟合，导致语言漂移（language drift）和多样性降低（reduced diversity）的问题。

rare-token identifier

不同于 textual inversion 中直接新增一个 token，作者设计了另一种触发词的方式：a [identifier] [noun] 。其中 noun 是目标 subject 具体属于的类别，比如上图的话，就可以选择 dog。而 identifier，是一个独一无二的触发词，代表了这个特定的 subject。这里的 identifier 应该是没有任何语义的，它只是一个标识符而已，所以为了避免选择的 identifier 被 tokenizer 拆成存在很强先验语义的 token，作者的做法是从词表中选择一些最少出现的词，来组成 identifier。

class-specific prior preservation loss

如果只使用 3-5 张目标 subject 的图片来微调整个 t2i 模型，难以避免的会出现过拟合的现象。具体可以描述为两个问题：语言漂移和多样性降低，就是说在这只特定的小狗的图片上微调之后，即使 prompt 里没有指定这只特定的小狗，只是指定了 dog，想生成任意的狗，模型也大概率会生成这只特定的狗；而且，模型生成结果的多样性会降低，生成结果看起来都差不多。

为了解决这个问题，作者想到扩充一些任意样貌 subject （比如任意狗）的图像数据。如何扩充数据呢？我们现成不就有一个生成任意 subject 的 t2i 通用模型吗，直接用它来生成就好了。作者将这称为 class-specific prior prservation。具体的，使用原本的通用 t2i 模型生成一组类别中任意 subject 的图片，即为类别先验 prior，然后在微调特定目标 subject。整体目标函数：
$${\mathbb{E}}_{\mathbf{x},\mathbf{c},ϵ,{ϵ}^{\prime },t}[w_t||{\hat{\mathbf{x}}}_{\theta }({\alpha }_t\mathbf{x}+{\sigma }_tϵ,\mathbf{c})-\mathbf{x}|{|}_2^2]+\lambda w_{t^{\prime }}||{\hat{\mathbf{x}}}_{\theta }({\alpha }_t{\mathbf{x}}_{\text{pr}}+{\sigma }_{t^{\prime }}{ϵ}^{\prime },{\mathbf{c}}_{\text{pr}})-{\mathbf{x}}_{\text{pr}}|{|}_2^2$$
其中前一项是在 3-5 张特定目标 subject 上进行微调的目标函数，而后一项就是 class-specific prior preservation loss，在生成的类别先验图片上的目标函数。

注意，DreamBooth 原文虽然是全参数微调，但是也可以结合 LoRA，指定某些层进行部分微调。

总结

DreamBooth 是 subject-driven 可控生成方向的一篇重要工作，后面各种 LoRA 的训练基本都使用 DreamBooth 的训练方法。