利用自身内部动态引导扩散Transformer

通过引入中间层监督和内部引导机制，显著提升扩散Transformer的生成质量和训练效率

前置知识

扩散模型

扩散模型是一类基于概率的生成模型，通过逐步添加噪声将数据转换为高斯噪声，然后学习逆向去噪过程从噪声恢复原始数据。其核心是学习从时间t的噪声状态xt预测原始数据x*的去噪函数Dθ(xt, t)，通过最小化去噪目标函数E[x*~p(x)]E[ε~N(0,I)]||Dθ(xt, t) - x*||^2来训练模型。扩散模型具有强大的分布覆盖能力和高质量的生成结果。

本文基于扩散模型和Flow-matching框架，理解扩散模型的基础原理对掌握Internal Guidance方法的训练和采样机制至关重要。

Classifier-Free Guidance

分类器无关引导是一种在采样阶段提升生成质量的常用技术，通过在条件和无条件预测之间进行外插来获得引导分数。具体公式为引导分数等于条件预测加上w乘以条件预测与无条件预测的差值，其中w是引导强度参数。虽然CFG能够显著提升生成质量，但过高的引导系数会导致样本过度简化或失真，并降低生成多样性。

CFG是目前扩散模型的标准引导方法，本文的Internal Guidance与CFG具有互补性，理解CFG的原理和局限性有助于理解IG的创新点和优势。

Flow Matching

Flow Matching是一种生成建模框架，通过在有限时间间隔内插值噪声和数据来定义生成过程，与前向随机微分方程扩散模型不同。Flow Matching和扩散模型都从高斯噪声开始，通过随机过程xt等于αt乘以x*加上σt乘以ε逐渐转换为数据样本，其中αt和σt分别是递减和递增的时间函数。

本文的工作基于Flow Matching和扩散模型的统一视角，SiT模型正是采用Flow Matching训练，理解Flow Matching的数学基础有助于理解本文的训练目标设计。

Transformer架构

Transformer是一种基于自注意力机制的深度学习架构，由多头注意力层、前馈神经网络和层归一化组成。在扩散模型中，DiT和SiT采用纯Transformer架构替代传统的UNet，通过可扩展的注意力机制实现高效的高维数据生成，具有出色的可扩展性和性能表现。

本文的Internal Guidance方法专门针对扩散Transformer设计，利用Transformer的中间层输出进行引导，理解Transformer的架构和层次结构是理解IG方法实现的基础。

研究动机

现有扩散模型在生成高质量图像时面临两个主要问题：第一，由于训练数据和训练量的限制，模型难以充分学习覆盖低概率区域的数据分布，导致在这些区域生成高质量图像的能力不足。第二，标准的Classifier-Free Guidance虽然能够引导样本向高概率区域移动，但过高的引导系数会使采样轨迹超出条件分布的期望范围，导致过度简化或失真的样本，同时降低生成多样性。替代的AutoGuidance方法通过用模型的劣化版本来引导模型本身，虽然能够在保持多样性的同时提升生成质量，但需要精心设计退化策略、额外的训练或额外的采样步骤，这些限制使得这些方法难以在大规模图像生成器中实际应用。

本文的目标是本文的目标是提出一种简单有效的采样引导策略Internal Guidance，通过利用扩散Transformer内部的中间层输出在训练过程中引入辅助监督，在采样过程中利用中间层和深层输出的关系进行外插引导。IG旨在实现类似AutoGuidance的效果，即在保持生成多样性的同时提升生成质量，但不需要任何额外的采样步骤、额外的训练或复杂的退化策略设计。此外，IG还希望能够与CFG和guidance interval等方法结合使用，进一步提升生成性能，并为大规模扩散模型的应用提供通用的增强组件。

与已有工作不同的是，本文的独特切入角度在于完全避免依赖外部设计的退化模型或预训练表示模型，而是直接利用模型自身的内部动态进行引导。与现有的AutoGuidance方法需要专门训练劣化版本模型不同，IG通过在训练过程中引入中间层辅助监督自然地获得劣化版本的输出，这种方法是即插即用的，不需要针对特定实验设置进行专门训练。与基于自监督学习正则化的方法相比，IG采用了最简单的辅助监督方法，却达到了 comparable 甚至更好的收敛性能。这种以模型自身引导自身的思想提供了一种全新的视角来解决扩散模型的引导问题。

核心方法

Internal Guidance方法包含两个核心组件：训练阶段的中间层辅助监督和采样阶段的内部引导。在训练阶段，我们在网络的中间层后额外定义一个输出层，引入辅助监督损失，使模型能够从中间层表示产生较弱的生成输出。在采样阶段，利用中间层输出Di和深层输出Df之间的关系进行外插引导，通过公式Dw等于Di加上w乘以Df减去Di获得引导后的输出。这种方法本质上相当于利用模型的中间层输出作为劣化版本來引导最终的深层输出，类似于AutoGuidance但完全避免了复杂的退化策略设计和额外训练。IG是一个即插即用的方法，可以无缝集成到现有的扩散Transformer中，几乎不需要额外的计算成本。

IG的核心创新点在于完全利用模型自身的内部动态进行引导，而不依赖任何外部的退化模型或预训练表示。具体而言，通过在训练过程中引入中间层辅助监督，自然获得一个相对于最终输出较弱的版本，然后在采样过程中利用这个弱输出来引导强输出。这种设计与AutoGuidance有本质区别：AutoGuidance需要专门训练一个劣化版本的模型，而IG通过单次训练自然获得引导所需的不同强度输出。这种设计不仅简化了实现，还避免了复杂的超参数调优和额外的计算开销。此外，IG还可以与CFG、guidance interval等方法无缝结合，提供更灵活的引导策略组合。

方法步骤详情

IG方法包含训练和采样两个阶段。训练阶段：首先在网络的中间层如第4层或第8层后额外定义一个输出层，得到中间输出Di和最终输出Df。然后定义中间损失Linter等于Di减去x*的范数平方和最终损失Lfinal等于Df减去x*的范数平方，总损失为L等于Lfinal加上λ乘以Linter，其中λ是平衡两个损失的权重参数。训练时模型同时学习从中间层和最终层预测原始数据。采样阶段：在采样过程中，同时获取中间层输出Di和深层输出Df，然后通过外插公式Dw等于Di加上w乘以Df减去Di获得引导后的输出，其中w是引导强度参数。还可以引入引导区间，使引导只在特定的噪声级别范围内生效。

技术新颖性

IG的技术新颖性体现在多个方面：首先，它提出了完全基于模型内部动态的引导范式，避免了对外部退化模型的依赖，这是对AutoGuidance思想的简化但更为有效的实现。其次，IG证明了简单的中间层辅助监督就能达到与复杂自监督学习正则化方法相当的收敛性能，这挑战了需要复杂设计才能提升训练效率的常规认知。第三，IG首次系统性地分析了内部引导与CFG的兼容性，以及内部引导的独特引导区间特性，即与CFG的引导区间相反，IG在高噪声和中噪声区间应用引导更有效。最后，IG还启发了新的训练加速方法设计，通过将引导信号融入损失函数，实现了比基于预训练表示模型的方法更快的收敛速度。

**Figure 2**: The overall framework of our proposed Internal Guidance. We introduce an additional auxiliary supervision loss during training, and utilize the intermediate layer outputs during sampling process to guide the final outputs.

**Figure 3**: A fractal-like 2D distribution with two classes indicated with gray and orange regions.

实验结果

实验结果表明IG在多个基准模型上都取得了显著性能提升。在SiT-XL/2上，使用IG后仅用80个训练epochs就达到FID等于5.31，超过了 vanilla SiT-XL在1400个epochs和REPA在800个epochs的性能。继续训练到800个epochs后，SiT-XL/2加IG达到FID等于1.75无CFG和FID等于1.46结合CFG和guidance interval。在LightningDiT-XL/1上，使用IG后在60个epochs就达到FID等于2.42，超过了 vanilla LightningDiT-XL/1在64个epochs的性能。在680个epochs后达到FID等于1.34无CFG，结合CFG后达到当前最先进的FID等于1.19。消融实验显示中间层监督放在前几层如第4层最有效，λ系数在0.5以下时性能稳定，IG与CFG结合使用时可以使用更低的IG系数获得更好的生成质量。此外，IG的计算开销极小，参数量仅增加0.44%，FLOPs增加0.01%，采样延迟仅增加0.16%。

**Table 1**: Auxiliary supervision loss on different intermediate layers on ImageNet 256×256 with internal guidance.

**Table 2**: Convergence analysis of the SiT-B/2 model on ImageNet 256×256 without internal guidance and classifier-free guidance.

**Table 5**: Class-conditional performance on ImageNet 256×256. SiT + IG reaches an FID of 1.75 in and further achieves an FID of 1.46 with CFG.

**Table 3**: The IG coefficient and the corresponding guidance interval analysis on ImageNet 256×256.

**Figure 1**: Visualization Results. We visualize our latent diffusion system with proposed IG together with LightningDiT-XL trained on ImageNet 256 × 256 resolution.

**Figure 4**: The combination of IG and CFG can yield a higher FID value, which is superior to simply apply CFG.

**Figure 5**: Scalability of IG. The relative improvement of IG over the vanilla model becomes increasingly significant as the model size grows.

**Figure 6**: Internal guidance inspires new training acceleration methods.

**Figure 7**: The combination of IG and CFG can produce better visual quality.

查看结构化数据

任务	指标	本文	基线	提升
ImageNet 256×256 class-conditional image generation	FID (lower is better)	1.19 (LightningDiT-XL/1+IG+CFG, 680 epochs)	1.28 (RAE/DiTDH, 800 epochs)	约7%提升，达到当前最先进水平
ImageNet 256×256 class-conditional image generation (without CFG)	FID (lower is better)	1.34 (LightningDiT-XL/1+IG, 680 epochs)	2.17 (LightningDiT-XL/1, 800 epochs)	约38%提升，显著优于基线
ImageNet 256×256 class-conditional image generation	Inception Score (higher is better)	269.0 (LightningDiT-XL/1+IG+CFG, 680 epochs)	295.3 (LightningDiT-XL/1+CFG, 800 epochs)	在显著降低FID的同时保持较高的IS
ImageNet 512×512 class-conditional image generation	FID (lower is better)	1.78 (SiT-XL/2+IG, 60 epochs)	2.62 (SiT-XL/2, 600 epochs)	训练效率提升10倍，性能提升32%

局限与改进

作者承认的局限性包括：当模型未完全训练时，过高的IG或AutoGuidance系数可能会引入新的离群值，因为更强的输出无法准确指向高概率区域。此外，虽然IG在多个基准模型上表现良好，但其最优的中间层监督位置和引导参数可能因模型架构和规模而异，需要一定的调优。我自己观察到的局限性包括：IG方法主要针对类条件生成进行了验证，在文本到图像等复杂条件生成任务上的效果尚未充分探索。另外，虽然IG与CFG的结合效果良好，但这种组合可能会增加超参数调优的复杂度，需要同时调整IG系数和CFG系数。IG在更大规模模型如超过XL尺寸上的可扩展性也值得进一步研究。

独立分析的弱点

独立分析的弱点包括：第一，IG的最优中间层监督位置存在模型规模依赖性，小模型如B/2在第4层效果最好，而大模型如XL需要放在第8层，这增加了应用时的调优成本。第二，IG在高模型系数和未完全训练场景下可能产生新的离群值，限制了在早期训练阶段的直接应用。第三，虽然IG避免了额外的采样步骤，但采样过程中需要同时计算中间层和深层输出，可能会轻微增加采样时的计算开销，尽管实验显示开销很小。第四，IG的引导区间特性与CFG相反，CFG在低噪声区间有效，IG在高噪声和中噪声区间有效，这种差异可能使得两者的最优组合策略更加复杂。改进方向包括：开发自适应的中间层监督位置选择策略、研究训练阶段逐步增加引导系数的curriculum、探索更高效的引导区间调度策略。

未来方向

作者提出的未来研究方向包括：进一步探索基于IG洞察的新训练加速方法设计，通过将引导信号更深入地融入损失函数来实现更快的收敛。基于成果可延伸的研究方向包括：将IG扩展到其他生成任务如文本到图像生成、视频生成和3D生成；研究IG与其他引导技术如Self-Attention Guidance、Smoothed Energy Guidance的结合；探索IG在不同扩散框架如Consistency Models、Rectified Flow中的适用性；开发自动化的超参数搜索策略来找到最优的中间层监督位置和引导参数；研究IG在多模态生成任务中的应用潜力；探索将IG的思想应用到其他类型的生成模型如GANs、VAEs中。

复现评估

论文的复现性较好，作者提供了详细的实现细节和超参数设置。作者基于开源的DiT、SiT和LightningDiT实现，使用混合精度训练、梯度裁剪和FlashAttention等技术来加速训练和节省GPU内存。数据预处理使用ImageNet-1K，图像预处理为256乘256分辨率，通过Stable Diffusion VAE或VA-VAE编码为潜在表示。计算资源方面，大模型使用NVIDIA A6000 pro 96GB GPU，小模型使用NVIDIA 4090 24GB GPU。论文提供了详细的超参数表格，包括架构设置、优化参数、训练目标、采样器和IG相关的参数。然而，论文没有明确说明代码是否开源，这可能影响完全复现的难度。整体而言，复现难度中等，需要较强的计算资源和大模型需要96GB显存GPU和扎实的深度学习实现经验。

利用自身内部动态引导扩散Transformer Guiding a Diffusion Transformer with the Internal Dynamics of Itself