动态分块扩散变换器：通过自适应计算实现弹性推理的视觉生成 Dynamic Chunking Diffusion Transformer
通过学习路由的动态分块机制替换固定patch化，实现扩散变换器的自适应计算分配
前置知识
Diffusion Transformer (DiT)
DiT是将变换器架构应用于扩散模型的架构，通过固定patch化操作将输入图像转换为固定网格的token序列。在每个去噪步骤中，模型对这个固定长度的token序列进行处理。标准DiT使用adaLN-Zero进行类条件控制，采用线性噪声调度，在Stable Diffusion VAE的潜空间中操作。这种固定token化方式导致不同区域（平滑背景、高频细节）和不同时间步（噪声大的早期、细节精炼的晚期）都分配相同的计算预算。
本文的核心改进就是针对DiT的固定patch化机制，理解DiT的基本架构和工作原理是理解DC-DiT创新点的前提。
Classifier-Free Guidance (CFG)
CFG是扩散模型中广泛使用的条件生成技术。在每个采样步骤中，模型同时执行条件分支（以类别标签为条件）和无条件分支（无类别信息），然后通过引导系数将两个预测结合：$\epsilon = \epsilon_{\text{uncond}} + w \cdot (\epsilon_{\text{cond}} - \epsilon_{\text{uncond}})$，其中$w$是引导强度。这种双分支评估方式导致计算量翻倍。
DC-DiT提出了Lite-CFG，通过给条件分支和无条件分支分配不同的尾丢弃比例来优化CFG的计算效率，这是论文的重要创新之一。
Token Merging (ToMe)
Token Merging是一种推理时加速技术，通过合并冗余token来减少序列长度。它基于token间的相似度，将相似的token合并为一个代表token，从而减少后续层需要处理的token数量。类似的技术还包括token剪枝（直接丢弃不重要的token）和注意力稀疏化（限制每个token只关注部分键值对）。
这些技术代表了当前提升扩散变换器效率的主要方向，DC-DiT的创新在于在token化阶段就实现自适应，而非在固定token序列上进行后处理。
Elastic Inference（弹性推理）
弹性推理是指单个训练好的模型检查点可以在多个计算预算下进行评估，而无需重新训练。通过调整推理时的token丢弃比例，可以在质量和效率之间实现平滑的权衡曲线。这类似于神经网络中的宽度或深度可调性，但在token序列长度维度上实现。
这是DC-DiT的核心特性之一，通过尾丢弃机制实现，使得一个模型检查点可以适应不同的硬件资源和实时性需求。
研究动机
当前的扩散变换器（DiT）依赖固定的patch化操作，将每个图像转换为相同数量的token网格。这种设计使得token数量成为架构常数，忽略了图像生成中的两个自然自适应性：（1）不同空间区域包含不同细节量——平滑背景区域、高频物体边界、纹理区域的信息密度差异巨大；（2）扩散轨迹从粗到细的结构——噪声大的早期时间步与细节精炼的晚期时间步对计算资源的需求截然不同。现有的动态计算方法（如DyDiT的token剪枝、Token Merging的token合并、TeaCache的时间步缓存）虽然减少了计算冗余，但都局限于固定的token序列内部操作。这些方法在固定token序列形成后才进行适应性调整，无法从根本上解决token序列本身静态定义的问题。
本文的目标是本文的目标是开发一种能够在token化阶段就实现自适应的扩散变换器架构。具体而言，DC-DiT旨在：（1）学习根据图像内容动态分配token预算——对可预测区域压缩更多token，对细节区域保留更多token；（2）根据去噪时间步调整计算分配——早期噪声大的步骤使用更少token，后期细节精炼步骤使用更多token；（3）通过多预算训练实现弹性推理，使得单个检查点可以在多个计算预算下评估；（4）兼容现有的加速技术，如DyDiT和TeaCache，实现额外的计算节省。
与已有工作不同的是，DC-DiT的独特切入角度在于将自适应性前移到token化阶段，而不是在固定token序列上进行后处理。现有的动态计算方法（如DyDiT、DiffCR、SparseDiT）在固定的patch网格内进行token选择或合并，而DC-DiT通过学习encoder-router-decoder架构，从一开始就决定保留哪些token。这种设计使得token化本身变得可学习，通过端到端的扩散训练来优化，无需额外的监督信号（如边界检测、分割标注）。此外，DC-DiT的路由器基于局部空间可预测性：难以从邻域预测的token被分配高边界概率，而局部冗余token被分配低概率。这种基于可预测性的路由机制自然产生了有意义的空间分割和时间步自适应压缩调度。
核心方法
DC-DiT的核心思路是将DiT的固定patch化替换为一个学习的encoder-router-decoder支架，使token化本身变得可学习。整体流程是：首先，一个各向同性编码器混合局部信息以支持路由决策；然后，路由器预测每个token的边界概率——基于局部空间可预测性，难以从邻域预测的token获得高概率；弹性分块层根据边界概率保留一个缩短的token序列；这个缩短的序列被送入DiT主干网络处理；去分块层恢复原始分辨率；最后解码器生成扩散预测。关键创新在于：（1）路由器学习在空间上分配更多token给细节区域、更少token给平滑背景；（2）在时间上，早期噪声大的步骤使用更少token，后期步骤使用更多token；（3）通过多预算训练实现弹性推理。
DC-DiT与已有方法的本质区别在于将自适应性前移到token化阶段。现有的动态计算方法（DyDiT的token剪枝、Token Merging的token合并）都在固定的token序列上操作，而DC-DiT通过学习路由器从一开始就决定保留哪些token。路由器的设计基于局部空间可预测性：编码器混合邻域信息后，难以预测的token（如边缘、纹理）获得高边界概率，可预测的token（如平滑背景）获得低概率。这种设计无需外部监督（如分割标注），完全通过扩散训练目标学习。此外，路由器的边界概率提供了token的重要性排序，使得尾丢弃机制能够实现弹性推理——单个检查点可以在多个计算预算下评估。多预算训练通过在训练时随机采样不同的尾丢弃比例，避免了训练-测试不匹配。
方法步骤详情
DC-DiT的方法包含以下步骤：（1）编码阶段：各向同性编码器将输入token序列 $X \in \mathbb{R}^{B \times L \times d}$ 处理为特征 $X_{\text{enc}} \in \mathbb{R}^{B \times L \times d}$，编码器使用卷积残差块混合局部空间信息，中间隐藏宽度为主干维度的四分之一以减少开销。（2）路由阶段：路由器将编码序列重塑为 $H \times W$ 特征网格，投影到瓶颈表示 $z$，使用轻量级 $3 \times 3$ 卷积预测器从局部上下文估计每个特征得到 $\hat{z}$，残差 $r = z - \hat{z}$ 衡量局部不可预测性，卷积评分器 $f(\cdot)$ 将残差映射为logit，边界概率 $p_i = \sigma(f_{\text{score}}(r)_i)$。（3）弹性分块：通过阈值 $p_i > 0.5$ 生成硬边界掩码，使用直通估计器（STE）保持可微性——前向传播使用硬掩码，反向传播梯度通过连续概率 $p_i$ 流动。尾丢弃机制按 $p_i$ 排序保留的token，丢弃最低 $\rho|B|$ 个token，其中 $\rho$ 是用户指定的丢弃比例。（4）主干处理：DiT主干在缩短的token序列上处理，添加基于原始2D网格位置的正弦位置嵌入。（5）去分块：通过空间平滑和插回映射恢复原始分辨率——高置信度token保留原始特征，低置信度token与邻近token混合，每个原始位置接收最近边界token的表示。（6）解码：解码器生成扩散预测，添加编码器输出的残差连接以保留细粒度空间信息。
技术新颖性
DC-DiT的技术新颖性体现在多个层面。首先，这是首个将自适应token化应用于扩散变换器的工作，通过encoder-router-decoder支架实现端到端学习的token化。其次，路由器基于局部空间可预测性的设计是独特的——通过测量token与邻域预测的残差来决定保留哪些token，无需外部监督。第三，弹性推理机制通过尾丢弃和多预算训练实现，使得单个检查点可以在多个计算预算下评估，这在扩散模型中是新颖的。第四，Lite-CFG通过给条件分支和无条件分支分配不同的计算预算，优化了CFG的效率。第五，DC-DiT可以从预训练的DiT检查点进行轻量级微调（upcycling），保留预训练的好处同时获得自适应token化能力。最后，DC-DiT与现有加速技术（DyDiT、TeaCache）正交，可以组合使用获得额外收益。
实验结果
DC-DiT在类条件ImageNet生成任务上展示了显著的改进。在256x256分辨率下，DC-DiT-S在无CFG设置中将FID从68.40降低到58.43（-14.6%），同时FLOPs从3.02增加到3.06；在尾丢弃0.5时，FLOPs降低到2.14（-29.1%），FID为66.56。DC-DiT-B在尾丢弃0.2时FID从43.47降低到33.91（-22.0%），FLOPs为11.94；在尾丢弃0.5时FLOPs降低到8.55（-25.4%），FID为44.64。DC-DiT-L在尾丢弃0.1时FID从23.33降低到22.32（-4.3%），FLOPs为39.61；在尾丢弃0.3时FLOPs降低到31.86（-20.8%），FID为23.94。在Lite-CFG（CFG=1.25）设置下，DC-DiT-S在尾丢弃0.9时FLOPs从6.04降低到4.73（-21.7%），FID从54.98降低到38.99（-29.1%）。DC-DiT-B在尾丢弃0.9时FLOPs从22.93降低到18.38（-19.8%），FID从29.72降低到18.49（-37.8%）。DC-DiT-L在尾丢弃0.9时FLOPs从80.41降低到52.88（-34.2%），FID从12.59降低到10.17（-19.2%）。DC-DiT-XL在尾丢弃0.9时FLOPs从118.13降低到76.13（-35.6%），FID从9.45降低到8.80（-6.9%）。在512x512分辨率下，DC-DiT-XL在Lite-CFG设置中FLOPs从522.50降低到330.00（-36.8%），FID从11.04增加到12.50。路由器学习到有意义的时空压缩模式：空间上，对象边缘、纹理等高细节区域获得高边界概率；时间上，早期噪声步骤保留更少token，后期步骤保留更多token。DC-DiT可以与DyDiT和TeaCache组合使用获得额外计算节省。在Z-Image上进行upcycling后，DC-Z-Image在尾丢弃30%时将FLOPs从5901降低到4211，DPG-Bench从86.82略微提升到87.64；在尾丢弃60%时FLOPs进一步降低到2570，DPG-Bench为86.32。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet 256x256 类条件生成（无CFG） | FID-50K | DC-DiT-S: 58.43 | DiT-S/2: 68.40 | -14.6% |
| ImageNet 256x256 类条件生成（Lite-CFG, CFG=1.25） | FID-50K | DC-DiT-B: 18.49 | DiT-B/2: 29.72 | -37.8% |
| ImageNet 256x256 类条件生成（Lite-CFG, CFG=1.25） | TFLOPs/img | DC-DiT-L: 52.88 | DiT-L/2: 80.41 | -34.2% |
| ImageNet 512x512 类条件生成（Lite-CFG, CFG=1.25） | TFLOPs/img | DC-DiT-XL: 330.00 | DiT-XL/2: 522.50 | -36.8% |
| Z-Image文本到图像生成（upcycling） | DPG-Bench | DC-Z-Image (tail drop 30%): 87.64 | Z-Image: 85.97 | +1.9% |
| Z-Image文本到图像生成（upcycling） | TFLOPs/img | DC-Z-Image (tail drop 60%): 2570 | Z-Image: 5603 | -54.1% |
局限与改进
尽管DC-DiT取得了显著成果，但仍存在一些局限性。首先，DC-DiT的encoder-router-decoder支架引入了额外的参数开销——例如DC-DiT-XL比DiT-XL多14M参数（689M vs 675M），虽然相对较小但并非零成本。其次，encoder和decoder的操作在序列压缩前后进行，其计算开销不可忽略，论文使用主干维度的四分之一作为中间隐藏宽度来减少开销，但这可能限制了编码器和解码器的表达能力。第三，多预算训练需要扩展训练步数以匹配基线的计算预算，因为尾丢弃减少了每次训练步骤的有效FLOPs，这增加了训练成本。第四，DC-DiT在XL模型尺度上虽然使用更低FLOPs达到了竞争性性能，但FID改进不如S、B、L尺度显著，可能暗示在更大模型上自适应token化的收益递减。第五，论文仅在类条件ImageNet生成和Z-Image上进行评估，尚未在更大规模的文本到图像生成基准（如COCO、PartiPrompts）上验证。此外，路由器的边界概率阈值0.5是固定的，最优阈值可能因数据集、分辨率、模型规模而异。最后，虽然论文展示了与DyDiT和TeaCache的兼容性，但未评估与更多动态计算方法（如Token Merging、SparseDiT）的组合效果。
独立分析的弱点
DC-DiT的几个弱点值得进一步分析和改进。首先，路由器的边界概率计算基于局部空间可预测性，但在高度纹理化或重复模式的区域（如草地、天空中的云），这种基于残差的方法可能无法准确区分重要和冗余token，可以考虑引入频率域分析或多尺度特征来改进路由决策。其次，尾丢弃机制按概率排序丢弃token，但这假设概率排序就是重要性排序，实际上可能存在低概率但对全局结构重要的token（如物体轮廓的稀疏部分），可以考虑引入任务感知的重要性度量。第三，去分块层的空间平滑使用高斯核，其带宽 $\sigma$ 是固定的，不同分辨率和内容可能需要不同的平滑程度，可以考虑学习自适应带宽。第四，多预算训练从固定集合 $\mathcal{R}$ 中采样 $\rho$，这种离散采样可能不如连续分布平滑，可以考虑使用重要性采样或课程学习策略。第五，论文未充分分析路由器的计算开销本身——路由器的卷积预测器和评分器增加了延迟，虽然FLOPs减少但实际加速比可能不如论文所示的FLOPs减少比例。
未来方向
基于DC-DiT的成果，可以探索多个未来研究方向。首先，将DC-DiT扩展到视频生成领域——视频的时间维度提供了额外的压缩机会，可以学习在时间和空间维度上同时进行自适应token化。其次，探索DC-DiT在3D生成任务中的应用，如NeRF或点云生成，其中3D结构的空间变化可能从自适应token化中受益。第三，将DC-DiT与更大规模的文本到图像模型（如SD3、FLUX）结合，在实际应用中验证其效果。第四，研究路由器的可解释性——边界概率图可以用于理解模型对图像内容的理解，可能应用于图像编辑、分割、显著性检测等下游任务。第五，探索更细粒度的自适应机制，如在不同注意力头或不同网络深度上使用不同的路由策略。第六，研究DC-DiT与其他效率技术（如量化、蒸馏）的组合，实现更极致的效率提升。第七，将自适应token化思想应用于其他模态的扩散模型，如音频生成、3D形状生成等。
复现评估
论文提供了较为详细的复现信息。代码方面，虽然论文未明确提及开源计划，但方法描述足够详细，包括架构细节（编码器使用卷积残差块、路由器的边界概率计算、尾丢弃机制）、训练超参数（全局batch size 256、AdamW优化器、学习率 $1 \times 10^{-4}$、比率损失权重 $\lambda=0.03$、目标压缩比 $N=4$）、以及多预算训练的尾丢弃比例集合 $\{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6\}$。数据方面，实验使用ImageNet数据集，这是公开可用的。算力方面，论文在AMD Instinct MI325X和MI300X GPU上训练，对于S、B、L、XL多个模型尺度进行了实验，训练成本与对应的DiT基线匹配。复现难度中等——主要挑战在于实现序列打包和FlashAttention的变量长度注意力支持，以及正确实现STE的梯度传播。总体而言，方法的核心思想清晰，技术细节充分，具备良好的可复现性。
论文图表
