复杂度平衡的扩散分割

基于de Boor等分布原理，将扩散时间轴划分为等复杂度片段并用专门子网络建模

前置知识

扩散模型 (Diffusion Models)

扩散模型是一类生成模型，通过学习将数据逐步添加噪声直到变成纯噪声，再学习逆向过程从噪声中恢复数据。在连续时间 formulation 中，使用神经网络预测每个时间步的速度场，通过ODE积分实现采样。这个框架将生成问题转化为学习从噪声分布到数据分布的连续路径。

本文的核心是优化扩散模型的时间轴分配，因此必须理解扩散模型如何在不同时间步建模从噪声到数据的转换过程，以及为什么单一网络难以适应不同时间步的复杂度差异。

Flow Matching / 速度预测

Flow Matching 是一种统一的生成建模框架，通过学习数据分布与简单噪声分布之间的连续路径速度场来训练模型。给定时间 $t \in [0,1]$ 和中间状态 $x_t$，神经网络学习预测真实速度 $u(x_t, t)$，优化目标是最小化预测速度与真实速度的均方误差 $\mathcal{L} = \mathbb{E}_{t,x_0,x_1} \|v_\theta(x_t, t) - u(x_t, t)\|^2$。采样时通过积分这个速度场从噪声生成数据。

本文从函数近似的角度重新审视扩散模型，将生成过程视为学习速度场近似问题，因此需要理解Flow Matching的数学表述。

de Boor 等分布原理

de Boor 等分布原理是逼近理论中的经典结果，用于解决最优域划分问题。核心思想是：在有限资源约束下，为了最小化最大逼近误差，应该将目标函数的定义域划分为若干子区间，使得每个子区间承担的逼近负担相等。数学上，给定监控函数 $m(t)$ 量化局部复杂度，最优划分点 $\{t_i\}$ 满足 $\int_{t_{i-1}}^{t_i} m(t) dt = \frac{1}{N} \int_0^1 m(t) dt$，其中 $N$ 是区间数量。

这是本文的理论基础，作者将时间轴分割问题框架化为域分解问题，并应用de Boor原理来推导最优划分策略。

Dirichlet 能量

Dirichlet 能量是函数空间的重要概念，用于量化函数的"粗糙度"或空间变化剧烈程度。对于向量场 $v(x)$，其Dirichlet能量定义为 $E_D(v) = \frac{1}{2} \int_{\mathbb{R}^d} \|\nabla_x v(x)\|^2 dx$，即梯度的 $L^2$ 范数积分。高Dirichlet能量意味着函数变化剧烈，需要更多的模型容量来精确逼近。在频域中，Dirichlet能量通过Parseval恒等式与频谱能量相关，可以用于估计函数的谱复杂度。

本文提出的第一个监控函数基于Dirichlet能量，用于量化流场的空间复杂度，这是推导复杂度平衡分割的关键技术组件。

路径加速度

路径加速度是几何建模中常用的复杂度度量，指采样轨迹对时间的二阶导数。给定速度场 $v_t(x)$，轨迹定义为 $\frac{dx_t}{dt} = v_t(x_t)$，加速度为 $\frac{d^2x_t}{dt^2}$。与速度不同，加速度能够过滤掉恒定速度的位移，只捕获路径的"弯曲度"或非线性程度。在曲线逼近理论中，高阶导数的最大值与逼近误差直接相关。

本文提出的第二个监控函数基于路径加速度，用于量化采样轨迹的几何复杂度，实验证明它比Dirichlet能量更直接地反映采样质量。

研究动机

现有连续时间生成模型采用单一神经网络架构处理整个去噪过程，这在计算上是低效的。生成过程需要穿越从各向同性噪声到复杂数据分布的显著不同信号区间，从粗略结构形成到高频细节细化。标准做法是扩大模型规模来提升性能，但这意味着在每个时间步都部署大规模网络，而实际上没有单独的去噪区间需要如此大的容量。更高效的替代方案是按时间分布容量，训练多个专门网络分别负责去噪过程的不同阶段，但现有的划分方法通常依赖启发式分割或计算昂贵的搜索过程，缺乏有原则的容量分布准则，限制了效率和通用性。

本文的目标是本文的目标是提出一个有原则的时间容量分配框架，将扩散时间轴划分为等逼近负担的片段，为更难以建模的区域分配更多表示能力。通过从函数近似理论的角度重新审视问题，应用de Boor等分布原理，设计无需昂贵搜索的最优时间分割方案，在不增加每步推理成本的情况下提升生成质量。

与已有工作不同的是，本文的独特切入角度是将时间轴分割问题框架化为逼近理论中的域分解问题。不同于现有工作依赖启发式分割（如均匀分割）或计算昂贵的搜索过程（需要训练多个大规模模型），本文提出从函数近似的误差边界出发，使用可计算的监控函数量化局部复杂度，通过等分布原则直接推导最优分割点。这种方法既具有理论保证，又避免了昂贵的计算开销，将抽象的逼近理论转化为可实践的容量分配工具。

核心方法

CBS的核心思想是：扩散生成过程在不同时间段的复杂度差异很大，应该为复杂度高的区域分配更多模型容量。作者从函数逼近理论出发，将时间轴视为一维域，应用de Boor等分布原理进行最优划分。关键步骤包括：先训练一个轻量辅助网络估计复杂度曲线，使用两个互补的监控函数（基于流场Dirichlet能量的空间度量，和基于采样轨迹加速度的几何度量）量化局部复杂度，然后将时间轴划分为等复杂度片段，最后为每个片段训练专门的子网络。这种方法实现了容量扩展与推理成本的解耦。

核心创新点在于将抽象的逼近理论转化为可实践的监控函数。作者推导了两个可计算的复杂度度量：一是基于Barron定理的Dirichlet能量监控函数，通过随机迹估计器高效计算Jacobian-Vector Products，避免高维梯度计算；二是基于路径加速度的几何监控函数，直接测量采样轨迹的非线性程度。这两个监控函数使得能够无需昂贵搜索就估计最优时间分割点，将传统的启发式或经验搜索方法替代为有原则的数学框架。与现有方法相比，CBS不需要训练多个候选分割方案，而是直接从复杂度分析出发确定分割点，大大降低了计算开销。

方法步骤详情

CBS的完整步骤如下：第一步，训练一个轻量级辅助网络 $v_\theta^{aux}$，在完整时间区间 $t \in [0,1]$ 上使用标准速度预测目标，这个网络只需在小数据集上训练少量迭代；第二步，使用这个辅助网络生成 $K$ 条采样轨迹 $\{x_t^k\}_{k=1}^K$，用于估计监控函数 $m(t)$；第三步，对于Dirichlet能量监控，在每个时间网格点使用随机迹估计器计算 $E_D(v_t)$，避免显式计算高维Jacobian；第四步，对于路径加速度监控，通过有限差分计算 $m(t) = \frac{1}{K} \sum_{k=1}^K \|v_{t+\Delta t}(x_{t+\Delta t}^k) - v_t(x_t^k)\|$；第五步，应用de Boor等分布原理，计算监控函数的累积和，选择将总监控值均匀分割的时间点作为最终分割边界 $\{t_i\}$；第六步，为每个时间片段 $[t_i, t_{i+1}]$ 训练专门的子网络，每个网络只在其指定的时间区间内采样 $t$ 值；第七步，推理时根据当前时间步自动切换到对应的子网络，完成生成过程。

技术新颖性

技术新颖性体现在多个方面：首先，这是首个将逼近理论中的de Boor等分布原理应用于扩散模型时间轴分割的工作，为容量分配提供了坚实的数学基础；其次，提出了两个可计算的监控函数，特别是将Barron定理的谱复杂度边界通过Dirichlet能量转化为可计算的度量，以及使用路径加速度直接量化采样轨迹的几何复杂度；第三，证明了复杂度估计对架构规模和训练时长具有鲁棒性，即使使用在10%数据上训练的小网络也能得到相同的分割边界，这意味着时间分割可以在预训练阶段以可忽略的计算开销完成；第四，方法具有架构无关性，在SiT、JiT、UNet等不同架构上都取得了显著提升。

实验结果

核心发现是CBS在多个架构和数据集上一致提升了生成质量，且不增加每步推理成本。在ImageNet-256上使用SiT-XL/2，CBS的FID从均匀分割的17.97降至4.03，Inception Score从73.85提升至195.72，Precision从0.63提升至0.80，提升幅度约35%。CBS在更大的模型上表现更好，SiT-XL/2的提升显著大于SiT-S/2和SiT-B/2，表明复杂度驱动的分割在模型容量更大时优势更明显。当使用Classifier-Free Guidance（CFG）时，CBS的提升更加显著，说明专门的子网络能更好地处理CFG引入的复杂空间梯度。可扩展性实验显示，随着网络数量 $N$ 从1增加到4，FID持续下降（N=1时34.84，N=4时29.33），IS持续上升（N=1时41.53，N=4时55.03），验证了通过更细粒度的时间分割缓解局部容量瓶颈的有效性。最优性验证实验表明，扰动分割边界后性能下降，证实基于复杂度的分割接近最优解。监控函数比较发现，路径加速度监控函数略优于Dirichlet能量，作者认为这是因为它更直接地测量最终采样质量。效率分析表明，即使使用在10%数据上训练的小网络，也能得到与完整训练的SiT-XL/2几乎相同的复杂度曲线和分割边界，证明时间分割估计的计算开销可以忽略不计。方法在像素空间的ImageNet-64和CIFAR-10上也有效，验证了架构无关性。

**Table 1**: Quantitative evaluation of SiT on ImageNet-256.

**Table 2**: Quantitative evaluation of JiT-B/4 on ImageNet-64.

**Table 3**: Unconditional generation on CIFAR-10.

**Table 4**: Scaling CBS across multiple networks.

**Table 5**: Empirical validation of boundary optimality on SiT-B/2.

**Table 6**: Comparison of Monitor Functions.

查看结构化数据

任务	指标	本文	基线	提升
ImageNet-256图像生成	FID (with CFG)	4.03 (SiT-XL/2)	6.53 (SiT-XL/2 monolithic)	38.3%降低
ImageNet-256图像生成	Inception Score (with CFG)	195.72 (SiT-XL/2)	162.23 (SiT-XL/2 monolithic)	20.6%提升
ImageNet-256图像生成	FID (without CFG)	15.81 (SiT-XL/2)	18.04 (SiT-XL/2 monolithic)	12.4%降低
ImageNet-64像素空间生成	FID (with CFG)	13.93 (JiT-B/4)	16.41 (JiT-B/4 monolithic)	15.1%降低
CIFAR-10无条件生成	FID	2.72 (UNet)	3.55 (UNet monolithic)	23.4%降低
时间分割可扩展性	FID下降 (N=1到N=4)	从34.84降至29.33	单网络基线	15.8%降低

局限与改进

作者明确指出本文专注于时间轴，在时域上达到了接近最优的解决方案，但相信应该扩展到其他形式的分割。空间分割（在令牌之间路由网络）是一个潜在更强的方向，但推导适当的监控函数是一个挑战。从技术角度观察，方法需要训练一个辅助网络来估计复杂度，虽然开销可以忽略，但仍然是一个额外的训练步骤。监控函数（特别是Dirichlet能量）依赖于随机估计，可能在高维空间中引入方差。方法主要在ImageNet等大规模数据集上验证，在小数据集或分布偏移情况下的鲁棒性需要进一步研究。论文中主要评估了FID、IS等指标，对计算效率（训练时间、推理速度）的详细分析相对有限。

独立分析的弱点

独立分析的主要弱点包括：首先，CBS需要预先确定分割数量 $N$，论文中主要使用 $N=3$，但不同数据集和架构可能需要不同的最优 $N$ 值，缺乏自动确定 $N$ 的方法。其次，监控函数的选择对性能有影响，虽然路径加速度略优于Dirichlet能量，但两者都在不同架构上表现稳定，缺乏理论分析指导何时使用哪个监控函数。第三，方法目前主要在Transformer架构上验证，在更复杂的模型（如现代大规模扩散模型）上的适用性需要进一步验证。第四，等分布原则假设所有子网络具有相同的建模能力，但在实践中可能希望根据局部复杂度分配不同的网络大小，这需要扩展理论框架。改进方向包括：开发自动确定最优 $N$ 的方法；研究混合监控函数或自适应选择监控函数；扩展到非均匀容量分配；在更多架构和数据集上验证；分析对训练数据规模和分布变化的鲁棒性。

未来方向

作者提出的未来方向是扩展到空间分割，在令牌级别路由网络，这预期是更强的方向但需要推导适当的监控函数。基于本文成果的可延伸方向包括：将等分布原则应用于其他生成模型（如GANs、VAEs）的容量分配；研究层次化分割，先进行粗粒度时间分割，再在每个片段内进行细粒度空间或特征分割；探索动态分割，在训练过程中根据学习进度调整分割边界；应用监控函数思想到模型压缩和蒸馏，识别并保留关键的网络组件；研究复杂度监控函数与其他模型设计原则（如MoE路由、动态计算图）的结合；扩展到多模态生成，处理不同模态的复杂度差异；分析复杂度平衡分割与样本质量、多样性、速度的trade-off关系。

复现评估

论文提供了项目页面 https://noamissachar.github.io/CBS/，但未明确说明是否开源代码。实验使用标准数据集（ImageNet-256、ImageNet-64、CIFAR-10）和标准架构（SiT、JiT、UNet），便于复现。论文声称辅助网络只需在10%数据上训练少量迭代就能得到可靠的分割边界，这降低了复现的计算门槛。然而，实验的具体超参数（如学习率、批量大小、训练轮数等）主要在附录C中描述，可能增加复现的难度。论文的主要发现（复杂度平衡分割优于均匀分割）具有很强的理论基础和实验支撑，复现应该相对直接。对于大规模实验（如SiT-XL/2），可能需要显著的计算资源，但方法的核心思想可以在较小的模型上验证。

复杂度平衡的扩散分割 Complexity-Balanced Diffusion Splitting