贝叶斯扩散模型泛化相变的精确信息论

贝叶斯扩散模型泛化相变的精确信息论

An exact information theory of generalization phasetransitions in Bayesian diffusion models

https://arxiv.org/pdf/2607.08041

摘要

扩散模型如何规避维度灾难，从有限的训练集中学习高维空间上的复杂分布，而不是记住它，仍然是一个基本的谜团。为了解决这个问题，我们引入了可解析处理的贝叶斯信息受限扩散（BIRD）模型，其中每个像素观察关于噪声数据的受限信息。BIRD模型通过使用贝叶斯后验推断哪个过去的训练样本产生了其当前的受限观察，从而在时间上反转扩散。这个模型类推广了使用空间局部信息限制的现有解析扩散模型。我们表明，在训练早期，跨不同架构（如UNets和DiTs），空间局部BIRD模型紧密近似训练过的扩散模型。在对数据分布的最小假设下，我们在训练数据量、反向生成过程时间和信息限制量的联合空间中，识别出记忆和泛化之间的信息论相边界：当其受限噪声观察与训练数据之间的互信息超过训练点数量的对数时，BIRD模型进行记忆，否则进行泛化。跨越一系列数据集的实验证实了我们对转变位置的理论预测。我们发现生成在记忆的边缘进行：空间局部BIRD模型和早期训练的扩散模型都通过随着时间的推移越来越限制信息来追踪记忆-泛化相边界。总的来说，我们的结果揭示了信息限制在生成式AI中规避维度灾难的基本作用。

1 引言与相关工作

值得注意的是，生成式AI现在可以用可处理数量的训练数据学习高维空间上的复杂分布。例如，扩散模型[1–3]稳健且一致地泛化：当在完全不相交的数据子集上训练时，当数据子集大小适度大时，两个独立的模型将收敛到相同的通用模型[4]。低于这个适度的数据阈值，这两个模型可能反而记住它们自己的训练数据。然而，朴素地，维度灾难会表明，对于任意数据分布，从记忆过渡到泛化应该需要不可处理的（例如，关于环境维度呈指数级）数据量[5]。这引发了一个基本的谜团，即扩散模型从记忆到泛化的转变何时发生，以及为什么它不需要那么多训练数据。已经采取了几种理论方法来探索这个谜团。这些方法通常涉及分析简化扩散模型在简化数据集上的行为。扩散模型反转将数据转换为噪声的前向扩散过程。存在这个过程的一个精确数学反转，它从噪声开始并遵循有限训练集的理想或经验得分函数（附录B，[6]）。然而，这种经验得分模型需要关于维度呈指数级的数据量来避免记忆[6, 7]。虽然几篇论文探讨了数据分布的低维性可能拯救经验得分模型泛化的假设[8–10]，但广泛的实验工作表明，在现实建模条件下，经验得分模型在现实维度的数据集上无法泛化[11–17]。

经验得分模型无法以可处理的训练数据量实现泛化，这表明扩散模型中成功的泛化可能源于对归纳偏置的限制，这些偏置会阻碍对经验得分函数的学习，例如平滑性[12, 18]、线性[19, 20]、低维性[21, 22]、几何自适应基[4]、早期训练[23–26]以及模型参数化/架构[27–31]。然而，这些工作中没有一项推导出一个高度准确的预测理论，来描述训练后的非线性扩散模型在泛化阶段在各种基准数据集上实际收敛于何种状态。实现这一目标的此类理论之一出现在文献[13, 14]中，该理论提出了局部得分模型（LS），这是一种受局部性约束的最优扩散模型，其中每个像素只能使用图像的局部空间邻域来计算得分。文献[13]还引入了等变（ES）和（边界破缺的）等变-局部得分模型（ELS），这些模型引入了（平移）群等变性及其通过边界破缺的额外归纳偏置。特别是，文献[13]表明，此类解析模型能够以极高的定量精度，逐案预测仅含卷积的训练扩散模型的图像输出。随后，文献[32]提供了一种通过最优线性去噪获取LS模型局部尺度和图块几何结构的方法。有趣的是，这些研究发现，图像生成过程从较大的粗糙空间尺度开始，逐步过渡到较小的精细尺度。此后，局部泛化机制也在后续研究中得到了进一步探讨[33–38]。重要的是，尽管文献[13]为处于泛化阶段深处的卷积扩散模型开发了一种预测性解析理论，但该工作及其他研究仍留下了一个未解之谜：对于更强大的模型，在复杂的基准数据集上，从记忆到泛化的转变究竟在何时何地发生。普遍地回答这一问题极具挑战性，因为先验上它可能联合依赖于架构带来的显式偏置、训练动态带来的隐式偏置、数据分布的结构、生成过程中的时间或噪声水平，以及当然还有数据量。为了在理论上理解这一问题，本文引入了一类通用的扩散模型，我们称之为贝叶斯信息受限扩散（BIRD）模型。该模型受文献[13, 14]中等变/局部得分模型的启发，并涵盖了所有这些模型。BIRD模型巧妙地结合了直观吸引力、预测能力和理论可处理性，为理解记忆-泛化转变提供了重要的概念性见解。我们的主要贡献如下：

（1）我们引入了一个直观易懂的BIRD模型通用概念，其中每个图像像素x可被视为一个智能体，在反向生成过程的时间t观测关于噪声图像φt的受限信息Cx,t，并利用训练数据上的最优贝叶斯推断来逆转前向扩散过程（图1）。这构成了一个从有限训练集中获得的定义明确的扩散模型，无需任何显式的架构假设或学习动态。

（2）我们证明了空间局部BIRD模型（其中受限观测Cx,t对应于x周围的局部图像块）能够对许多神经扩散模型在训练早期学到的内容进行极其准确的预测：在训练初期，针对包括带自注意力的UNet和扩散Transformer（DiT）在内的多种架构，它们能够逐案预测神经扩散模型的单幅图像输出，中位决定系数r²高达0.85至0.93。这一结果超越了文献[13]的结论，后者仅证明了对小型CNN具有类似的高预测性。

（3）理论上，在对数据分布的极弱假设下，我们为一般BIRD模型在训练数据量、时间（或反向过程中的噪声水平）和信息限制量构成的联合空间中，确立了一个精确的记忆-泛化相变位置信息论判据：若在前向扩散过程中，观测值Cx,t与真实数据分布之间的互信息大于（小于）数据量的对数，则BIRD模型将发生记忆（泛化），因为对底层图像的贝叶斯猜测将过于容易（困难）。我们在多种数据集上通过实验证实了理论预测的相边界位置。这些结果表明，只需使数据量相对于互信息呈指数增长即可避免记忆，并突显了信息限制在促进泛化中的根本作用。

（4）聚焦于空间局部BIRD模型，其中信息限制通过每个像素观测尺度为L的局部图块来实现，我们利用多种信息论分析，揭示了在反向生成时间t（或噪声水平σt）与图块尺度L构成的联合空间中，记忆-泛化相变边界的形状。我们发现了一个随时间变化的临界图块尺度Lc(t)，在反向生成过程中，随着时间t和噪声σt从t=1减小至t=0，该尺度单调递减。对于任意t，若L(t) > Lc(t)（或L(t) < Lc(t)），则BIRD模型将发生记忆（或泛化）。最优生成与去噪过程沿着这一相变边界推进：本质上，成功的生成是在记忆的边缘进行的（图1b）。

综上所述，这些结果为一类通用的扩散模型中记忆与泛化之间的相变提供了根本性的见解，并突显了信息限制在允许以相对适度的数据量实现泛化方面所发挥的作用。下文我们将对结果进行总结，更详尽的阐述可见附录（例如，参见附录A以获取所有数学符号与背景的总结）。

2 整体框架

2.1 贝叶斯信息受限扩散（BIRD）模型

对于最优去噪器而言，贝叶斯猜测游戏不幸的简易性提出了一个问题：人们是否可以获得替代的贝叶斯扩散模型，这些模型既保留了解析可处理性和概念简洁性，又能实现泛化，从而表现得更像训练后的神经网络扩散模型。我们要证明，通过结合两种关键方法，我们可以同时实现理论可处理性和现实泛化：1）通过限制模型用于猜测训练图像的图像信息，使贝叶斯猜测游戏变得更难；以及 2）允许图像的不同部分使用不同的受限信息片段。我们将此类模型称为贝叶斯信息受限扩散（BIRD）模型，其灵感源自并推广了 [13, 14]。

接下来我们在第 3 节中展示，空间局部 BIRD 模型具有高度的相关性，因为它们很好地描述了许多扩散模型架构的早期学习阶段。在此之后，我们在第 4 节中展示，所有 BIRD 模型的一般类别提供了一个非常有用的抽象，即针对所有此类 BIRD 模型，基本上对于任何数据分布，都可以推导出记忆-泛化相变的一般理论。

3 BIRD 模型描述扩散模型的早期学习阶段

BIRD 模型理论与神经实验之间的这种强烈一致性，促使我们将 BIRD 模型用作一个理论实验室，以研究记忆-泛化相变的本质。我们将在第 4 节中阐明这一理论。

4 BIRD 模型中的记忆-泛化相变

为了实际定位记忆与泛化之间的相变位置，我们在附录 D.5 中证明，对于具有次指数尾（sub-exponential tails）的后验分布，在大数据集尺寸和高观测维度极限下，

5 BIRD 模型如何规避维度灾难

我们现在解释，对于尺度不变的自然图像，空间局部的 BIRD 模型如何能够有效地去噪，同时避免记忆，即使对于大图像也是如此。

5.1 记忆-泛化转变处的临界图块尺度

5.2 有效线性去噪的谱长度尺度

6 讨论

利用信息论，我们能够为高度通用的 BIRD 模型类别以及非常通用的数据分布类别获取记忆-泛化相图。虽然 BIRD 模型通过信息限制引入了归纳偏置，但我们理论的一个显著局限性是它无法捕捉由于架构选择或学习动态可能产生的其他归纳偏置。例如，BIRD 模型无法表达在训练后的模型中可能出现的对线性的偏置。然而，我们尚未发现任何扩散模型理论能够精确描绘具有任意架构和数据分布的记忆-泛化相变边界。所有现有理论要么使用简化的架构，要么使用简单的数据模型。因此，BIRD 模型提供了一个令人兴奋的新理论实验室，用于研究跨越许多可能数据分布的记忆和泛化。

原文链接：https://arxiv.org/pdf/2607.08041