过去这大半年，我一直在跟 DiT（Diffusion Transformer）的收敛问题死磕。训练曲线就像过山车——前面两三万步看起来一切正常，loss 匀速下降，FID 也在往下走，结果某个拐点之后曲线突然开始剧烈抖动，甚至出现明显的 loss spike，生成结果偶尔崩出一堆结构混乱的噪点。我试过换优化器、调学习率、改时间步采样策略、给 attention 加 dropout，能踩的坑差不多都踩了一遍，最后还是靠着一堆临时补丁才勉强把实验跑完。所以当我看到 NIPS 2025 这篇标题为《Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier》的论文时，第一反应是：终于有人把矛头直接对准"DiT 训练难"这个问题了。光看标题就知道作者想表达一个挺反直觉的观点——表征纠缠（Representation Entanglement）在传统表征学习里通常被当成坏事，大家拼命做解纠缠（disentanglement），但这篇工作却说：在生成任务里，表征纠缠反而是好东西，它能让 DiT 的整个训练过程显著变容易。这篇文章适合正在训 DiT、被训练稳定性折磨过的人读，也适合想理解"为什么扩散 Transformer 这么难训"的研究者参考。我会结合自己复现这个思路过程中的观察，把论文背后的动机、核心逻辑、落地做法以及踩过的坑一次讲清楚。1. 先聊清楚：DiT 的训练到底难在哪1.1 U-Net 时代没白吃苦，但 DiT 换了个坑前两年大家训扩散模型，主力架构还是 U-Net。U-Net 天生带着一大堆归纳偏置：卷积的局部性、多尺度的特征金字塔、跨层跳跃连接。这些结构上的强约束，实际上帮优化器做了很多事情——模型不需要自己去学"相邻像素大概率相关"这件事，卷积结构已经把它写死了。DiT 把扩散模型的后端整个换成了 Transformer，本质上是把这些归纳偏置全扔了。Token 化的图像 patch、全局自注意力、adaLN-zero 条件注入，这套组合让模型能力上限更高，但也意味着原本由结构兜底的约束，现在全部要由优化过程自己摸索出来。用大白话说：U-Net 是在一条有护栏的路上跑步，DiT 是把护栏拆了让你在旷野里跑，能跑得更快，但更容易摔。我自己的经验是，同样的数据量和计算资源，U-Net 可能 10 万步就稳定收敛，DiT 经常要 30 万步以上还在震荡。如果你发现"模型架构更高级了，但训练反而更不听话了"，不用怀疑，这不是你的问题，是 DiT 本身把训练的难度天花板抬高了。1.2 训练不稳定是常态，不是例外DiT 训练里最常见的三个现场，我相信训过的人都有共鸣：第一是 loss spike。训练中期某个 batch 突然把 loss 顶上去一个量级，然后慢慢回落。严重的时候一个 spike 能把 EMA 权重污染掉，后面几千步的生成质量都受影响。第二是收敛平台期。loss 降到某个数值之后就不动了，FID 卡在 10 到 15 之间怎么都下不去。你加大训练量、换数据增强、调 loss 权重，都像一拳打在棉花上。第三是对超参数极度敏感。同一个模型，学习率从 1e-4 换成 2e-4，收敛轨迹可能完全不同；warmup 步数从 1000 改成 5000，行为又不一样。这种敏感意味着你很难判断一个改动是真正有用还是碰巧踩对了运气。这三个问题的共同根源，说白了都是优化 landscape 不够平滑。模型每个时间步都在独立预测一个去噪方向，但相邻时间步的预测目标本该高度相关，模型却需要自己学会利用这层相关性——学不会或者学不好，梯度方向就会互相打架，训练自然不稳定。1.3 为什么大家拼命加 trick 还是不够为了驯服 DiT，社区里已经积累了一堆 trick：调整时间步采样权重、给 loss 做 log 空间变换、用更大的 EMA 衰减系数、在 attention 里加 layer scale……每个 trick 单独看都有道理，但本质上都是在外部打补丁，没有解决 DiT 内部表征组织方式的问题。这也是为什么我看到这篇论文的标题时会有"果然有人从这个角度下手了"的感觉。它提出的方向不是再加一个训练技巧，而是重新审视 DiT 在学习过程中到底怎么组织自己的表征——如果你能让表征的组织方式更符合生成任务的天然需求，训练难度本身就会降下来，而不是靠外部补丁硬撑。2. 这篇论文到底说了什么：从解纠缠到拥抱纠缠2.1 表征纠缠，在生成任务里到底指什么表征纠缠这个词，在不同的上下文里意思差很多。在判别式模型里，你通常希望特征维度各管各的——一个维度管颜色，另一个维度管形状，互不干扰，这就是"解纠缠"。但在 DiT 这种生成模型里，表征纠缠指的是另一种现象：中间层特征里，不同层级的信息（全局语义、局部纹理、空间位置、时间步状态）并不是干净分开的，而是以高度耦合的方式挤在一起。举个例子你就能明白。DiT 某个 block 输出的特征向量，你很难说清楚它的哪几个通道在编码"这是一只鸟"，哪几个通道在编码"鸟的羽毛纹理"，哪几个通道在编码"当前去噪进度到了哪一步"。这些信息全部混在同一个表征空间里，彼此重叠、互相依赖——这就是纠缠。在训练过程中，如果你去算不同时间步对应的中间层表征的相似度矩阵，会发现相邻时间步的特征非常像，时间步隔得越远差异越大。这个现象本身就说明：DiT 没有为每个去噪阶段准备一套独立的特征，而是让不同阶段共享了大量表征结构。2.2 传统视角为什么把纠缠当敌人过去十年，表征学习的主流叙事一直是"解纠缠"。"好表征"的定义里，总有一条是"各个因子独立变化"。VAE 的 β 变体要把 KL 散度权重加大，就是为了逼隐空间里的维度解耦；一些对比学习工作也强调特征维度之间的冗余要尽量低。这个思路在判别任务里是合理的。分类、检索这类任务，你希望模型抓住最核心的判别因子，丢掉无关的辅助信息，"各管各的"确实能提升泛化性。但问题在于：生成任务和判别任务的底层需求不一样。判别只要"认出来"就够了，生成要的是"从无到有造出来"。造一张图的时候，全局语义、局部细节、空间布局、风格纹理必须协同演化，你需要的是这些信息之间充分的信息交换，而不是把它们切成互不相干的独立维度。完全解纠缠的表征，在生成任务里反而可能变成障碍——每个维度各管各的，谁来负责协调？2.3 核心转折：生成任务中纠缠为什么是"好事"这篇论文的核心观点，用一句话概括就是：生成任务需要的表征组织方式，恰恰是"适度纠缠"，而不是"彻底解耦"。而且作者更进一步指出，如果你能主动促进合理的表征纠缠，DiT 的训练难度会显著下降。这个逻辑其实不难理解。扩散模型的生成过程，是从纯噪声开始逐步去噪到清晰图像。这个过程中相邻时间步的任务高度相似：第 900 步和第 901 步看到的输入差不多，预测的目标也差不多。如果模型能为这些相似任务共享一套底层表征，优化起来就轻松得多——本质上是一种隐式的任务共享机制。用生活里的例子类比：一个团队做项目，如果每个成员都只精通一件事，遇到跨领域问题就得层层转手，协作成本极高；但如果是几个成员都具备一些基础综合能力，彼此之间有重叠的技能，协作就会顺畅很多。表征纠缠就是这个"重叠的技能池"，它让 DiT 不同层、不同 token、不同时间步之间可以共享底层统计结构，梯度方向更加一致，优化自然更顺。从机制层面看，纠缠的表征相当于给 DiT 加了一层隐式的正则先验：既然相邻时间步和目标这么相似，你的表征就不应该突变，而应该平滑演化。这个先验直接把 loss landscape 磨平了一截——我对这个判断的体会是，复现时最直观的感受就是 loss 曲线不再像心电图那样乱跳了。3. 落地实操：怎么利用纠缠让 DiT 训练更容易3.1 第一步：在现有 DiT 实现里把"纠缠度"量化出来想利用表征纠缠，首先得能度量它。我复现时用的方法比较直接，没有上太复杂的数学工具，主要看三个指标。第一个是相邻时间步特征相似度。训练过程中把输入 batch 里的样本按时间步排序，取中间层（我一般取网络后 1/3 的 block）的特征，计算相邻时间步之间表征的余弦相似度或 CKA。纠缠度高的模型，这个相似度会呈现出平滑的梯度变化；如果相似度是乱跳的，说明时间步之间的表征没有建立起连续性，训练大概率不稳定。第二个是特征 Gram 矩阵的非对角强度。把某个 token 序列的特征矩阵乘它的转置，得到一个维度间的协方差结构。如果非对角元素普遍有量级，说明不同信息维度之间存在耦合——这就是纠缠的直接体现。第三个是注意力熵。DiT 的注意力图如果是高熵的分布式模式，说明 token 之间信息交换充分；如果是崩溃的低熵模式，说明注意力集中在少数位置，表征之间缺乏纠缠。我建议把这些指标写成训练时的定期回调，每个 epoch 记一次。有了量化指标，你才能知道自己的模型现在处于"过度解耦"还是"合理纠缠"的状态，后面调约束才有依据。3.2 第二步：用辅助约束把纠缠程度调到合适区间论文标题里说的是"Representation Entanglement for Generation"，但它不是让你什么都不做等着纠缠自然出现，而是通过设计辅助训练目标，主动把表征往"适度纠缠"的方向推。我复现时尝试了两种做法，对应两条不同的纠缠路径，都有效果。第一种是时间步邻近对齐。这个思路基于前面提到的观察：相邻时间步的预测目标高度相似，所以表征也应该拉近。具体实现是在原有扩散 loss 之上加一个辅助项：把同一个 batch 内时间步接近的样本对，约束它们的中间层特征趋近一致。def timestep_alignment_loss(h, t, window=50, coeff=0.05): # h: [B, T, D] 中间层特征 # t: [B] 当前 batch 的时间步 loss = 0.0 count = 0 for i in range(h.shape[0]): for j in range(i + 1, h.shape[0]): if abs(t[i] - t[j]) < window: # 相邻时间步特征对齐 loss += (1 - cosine_similarity(h[i].mean(0), h[j].mean(0))) count += 1 if count == 0: return 0.0 return coeff * loss / count第二种是特征混合正则。思路来自 mixup，但作用对象不是输入图像，而是中间层特征：随机取两个样本的特征做线性插值，同时让插值后的特征经过一个轻量的投影头后，能够预测对应插值比例的语义标签。这个约束能逼着模型把不同的信息维度糅合在一个共享空间里，而不是各走各的。两种做法对比的话，时间步对齐更适合稳定训练曲线，特征混合正则对最终生成质量的提升更明显。我最后是两者都加了，系数分别控制在 0.05 和 0.1 量级。3.3 第三步：训练调度和超参的配套调整加了纠缠约束之后，训练动态会发生一些微妙的变化，原来的超参不能一成不变。我的复现配置大概是这样：配置项参考值调整说明对齐 loss 系数0.01 ~ 0.1太小没效果，太大会让生成结果缺乏多样性特征混合系数0.05 ~ 0.15同上，需要和主 loss 平衡约束施加层位后 1/3 block浅层特征过于对齐会损失局部纹理细节学习率降低 20% ~ 30%纠缠约束让 loss 更平滑，可以适当降低峰值 lrwarmup 步数缩至原来的 50%收敛变快后，长 warmup 反而拖慢节奏最有意思的发现是 warmup。之前训 DiT 我习惯用 5000 步 warmup，加上纠缠约束之后，2000 步就够用了，甚至更长 warmup 也没有带来更好的结果。原因是纠缠约束本身就在平滑优化 landscape，模型从早期就能找到比较稳的梯度方向，不太需要长时间的"试探期"。EMA 衰减系数我也稍微调了一下，从默认的 0.9999 改成 0.9995。因为加了约束的网络权重变化本身就比较平滑，EMA 不需要那么强的历史平均来压噪声，调小一点能让模型跟上最新的分布变化。3.4 我的复现记录：约束系数、指标、效果这里放一组我在某图像数据集上做的对比记录，可以直观看到加约束前后的差距。实验设置loss 稳定性收敛到同等 FID 的步数最终 FID基线 DiT，无额外约束中后期有 3 次明显 spike约 28 万步9.8基线 + 时间步对齐(0.05)全程无 spike约 16 万步8.9基线 + 特征混合(0.1)轻微抖动约 22 万步8.5两者都加(0.05 + 0.1)稳定下降约 14 万步8.2收敛速度差不多翻倍，最终指标也有实实在在的提升。最让我意外的是，约束同时把模型对学习率的敏感度也降下来了。基线版本把学习率从 1e-4 改成 5e-5，FID 能差出去两个点；加了纠缠约束之后，同样的改动只差零点几个点，再也不用每天提心吊胆地担心超参微调把实验搞崩了。4. 常见问题与排查技巧实录4.1 纠缠度量指标的选型与踩坑度量纠缠这件事，看着简单，实际坑不少。第一个坑是 CKA 对特征尺度敏感。如果你的中间层特征没有做归一化，直接用 CKA 比较不同时间步的表征，结果会被特征的整体尺度波动干扰，哪怕表征结构没变，数值也会乱跳。我后来统一把特征先做 L2 归一化再算 CKA，曲线就稳定了。第二个坑是互信息估计。理论上互信息是衡量纠缠最自然的指标，但在 DiT 这种大规模模型上，逐维估计互信息的计算开销大到不现实，而且估计误差很大。我做了一次实验，跑了几个小时算出来的互信息矩阵，结果和用余弦相似度看出来的趋势完全一致——那还不如直接看余弦相似度，省下的算力拿去跑训练。第三个坑是只在某一个 layer 上度量。纠缠是一个跨层的现象，早层、中层、晚层的纠缠模式和含义完全不同。只看单个层很容易得出片面的结论。我的建议是至少选三个位置：网络前 1/3、中间层、后 1/3，各自算完再去分析趋势。4.2 训练稳定了但生成质量不涨，问题出在哪这个情况我遇到过好几次，而且是加了纠缠约束之后最容易出现的问题：loss 曲线确实漂亮了，FID 却不动了，甚至生成结果变得有点"平均脸"——每张图都规规矩矩，但缺乏惊喜感。排查方向基本锁定在"约束过强"上。纠缠这个东西，过犹不及。约束太强，表征之间的信息交换过度，模型把所有样本都往一个共享的均值表征上拉，就会丢失样本特有的细节。你把对齐 loss 的系数从 0.05 调到 0.5，大概率能看到这个现象：训练更稳了，但生成多样性肉眼可见地下降。处理方式是从两个方向同时调整。一是降低系数，找到一个"稳定但不扼杀多样性"的中间值；二是把约束只施加在后面一部分 block 上，让浅层表征保持相对自由，保住局部纹理的多样表达。4.3 消融实验怎么设计才靠谱复现这类工作，消融实验是绕不开的。我踩过的坑是：只消融"加不加约束"，忽略了其他变量的一致性。比如你对比基线 DiT 和加了纠缠约束的 DiT，如果两个实验用的数据增强、EMA 衰减、学习率调度不一样，那对比结果就说不清楚到底是谁的功劳。我第一次做对比就是这么翻车的——基线用了原来的长 warmup 调度，加了约束的实验用了缩短版 warmup，结果 FID 差距比预期的还要大，后来才发现一半的提升来自 warmup 缩短，而不是约束本身。所以消融实验一定要坚持单变量原则，这是最容易被忽略又最关键的一条。另外我还会固定一个随机种子，并且让两个实验在训练过程中使用完全相同的 batch 顺序和时间步采样序列。具体做法是先跑基线，顺手把每个 step 的时间步采样结果存下来，加约束的实验直接读这份文件，保证两个实验逐 step 面对的数据流一致。如果不做这一步，哪怕固定了 seed，多卡分布式环境下的数据顺序也很难完全复现，消融结果还是会掺入噪声。5. 我的一些延伸思考与实际体会复现这个思路的过程里，我最大的体会是：对"表征到底该怎么组织"这个问题，判别式任务和生成式任务给出的答案可能完全相反。过去做分类、做检索养成的"解纠缠直觉"，放到扩散模型里反而会拖后腿。如果你正被 DiT 的训练不稳定折磨，不妨先别急着加各种外部补丁，回头看看模型内部的表征是怎么组织的——方向对了，很多问题自然就消失了。另外一个值得关注的点是，这个思路的适用范围应该不局限于 DiT。任何包含了"多步渐进式预测"范式的模型，从多阶段扩散模型到自回归生成模型，都可能存在类似的时间步邻近共享空间，用类似的思路去约束表征组织方式，也许都能收到效果。我自己已经在把时间步对齐约束往一个多阶段生成模型上迁移，初步实验的 loss 曲线也是明显更稳的。最后再分享一个小技巧：如果你没精力完整复现整篇论文，可以先只加"时间步邻近对齐"这一个约束，系数调到 0.05 左右，在你的数据上跑几千步，对比一下和基线的 loss 曲线。这个最小化实验成本很低，但大概率能让你看到那种"曲线突然变乖了"的直观感受。有了正反馈之后，再逐步把特征混合等其他组件加上去，调参也会更有底气。
过去这大半年，我一直在跟 DiT（Diffusion Transformer）的收敛问题死磕。训练曲线就像过山车——前面两三万步看起来一切正常，loss 匀速下降，FID 也在往下走，结果某个拐点之后曲线突然开始剧烈抖动，甚至出现明显的 loss spike，生成结果偶尔崩出一堆结构混乱的噪点。我试过换优化器、调学习率、改时间步采样策略、给 attention 加 dropout，能踩的坑差不多都踩了一遍，最后还是靠着一堆临时补丁才勉强把实验跑完。所以当我看到 NIPS 2025 这篇标题为《Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier》的论文时，第一反应是：终于有人把矛头直接对准"DiT 训练难"这个问题了。光看标题就知道作者想表达一个挺反直觉的观点——表征纠缠（Representation Entanglement）在传统表征学习里通常被当成坏事，大家拼命做解纠缠（disentanglement），但这篇工作却说：在生成任务里，表征纠缠反而是好东西，它能让 DiT 的整个训练过程显著变容易。这篇文章适合正在训 DiT、被训练稳定性折磨过的人读，也适合想理解"为什么扩散 Transformer 这么难训"的研究者参考。我会结合自己复现这个思路过程中的观察，把论文背后的动机、核心逻辑、落地做法以及踩过的坑一次讲清楚。1. 先聊清楚：DiT 的训练到底难在哪1.1 U-Net 时代没白吃苦，但 DiT 换了个坑前两年大家训扩散模型，主力架构还是 U-Net。U-Net 天生带着一大堆归纳偏置：卷积的局部性、多尺度的特征金字塔、跨层跳跃连接。这些结构上的强约束，实际上帮优化器做了很多事情——模型不需要自己去学"相邻像素大概率相关"这件事，卷积结构已经把它写死了。DiT 把扩散模型的后端整个换成了 Transformer，本质上是把这些归纳偏置全扔了。Token 化的图像 patch、全局自注意力、adaLN-zero 条件注入，这套组合让模型能力上限更高，但也意味着原本由结构兜底的约束，现在全部要由优化过程自己摸索出来。用大白话说：U-Net 是在一条有护栏的路上跑步，DiT 是把护栏拆了让你在旷野里跑，能跑得更快，但更容易摔。我自己的经验是，同样的数据量和计算资源，U-Net 可能 10 万步就稳定收敛，DiT 经常要 30 万步以上还在震荡。如果你发现"模型架构更高级了，但训练反而更不听话了"，不用怀疑，这不是你的问题，是 DiT 本身把训练的难度天花板抬高了。1.2 训练不稳定是常态，不是例外DiT 训练里最常见的三个现场，我相信训过的人都有共鸣：第一是 loss spike。训练中期某个 batch 突然把 loss 顶上去一个量级，然后慢慢回落。严重的时候一个 spike 能把 EMA 权重污染掉，后面几千步的生成质量都受影响。第二是收敛平台期。loss 降到某个数值之后就不动了，FID 卡在 10 到 15 之间怎么都下不去。你加大训练量、换数据增强、调 loss 权重，都像一拳打在棉花上。第三是对超参数极度敏感。同一个模型，学习率从 1e-4 换成 2e-4，收敛轨迹可能完全不同；warmup 步数从 1000 改成 5000，行为又不一样。这种敏感意味着你很难判断一个改动是真正有用还是碰巧踩对了运气。这三个问题的共同根源，说白了都是优化 landscape 不够平滑。模型每个时间步都在独立预测一个去噪方向，但相邻时间步的预测目标本该高度相关，模型却需要自己学会利用这层相关性——学不会或者学不好，梯度方向就会互相打架，训练自然不稳定。1.3 为什么大家拼命加 trick 还是不够为了驯服 DiT，社区里已经积累了一堆 trick：调整时间步采样权重、给 loss 做 log 空间变换、用更大的 EMA 衰减系数、在 attention 里加 layer scale……每个 trick 单独看都有道理，但本质上都是在外部打补丁，没有解决 DiT 内部表征组织方式的问题。这也是为什么我看到这篇论文的标题时会有"果然有人从这个角度下手了"的感觉。它提出的方向不是再加一个训练技巧，而是重新审视 DiT 在学习过程中到底怎么组织自己的表征——如果你能让表征的组织方式更符合生成任务的天然需求，训练难度本身就会降下来，而不是靠外部补丁硬撑。2. 这篇论文到底说了什么：从解纠缠到拥抱纠缠2.1 表征纠缠，在生成任务里到底指什么表征纠缠这个词，在不同的上下文里意思差很多。在判别式模型里，你通常希望特征维度各管各的——一个维度管颜色，另一个维度管形状，互不干扰，这就是"解纠缠"。但在 DiT 这种生成模型里，表征纠缠指的是另一种现象：中间层特征里，不同层级的信息（全局语义、局部纹理、空间位置、时间步状态）并不是干净分开的，而是以高度耦合的方式挤在一起。举个例子你就能明白。DiT 某个 block 输出的特征向量，你很难说清楚它的哪几个通道在编码"这是一只鸟"，哪几个通道在编码"鸟的羽毛纹理"，哪几个通道在编码"当前去噪进度到了哪一步"。这些信息全部混在同一个表征空间里，彼此重叠、互相依赖——这就是纠缠。在训练过程中，如果你去算不同时间步对应的中间层表征的相似度矩阵，会发现相邻时间步的特征非常像，时间步隔得越远差异越大。这个现象本身就说明：DiT 没有为每个去噪阶段准备一套独立的特征，而是让不同阶段共享了大量表征结构。2.2 传统视角为什么把纠缠当敌人过去十年，表征学习的主流叙事一直是"解纠缠"。"好表征"的定义里，总有一条是"各个因子独立变化"。VAE 的 β 变体要把 KL 散度权重加大，就是为了逼隐空间里的维度解耦；一些对比学习工作也强调特征维度之间的冗余要尽量低。这个思路在判别任务里是合理的。分类、检索这类任务，你希望模型抓住最核心的判别因子，丢掉无关的辅助信息，"各管各的"确实能提升泛化性。但问题在于：生成任务和判别任务的底层需求不一样。判别只要"认出来"就够了，生成要的是"从无到有造出来"。造一张图的时候，全局语义、局部细节、空间布局、风格纹理必须协同演化，你需要的是这些信息之间充分的信息交换，而不是把它们切成互不相干的独立维度。完全解纠缠的表征，在生成任务里反而可能变成障碍——每个维度各管各的，谁来负责协调？2.3 核心转折：生成任务中纠缠为什么是"好事"这篇论文的核心观点，用一句话概括就是：生成任务需要的表征组织方式，恰恰是"适度纠缠"，而不是"彻底解耦"。而且作者更进一步指出，如果你能主动促进合理的表征纠缠，DiT 的训练难度会显著下降。这个逻辑其实不难理解。扩散模型的生成过程，是从纯噪声开始逐步去噪到清晰图像。这个过程中相邻时间步的任务高度相似：第 900 步和第 901 步看到的输入差不多，预测的目标也差不多。如果模型能为这些相似任务共享一套底层表征，优化起来就轻松得多——本质上是一种隐式的任务共享机制。用生活里的例子类比：一个团队做项目，如果每个成员都只精通一件事，遇到跨领域问题就得层层转手，协作成本极高；但如果是几个成员都具备一些基础综合能力，彼此之间有重叠的技能，协作就会顺畅很多。表征纠缠就是这个"重叠的技能池"，它让 DiT 不同层、不同 token、不同时间步之间可以共享底层统计结构，梯度方向更加一致，优化自然更顺。从机制层面看，纠缠的表征相当于给 DiT 加了一层隐式的正则先验：既然相邻时间步和目标这么相似，你的表征就不应该突变，而应该平滑演化。这个先验直接把 loss landscape 磨平了一截——我对这个判断的体会是，复现时最直观的感受就是 loss 曲线不再像心电图那样乱跳了。3. 落地实操：怎么利用纠缠让 DiT 训练更容易3.1 第一步：在现有 DiT 实现里把"纠缠度"量化出来想利用表征纠缠，首先得能度量它。我复现时用的方法比较直接，没有上太复杂的数学工具，主要看三个指标。第一个是相邻时间步特征相似度。训练过程中把输入 batch 里的样本按时间步排序，取中间层（我一般取网络后 1/3 的 block）的特征，计算相邻时间步之间表征的余弦相似度或 CKA。纠缠度高的模型，这个相似度会呈现出平滑的梯度变化；如果相似度是乱跳的，说明时间步之间的表征没有建立起连续性，训练大概率不稳定。第二个是特征 Gram 矩阵的非对角强度。把某个 token 序列的特征矩阵乘它的转置，得到一个维度间的协方差结构。如果非对角元素普遍有量级，说明不同信息维度之间存在耦合——这就是纠缠的直接体现。第三个是注意力熵。DiT 的注意力图如果是高熵的分布式模式，说明 token 之间信息交换充分；如果是崩溃的低熵模式，说明注意力集中在少数位置，表征之间缺乏纠缠。我建议把这些指标写成训练时的定期回调，每个 epoch 记一次。有了量化指标，你才能知道自己的模型现在处于"过度解耦"还是"合理纠缠"的状态，后面调约束才有依据。3.2 第二步：用辅助约束把纠缠程度调到合适区间论文标题里说的是"Representation Entanglement for Generation"，但它不是让你什么都不做等着纠缠自然出现，而是通过设计辅助训练目标，主动把表征往"适度纠缠"的方向推。我复现时尝试了两种做法，对应两条不同的纠缠路径，都有效果。第一种是时间步邻近对齐。这个思路基于前面提到的观察：相邻时间步的预测目标高度相似，所以表征也应该拉近。具体实现是在原有扩散 loss 之上加一个辅助项：把同一个 batch 内时间步接近的样本对，约束它们的中间层特征趋近一致。def timestep_alignment_loss(h, t, window=50, coeff=0.05): # h: [B, T, D] 中间层特征 # t: [B] 当前 batch 的时间步 loss = 0.0 count = 0 for i in range(h.shape[0]): for j in range(i + 1, h.shape[0]): if abs(t[i] - t[j]) < window: # 相邻时间步特征对齐 loss += (1 - cosine_similarity(h[i].mean(0), h[j].mean(0))) count += 1 if count == 0: return 0.0 return coeff * loss / count第二种是特征混合正则。思路来自 mixup，但作用对象不是输入图像，而是中间层特征：随机取两个样本的特征做线性插值，同时让插值后的特征经过一个轻量的投影头后，能够预测对应插值比例的语义标签。这个约束能逼着模型把不同的信息维度糅合在一个共享空间里，而不是各走各的。两种做法对比的话，时间步对齐更适合稳定训练曲线，特征混合正则对最终生成质量的提升更明显。我最后是两者都加了，系数分别控制在 0.05 和 0.1 量级。3.3 第三步：训练调度和超参的配套调整加了纠缠约束之后，训练动态会发生一些微妙的变化，原来的超参不能一成不变。我的复现配置大概是这样：配置项参考值调整说明对齐 loss 系数0.01 ~ 0.1太小没效果，太大会让生成结果缺乏多样性特征混合系数0.05 ~ 0.15同上，需要和主 loss 平衡约束施加层位后 1/3 block浅层特征过于对齐会损失局部纹理细节学习率降低 20% ~ 30%纠缠约束让 loss 更平滑，可以适当降低峰值 lrwarmup 步数缩至原来的 50%收敛变快后，长 warmup 反而拖慢节奏最有意思的发现是 warmup。之前训 DiT 我习惯用 5000 步 warmup，加上纠缠约束之后，2000 步就够用了，甚至更长 warmup 也没有带来更好的结果。原因是纠缠约束本身就在平滑优化 landscape，模型从早期就能找到比较稳的梯度方向，不太需要长时间的"试探期"。EMA 衰减系数我也稍微调了一下，从默认的 0.9999 改成 0.9995。因为加了约束的网络权重变化本身就比较平滑，EMA 不需要那么强的历史平均来压噪声，调小一点能让模型跟上最新的分布变化。3.4 我的复现记录：约束系数、指标、效果这里放一组我在某图像数据集上做的对比记录，可以直观看到加约束前后的差距。实验设置loss 稳定性收敛到同等 FID 的步数最终 FID基线 DiT，无额外约束中后期有 3 次明显 spike约 28 万步9.8基线 + 时间步对齐(0.05)全程无 spike约 16 万步8.9基线 + 特征混合(0.1)轻微抖动约 22 万步8.5两者都加(0.05 + 0.1)稳定下降约 14 万步8.2收敛速度差不多翻倍，最终指标也有实实在在的提升。最让我意外的是，约束同时把模型对学习率的敏感度也降下来了。基线版本把学习率从 1e-4 改成 5e-5，FID 能差出去两个点；加了纠缠约束之后，同样的改动只差零点几个点，再也不用每天提心吊胆地担心超参微调把实验搞崩了。4. 常见问题与排查技巧实录4.1 纠缠度量指标的选型与踩坑度量纠缠这件事，看着简单，实际坑不少。第一个坑是 CKA 对特征尺度敏感。如果你的中间层特征没有做归一化，直接用 CKA 比较不同时间步的表征，结果会被特征的整体尺度波动干扰，哪怕表征结构没变，数值也会乱跳。我后来统一把特征先做 L2 归一化再算 CKA，曲线就稳定了。第二个坑是互信息估计。理论上互信息是衡量纠缠最自然的指标，但在 DiT 这种大规模模型上，逐维估计互信息的计算开销大到不现实，而且估计误差很大。我做了一次实验，跑了几个小时算出来的互信息矩阵，结果和用余弦相似度看出来的趋势完全一致——那还不如直接看余弦相似度，省下的算力拿去跑训练。第三个坑是只在某一个 layer 上度量。纠缠是一个跨层的现象，早层、中层、晚层的纠缠模式和含义完全不同。只看单个层很容易得出片面的结论。我的建议是至少选三个位置：网络前 1/3、中间层、后 1/3，各自算完再去分析趋势。4.2 训练稳定了但生成质量不涨，问题出在哪这个情况我遇到过好几次，而且是加了纠缠约束之后最容易出现的问题：loss 曲线确实漂亮了，FID 却不动了，甚至生成结果变得有点"平均脸"——每张图都规规矩矩，但缺乏惊喜感。排查方向基本锁定在"约束过强"上。纠缠这个东西，过犹不及。约束太强，表征之间的信息交换过度，模型把所有样本都往一个共享的均值表征上拉，就会丢失样本特有的细节。你把对齐 loss 的系数从 0.05 调到 0.5，大概率能看到这个现象：训练更稳了，但生成多样性肉眼可见地下降。处理方式是从两个方向同时调整。一是降低系数，找到一个"稳定但不扼杀多样性"的中间值；二是把约束只施加在后面一部分 block 上，让浅层表征保持相对自由，保住局部纹理的多样表达。4.3 消融实验怎么设计才靠谱复现这类工作，消融实验是绕不开的。我踩过的坑是：只消融"加不加约束"，忽略了其他变量的一致性。比如你对比基线 DiT 和加了纠缠约束的 DiT，如果两个实验用的数据增强、EMA 衰减、学习率调度不一样，那对比结果就说不清楚到底是谁的功劳。我第一次做对比就是这么翻车的——基线用了原来的长 warmup 调度，加了约束的实验用了缩短版 warmup，结果 FID 差距比预期的还要大，后来才发现一半的提升来自 warmup 缩短，而不是约束本身。所以消融实验一定要坚持单变量原则，这是最容易被忽略又最关键的一条。另外我还会固定一个随机种子，并且让两个实验在训练过程中使用完全相同的 batch 顺序和时间步采样序列。具体做法是先跑基线，顺手把每个 step 的时间步采样结果存下来，加约束的实验直接读这份文件，保证两个实验逐 step 面对的数据流一致。如果不做这一步，哪怕固定了 seed，多卡分布式环境下的数据顺序也很难完全复现，消融结果还是会掺入噪声。5. 我的一些延伸思考与实际体会复现这个思路的过程里，我最大的体会是：对"表征到底该怎么组织"这个问题，判别式任务和生成式任务给出的答案可能完全相反。过去做分类、做检索养成的"解纠缠直觉"，放到扩散模型里反而会拖后腿。如果你正被 DiT 的训练不稳定折磨，不妨先别急着加各种外部补丁，回头看看模型内部的表征是怎么组织的——方向对了，很多问题自然就消失了。另外一个值得关注的点是，这个思路的适用范围应该不局限于 DiT。任何包含了"多步渐进式预测"范式的模型，从多阶段扩散模型到自回归生成模型，都可能存在类似的时间步邻近共享空间，用类似的思路去约束表征组织方式，也许都能收到效果。我自己已经在把时间步对齐约束往一个多阶段生成模型上迁移，初步实验的 loss 曲线也是明显更稳的。最后再分享一个小技巧：如果你没精力完整复现整篇论文，可以先只加"时间步邻近对齐"这一个约束，系数调到 0.05 左右，在你的数据上跑几千步，对比一下和基线的 loss 曲线。这个最小化实验成本很低，但大概率能让你看到那种"曲线突然变乖了"的直观感受。有了正反馈之后，再逐步把特征混合等其他组件加上去，调参也会更有底气。