∞-Diff: Infinite Resolution Diffusion with Subsampled Mollified States

目录

0. 摘要

1. 简介

2. 生成式神经场

3. 无限维扩散模型

3.1 平滑化（Mollificaition）

3.2 无限维度平滑扩散

4. 参数化扩散过程

4.1 神经算子

4.2 多尺度架构

4.3 高效稀疏算子

5. 实验

6. 讨论

7. 结论

0. 摘要

这篇论文介绍了 ∞-Diff，这是一种定义在无限维希尔伯特（Hilbert）空间中的生成扩散模型，能够对无限分辨率的数据进行建模。通过在随机采样的坐标子集上进行训练，并仅在这些位置对内容进行去噪，我们学习了一个用于任意分辨率采样的连续函数。与之前基于神经场的无限维度模型使用需要潜在压缩（latent compression）的逐点函数不同，我们的方法采用了非局部积分算子（non-local integral operator）在希尔伯特空间之间进行映射，从而实现空间上下文的聚合。这是通过一种高效的多尺度函数空间架构实现的，该架构直接作用于原始稀疏坐标，并配合平滑扩散（mollified diffusion）过程以消除不规则性。通过在高分辨率数据集上的实验，我们发现即使在 8 倍下采样率下，我们的模型仍能保持高质量的扩散效果。这带来了显著的运行时间和内存节省，生成的样本具有更低的 FID 分数，并且在超越训练分辨率的同时保留细节。

1. 简介

深度神经网络通常假设数据可以用固定的均匀网格表示，但底层信号往往是连续的。因此，这些方法在分辨率增加时扩展性较差。神经场（Neural fields）（Xie等人, 2022; Sitzmann等人, 2020; Mildenhall等人, 2021）解决了这个问题，通过直接将数据表示为从坐标到强度（如像素值）的映射，使得参数化和内存/运行时间独立于分辨率，从而可以在通常无法适应内存的数据上进行训练。基于神经场的生成模型（Dupont等人, 2022a;b; Bond-Taylor和Willcocks, 2021; Du等人, 2021）被开发出来以利用这些特性。这些方法在坐标之间本质上是独立的，因此通过压缩的潜在向量来为网络提供全局信息。然而，这些方法生成的样本质量显著低于有限维度的生成模型。

2. 生成式神经场

神经场（Xie等人, 2022; Mildenhall等人, 2021; Sitzmann等人, 2020）是一种用于连续数据表示的方法，它将坐标 c 映射到值 v（如 RGB 强度），即 f_θ(c) = v。这使得表示数据所需的内存与其分辨率解耦。

【刚开始的直觉：分辨率越大，像素点越多，坐标应该也越多。所以内存与分辨率解耦是反直觉的。

实际上，在使用神经场时，通常将坐标归一化到 [-1,1]，然后通过一个定义在该区域的连续函数表示值 v：f_θ(c) = v。

】

映射函数 fθ 通常是一个 MLP 网络，通过最小化在每个坐标上对应于真实值的重建损失来进行优化。fθ 的局部性质允许通过在坐标子集上评估 fθ 来蒙特卡洛（Monte-Carlo）的近似损失，从而可以在比内存能容纳的更高分辨率数据上进行训练。由于 fθ 在每个坐标上是独立的，无法像卷积 / transformer 那样在多个点上进行转换，为了表示函数空间的方法通常还会以描述单个数据点的压缩潜在向量 z 为条件，即 fθ(c, z)。

• Dupont 等人（2022a）首先使用元学习（meta-learning）将数据集压缩到潜在条件神经场（latent conditional neural fields）中，然后使用 DDPM（Ho等人, 2020）或标准化流（Normalizing Flow）（Rezende和Mohamed, 2015）来近似潜在空间的分布。

• Bond-Taylor 和 Willcocks（2021）通过单步梯度步骤构建了一个类似 VAE 的模型来获取潜在变量。

• Zhuang等人（2023）设计了一个扩散模型，使用少量坐标子集来提供上下文。

• 最后，一些方法使用超网络来输出神经场的权重：Dupont 等人（2022b）将超网络定义为对抗框架中的生成器，而 Du 等人（2021）则使用流形学习来表示潜在空间。

3. 无限维扩散模型

在本节中，我们将扩散模型扩展到无限维，以便通过在训练中对坐标进行下采样来允许更高分辨率的数据训练，并允许在任意分辨率下进行训练/采样。我们认为将条件神经场应用于扩散模型是有问题的，因为需要压缩到潜在向量，这增加了复杂性和出错的机会，而去噪函数应为无需压缩的非局部积分算子。许多同期研究也开发了无限维度的扩散模型，包括 Kerrigan等人（2023），Lim等人（2023），和Franzese等人（2023）；我们建议也阅读这些研究，它们在理论处理上走得更远，而我们的研究更侧重于设计和实际扩展性。

为此，我们将扩散状态空间限制在希尔伯特空间 H 中，其元素 x ∈ H 是函数，例如 x: R^n → R^d。希尔伯特空间配备了内积 〈·, ·〉 和相应的范数 ||·||_H。为了简化，我们考虑 H 是从 [0, 1]^n 到 R^d 的 L2 函数空间，尽管以下部分可以应用于其他空间。因此，H 中的一个点可以表示图像、音频信号、视频、3D 模型等。可以用其特征函数 ˆμ（Da Prato和Zabczyk, 2014）在 H 中定义高斯测度（Gaussian measure） μ：

均值 m 位于希尔伯特空间 H 中，m∈H，协方差算子（C：H → H） 是自伴的（self-adjoint）（记作 C = C*），非负的，并且是迹类的（trace-class）（Kukush, 2020）

对于具有分布 μ 的高斯随机元素 x，有 x∼N(m,C)。Radon-Nikodym 定理表明，对于一个绝对连续于基准测度 μ 的测度 v，存在一个密度函数：例如，两高斯分布之间的密度由 Minh (2021) 给出；在函数扩散模型的背景下，请参见 Lim 等人 (2023) 和 Kerrigan 等人 (2023) 的详细讨论。

3.1 平滑化（Mollificaition）

在定义无限维度上的扩散时，似乎使用前向过程中的白噪声是很自然的，即每个坐标是一个独立同分布的高斯随机变量：即 N(0,C_I)，其中 C_I(z(s),z(s′)) = δ(s−s′)，使用 Dirac delta 函数 δ。然而，这种噪声并不位于希尔伯特空间 H 中（Da Prato 和 Zabczyk, 2014），因为它不满足迹类要求。取而代之的是，通过对应于线性算子 T 平滑核（mollifier kernel） k(s)>0 与白噪声的卷积，从而获得希尔伯特空间 H 中的高斯噪声，即 N(0,T·T*)，将白噪声平滑到 H 中（Higdon, 2002）。为了确保核和噪声之间的一一对应，k 必须满足

使得 T·T* 自伴且非负。考虑 k 是具有平滑参数 l>0 的高斯核，h=Tx 给出如下：

3.2 无限维度平滑扩散

为了在希尔伯特空间 H 中制定扩散模型，我们必须指定转移分布。然而，数据点 x 的不规则性可能会影响模型的稳定性，使其无法在不同的下采样率/分辨率之间进行泛化。通过仔细的超参数调整可以缓解这一问题，或者在我们的情况下，也可以对 x 进行平滑处理（就像之前的噪声平滑一样）。虽然这种需求取决于 x 的性质，但我们为了完整性而包含了这一处理。首先，我们定义边际分布

由此我们可以推导出后验分布的闭式（closed form）表示（证明见附录 B.1），即

将反向转移定义为

然后我们可以参数化 μ_θ: H×R→H 来直接预测 x0。损失函数在等式 (3) 中可以扩展到无限维（Pinski等, 2015）。然而，由于 Ho 等人(2020) 发现预测噪声可以提高图像质量，因此我们参数化 μθ 以预测 ξ ∼ N(0,TT*)，并通过重写损失函数来实现这一目标：

由于 T^{-1} 不影响最小值，我们遵循 Ho 等人 (2020) 的做法，使用简化的损失函数

Kerrigan 等人 (2023) 的同期工作表明，在无限维的极限情况下，损失函数仅对特定的 ~β_t 选择才是有限的，而 Lim 等人 (2023) 则发现了类似的情况，只有在对 μ_θ 的特定参数化下才成立；然而，由于损失函数是通过蒙特卡洛近似的，因此在实践中这并不是问题。

数据平滑化。通过对训练数据 x0 进行平滑处理以确保其规律性，生成的样本也具有相似的规律性。直接预测 x0 将给出原始数据的估计值，但通过预测 ξ 我们只能对 Tx_0 进行采样。然而，在具有适当边界条件的高斯平滑核的情况下，如果我们考虑 x(c) 的傅里叶变换 ^x(ω)，则 T^{-1} 的存在是明确的，然后高斯卷积可以定义为

因此 Tx 在任何傅里叶可变换的函数类上都是一对一的，Tx 的有界性确保了唯一性，因此也是可逆的（John, 1955）。显式地，逆变换由

给出（Hummel 等人, 1987）。然而，逆变换的条件较差，微小的变化（例如浮点误差）会破坏平滑性（Hummel 等人, 1987）。在这种情况下，维纳滤波器可用作近似逆变换，定义为

其中 ϵ\epsilonϵ 是逆信噪比的估计值（Biemond 等人, 1990）。

4. 参数化扩散过程

为了在希尔伯特空间中建模去噪函数，有些特性是可学习函数类必须满足的，以便允许在无限分辨率数据上进行训练：

1. 能够接受位于任意坐标的点作为输入。

2. 能够推广到与训练时不同数量的输入点，在规则网格上采样。

3. 能够捕捉全局和局部信息。

4. 能扩展到大量输入点，即在运行时间和内存方面是高效的。

最近的扩散模型通常使用包含卷积编码器和解码器的 U-Net（Ronneberger 等人, 2015），并在分辨率之间通过跳跃连接有效地捕捉全局和局部信息。不幸的是，U-Net 在固定网格上工作，使其不适合我们的需求。然而，我们可以从中获得灵感，构建一个满足这些期望特性的架构。

4.1 神经算子

神经算子（Li 等人, 2020；Kovachki 等人, 2021）是一种框架，旨在通过学习直接将 PDE 参数映射到解决方案，在单步中高效地求解偏微分方程。然而，更广泛地讲，它们能够学习两个无限维函数空间之间的映射，使其适合于参数化无限维的扩散模型。

令 X 和 S 是可分离的 Banach 空间，分别表示有噪声和去噪数据的空间；一个神经算子是一个映射 F_θ: X→S。由于 x∈X 和 s∈S 都是函数，我们只能访问逐点评估。令

为域 D=[0,1]^n 的 m 点离散化（即 c 是 m 个坐标），并假设我们有观测值 x(c)∈R^(m×d)。为了实现离散化不变性，神经算子可以在任意 c∈D 处进行评估，从而允许在离散化之间传递解决方案。每一层都通过一个非局部积分核算子 K(x;ϕ) 构建，由神经网络 κ_ϕ 参数化，在空间上聚合信息。

深度网络可以类似于传统方法，通过堆叠带有非线性激活函数的线性算子层来构建，即 v0→v1→⋯→vL，其中 v_l → v_(l+1) 定义为

其中 v0 = x 为输入，vl 为激活，vL = s 为输出，逐点线性变换 W: R^d→R^d，激活函数 σ: R→R。一个例子是傅里叶神经算子（Fourier Neural Operator，FNO）（Li et al., 2021），

其中 G 是傅里叶变换，R_ϕ 是傅里叶空间中的学习变换。当坐标位于规则网格上时，可以使用快速傅里叶变换，使得 FNO 快速且可扩展。

4.2 多尺度架构

虽然有满足所有要求属性（1-4）的神经算子存在，例如 Galerkin 注意力（Cao, 2021）（无 softmax 线性注意力）和 MLP-Mixers（Tolstikhin et al., 2021），但由于内存成本高，扩展到大规模坐标仍然具有挑战性。相反，我们设计了一个受 U-Net 启发的多尺度架构（图 4），分别聚合局部和全局信息。在连续设置中，有两种主要的下采样方法：（1）选择坐标的子集（Wang 和 Golland, 2022）和（2）将点插值到规则间隔的网格（Rahman et al., 2022）。我们发现，通过反复应用（1），在非常少的点上近似非均匀间隔网格上的积分算子时，表现和泛化效果都不好，可能是由于高方差。另一方面，虽然使用规则网格去除了一些稀疏性特性，但方差问题要小得多。因此，我们采用一种混合方法，将稀疏算子应用于原始不规则采样数据的局部区域；之后将点插值到规则网格，然后应用基于网格的架构来聚合全局信息；如果规则网格的维度足够高，这种组合应该足够。虽然可以使用 FNO（Li et al., 2021; Rahman et al., 2022）架构，但我们用密集卷积（Nichol 和 Dhariwal, 2021）取得了更好的效果，并使用稀疏算子进行分辨率变化。

4.3 高效稀疏算子

在稀疏层面，我们使用卷积算子（Kovachki et al., 2021），发现它们比 Galerkin 注意力更具表现力，且通过多尺度架构提供全局上下文。这通过一个限制在每个坐标局部区域 N(c) 的平移不变核来定义，

由于深度核在大核的参数效率更高（尤其是对连续参数化核而言），我们限制 κ 为深度核，并发现它们在使用较少采样坐标进行训练时更能够泛化；尽管规则卷积和深度卷积的稀疏率相同，但由于规则卷积中有大量的值，值之间有更多的空间相关性。当使用非常大量的采样坐标时，完全连续卷积在内存使用和运行时间方面非常不实用。然而，在实际中，图像是在离散网格上获取和存储的。因此，通过将图像视为高维但离散的实体，我们可以利用高效的稀疏卷积库（Choy et al., 2019；Contributors, 2022），使内存使用和运行时间更加合理。具体来说，我们使用了 TorchSparse（Tang et al., 2022），并对其进行修改以允许深度卷积。Wang 和 Golland (2022) 提议使用低差异坐标序列来近似积分，因为它们具有更好的收敛率。然而，我们发现均匀采样点更有效，可能是因为高频细节更容易捕捉。

5. 实验

6. 讨论

有许多有趣的方向可以改进我们的方法，包括更强大/更高效的神经算子、更高效的稀疏方法、更好的积分近似以及改进的 UNet 设计（Williams等人，2023）。我们已经证明了扩散模型可以在 8 倍采样率下进行训练，因此我们认为还有大量的性能提升空间。此外，最近的研究也很有趣，它们通过在反向过程中迭代上采样来加速扩散采样，这需要为每个分辨率配备一个独立的模型（Jing等人，2022；Zhang等人，2023）；而我们的方法的分辨率不变性允许使用一个单一模型来实现这一点。

最近的扩散进展也与我们的方法互补，这些包括一致性模型（Song等人，2023）、随机插值（Albergo等人，2023）、薛定谔桥（De Bortoli等人，2021）、临界阻尼扩散（Dockhorn等人，2022）、架构改进（Hoogeboom等人，2023）以及更快的求解器（Lu等人，2022）。与我们的平滑扩散类似，模糊化已被用于改进扩散（Rissanen等人，2023；Hoogeboom和Salimans，2023）。类似于 GASP（Dupont等人，2022b），其他神经场 GAN 方法也存在，例如 CIPS（Anokhin等人，2021）和 Poly-INR（Singh等人，2023），但这些方法使用卷积判别器，需要在固定网格上使用所有坐标，因此无法扩展到无限分辨率。此外，神经过程（Garnelo等人，2018；Dutordoir等人，2022）也与此相关，它们学习类似于高斯过程的函数分布，然而，这些方法解决的是条件推理问题，而我们构建的是一个无条件生成模型，用于处理复杂得多的数据。

与本工作同时进行的，其他论文独立提出了无限维度的扩散模型（Lim等人，2023；Franzese等人，2023；Hagemann等人，2023；Zhuang等人，2023；Kerrigan等人，2023；Baldassari等人，2023；Pidstrigach等人，2023），这些方法与我们的方法互补，并在许多方面有所不同。我们的工作主要关注高效建模复杂高维数据所需的实际开发，而这些论文则更多地关注理论基础，通常仅应用于简单数据（如高斯混合和 MNIST）。特别值得注意的是，Kerrigan等人（2023）也在希尔伯特空间中开发了扩散模型，他们在基础理论上走得更远，包括讨论获得良好定义模型的要求，并考虑了不同的函数空间；Lim等人（2023）开发了作为 SDE 定义的无限维度扩散；Franzese 等人（2023）则证明了反向 SDE 的存在性。与我们的工作不同，这些方法利用了条件神经场或在统一网格的坐标上操作，而我们的方法在原始稀疏数据上操作，允许更好的扩展性。与本工作在扩展性上最接近的是扩散概率场（Zhuang等人，2023），该方法独立地去噪坐标，使用小坐标子集作为上下文；但这比我们的方法更具限制性，并且其分辨率远小于我们的方法（最大64×64）。

7. 结论

总之，我们发现我们的无限维度希尔伯特空间扩散模型，其过渡密度由非局部积分算子表示，能够生成高质量的任意分辨率样本。尽管在训练期间仅观察到少量像素子集，但样本质量显著优于先前的无限维度生成模型，并且与同时训练所有像素的最先进有限维度模型竞争。虽然之前的无限维度方法使用的是潜在条件神经场，我们的发现表明，直接在原始数据上操作的稀疏神经算子是一种有能力的替代方案，提供了显著的优势，因为它消除了逐点约束，因此无需潜在压缩。未来的工作将受益于改进的神经算子，这些算子能够在更大的稀疏性下有效地操作，以提高我们方法的效率并实现进一步的扩展。

论文地址：https://arxiv.org/abs/2303.18242

项目页面：https://github.com/samb-t/infty-diff

进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群
加 VX 群请备注学校 / 单位 + 研究方向