Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions

技术摘要：扩散模型在低维多模态分布学习中的统计最优性

1. 问题表述

基于分数的扩散模型在生成高维数据方面已展现出实证成功，特别是那些表现出低维和多模态结构的数据。然而，现有的关于其统计效率的理论保证往往受困于“维数灾难”，依赖于强正则性假设，如一致有界的密度、全局平滑的分数函数或对数凹性。这些假设未能捕捉内在的低维结构，并排除了重要的多模态分布（其密度在模态之间消失）。

本文研究了学习支持于**低维子空间并集（UoS）**上的分布的样本复杂度。具体而言，假设目标分布 $p^\star$ 支持于 $M$ 个线性子空间 $\{V_i\}_{i=1}^M$ 的并集上，其内在维度为 $k_i$，其中最大内在维度为 $k = \max_i k_i$。在每个子空间内，假设分布是次高斯的。目标是确定扩散模型能否实现由内在维度 $k$ 而非环境维度 $d$ 决定的样本复杂度，且无需施加限制性的平滑或有界密度假设。

2. 方法论

作者提出了一个以基于核的正则化分数估计量为核心的理论框架，用于高斯平滑分布 $p_t = p^\star * \mathcal{N}(0, tI_d)$。该方法主要分两个阶段进行：

A. 分数分解与估计策略
核心见解是对分数函数 $s^\star_t(x) = \nabla \log p_t(x)$ 进行分解。在 UoS 假设下，分数具有混合类型的分解形式：
$$s^\star_t(x) = \sum_{i=1}^M w_t(i, x) \cdot s_t(i, x)$$
其中 $w_t(i, x)$ 表示 $x$ 源自子空间 $V_i$ 的后验概率，而 $s_t(i, x)$ 是与该子空间相关的分数分量。关键在于，$s_t(i, x)$ 允许进行法向 - 切向分解：
$$s_t(i, x) = -\frac{1}{t}(x - \text{proj}_i(x)) + A_i s^{\text{low}}_t(i, A_i^\top x)$$
此处，法向分量是一个闭式高斯项，而切向分量 $s^{\text{low}}_t$ 是定义在子空间上的 $k_i$ 维分数函数。

B. 估计量构建
提出的估计量 $\hat{s}_t$ 通过两步过程构建（为理论清晰起见使用了样本划分）：

1. 子空间恢复：使用部分样本恢复子空间 $\{V_i\}$ 和一个分类函数 $c(\cdot)$，该函数将点分配给子空间。这被视为一个子空间聚类问题。
2. 分量估计：
   • 混合权重：使用核密度估计量来估计总密度 $p_t(x)$ 和部分密度 $q_t(i, x)$（来自子空间 $i$ 的质量）。权重估计量为 $\hat{w}_t(i, x) = \hat{q}_t(i, x) / \hat{p}_t(x)$。
   • 分数分量：对于每个子空间，使用带有特定正则化方案的核密度估计量来估计低维分数 $s^{\text{low}}_t$。该估计量涉及：
     • 一个代入估计量 $\nabla \hat{g}_t / \hat{g}_t$。
     • 一个阈值函数 $\psi$，在分母较小的低密度区域将分数设为零（防止不稳定性）。
     • 一个截断算子，用于限制估计量的范数。

C. 采样过程
将估计的分数代入反向随机微分方程（SDE）采样器（算法 1），并设置早期停止时间 $\tau$ 以避免具有挑战性的小时间区间。

3. 主要贡献与结果

A. 分数估计误差界
本文建立了所提出的分数估计量的有限样本 $L_2$ 误差界。假设子空间恢复精确，期望平方误差满足：
$$\mathbb{E}[\|\hat{s}_t(X) - s^\star_t(X)\|_2^2] = \tilde{O}\left( \frac{1}{n} \left( \frac{1}{t} + \frac{1}{t^{(k \vee 2)/2 + 1}} \right) \right)$$
其中 $n$ 是样本量，$k$ 是内在维度。该速率取决于 $k$ 而非环境维度 $d$。

B. 采样复杂度与收敛性
基于分数估计保证，作者证明扩散采样器在 1-Wasserstein 距离（$W_1$）下实现 $\epsilon$-精度样本所需的样本量最多为：
$$\tilde{O}(\epsilon^{-(k \vee 2)})$$
该收敛速率：

• 匹配极小极大最优性：它与学习 $k$ 维分布的极小极大最优速率一致。
• 打破维数灾难：依赖关系在于内在维度 $k$，而非环境维度 $d$。
• 放宽假设：该理论仅要求分布在每个子空间上是次高斯的。它不要求：
  • 分数或密度的平滑性。
  • 对数凹性。
  • 远离零的一致有界密度（允许存在良好分离的多模态分量）。

C. 数值验证
在 $\mathbb{R}^{48}$ 中支持于 128 个子空间（内在维度 $k=3$）的分布上的合成实验证实，经验 $L_2$ 分数估计误差的缩放符合内在维度，与理论预测一致，而非环境维度。

4. 意义与主张

本文声称提供了严格的理论依据，解释了扩散模型在复杂、高维、多模态设置中的实证成功。其主要意义在于证明扩散模型可以统计地适应内在的低维结构，而无需依赖困扰先前理论分析的强正则性假设。

作者强调，虽然基于核的估计量主要是一种理论证明工具，而非神经网络（NN）估计量的实际替代方案，但结果确立了可实现的基本统计极限。这为未来对基于 NN 的分数估计量的分析提供了具体的目标和结构路线图，指出最优性的关键在于捕捉低维法向 - 切向分解，并在无密度下界约束的情况下处理多模态性。

该工作承认，误差界前缀中关于 $d$ 的线性依赖可能是证明技术的产物，而将这些结果扩展到噪声环境及实用的基于 NN 的估计量仍是未来工作的开放方向。