ComboStoc：面向扩散生成模型的组合随机性 ComboStoc: Combinatorial Stochasticity for Diffusion Generative Models

训练阶段流程如下（以 SiT-XL/2 图像为例，3D 形状结构相同）：① 采样一批数据 $x_1\in\mathcal{D}$ 与高斯噪声 $z\sim\mathcal{N}(0,1)$；② 按配置生成张量 $\mathbf{t}$——unsync_none 为 $(N,1,1,1)$、unsync_patch 为 $(N,C,H,W)$、unsync_vec 为 $(N,C,1,1)$、unsync_all 为 $(N,C,H,W)$，ImageNet 上为稳定起见只把 50% 样本异步化、50% 保持同步（当作 curriculum），PartNet 因数据小则全部异步；③ 按 $\mathbf{x}_t=(1-\mathbf{t})\odot z+\mathbf{t}\odot x_1$ 计算张量化中间态；④ 把 $\mathbf{t}$ 送入新的 timestep embedding（图 5b）：先用 sine/cosine 把每个标量映到 4 维，再用 SiT 的 patchwise embedding 把局部 patch 池成 hidden dim $C_H$ 的 token，与类别标签一同通过 adaLN-zero 调制 Transformer block；⑤ 网络 $f_\theta(\mathbf{x}_t;c,\mathbf{t})$ 预测速度 $\hat{v}=x_1-z$；⑥ 计算 MSE 损失反传。推理阶段分两路——同步路径 Alg.1A 直接用 $\mathbf{t}^{(k)}=(t_k)\cdot\mathbf{1}$，与标准 SiT 完全一致，250 步 SDE 积分；异步路径 Alg.1B 给定 mask $m\in[0,1]^{\text{shape}(x)}$，初始点 $x^{(0)}=(1-m)\odot z+m\odot x_1$，每条维度用独立 $t^{(0)}_i=m_i$ 起步，步长 $\Delta t=(1-m_i)/K$，积分 $K$ 步；为避免异步下预测速度 $\hat{v}=x_1-z$ 仍走对角线、导致 Eq.9 的 off-diagonal 偏差，额外在每步加上补偿项 $v_{cmpn}=\delta(x_t)$ 的反方向梯度（式 10），把 off-diagonal 漂移沿势能 $\Phi=\frac{1}{2}\|\delta_t\|^2$ 的负梯度拉回对角。

技术新颖性可拆为四点：(i) 路径空间偏置的形式化证明——通过数值积分 + 梯度闭式推导 $(x_1-x)\cdot\nabla\rho>0$，把经验现象（FID 收敛慢、3D 生成失败）落地到可量化的密度公式 Eq.2-4 上，让后续工作有了直接复用的诊断工具。(ii) 时间步张量化——公式改动仅一行，但要求 timestep embedding 同步重构，作者设计了两层 MLP+sine/cosine+patch pooling 的轻量模块，使得 GFlops 仅从 237.34 升至 352.46、推理延迟几乎不变（48ms vs 49ms SiT vs 50ms DiT，表 8）。(iii) 异步推理与现有控制方法的统一——传统 inpainting 是二值 mask，ComboStoc 把 mask 推广到连续 $m\in[0,1]$，自然支持图 12 的「中心 $\approx 0.85$ 渐变到 0」软修复、图 13 的「空间 + 通道联合 mask」探索 VAE latent 语义，更支持图 14/15 的 3D 形状补全/装配，不需要任何任务特定训练。(iv) 正交性——ComboStoc 只改输入分布、不改网络结构或损失函数，因此理论上与 REPA/DeepFlow/RAE 等加速方法正交，为未来组合留下了空间。

独立分析的弱点与改进方向：(1) **Batch mixing 是次优策略**——文末 footnote 承认「ImageNet 上 50% 同步 50% 异步」是经验做法，初步测试显示全部异步配合某种 curriculum 可能更好，未来可探索动态比例调度（如前期 100% 异步稳定后逐步混入同步样本）。(2) **Embedding 容量可能欠拟合**——unsync_none（仅 $N\times 1\times 1\times 1$ 标量 timestep）反而比 baseline SiT 略差，作者归因为新 embedding 模块的 $C_C=4$ 太窄；改进方向是让 $C_C$ 可学习、并对不同维度用独立的可学习偏置。(3) **GFlops 上升 48%**——多出来的代价全在 conditioning tensor 的 MLP 上，可考虑把 patchwise embedding 与 image embedding 共享权重，或只在最后几层注入 $\mathbf{t}$。(4) **补偿项的解析形式过强假设**——off-diagonal drift 最小化依赖 $\delta(x_t)=x_t-x_1-\text{proj}(x_t-x_1)$ 的解析梯度，假设 $z,x_1$ 在推理时已知，实际异步 inpainting 下 $x_1$ 的未知维度会让补偿失效，应改为网络预测补偿量。(5) **结构化 3D 形状的旋转未支持**——Fig.15 把部件旋转作为已知输入而非生成目标，限制了装配场景的实用性；可通过把旋转矩阵拆为 6D continuous representation 纳入 $\mathbf{t}$ 张量化来补齐。(6) **多样性 vs 质量的 Pareto 仍可优化**——Tab.5 显示 ComboStoc COV/MMD 略逊于 StructRe，提示在保持均匀采样优势的同时引入部分部件层次信息可能进一步压缩差距。

作者明文给出的未来方向：(i) 把 ComboStoc 与层次化生成（如 StructRe 的 part rewriting、StructureNet 的 part tree）结合，扬结构正则 + 组合采样之长；(ii) 探索 ComboStoc 在分子对接、蛋白折叠等「强组合结构 + 科学意义」领域的应用（Corso 2022、Wu 2024、Yim 2024），这些任务的解空间维度与属性数比图像和 3D 形状更高、更结构化。论文作者强调 ComboStoc 与 REPA、DeepFlow、RAE 等加速方法「正交」，因此最直接的延伸是把张量化 timestep 与表征蒸馏、深度速度监督、预训练 encoder latent 结合，可预期在 ImageNet FID 上进一步压缩收敛步数。基于结果的衍生方向：(iii) 把分级控制扩展到视频——给不同时间帧分配不同 $t_0$，实现「先决定远景再生成近景」的长视频生成；(iv) 把 mask $m\in[0,1]$ 与 text prompt 联合编码，实现「在某些区域保留参考图、其它区域遵循 prompt」的局部编辑；(v) 把异步推理用于 inpainting 之外的图像翻译（image-to-image）任务，迁移时只需把 $x_1$ 换为条件图；(vi) 探索跨模态应用，把 timestep 张量化扩展到 3D occupancy / NeRF 等更结构化的表示。

复现评估总体良好：(1) **代码开源**——GitHub 仓库 https://github.com/Xrvitd/ComboStoc 已发布，PyTorch 实现，依赖 numpy 和 pytorch 的 broadcast semantics 即可实现张量化 timestep，无需修改核心网络结构。(2) **数据**——图像实验直接用公开 ImageNet 1.3M (256×256 中心裁剪)；3D 实验用 PartNet (18K shapes)，并复用了 Wang et al. 2025 公开的 pretrained point cloud VAE + transformer VAE，避免重新训练部件编码器。(3) **算力门槛偏高**——图像主实验 4×H100 GPU 训练 800K 迭代 7.5 天；3D 实验 4×A100 GPU 训练 1.5K epoch 约 3 天；小数据消融用 8×H20 训练 50K 步。一般研究者复现 Tab.1 主结果需要约 30×H100·天的算力。(4) **实现细节充分**——Alg.1 给出同步/异步推理伪代码，图 5 给出新的 timestep embedding 模块结构，Eq.5/10/11 给出训练/补偿/推理公式；超参（学习率 1e-4、AdamW、batch 256 image / 16 3D、CFG 范围 1.0/1.5/4.0、patch 2×2 image、256 parts 3D）均在 Sec.5.1 列出。(5) **难度**——中等偏上：核心改动只有一行公式与一个 embedding 模块，但异步推理的补偿项涉及 ODE/SDE 积分细节、GFlops 优化需要 GPU kernel 经验，建议先复现 unsync_none baseline（最简单）再逐步打开异步与分级控制。