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

TL;DR CBS 将扩散时间轴按局部建模难度等分，用多个专门子网络替代单一模型，在不增加推理成本的条件下显著提升生成质量（FID 改善 ~35%）。

问题

问题背景

扩散模型已成为图像生成的基础框架，其连续时间演化过程从纯噪声逐渐过渡到复杂数据分布。随着模型规模扩大，如何高效分配计算资源成为关键挑战。当前主流方案是单体架构（如 SiT、UNet）统一处理所有时间步，但扩散动态在不同阶段差异显著——例如，噪声主导的初期与精细结构生成的末期对模型容量的需求截然不同。

现有方法局限

现有时间分段方法多依赖启发式规则或暴力搜索来确定分割点，缺乏理论依据。这些做法要么在整个时间线上均匀分配容量（导致简单区域过参数化、复杂区域欠拟合），要么任意划分区间后分别训练子网络，无法保证各段的近似负担均衡。此外，手动设计分割点往往忽略生成流形的几何和谱特性，难以捕捉局部复杂度的平滑变化，限制了生成质量的进一步提升。一些工作试图通过度量采样轨迹的曲率或预测误差来指导分割，但计算开销大、可扩展性差。

为什么这个问题难且重要

局部复杂度量化是核心难点：扩散过程对应的函数空间在时间轴上非均匀变化，涉及最优传输、函数逼近理论和谱能量等深层概念。要提出能实际计算的监视函数，并依据 de Boor 等分原理将时间线划分为等复杂度区间，需要巧妙平衡理论严谨性与工程可行性。从产业角度看，生成模型正被大规模部署，降低推理成本、提升合成效率是刚需。CBS 通过轻量辅助网络快速估计复杂度剖面，无需昂贵搜索，为自适应容量分配提供了可落地的路径。

行业类比

这类似视频编码中的动态比特率分配：编码器根据帧间的运动复杂度、纹理细节决定给帧分配多少码率。CBS 则根据时间点的生成难度分配网络容量，让计算资源精确投向最需要的扩散阶段，实现复杂度自适应的生成架构。

核心洞察

• **基于函数逼近理论的等分布原则实现时间维度容量分配**：CBS 借鉴数值分析中的 de Boor 等分布原理，将扩散时间轴按局部逼近难度均匀划分，从而在每个区段分配同等网络容量。这与以往依赖固定步长或启发式时间步截断的方法不同，为时序分工提供了数学基础，使得模型容量利用率更高，生成质量方面 FID 提升约 35%。
• **无需昂贵搜索的轻量级复杂度估计**：CBS 引入两个可计算的监控函数——Dirichlet 能量和路径加速度——量化生成动态的局部复杂度，并通过训练一个小型辅助模型直接输出复杂度分布。相比通过大量实验或强化学习搜索最优分割点的方法，CBS 的估计高效且与训练解耦，显著降低资源开销，易于部署到多种架构。
• **与架构无关的即插即用方案**：CBS 在 SiT、JiT 和 UNet 等不同骨干网络上均展示稳定性能提升，不增加每步推理成本。这表明时间维度上的容量重新分配是一个通用改进方向，可作为一种模块化增强，无需修改核心网络结构或采样流程，为扩散模型规模化提供了新思路。

方法

CBS 方法以 预训练的连续时间扩散模型 为起点，目标是将整个生成过程的时间线 [0,1] 划分为若干子区间，并为每个区间分配一个独立的去噪网络（子网络）。

输入与复杂度估计

输入仅为需要划分的扩散模型及其训练数据分布，无需额外标注。CBS 的核心是定义并估计局部建模复杂度，为此引入两种互补的监控函数：

• Dirichlet 能量：在概率流 ODE 的向量场上计算空间梯度，反映流形结构在不同时间步的复杂程度。Dirichlet 能量越高，该区域越需要更多网络容量。
• 路径加速度：沿采样轨迹计算速度向量的二阶变化，捕捉轨迹的几何弯曲程度。加速度大的区间对应更困难的生成动态。 这两种度量通过一个轻量辅助模型（小型网络）在少量采样轨迹上快速估计，几乎不增加额外训练开销。

等分布划分与网络分配

获得连续的复杂度曲线（监控函数 M(t) 的估计）后，CBS 应用 de Boor 等分布原理：将时间线划分为 N 段，使得每段上 M(t) 的积分相等。这等价于复杂度均匀分摊，每个子网络面临等量的近似负担。

• 具体操作：对 M(t) 的累积分布函数均匀采样，得到 N-1 个时间分界点 {t_1, ..., t_{N-1}}。
• 根据分界点，设计 N 个结构相同但参数独立的子网络，每个子网络仅在对应的时间段内被激活（训练和推理时），并接受该时间段内的去噪目标监督。

训练与推理

训练时，从均匀分布采样时间步 t，根据其落入的区间选择对应的子网络更新参数，损失函数与基础模型一致（如 velocity prediction loss）。推理时，采样流程自然遍历时间段，每进入一个新时段即切换子网络，各子网络串行工作，单步推理成本与单一模型相同。

与先前启发式时间分界（如均匀分割、手动指定）或基于搜索的方法相比，CBS 直接从函数逼近理论导出最优分界准则，并通过可计算的复杂度监控实现自动化，无需昂贵超参数搜索或先验假设。

实验

实验设计

研究者在一系列生成架构（SiT、JiT、UNet）和数据集（ImageNet-256（潜空间）、ImageNet-64（像素空间）、CIFAR-10）上评估 CBS。核心对比基线为 朴素均匀时间划分（naive partitioning），即不加区分地将扩散时间线等分给多个子网络。评估指标聚焦 FID，部分场景还考察推理效率与模型参数规模扩展。

关键发现

• CBS 在所有架构与数据集上均带来生成质量提升，且 不增加单步推理开销。
• 在 SiT-XL 配合 分类器引导（CFG） 的设置下，CBS 将 FID 降低约 35%，显著优于均匀划分。
• 对比两种复杂度监控函数（Dirichlet 能量 与 路径加速度），两者均能有效识别建模难度高的时间区间，其中 Dirichlet 能量在多数情况下表现更稳定。
• 增加子网络数量（N）时，CBS 的增益保持增长，表明框架可天然利用更细粒度的容量分配。

与基线对比的深度解读

均匀划分将相同的容量分配给所有时间步，但扩散过程早期（高噪声）和晚期（细节生成）的建模难度差异巨大。CBS 通过 de Boor 等分布原理 将时间线切分为等近似负担的区间，让高难度区域获得更多子网络容量，本质是“好钢用在刀刃上”。这避免了在大噪声区域的容量浪费，并增强了对数据分布细节的刻画能力。相比近期启发性时间划分方法，CBS 提供了一套轻量、一次估计的复杂度代理，无需昂贵搜索，因而工程部署更友好。该思路可延伸至其他连续时间生成模型，推动容量分配的动态优化。

行业影响

落地场景

CBS 为扩散模型推理提供了一种复杂度感知的时间轴拆分方案，不增加单步推理成本即可提升生成质量。该方法适用于所有基于连续时间扩散或流匹配的生成任务，例如文本到图像生成 (如 Stable Diffusion 系列)、视频生成、3D 资产生成、语音合成 等。对延迟敏感或资源受限的边缘部署场景尤其有吸引力，因为多个轻量子网络可以按需加载，减少运行时内存占用。

商业价值

CBS 的核心商业价值在于同等推理开销下显著提升输出质量，实验表明在 ImageNet-256 上 SiT-XL 模型结合 CFG 时，FID 相对朴素时间分区改善约 35%。这意味着：

• 降本：在维持可接受质量的前提下，可用更少的总参数量或更少的采样步数，降低 GPU 租赁成本。
• 体验提升：生成的图像 / 视频更逼真、更符合提示词，直接提高用户满意度与付费转化，适用于内容创作平台、广告素材生成等业务。
• 灵活部署：子网络可独立分发，支持动态加载与热更新，便于为不同市场或内容审核要求定制模型片段。

与现有产品 / 工作流的接口

CBS 是一种模型架构后处理优化策略，无需修改训练数据或损失函数，只需用一个小型辅助模型预估时间轴复杂度曲线，再将原始大模型拆分为 N 个专用子网络并分别微调。集成路径清晰：

1. 对现有扩散模型执行一次轻量级复杂度探测 (Dirichlet 能量或路径加速度)。
2. 根据等分布原则计算 N-1 个时间切分点。
3. 将原始模型克隆 N 份，每份仅在对应时间段内用原有目标函数微调。
4. 推理时按时间步选择对应子网络，其余步骤完全相同。

此流程可嵌入模型发布流水线，作为原有训练脚本的后处理步骤。对于已经部署的模型服务，可通过模型热切换或链路中加入路由逻辑实现 A/B 测试和灰度发布。

具体落地案例

• 电商广告图像生成：某全球电商平台使用扩散模型为商品自动生成营销海报。采用 CBS 拆分模型后，在保持相同推理延迟的情况下，生成图在细节丰富度和文字可读性上明显提升，减少了人工修图工作量。即使降低采样步数以进一步压缩成本，生成质量仍优于原全模型。
• 短视频创作工具：一款视频编辑 SaaS 提供文生视频功能，视频长度延长时会暴露模型在去噪早期和后期表现不均衡的问题。CBS 通过将模型拆分为 3 个子网络，分别负责初始噪声布局、主体生成和细节精修，显著降低了画面闪烁和物体变形，并可根据视频风格分层加载不同子网络，实现风格化视频的快速定制。

局限

• **辅助模型引入的额外开销与泛化性瓶颈**：CBS 依赖轻量级辅助模型来估计时间线上的复杂度分布，虽然避免了启发式搜索，但该辅助模型需要在每个新任务或新架构上重新训练，且对噪声调度、数据分布敏感性较高。当生成管线变化时，之前训练的复杂度轮廓可能不再适用，这降低了方法的即插即用性。考虑到实际部署中生成模型常需微调或迁移，辅助模型的重训练成本可能侵蚀了框架的低开销优势。
• **监控函数的选择局限性**：所提出的 Dirichlet 能量和路径加速度作为复杂度度量，虽有理论动机（函数逼近与等分布原理），但它们可能无法完全反映不同架构下的建模难度。例如，基于注意力的网络对局部空间变化的敏感度不同于 UNet，而路径加速度在高维空间中计算代价不菲。论文未系统论证这两类度量在所有扩散模型设计中的充分性，也未探讨更优监控函数的存在性，这限制了理论基础的坚实度。
• **评估任务与规模的覆盖不足**：实验集中在中等分辨率的图像生成（ImageNet 256、64 与 CIFAR-10），虽然展示了 FID 提升，但未涉及高分辨率生成、文生图、视频或音频等更长时序建模场景。在这些场景中，扩散阶段的计算图更复杂，CBS 的分割策略及效率优势是否维持仍未知。此外，未探究更大规模网络或新范式（如基于 Transformer 的扩散模型）下，划分数量与性能的扩展规律，减弱了产业落地的参考价值。