∞-Brush: Controllable Large Image Synthesis with Diffusion Models in Infinite Dimensions
目录
0. 摘要
1. 简介
2. 相关工作
4. 提出的方法
4.1 函数空间中的条件扩散模型
4.2 带有跨注意力神经算子的条件去噪器
5. 实验
6. 限制
7. 结论
0. 摘要
将复杂、特定领域的信息合成高分辨率图像在生成建模中仍然是一个重大挑战，尤其是在如数字病理学和遥感等大图像领域的应用中。现有方法存在关键局限性：像素或潜在空间中的条件扩散模型在超过其训练分辨率时会失去保真度，而且图像尺寸越大，计算需求越显著增加。基于小块的拼接方法虽然具有计算效率，但由于过度依赖局部信息，无法捕捉远距离的空间关系。在本文中，我们引入了一种新的无限维度条件扩散模型，∞-Brush，用于可控的大图像合成。我们提出了一种跨注意力神经算子，以实现函数空间中的条件化。我们的模型克服了传统有限维度扩散模型和基于小块的方法的限制，提供了更好的扩展性，并且在保持细节的同时，更好地保留了全局图像结构。据我们所知，∞-Brush 是第一个在函数空间中的条件扩散模型，能够在可控条件下合成任意分辨率、最高达 4096 × 4096 像素的图像。
1. 简介
目前的最先进方法在可控大图像生成方面仍然存在显著局限性。这些方法大致分为两类：
第一类方法直接在有限的潜在或像素空间中使用条件扩散模型，受设计限制，这些方法只能在其训练时的分辨率下生成图像。例如，SDXL [26] 和 Matryoshka Diffusion [11] 可以生成分辨率高达 1024 × 1024 像素的图像。尽管这些方法在一定程度上表现出色，但无法在更高分辨率下生成图像而不损失质量或保真度。此外，随着分辨率的增加，训练和运行这些模型所需的计算资源呈指数增长，这使得处理更大尺寸的图像过程变得越来越低效。
第二类策略由 MultiDiffusion [1] 引入，并由 Graikos 等人 [10] 进行了改进，涉及一种基于小块的生成方法，将大图像生成分割成较小的部分。该技术通过对大图像的局部区域进行训练，并使用扩展算法进行大图像合成。虽然这种方法在计算上更为高效，并且能够生成足够逼真的大图像，但它在捕捉远距离空间依赖性方面表现不足（如补充材料中讨论）。这一限制源于对局部信息的高度依赖，因为每个小块的生成主要受局部条件的影响，而不会受到远处小块信息的影响。
上述方法都在有限的图像或潜在空间中操作，无法在生成过程中显著超出训练图像的尺寸。因此，必须直接在整个大图像上训练模型，导致难以克服的计算成本。最近，Bond-Taylor 等人 [2] 通过将图像表示为希尔伯特空间 H 中的函数，展示了可以在固定尺寸输入上训练的同时合成任意大的图像。然而，他们的模型无法进行条件化，这对于在下游应用中高效利用模型（如数据增强）是必要的。
2. 相关工作
无限维度的扩散模型。Kerrigan 等人 [19] 首次提出了将扩散模型应用于功能性数据的概念，开创了生成模型可以超越有限维度空间限制的想法。基于无限维度扩散的思想，Lim 等人 [22] 和 ∞-Diff [2] 专门讨论了在函数空间中表示的图像生成。然而，这些方法无法进行条件化以实现可控的图像生成。据我们所知，我们的 ∞-Brush 结合了一种新的跨注意力神经算子，是第一个用于可控大图像合成的无限维度条件扩散模型。
4. 提出的方法
我们提出了一种在函数空间 H 中的新的条件扩散模型。在前面的背景介绍的基础上，我们现在构建我们条件扩散模型在无限维度中的前向和逆向过程以及训练目标。此外，我们提出了一种新颖的架构，利用带有跨注意力神经算子的条件去噪器来参数化去噪过程。
4.1 函数空间中的条件扩散模型
在图像生成的背景下，我们通过对每个图像采样 N 个坐标，将函数 u_j 在网格
上离散化，这会导致非平滑的输入空间。为了实现更平滑的函数表示，应用平滑算子 A: H→H（例如截断的高斯核）来近似函数空间 H 中的粗糙输入。
前向过程。我们在无限维度中的条件扩散模型的前向过程等价于函数空间中的无条件扩散模型，它逐渐将概率测度 Q_0 = Q_data 加扰到高斯测度 N(m,C)，并允许在任意时间步 t 进行采样：
逆向过程。我们使用变分方法来近似后验测度，用 H 空间上的变分测度族并结合条件嵌入 e 来控制生成过程。我们用条件高斯测度来建模基础的后验测度 Q(u_{t−1}∣ut)：
命题 1（学习目标）。函数空间中的条件扩散模型的交叉熵具有变分上界：
证明。请参阅附录 A 中的完整证明。
要计算概率测度之间的 KL 散度 KL(Q || P)，我们需要利用 KL 散度的测度理论定义，如下列引理中所述【6】。（见原论文）
4.2 带有跨注意力神经算子的条件去噪器
我们的 ∞-Brush 利用了一种分层去噪器架构，包括一个用于高效捕捉细节的稀疏层次和一个用于全局信息的网格层次（图 2）。我们通过随机选择坐标子集
来离散化噪声函数 u∈H 和去噪函数 s∈H。在稀疏层次上，我们依次应用稀疏神经算子、我们的跨注意力神经算子和自注意力机制对函数进行逐点评估。
普通注意力机制的计算复杂度相对于序列长度（即函数样本数量，在此情况下为 N）为二次 O(N^2·d)，相对于其维度 d 为线性。对于在无限维度中学习算子，N 可以达到数百万个点（例如在生成 4096 × 4096 的图像时，N≈1600万）。我们通过提出一种计算复杂度相对于 N 为线性的跨注意力神经算子来解决这个问题。
具体来说，在跨注意力神经算子中，假设我们有 L 个条件嵌入
在我们的 ∞-Brush 中，L=3，分别表示扩散时间步嵌入 t、条件嵌入 e 和坐标嵌入 c。首先，我们计算查询 Q = (q_i)、键 K_l = (k^l_i) = Y_l·W_k 和值 V_l = (v^l_i) = Y_l·W_v，然后将所有 qi 和 ki 归一化为 ~q_i = softmax(qi) 和 ~k_i = softmax(ki)。最后，跨注意力公式为：
其中 α 是归一化系数。与普通注意力的主要区别在于我们首先将逐点向量 ~k^l_i 和 v^l_i 相乘，然后再与 ~q_t 进行点积。因此，公式 (15) 的复杂度为
相对于点数 N 是线性的。
（2020|ICML PMLR，线性 Transformer，核函数，RNN）Transformer 是 RNN
稀疏层次的输出通过 k 近邻法线性插值到规则间隔的网格，这是网格层次模型的输入。网格数据点传递给基于网格的、有限维度的 UNO 架构【2, 22】，该架构用于聚合全局信息。UNO 架构基于广泛使用的 UNet 模型，该模型已被广泛研究用于条件有限维度扩散模型【30】。按照这方面的文献，我们在 UNet 去噪器的瓶颈处使用普通的跨注意力来整合网格层次的条件信息。在实验中，我们展示了由于网格层次的粗略插值不能完全表示函数，因此需要在有限维度（网格）和无限维度（稀疏）层次都应用条件化，以获得高质量的结果。
5. 实验
6. 限制
尽管 ∞-Brush 生成的图像在全局结构一致性和细节保留方面表现更好，但在局部细节上并不优于其他方法。我们强调几个关键原因，这些原因可能阻碍了我们模型的性能。首先，我们的模型参数量最少，只有 [10] 模型的一半。我们预计随着更多研究集中于无限维度扩散模型，模型规模会扩大，性能也会随之提升，正如在常规有限扩散模型中观察到的那样。此外，SDXL 和 [10] 都利用了预训练模型进行初始化，而我们的模型是从零开始训练的，因为目前没有可用的无限维度预训练模型，这导致在较小数据集上的表现较差。
7. 结论
总之，∞-Brush 在条件大图像生成领域，尤其是对高分辨率和领域特定条件生成需求的应用中，迈出了必要的一步。本文展示了我们的方法有效解决了先前扩散模型固有的可扩展性限制，同时保留了对生成输出的高度控制。通过提出一种在函数空间中的新型条件扩散模型，并辅以跨注意力神经算子，我们不仅在图像的全局结构上实现了最先进的保真度，还在高分辨率图像中保持了可接受的细节，而不需要通常与此类任务相关的过度计算成本。在未来的工作中，我们计划设计本地神经算子以捕捉精细细节，并从有限维度扩散模型中转移知识，以实现更强大的初始化。
论文地址：https://arxiv.org/abs/2407.14709
项目页面：https://github.com/cvlab-stonybrook/infinity-brush
进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群
加 VX 群请备注学校 / 单位 + 研究方向
