Weights2wights: Interpreting the Weight Space of Customized Diffusion Models

导语

可控生成是图像生成领域的一个重要方向。从最基础的文本条件生成，到 ControlNet、IP-Adapter 等图像条件生成，再到各种概念定制化生成，扩散模型的可控生成技术越来越丰富。

概念定制化生成（Concept Customized Generation）是指给定一个特定的个体（比如我家的狗，而不是别的狗），然后要求模型结合其他条件（如文本 prompt）生成出该特定个体的图片。最常见的就是固定某个人物，生成出该人物的各种图片。概念定制化生成的方法也已经有很多，有经典的 Textual Inversion、DreamBooth、LoRA，也有最新的 ConsiStory、OMG 等方法，本专栏已有多篇详细的介绍。最常见的思路是：给定目标概念的 3-5 张图片作为训练集，在文生图模型上进行微调，并绑定到特定的触发词，在生成时使用触发词引导模型生成目标概念。

这种方法有不错的定制化生成效果，但是用起来比较麻烦，对于每个目标概念，首先收集需要多张目标概念，还要进行模型训练。更理想的方案是，让模型学习到概念定制化生成的能力，在推理生成时，使用单张目标概念图输入模型作为参考，模型能够直接生成目标概念。这相当于是实现了 zeroshot 的概念定制化生成。

本文介绍 UC Berkeley、Snap 和 Stanford University 联合提出的 Weights2weights 方法。使用一系列训练好的 DreamBooth LoRA 权重作为训练集，建模出了权重空间的低维流形子空间，称为 weights2weights 空间（w2w space）。基于这个 w2w 空间，作者提出了三种应用方式：Sampling、Inversion 和 Editing。以下将详细介绍。

方法

基础介绍

本文提出的 Weights2weights 方法是基于三个紧密相关的基础方法：LDM、DreamBooth、LoRA。我们首先简要介绍这三个基础方法。

LDM

LDM（Latent Diffusion Model）通过引入 VAE 将扩散模型做在了隐空间，从而降低训练和推理成本。并通过交叉注意力将文本 prompt 条件注入到去噪扩散模型中。LDM 的训练目标可写为如下条件去噪的形式：
$${\mathbb{E}}_{\mathbf{x},\mathbf{c},ϵ,t}[w_t||ϵ-{ϵ}_{\theta }({\mathbf{x}}_t,\mathbf{c},t)|{|}_2^2]$$
其中 ${ϵ}_{\theta }$ 就是去噪 UNet，$\mathbf{c}$ 是文本条件 embedding，$w_t$ 是关于时间步 $t$ 的加权函数。

DreamBooth

DreamBooth 是概念定制化生成的经典之作，其创新点有二：一是使用词表中无意义的低频词作为触发词（如 sks、mnsd），二是提出了 prior preservation loss，引入原模型生成的目标概念同类别其他图像到训练数据中，作为一个正则项，避免模型过拟合到目标概念上。DreamBooth 在本专栏已有介绍。其目标函数可写为以下两项：
$${\mathbb{E}}_{\mathbf{x},\mathbf{c},ϵ,t}[w_t||ϵ-{ϵ}_{\theta }({\mathbf{x}}_t,\mathbf{c},t)|{|}_2^2+\lambda w_t^{\prime }||{ϵ}^{\prime }-{ϵ}_{\theta }({\mathbf{x}}_t^{\prime },{\mathbf{c}}^{\prime },t)|{|}_2^2]$$
其中第一项就是标准的 LDM 损失，第二项是使用原模型生成的目标概念同类别其他图像 ${\mathbf{x}}^{\prime }$ 和类别词 ${\mathbf{c}}^{\prime }$ 的 ppl 损失。

LoRA

DreamBooth 原文的训练方式是对整个模型进行全量微调，但是这样训练成本太高，并且参数空间太大。社区一般训练一个 LoRA 进行高效微调。LoRA（Low Rank Adaption）是一种更新低秩参数的高效微调方法。具体来说，原模型有参数矩阵 $W\in {\mathbb{R}}^{m\times n}$，LoRA 更新的是 $\Delta W=BA$，其中 $B\in {\mathbb{R}}^{m\times r},A\in {\mathbb{R}}^{r\times n},r<<\text{min}(m,n)$ 是分解的低秩矩阵。训练完成后，低质 LoRA 权重可以直接融合到模型中 $W^{\prime }=W+\alpha W$，其中 $\alpha \in \mathbb{R}$ 是加权系数。

构建 w2w 权重流形

我们的目标是建模 w2w 空间，从而让模型具有 zeroshot 概念定制化生成的能力。

首先要构建数据集。不同于其他模型，w2w 的数据集不是图片和文本，而是模型权重。具体来说，作者使用 DreamBooth 方法对 LDM 进行了概念定制化微调，训练了 $N$ 个概念定制化模型。为了降低参数空间的维度，这里使用了 LoRA 方法进行微调训练。训练完成后，将 LoRA 权重矩阵拉直，就得到了 $N$ 组 $d$ 维的参数向量 ${\theta }_i\in {\mathbb{R}}^d,i=1,\dots ,N$。这样我们就得到了模型权重数据集 D = { θ 1 , θ 2 , … , θ N } \mathcal{D}=\{\theta_1,\theta_2,\dots,\theta_N\} D={θ1​,θ2​,…,θN​}。

然后，我们要建模权重空间的流行子空间。假设我们的权重参数 ${\theta }_i\in {\mathbb{R}}^d$ 分布在一个低维流形上，并且该流形是 $d$ 维权重空间一个子集。之前已经有工作发现，高维的概念可以被编码到表征的线性子空间，我们将这个流形建模为一个 $m$ ($m<d$) 维线性子空间，称之为 w2w 空间。我们使用基向量 w = { w 1 , … , w m } ,   w i ∈ R d \mathbf{w}=\{w_1,\dots,w_m\},\ w_i\in\mathbb{R}^d w={w1​,…,wm​}, wi​∈Rd 的线性组合来表示这个子空间中的点。实际中，作者对 $N$ 个权重采用了 PCA 进行降维，保留前 $m$ 个主要分量，来得到这个 $m$ 个基向量。

Sampling、Editing、Inversion

基于构建好的 w2w 空间，作者尝试了三种应用方式，分别是 Sampling、Editing 和 Inversion。

Sampling

首先，我们可以对这个空间进行采样。空间由训练好的概念定制化生成 LoRA 权重组合而成，其采样出的新的权重向量可以定制化生成一个新的概念。具体来说，我们可以采样出一组系数 { β 1 , … , β m } \{\beta_1,\dots,\beta_m\} {β1​,…,βm​} 来基于基向量组表示一个概念定制化 LoRA 权重，其中 ${\beta }_k$ 是从均值为 ${\mu }_k$，标准差为 ${\sigma }_k$ 的正态分布中采样出来的。这里的 ${\mu }_k,{\sigma }_k$ 是在所有训练模型的第 $k$ 个 PCA 主要分量上计算出来的。

笔者自己在官方 demo 上 sample 权重出的图中的人脸 ID 不是很一致呢

Editing

与 GAN Latent Inversion 类似的，我们可以找到一个方向 $\mathbf{n}\in {\mathbb{R}}^d$，来定义一个将嵌入在模型权重中的某种二元属性（如男/女，老/幼）分隔开来的超平面。这里假设模型编码的概念属性二元标签是已有的。我们就可以使用模型权重作为数据，基于这些标签，训练线性分类器，强加权重空间中的分离超平面。给定一个由权重 $\theta$ 参数化的概念，我们可以通过沿着方向 $\mathbf{n}$ 行进来操作单一属性，该方向与分离超平面正交：${\theta }_{\text{edit}}=\theta +\alpha \mathbf{n}$。

Inversion

传统的生成模型 Inversion 是指给定一张图片，找到一个对应的 latent 向量，将其输入到模型中可以重构出原始图片。这相当于在学习到的数据流形上找到输入的映射。有了 w2w 空间，我们建模出了一个数据流形，它是模型权重而不是图像。借鉴隐层优化方法的做法，作者这里提出了一种基于梯度的 inversion 方法，将单个概念从图像中 invert 到我们构建的空间。给定一张图片 $\mathbf{x}$，有受限去噪目标：
$${\text{max}}_{\theta }{\mathbb{E}}_{\mathbf{x},\mathbf{c},ϵ,t}[w_t||ϵ-{ϵ}_{\theta }({\mathbf{x}}_t,\mathbf{c},t)|{|}_2^2],\text{s.t.}\theta \in w2w$$
我们通过优化一组基向量系数 { β 1 , β 2 , … , β m } \{\beta_1,\beta_2,\dots,\beta_m\} {β1​,β2​,…,βm​}，而不是权重本身，来将模型权重限制在 w2w 空间内。这里就需要再用 DreamBooth 中的 ppl 损失来作为正则了，因为 w2w 空间本就是有 DreamBooth LoRA 的权重构建出来的，已经继承了泛化能力。

总结

不同于 DreamBooth、Textual Inversion 对每个新的概念都要进行训练，weights2weights 通过构建 w2w 空间，具有了概念定制化生成的元能力，相当于可以进行 zeroshot 的概念定制化生成。