随着扩散 Transformer (DiT) 的采用，生成式 AI 的格局发生了巨大变化。像 Stable Diffusion 3 和 Sora 这样的模型表明，用 Transformer 架构取代传统的 U-Net 骨干网络可以带来可扩展的高保真结果。然而，这种性能的提升伴随着高昂的计算成本。

当前的扩散模型在一个静态范式下运行: 它们为去噪过程的每一步分配固定且大量的计算资源。无论模型是在解析构图的模糊轮廓，还是在细化猫毛的纹理，它都会消耗相同数量的 FLOPs (浮点运算次数) 。

本篇博文探讨 FlexiDiT , 这是一个挑战这种低效性的新颖框架。通过将 DiT 模型“灵活化 (flexifying) ”，研究人员找到了一种动态调整每一步计算预算的方法，在不牺牲质量的情况下实现了巨大的效率提升——图像生成效率提升超过 40%，视频生成高达 75%。

直觉: 并非每一步都生而平等

为了理解 FlexiDiT 为何有效，我们首先需要看看扩散 Transformer 是如何处理图像的。

在标准的视觉 Transformer (ViT) 或 DiT 中，图像并不是作为像素网格处理的。相反，它被切分成一个个 patch (例如 \(2 \times 2\) 或 \(4 \times 4\) 像素) 。这些 patch 被展平并投影为嵌入 (embeddings) ，从而形成一个 token 序列。

如上所示，patch 大小 (\(p\)) 决定了序列长度 (\(N\))。如果将 patch 大小加倍，token 的数量就会减少四倍。由于 Transformer 中自注意力机制的计算复杂度随着 token 数量呈二次方增长 (\(O(N^2)\))，减少 token 数量会带来巨大的加速。

扩散的频谱现实

扩散模型通过迭代去除噪声来生成图像。然而，在此过程中恢复的“信号”性质是变化的。

1. 早期步骤 (高噪声) : 模型关注低频细节——全局结构、布局和大形状。
2. 晚期步骤 (低噪声) : 模型关注高频细节——纹理、锐利边缘和细微噪声。

FlexiDiT 背后的研究人员通过对单个扩散更新步骤应用低通和高通滤波器，凭经验证明了这一点。

这里的关键见解是 早期去噪步骤关注低频信息。 你不需要细粒度的 token 网格来决定天空和地面的位置。你可以用更粗糙的 token (更大的 patch) 达到同样的效果。相反，后续步骤需要高分辨率的 token 来生成清晰的细节。

为早期的粗略步骤使用“强大”的模型 (小 patch，高计算量) 是一种资源浪费。FlexiDiT 解决了这个问题，它允许模型在“弱” (大 patch) 和“强大” (小 patch) 模式之间即时切换。

FlexiDiT 框架

FlexiDiT 的目标是让预训练的 DiT 变得“灵活”，意味着它可以处理不同 patch 大小的输入。这使我们能够使用 动态推理调度器 : 在早期步骤使用大 patch，在晚期步骤使用小 patch。

研究人员提出了两种主要方法来实现这一点，具体取决于你是想微调整个模型还是保持原始权重冻结。

方法 1: 共享参数 (全量微调)

如果你能访问原始训练数据和资源，你可以微调整个模型。核心 Transformer 块保持不变——从技术上讲，Transformer 可以处理任何长度的序列。挑战在于输入和输出层，它们与特定的空间维度相关联。

为了解决这个问题，作者调整了 分词 (嵌入) 和 反分词 层。

如上图 (左) 所示，模型为新的 patch 大小引入了新的嵌入权重。

1. 初始化: 新权重通过伪逆投影 (双线性插值) 从预训练权重初始化。这确保了模型在开始时就对图像结构有良好的理解。
2. 位置编码: 插值以匹配新的网格大小。
3. 训练: 模型被微调以使用随机选择的 patch 大小去噪图像。

这使得同一组 Transformer 权重能够处理不同长度的序列。

方法 2: LoRA (参数高效微调)

在许多情况下，全量微调太昂贵或者无法获取原始数据集。为此，作者利用了 低秩适应 (LoRA) 。

他们不重新训练庞大的 DiT 块，而是冻结原始权重。然后，他们专门为新的“弱”patch 大小添加小型的、可训练的 LoRA 矩阵。

在使用“弱”步骤 (大 patch) 进行推理期间，LoRA 适配器处于激活状态。在使用“强大”步骤 (原始 patch 大小) 期间，适配器被停用，确保模型完全保留其原始的高质量行为。

知识蒸馏

为了提高训练效率，作者使用原始的“强大”模型作为教师。他们训练“弱”模型来模仿强大模型的预测。这是通过最小化强大模型 (\(\epsilon_{\theta}(\dots; p_{\text{powerful}})\)) 和弱模型 (\(\epsilon_{\theta}(\dots; p_{\text{weak}})\)) 的噪声预测之间的距离来定义的:

\[ \mathbb { E } _ { t , \mathbf { x } _ { t } } \Vert \epsilon _ { \theta } ( \mathbf { x } _ { t - 1 } \vert \mathbf { x } _ { t } ; \boldsymbol { p } _ { \mathrm { p o w e r f u l } } ) - \epsilon _ { \theta } ( \mathbf { x } _ { t - 1 } \vert \mathbf { x } _ { t } ; \boldsymbol { p } _ { \mathrm { w e a k } } ) \Vert _ { 2 } . \]

这确保了弱模型学会尽可能逼近强大模型的输出，而不是从头开始学习。

推理: 动态调度器

一旦模型被训练为处理多种 patch 大小，推理策略就很简单了。

我们定义一个时间表，其中前 \(T_{\text{weak}}\) 步使用大 patch 大小 (低计算量) ，其余步骤使用原始的小 patch 大小 (高计算量) 。

• 弱模型: Patch 大小 \(4 \times 4\)。序列长度比像素少 \(16 \times\)。速度快。
• 强大模型: Patch 大小 \(2 \times 2\)。序列长度比像素少 \(4 \times\)。速度慢。

通过简单地调整 \(T_{\text{weak}}\)，我们可以在最高质量和最快速度之间进行调节。

左边的图表 (图 6) 是最关键的证据。它显示了相对于使用的 FLOPs 百分比的 FID (Fréchet Inception Distance，越低越好) 。

• 在 100% FLOPs (基线) 时，FID 约为 2.25。
• 在 ~60% FLOPs 时，FID 仍然 约为 2.25。

这意味着标准扩散模型中大约 40% 的计算是多余的 , 可以卸载到更弱、更快的模型上，而不会有任何感知上的质量损失。

实验与结果

作者在类别条件生成 (ImageNet) 、文本到图像 (T2I) 和视频生成上测试了 FlexiDiT。

文本到图像性能

对于文本到图像模型 (如 Emu 或 PIXART) ，结果依然强劲。作者绘制了 FID 与 CLIP 分数 (衡量图像与文本提示匹配程度的指标) 的对比图。

结果表明，FlexiDiT 模型 (各种彩色点) 与计算密集型的基线处于相同的帕累托前沿。简单来说: 你可以达到相同的对齐度和图像质量，但速度要快得多。

视频生成: 巨大的胜利

视频生成是 Transformer 二次复杂度真正令人头疼的地方。视频增加了时间维度，导致 token 数量激增。

FlexiDiT 在这里应用了同样的逻辑: 在早期步骤中，在空间 和 时间上都使用更大的 patch。

对于视频，节省是巨大的。即使在使用 仅 ~25% 的基线计算量 时，模型也能保持其 VBench 质量分数。这对于 AI 视频生成的经济性来说可能是一个改变游戏规则的因素。

不仅仅是理论上的 FLOPs

效率研究中一个常见的批评是，“理论上的 FLOPs”并不总是能转化为现实世界的延迟 (速度) 改善，因为存在内存带宽瓶颈或硬件效率低下的问题。

作者在 NVIDIA H100 GPU 上对 FlexiDiT 进行了分析以验证增益。

上面的图表显示，FlexiDiT (特别是弱模型配置) 保持了高 GPU 利用率。由于自注意力机制在长序列 (小 patch) 下会产生瓶颈，改用大 patch 实际上能更有效地利用硬件，这意味着延迟的减少是真实的，并且与 FLOPs 的减少成正比。

结论

FlexiDiT 重新审视了扩散模型的一个基本假设: 即每一步都需要相同数量的“脑力”。通过将计算预算与去噪过程的频谱性质相对齐——粗糙特征用粗略计算，精细特征用精细计算——作者解锁了显著的效率增益。

关键要点:

1. 动态 Patch: 在推理过程中改变 patch 大小允许可变的计算量。
2. 无质量损失: 你可以减少约 40% 的图像计算量和约 75% 的视频计算量，而几乎没有视觉退化。
3. 易于集成: 基于 LoRA 的方法允许以极少的训练数据 (约 5000 张图像) 和计算量 (原始训练成本的 <5%) 将其应用于现有的预训练模型。

随着模型不断扩展，像 FlexiDiT 这样优化 推理 阶段的技术对于使高保真生成式 AI 变得易于访问和负担得起至关重要。

---