ConsiStory Training-Free Consistent Text-to-Image Generation

ConsiStory: Training-Free Consistent Text-to-Image Generation

TL; DR：ConsiStory 无需训练，是在推理时通过魔改 Attention 等机制，保持物体的一致性，同时构图和背景多样性丰富且不会丢失 prompt 元素。

---

导语

现有的图像 concept 定制化生成方法一般需要微调训练（如 Textual Inversion、LoRA）来教会模型新的词。更严重的是，这些方法很难对齐 prompt 文本和绘制多主题。本文提出 ConsiStory，无需训练，通过共享预训练模型的内部激活值，就能做到 subject 一致的生成。作者引入了主题驱动的共享注意力层（subject-driven attention ）和基于对应关系的特征注入（correspondency-based feature injection）来促进图像间的主题一致性。并提出一种策略，在保持主题一致性的同时，丰富图像布局的多样性。并且，ConsiStory 可以自然地扩展到 multi-subject 的场景，并能实现常见物体的定制化生成。

先看一下效果，主要看两个点：一是看 subject（人物、耳机、猫）的一致性得到了保持，而是看不同的图对 prompt 的跟随能力，和构图、背景的多样性，都是不错的。

这篇实际中要用，写得细一点。

方法

目标是这样的：给定一批 prompt 来进行文生图，这批 prompt 中都包含了同一个 subject，在生成的一批结果中，要保证两点，一是这批图片中的指定 subject 看起来能保持一致，二是这批图片之间的布局、背景等足够多样。本文方法 ConsiStory 无需训练，是通过在 UNet 的 Attention 层共享模型的激活值，从而保持 subject 一致；通过区域掩码避免背景和其他物体受到影响，保持其多样性。

ConsiStory 方法包含三个主要部分。

• 首先是引入了一个 subject-driven self-attention（SDSA）机制，用于在一个 batch 内，相关的模型层激活值之间，共享特定 subject 之间的信息。

• 第二，观察到单独的 SDSA 会导致一个 batch 内生成结果的布局、背景等也都过于一致，缺乏多样性。为此提出 attention-dropout 和 blending query 的机制，来缓解这种模式崩塌。

• 第三，引入了一个 feature injection 机制，通过 DIFT 特征得到图像间的密集相关性图，将每张生成图片的特征与彼此映射，从而进一步对结果进行精修。

subject-driven self-attention

要保持一批生成图像之间的一致性，Tune-A-Video 等工作中的共享 KV 是一个常见的思路：扩展 unet 中的 self-attention 层，使得每张图片中的 query 同时 attend batch 内所有其他图片的 key 和 value，从而每张图片能够 attend 到彼此，共享视觉特征。在视频生成中，共享 KV 能提高生成出的视频的帧间一致性。然而，我们生成一批图片并保持物体一致性，与视频生成的要求有所不同。一是视频生成的所有帧都是完全共享同一个 prompt 的，而我们生成的一批图片的每个 prompt 都是不同的；二是视频生成要求不同帧间的布局和背景不同，我们则期望在不同 prompt 下，每张图片的布局和背景尽量多样。

很直觉的解决方法就是在 unet 生图时加上 mask，只在我们指定 subject 的位置共享 KV，因为我们只想要目标 subject 在一批图像中保持一致。具体地，本文参照之前的 MasaCtrl 等工作，unet 隐层去噪时，通过在所有 timesteps 间平均文本 prompt 条件的交叉注意力权重值，识别出可能会包含目标 subject 的 patch，形成特定 subject 的 mask，只在这些 mask 上共享 KV。用符号表示为：
$$\begin{gathered} K^{+}=[K_1\oplus K_2\oplus \cdots \oplus K_N]\in {\mathbb{R}}^{N\cdot P\times d_k} \\ {V}^{+}=[V_1\oplus V_2\oplus \cdots \oplus V_N]\in {\mathbb{R}}^{N\cdot P\times d_v} \\ {M}_i^{+}=[M_1\dots M_{i-1}\oplus 1\oplus M_{i+1}\dots M_N] \\ {A}_i^{+}=\text{softmax}(Q_i{K}^{{+}^{\top }}/\sqrt{d_k}+\log M_i^{+})\in {\mathbb{R}}^{P\times N\cdot P} \\ {h}_i=A_i^{+}\cdot V^{+}\in {\mathbb{R}}^{P\times d_v} \end{gathered}$$

其中 $i$ 是一批 $N$ 张图像的索引，$P$ 是 token 序列的长度。对于第 $i$ 张图片的 mask，与自己的自注意力肯定是全都保留（即式中 $1$），与其他图片的 mask 按照上述方式卡阈值计算。通过常见的 log 方式（$+\log M_i^{+}$）把被 mask 掉的地方的值设为负无穷，实现 null-out mask，从而设置注意力 mask。

enriching layout diversity

SDSA 已经能够恢复生图结果的 prompt 对齐度，但实际中结果的构图布局、位置姿势还是会比较相似。为了削弱 SDSA 共享 KV 的影响，提高生成结果的多样性，本文使用了两个策略：query blending 和 self-attention dropout。

query-blending 是指将原始的 query 和经过 SDSA 的 query 进行混合。记时间步 $t$ 的噪声为 $z_t$，模型会进行两次 forward，首先是对 $z_t$ 进行原始的去噪，把 unet 中自注意力层的 queries 都记录下来，记为 $Q_t^{vanilla}$ ，然后第二次 forward 使用 SDSA 再对 $z_t$ 进行去噪，在这次中使用两次 query 的插值，作为最终的 query：
$$Q_t^{\ast }=(1-v_t)Q_t^{SDSA}+v_t{Q}_t^{vanilla}$$
self-attention dropout 是指在每个 timestep，对 SDSA 中的 mask $M_i$ ，有一定概率会被置为 0，从而减弱不同图像之间的 KV 共享，使得构图布局能够更加多样丰富。而且，通过调整 dropout 的概率，我们还可以在 subject consistency 和 layout diversity 之间进行 trade-off。

feature injection

SDSA 已经能够提高 subject consistency，但可能在视觉细节方面还不够好。为了提高不同图像的对应位置处（如左眼）的细粒度的相似性，本文还引入了 feature injection 机制。

首先使用 DIFT 特征，在每个两两图像对 $I_t,I_s$ 之间，构建图像块的对应关系图，记为 $C_{t\to s}$ 。就是对于 $I_t$ 中的第 $p$ 个 patch，可以通过 $C_{t\to s}[p]$ 来映射到 $I_s$ 中与之最相关的 patch，如下图所示。实际中，对每张图像的第 $p$ 个 patch，我们会选取 batch 内与之最相似（DIFT 的余弦相似度最高）的另一张图的对应 patch
$$\text{src}(p)=\arg \underset{s\ne t}{\max }\text{similarity}(D_t[p],D_s[C_{t\to s}[p]])$$
然后将该 patch 的特征与找到的最相似的 patch 特征混合：
$${\hat{x}}_t^{out}=(1-\alpha )\cdot x_t^{out}+\alpha \cdot \text{src}(x_t^{out})$$
实际中，也是只有 mask 内的、且相似度要高于一定阈值的 patch 会去进行 feature injection。

anchor images and reusable subjects

ConsiStory 还有一个做法，就是可以选择一些锚点图像 anchor images。在有 anchor images 时，SDSA 就不是在所有的图像之间共享 KV 了，而是只与 anchor images 共享 KV；feature injection 也是只考虑将 anchor images 的特征作为 $\text{src}$ 。注意 anchor images 之间是可见的，但是不会受到非 anchor images 的影响。

anchor images 有几点好处：一是可以减少显存占用，因为共享 KV 的尺寸变小了；二是实验显示可以提高生图质量；最重要的是我们就可以在不同的 batch 内复用同一个 subject，生成更多同 subject 的图片。

multi-subject consistent generation

单图多物体的一致性生成一直是业界的难题。ConsiStory 也可以用于多 subject 的一致性生成，直接取多个 subject 的 mask 并集就可以了。最强的是，其他多物体一致性生成中，不同 subject 之间的信息泄露互相干扰的问题在 ConsiStory 中问题不大，因为注意力 softmax 的指数形式，它可以看作是抑制不相关 subject 之间的信息泄漏的 gate。 类似地，在 feature injection 中对对应图进行阈值处理会产生门控效应，防止信息泄漏。

总结

ConsiStory 无需训练，在生图 batch 内共享 KV 实现 subject 一致的生成，为避免背景受到影响，仅在 attention mask 内共享。使用混合 query 特征和 attention dropout 削弱 KV 共享的影响，从而丰富布局和背景的多样性。通过 DIFT 特征找到对应的 patch 位置，使用 feature injection 进行细粒度的优化。

无需训练，可以指定想要保持一致的物体或特征，可以单图参考也可以无参考，支持单图多物体一致性生成，是目前看起来多物体生图一致性问题比较理想的解决方案。