一个反直觉的发现

扩散Transformer的隐藏瓶颈：块间表示多样性如何决定生成质量

作者

快速导读：提出Weighted Diversity Score (WDS)量化DiT块间表示多样性，并基于此设计DiverseDiT++框架，通过长残差连接和表示多样性损失提升生成质量。

扩散Transformer（DiT）已成为视觉合成领域的核心架构，但其内部表示学习机制长期处于黑箱状态。现有方法如REPA依赖外部预训练编码器来对齐中间表示，虽然有效，却未能解释为何这种对齐能提升生成质量。DiverseDiT++论文首次从块间表示多样性的角度给出了系统性的回答。

论文的核心发现是：DiT中不同Transformer块学习到的表示应当保持足够的差异性，而这种块间多样性与最终生成质量之间存在强相关性。基于这一洞察，作者提出了Weighted Diversity Score (WDS)作为量化指标，并设计了DiverseDiT++框架，通过长残差连接和表示多样性损失来显式促进块间多样性，在多个视觉和科学任务上取得了显著的性能提升。

英文题目：DiverseDiT++: Quantifying, Analyzing, and Promoting Representation Diversity in Diffusion Transformers

论文出处：arXiv 每日论文精选 · arXiv:2608.03082

原始论文：PDF / 论文页面

这篇论文解决什么问题？

扩散Transformer将传统的U-Net骨干网络替换为纯Transformer架构，在ImageNet等基准上展现了优异的生成能力。然而，与CNN-based的扩散模型不同，DiT内部各Transformer块如何协同工作、学习到怎样的表示，一直缺乏系统性的理解。这种认知的缺失导致模型设计高度依赖经验调参，难以针对性地优化。

REPA（Representation Alignment）是近年来代表性的改进方法，它通过引入外部预训练编码器（如DINOv2）来对齐DiT中间层的表示，显著提升了生成质量。但REPA的成功机制并不清晰：是对齐本身带来了提升，还是对齐过程中意外触发了其他有益特性？论文通过CKA分析发现，REPA的有效性实际上源于它促进了不同块之间的表示多样性——对齐操作使得被对齐的块与其他块之间的表示差异增大。

这一发现引出了一个更根本的问题：块间表示多样性是否本身就是DiT有效学习的关键因素？如果是，能否不依赖外部模型，直接通过架构设计和训练目标来促进这种多样性？论文围绕这两个问题展开了系统性的探索。

值得注意的是，块间多样性的重要性并非DiT独有。在深度神经网络的理论研究中，不同层学习不同层次的特征是模型表达能力的来源。但在扩散模型的迭代去噪场景下，各块面对的是不同噪声水平的输入，其表示学习的动态过程更为复杂，需要专门的量化工具和优化策略。

核心创新

  • 首次系统分析DiT表示学习动态，揭示块间表示多样性是有效学习的关键因素。
  • 提出Weighted Diversity Score (WDS)指标，量化块间表示多样性，与log(FID)强相关（Pearson's r = -0.869）。
  • 设计DiverseDiT++框架，通过长残差连接和表示多样性损失显式促进块间多样性，无需外部模型。
  • 将方法泛化至蛋白质反向折叠和3D分子生成等科学任务，验证表示多样性作为通用原则。

方法概览

提出DiverseDiT++框架，包含两个核心组件：1) 长残差连接，将浅层输出注入深层以多样化块输入；2) 表示多样性损失，由正交损失、互信息最小化损失和特征分散损失组成，并采用自适应权重策略避免梯度冲突。

  • Weighted Diversity Score (WDS)：基于Centered Kernel Alignment (CKA)计算相邻块之间的表示相似性，然后对不同块对的相似性进行加权求和。权重设计使得相邻块之间的差异贡献更大，因为信息在相邻块间流动最为直接。WDS值越高，表示块间多样性越大。论文在97种不同训练设置下验证了WDS与log(FID)的Pearson相关系数达到-0.869，表明块间多样性确实与生成质量强相关。
  • 长残差连接（Long Residual Connections）：不同于标准DiT中仅有的块内残差连接，DiverseDiT++引入跨块的长跳跃连接，将浅层块的输出直接注入到深层块的输入中。这种设计使得深层块接收到的输入不再仅仅是前一块的输出，而是混合了来自多个浅层块的信息，从而自然地增加了不同块输入表示的多样性。具体实现上，论文采用可学习的加权融合方式，而非简单的恒等映射。
  • 表示多样性损失（Representation Diversity Loss）：由三个互补的损失项组成。正交损失鼓励不同块输出的特征向量相互正交，减少线性相关性；互信息最小化损失通过对比学习框架降低块间表示的统计依赖性；特征分散损失推动各块的特征分布在空间中占据不同的区域。三项损失的组合确保了多样性在多个层面上得到促进。
  • 自适应权重策略：直接联合优化三个损失项可能导致梯度冲突，使训练不稳定。论文设计了自适应权重机制，根据各损失项当前的梯度方向和大小动态调整权重，确保它们协同工作而非相互抵消。权重范围需要手动设定，但在此范围内是自动调节的。
  • 与外部方法的兼容性：DiverseDiT++的设计使其可以与REPA、SRA等依赖外部模型的方法叠加使用。实验表明，在REPA基础上添加DiverseDiT++的组件可以进一步提升性能，说明内部多样性促进和外部表示对齐是互补的优化方向。
  • 跨领域泛化设计：在蛋白质反向折叠任务中，长残差连接被适配到基于坐标的Transformer架构中，多样性损失则作用于氨基酸残基的中间表示。在3D分子生成任务中，多样性损失被应用于原子类型和坐标的联合表示空间。这些适配保持了核心思想的统一性，同时尊重了各领域的特定需求。
  • 计算开销控制：长残差连接引入的额外参数主要集中在融合权重的学习上，数量远小于主网络的参数量。多样性损失的计算主要在训练阶段进行，不增加推理时的计算负担。论文报告了具体的额外训练时间开销，表明其在实际应用中是可行的。
  • 层选择策略：并非所有块都需要参与长残差连接。论文通过实验发现，选择性地将浅层（如前1/3的块）连接到深层（如后1/3的块）效果最佳，中间块保持标准连接。这种稀疏连接设计在效果和效率之间取得了平衡。

逐图理解论文

一个反直觉的发现

一个反直觉的发现
Fig. 2. CKA representation similarities of models trained on various settings. We can observe that 1) the discrepancies between different blocks increases as training progresses; 2) aligning specific blocks significantly increases the dissimilarity between the corresponding block and other blocks; 3) aligning on more blocks with different pretrained encoders brings marginal performance improvements. Detailed quantitative results are provided in Table I.

图2展示了不同训练设置下DiT各块之间的CKA表示相似性矩阵。可以观察到：随着训练进行，不同块之间的差异增大（矩阵非对角线区域颜色变浅）；REPA对齐特定块后，该块与其他块的相似性显著降低。这直接启发了块间多样性的研究。

量化多样性与核心贡献

量化多样性与核心贡献
Fig. 4. WDS and log(FID) correlation for SiT-XL at 255×256 resolution.

图4展示了WDS与log(FID)的散点图，覆盖97种不同的训练设置。强负相关关系（Pearson's r = -0.869）直观地证明了块间多样性越高，生成质量越好。这一图是全文核心论点的实证基础。

如何验证与效果

如何验证与效果
Fig. 6. Generated samples from different training iterations. Images are sampled using the same seed, noise and class label. We use a classifier-free guidance scale of 4.0 during sampling.

图6展示了不同训练迭代次数下的生成样本对比。使用相同的噪声、种子和类别标签，DiverseDiT++在相同训练量下生成质量明显更优，直观验证了方法的加速收敛效果。

结论与意义

结论与意义
Fig. 5. Detailed diagram of our DiverseDiT++. Our method incorporates long residual connections to diversify input representations across blocks and a representation diversity loss to encourage blocks to learn distinct features.

图5详细展示了DiverseDiT++的架构。左侧为长残差连接的示意图，显示浅层输出如何注入深层；右侧为多样性损失的三个组件及其作用位置。该图是理解方法整体设计的关键。

实验如何设计？

  • 主要实验在ImageNet 256×256和512×512分辨率上进行，基础架构采用SiT-XL，同时也在REPA、MeanFlow等变体上验证了方法的通用性。
  • 训练配置包括400K迭代的无CFG设置和200/800 epoch的带CFG设置，覆盖了从快速验证到完全收敛的不同训练规模。
  • 评估指标包括FID、sFID、IS等标准生成质量指标，以及Precision和Recall以衡量生成样本的保真度和多样性平衡。
  • 科学任务实验包括蛋白质反向折叠（PDB数据集，评估序列恢复率和scRMSD）和3D分子生成（GEOM-Drug和QM9数据集，评估分子稳定性和能量指标）。
  • 消融实验系统性地移除了多样性损失的三个组件、长残差连接、自适应权重策略，以量化各组件的独立贡献。
  • 与DispLoss和SRA的兼容性实验验证了DiverseDiT++作为即插即用模块的灵活性。

关键结果与论文证据

  • WDS与log(FID)的Pearson相关系数为-0.869（第6页），表明块间表示多样性可以解释约75%的生成质量方差，是极为可靠的诊断指标。
  • 在SiT-XL上，DiverseDiT++在400K迭代无CFG设置下将FID从基线的15+降至12.38（第8页），在REPA-XL上进一步降至8.06，证明了方法的叠加效益。
  • 在200 epoch带CFG的完全训练设置下，DiverseDiT++在256×256分辨率上达到FID 1.49，800 epoch时进一步降至1.43（第9页），与同期最先进方法持平或更优。
  • 在512×512分辨率上，200 epoch训练即达到FID 1.97（第9页），验证了方法对高分辨率场景的适应性。
  • 在MeanFlow-XL/2上，DiverseDiT++将一步生成的FID降至2.94（第10页），表明多样性促进对少步采样同样有效。
  • 蛋白质反向折叠任务中，序列恢复率达到38.30%，scRMSD降至1.81（第11页），均优于基线方法。
  • 分子生成任务中，在GEOM-Drug数据集上分子稳定性达到97.75%（第11页），表明多样性损失有助于学习更符合物理约束的表示。
  • 消融实验显示，三个多样性损失组件各自独立贡献，但联合使用效果最佳；长残差连接单独使用即可带来显著提升，与多样性损失结合时产生协同效应（第11页）。

阅读时需要注意

  • WDS的计算依赖CKA，其敏感性在极大模型（如超过100层的Transformer）或非标准架构上尚未充分验证，可能需要在新的场景下重新校准。
  • 多样性损失的自适应权重范围需要手动设定，虽然论文给出了推荐范围，但在新任务上可能需要额外的调参成本。
  • 科学任务的评估仅覆盖了蛋白质和分子生成两个领域，在其他科学模态（如材料、气象）上的泛化性有待验证。
  • 长残差连接的设计目前基于均匀间隔的层选择策略，对于非均匀深度的架构可能需要重新设计连接模式。

关联工作

  • REPA（Yu et al., 2024）：通过外部DINOv2编码器对齐DiT中间表示，是本文最重要的对比基线和分析对象。论文揭示了REPA的有效性源于其意外地促进了块间多样性。
  • DispLoss（Chen et al., 2024）：提出分散损失来分离内部表示，但聚焦于单个块内的特征分散，未考虑块间关系。本文的多样性损失与之兼容，可叠加使用。
  • SRA（Kim et al., 2024）：通过师生模型对齐表示，无需外部模型。本文方法与其互补，联合使用可进一步提升性能。
  • CKA（Kornblith et al., 2019）：用于量化神经网络表示相似性的经典工具，本文在其基础上设计了WDS指标，将其从分析工具升级为诊断和优化工具。

发表评论