提出了一种统一的扩散框架, 可以在一个模型中适应与多模态数据集相关的所有分布. 它能够执行图像生成, 文本生成, 文本到图像生成, 图像到文本生成以及图像-文本对生成, 而无需额外开销, 另外还有一些洞察.

UniDiffuser

介绍

Diffusion高速进步, 文本->图片效果好, 但是仅允许执行单一任务, 希望生成各种类型的多模态内容, 使用任意作为条件. 举了个艺术家的例子, 感觉不太好就不放了.

文本到图像的生成可以被形式化为学习条件分布p(图像|文本), 一种经典的拟合所有相关分布的方法是隐式的, 首先学习联合分布, 然后通过额外的过程如马尔可夫链蒙特卡洛推断边缘和条件分布, 这在大规模多模态数据上是不可行的.

本文提出了一种基于扩散的框架, 在一个模型中明确拟合所有相关分布, 而不会引入额外的训练或推理开销. 他们的关键见解是, 学习所有分布的扩散模型可以统一为预测扰动数据中的噪声, 其中扰动水平即时间步对于不同的模态可以不同.

最后点了一下实验结果, 性能卓越.

背景

先是讲了diffusion的加噪, 去噪过程, 模型就是预测这个过程. 接着讲了条件生成的原理, 本质原理没变, 只是增加了一个条件. 最后讲了CFG, 让模型同时做无条件生成和有条件生成, 两者的结果按比例混合.

引出灵感: 既然有条件生成和无条件生成在底层逻辑上是一回事, 我们完全可以把图片, 文字, 各种条件全塞进同一个模型, 让它一次性把所有噪声都预测出来.

方法

对比着看, 我们知道在stable diffusion中, 我们将文本条件过一个CLIP升维, 然后通过交叉注意力做KV注入. 但是在它提出的UniDiffuser, 图像和文本的注意力是完全对称的, 将文本, 图像以及它们各自的时间步长, 全部统一为序列中的token, 送入一个纯Transformer中进行自注意力计算.

这个转换具体来说, 就是使用CLIP来提取文本特征, 输出是(77, 768), 增加一个额外的线性层(77, 64); 得到文本的潜在表示 $y_{t^y}$ , 还有图像的潜在表示 $x_{t^x}$ 后, 将他们分别展平, 并加上各自的时间步 $t^x$ 和 $t^y$ 的嵌入向量; 这些 Token在真正进入 U-ViT 的注意力模块之前, 都会通过各自的线性层统一映射到 1536 维, 形成一个超长的序列, 然后送入层叠的 Transformer Block 中, 在这个过程中, 图像和文本通过全局的 Self-Attention 自由地进行跨模态信息交互.

UniDiffuser

传统的生成模型通常只能处理单一任务, 而 UniDiffuser 的目标是设计一个单一模型, 能够同时捕获由图文对数据 $q(x_0, y_0)$ 决定的所有相关分布.

边缘分布(无条件生成): $q(x_0)$ (无条件图像生成)和 $q(y_0)$ (无条件文本生成).
条件分布(条件生成): $q(x_0|y_0)$ (文本生图)和 $q(y_0|x_0)$ (图生文本).
联合分布(联合生成): $q(x_0, y_0)$ (图文对同时生成).

作者指出, 使用扩散模型学习任何概率分布, 其本质都等价于在给定的扰动数据下, 去估计噪声的条件期望.

建模边缘分布 $q(x_0)$ 等价于估计 $\mathbb{E}[\epsilon^{x}|x_{t}]$ .
建模条件分布 $q(x_0|y_0)$ 等价于估计 $\mathbb{E}[\epsilon^{x}|x_{t},y_{0}]$ .
建模联合分布 $q(x_0,y_0)$ 等价于估计 $\mathbb{E}[\epsilon^{x},\epsilon^{y}|x_{t},y_{t}]$ .

作者发现, 上述所有任务的条件期望都可以完美地统一成一个通用的广义形式:

\mathbb{E}[\epsilon^{x},\epsilon^{y}|x_{t^{x}},y_{t^{y}}]

在这里, 图像的时间步长 $t^x$ 和文本的时间步长 $t^y$ 可以是各自独立且不同的, 通过巧妙地设置这两个时间步长的数值, 就能自由切换不同的生成任务.

无条件生成(边缘化): 将某一模态的时间步长设为最大值 $T$ (代表纯噪声，此时该模态的信息会被模型自动忽略). 例如, 设置 $t^y = T$ 时, 模型会忽略文本, 退化为图像无条件生成.
条件生成(显式条件): 将作为条件的那一模态的时间步长设为 $0$ (代表完全清晰的真实输入), 例如, 设置 $t^y = 0$ 时, 模型就变成了文本生图.
联合生成: 将两者的步长设为相同且同步变化的 $t$ (即 $t^x = t^y = t$ ), 用于同时生成相互对齐的图像和文本.

基于这种统一的视角, UniDiffuser 只需要训练一个联合噪声预测网络 $\epsilon_{\theta}(x_{t^{x}},y_{t^{y}},t^{x},t^{y})$ . 在训练时, 它独立, 均匀地对 $t^x$ 和 $t^y$ 进行采样, 并最小化一个极简的 $l_2$ 回归损失:

\mathbb{E}_{x_{0},y_{0},\epsilon^{x},\epsilon^{y},t^{x},t^{y}}||\epsilon_{\theta}(x_{t^{x}},y_{t^{y}},t^{x},t^{y})-[\epsilon^{x},\epsilon^{y}]||_{2}^{2}

Classifier-Free Guidance for Free

在推理时做CFG依然需要两次前向传播, 它这里的free是指它天然的适合将一个条件置空, 训练时完全不需要引入任何空标记. 就是不需要再训练侧做修改.

UniDiffuser on Images and Texts

把图文压缩进latent

图像使用sd的vae和clip分别提取重建特征和语义特征, 文本微调了GPT-2编码器, 压缩好的文本特征作为前缀, 让他逐字重新吐出来.

作者发现经过这样处理后, 图像和文本的潜在特征数值范围竟然差不多, 基本都服从正态分布, 落在 [−2,2] 的区间内. 这意味着不需要再做复杂的额外归一化, 就能直接把它们送进同一个网络里去训练.

用 Transformer 预测噪声

采用U-ViT的Transformer架构, 带有噪声的图像特征, 文本特征, 以及它们各自的时间步长, 全部被一视同仁的当成Token送进Transformer.

实验

baseline是Versatile Diffusion, 指标是FID和CLIP Score.

感受

它虽然提供了五种任务场景, 但是我觉得真正有用的场景依然只是t2I, 这篇在当时看可能没什么亮点, 但是这种架构是有很多优势的, 比如后面DreamZero搞具身的就从中得到了启发, 迁移到了video-action上面.

One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale

介绍

背景

方法

UniDiffuser

Classifier-Free Guidance for Free

UniDiffuser on Images and Texts

把图文压缩进latent

用 Transformer 预测噪声

实验

On this page

介绍

背景

方法

UniDiffuser

Classifier-Free Guidance for Free

UniDiffuser on Images and Texts

把图文压缩进latent

用 Transformer 预测噪声

实验