生成模型常见的条件融合方式

目前生成模型主要有4中常见的条件融合方式以实现可控生成：条件归一化层，Decoupled Cross-Attention，self-attention层进行融合，特征值逐元素求和。本文首先介绍下各种方法现，然后进行总结，最后提出一下展望。

1. 条件归一化层
   过去和现在的一些工作，会通过条件归一化层将一些条件（比如类别和文本）融合到生成模型以实现条件可控生成。假设生成模型某层的特征是$x\in R^{b\times hw\times c}$, $P$是输入条件，$f(P)$是条件特征提取网络，常见的范式如下：
   $$\begin{gathered} f(P)=\gamma ,\beta \\ y=\gamma \frac{x-\mu }{\sigma +ϵ}+\beta \end{gathered}$$, 其中$\mu ,\sigma$分别是$x$的均值和方差, $y$将要作为生成模型下一层的输入, $ϵ$是较小常量防止数值问题（为了方便，后文省略）。

   (1) Adaptive Instance Normalization (AdaIN)
   $$\begin{gathered} f(P)=\gamma ,\beta (\gamma \in R^{b\times 1\times c},\beta \in R^{b\times 1\times c}) \\ y=\gamma \frac{x-\mu }{\sigma }+\beta (\mu \in R^{b\times 1\times c},\sigma \in R^{b\times 1\times c}) \end{gathered}$$
   (2) 为解决条件归一化层（比如AdaIN和条件BatchNorm）缺失空间信息的问题, SPatially-Adaptive DEnormalization (SPADE, 2019)提出具有空间维度的$\mu ,\sigma$：
   $$\begin{gathered} f(P)=\gamma ,\beta (\alpha \in R^{1\times hw\times c},\beta \in R^{1\times hw\times c}) \\ y=\gamma \frac{x-\mu }{\sigma }+\beta (\mu \in R^{1\times 1\times c},\gamma \in R^{1\times 1\times c}) \end{gathered}$$
   
   （3） DiT为了实现恒等函数（the identity function）以加速大规模训练，提出adaLN-Zero block。具体地，其在adaLN（广泛用于GAN和Diffusion Unet）回归$\gamma ,\beta$的基础上,还会回归一个$\alpha$参数(这个$\alpha$会被0初始化以实现恒等函数)：
   $$\begin{gathered} f(P)=\gamma ,\beta ,\alpha (\gamma \in R^{b\times 1\times c},\beta \in R^{b\times 1\times c},\alpha \in R^{b\times 1\times c}) \\ y=\alpha \times MultiHeadSelfAttention(\gamma \frac{x-\mu }{\sigma }+\beta )(\mu \in R^{b\times hw\times 1},\sigma \in R^{b\times hw\times 1}) \end{gathered}$$
   

2. Decoupled Cross-Attention
   Decoupled Cross-Attention最初由IPAdapter(2023)提出，其在去噪网络crossattention layer旁边再引入一个cross attention分支，然后对二个cross attetnion的输出进行求和送给去噪网络下一层。具体地，假设生成模型某层的特征是$x\in R^{b\times hw\times c}$, $P$是cross attention的条件特征，$C$是新的cross attention层的条件特征，我们可以得到：
   $$y=\text{Softmax}(\frac{W_{q}x(W_{k}P{)}^{\top }}{\sqrt{d}})W_{v}P+\lambda \text{Softmax}(\frac{W_{q}x(W_k^{\prime }C{)}^{\top }}{\sqrt{d}})W_v^{\prime }C,$$
   其中$d$ 是归一化因子，$W_q$, $W_k$, $W_v$是去噪网络crossattention的投影矩阵, $W_k^{\prime }$, and $W_v^{\prime }$是新crossattention的投影矩阵，$y$将要作为生成模型下一层的输入, $\lambda$是权衡权重。
   
   对于没有cross attention层的模型比如FLUX(2024)，目前也有方法比如PuLID-FLUX-v0.9.0(2024)在self attention层旁引入cross attention分支。

3. self-attention层进行融合
   如果提取条件特征的网络是去噪网络的副本，条件特征则通常是在去噪网络的self-attetion层进行融合。
   (1) 假设生成模型层的特征是$x\in R^{b\times hw\times c}$, $P$是对应的条件特征：
   $$\begin{gathered} x^{\prime }=concat([x,P],axis=1) \\ y=SelfAttention(x^{\prime },x^{\prime },x^{\prime })[:,:hw,:], \end{gathered}$$
   其中concat是特征拼接函数，模型经过SelfAttention操作后会丢弃多余维度特征以保证y的维度和x的维度一致。该方法最初由AnimateAnyone(2023)提出。
   
   (2) 第一种方法有冗余特征，因此可以改进去噪网络的self-attention层的以实现更好的条件特征融合。具体地，假设生成模型层的特征是$x\in R^{b\times hw\times c}$, $P$是对应的条件特征：
   $$\begin{gathered} x^{\prime }=concat([x,P],axis=1) \\ y=SelfAttention(x,x^{\prime },x^{\prime })=\text{Softmax}(\frac{W_{q}x(W_k{x}^{\prime }{)}^{\top }}{\sqrt{d}})W_v{x}^{\prime }, \end{gathered}$$
   其中concat是特征拼接函数，$SelfAttention$的key和value将由$x^{\prime }$得到而不是$x$。
   (3) 为了不影响去噪模型的原本的能力，可以引入新的crossattention分支进行特征融合。具体地，假设生成模型层的特征是$x\in R^{b\times hw\times c}$, $P$是对应的条件特征：
   $$y=\text{Softmax}(\frac{W_{q}x(W_{k}x{)}^{\top }}{\sqrt{d}})W_{v}x+\lambda \text{Softmax}(\frac{W_{q}x(W_k^{\prime }P{)}^{\top }}{\sqrt{d}})W_v^{\prime }P,$$
   其中$d$ 是归一化因子，$W_q$, $W_k$, $W_v$是去噪网络self-attention投影矩阵, $W_k^{\prime }$, and $W_v^{\prime }$是新crossattention的投影矩阵，$y$将要作为生成模型下一层的输入, $\lambda$是权衡权重。

4. 特征直接逐元素求和
   (1) 同样假设生成模型某层的特征是$x\in R^{b\times hw\times c}$, $P\in R^{b\times hw\times c}$是提取的对应的条件特征。ControlNet(2023)等方法直接将$x$和$P$进行逐元素求和：
   $$y=x+\lambda P,$$其中 $y$将要作为生成模型下一层的输入, $\lambda$是权衡权重。
   

   (2) ControlNeXt(2024)为了解决训练开始阶段$x$和$P$特征分布不配备造成的收敛缓慢的问题，提出cross normalization，即使用$x$和$P$的均值和方差去平衡去噪网络和条件特征提取网络对最终输出的影响。具体地，假设生成模型某层的特征是$x\in R^{b\times hw\times c}$, $P\in R^{b\times hw\times c}$）是提取的对应的条件特征：
   $$y=x+\lambda (\frac{P-{\mu }_P}{{\sigma }_P}{\sigma }_x+{\mu }_x),$$其中${\mu }_x\in R^{b\times 1\times 1}$和${\sigma }_x\in R^{b\times 1\times 1}$是$x$均值和标准差，${\mu }_P\in R^{b\times 1\times 1}$和${\sigma }_P\in R^{b\times 1\times 1}$是$P$的均值和标准差， $y$将要作为生成模型下一层的输入, $\lambda$是权衡权重。

总结：
上文提到了条件归一化层，Decoupled Cross-Attention，self-attention层进行融合，特征值逐元素求和四种特征融合方式。显然地：
(1) 条件归一化层目前主要是用于一些简单的条件，比如文本和timestep。
(2) 特征值逐元素求和目前主要是用于保留条件的空间信息。
(3) Decoupled Cross-Attention和self-attention层进行融合目前主要是用于保留条件的语义和细节。在Cross-Attention层进行融合的话需要考虑到文本可控性和条件信息注入程度之间的权衡问题，即2个分支哪个分支占主导地位。

值得注意的是，尽管本文主要关注这四种特征融合方式，但是依然有其它的方式比如使用条件特征替换预训练模型的文本特征。

展望：
目前的条件融合方式主要是基于Diffusion Unet结构探索得到。对于FLUX这种将文本和图像耦合在一起的新兴模型，如何高效地将除文本外的额外信息融合到模型值得探索。