[2025] Just-in-Time: 彻底释放 Diffusion Transformer 的空间加速潜力

TL;DR

Diffusion Transformer (DiT) 虽然在图像质量上达到新高度，但其 Self-Attention 的二次方复杂度让高分辨率推断变得异常昂贵。本文提出的 Just-in-Time (JiT) 是一个无需重训练的加速框架。它利用“全局结构早于局部细节”的物理直觉，在生成早期只计算少数锚点 Token，后期再动态激活关键区域。在 FLUX.1 模型上，它实现了 7 倍加速且肉眼几乎无法分辨质量损失。

背景定位：为何我们还在“平等”对待每一个 Token？

在扩散生成的早期阶段，模型主要在构建轮廓和低频色彩。此时，对数千个 Token 进行全量 Self-Attention 是一种极大的冗余。现有的时间步压缩（Few-step Distillation）或缓存技术（Teacache）已经做得很好，但在“空间维度”上，我们依然缺乏一种既高效又不会产生空间伪影（Artifacts）的通用方案。

核心机制：SAG-ODE 与确定性微流

JiT 的核心贡献在于将生成过程看作一个可以“动态扩展示范”的常微分方程（ODE）求解过程。

1. 空间近似生成 ODE (SAG-ODE)

作者没有直接丢弃 Token，而是构建了一个 Augmented Lifter (Π)。

• 核心逻辑：在每一阶段，只从全量 Token 中选取一个子集 ${\Omega }_k$（锚点）。
• 物理直觉：Transformer 只处理这些锚点，而剩下 Inactive Token 的速度场（Velocity Field）则通过锚点的结果进行空间插值得到。
• 一致性保证：数学推导证明（见原文 Eq. 7），这种近似在锚点位置是完全精确的，保证了核心语义不偏移。

图 1: JiT 框架概览。a) SAG-ODE 演化；b) DMF 微流切换；c) 从粗到精的生成演化；d) 采样轨迹的动态资源分配。

2. 确定性微流 (DMF)

当从稀疏阶段转换到稠密阶段时，新注入的 Token 如果处理不当会导致图像“闪烁”或结构断裂。

• 解决方案：JiT 引入了一个限时（Finite-time）的 hitting ODE。它计算一个包含结构先验和噪声水平的目标状态 $y^{\star }$，强制新 Token 在极短的时间内“合流”到主轨迹上。这种设计完全抛弃了传统的、易产生模糊的上采样算子。

3. 重要性引导激活 (ITA)

哪里的 Token 更重要？JiT 通过计算速度场的局部方差来定义重要性图。方差大的地方通常意味着正在生成复杂的边缘或纹理，系统会优先在这些地方“恢复”全量计算。

实验战绩：接近 50-step 的完美复刻

在最先进的 FLUX.1-dev 模型上，JiT 的表现令人惊艳：

• 量化指标：在 7.07x 加速下，其 HPSv2.1（人类偏好得分）保持在 29.02，远高于其他加速方法（如 RALU 或 Teacache）。
• 文字渲染：文字渲染是 DiT 的“试金石”。实验显示，JiT 在 7 倍加速下依然能清晰渲染英文单词，而基线方法早已陷入混乱。

表 1: JiT 与各主流加速方法在不同加速比下的量化指标对比。

深度洞察

JiT 的成功在于它深刻理解了扩散模型流量分配（Flow Matching）的统计特性。

• Inductive Bias：它利用了图像天然的分级特性（Coarse-to-fine）。
• 无缝扩展：作者甚至将其应用到了 HunyuanVideo-1.5 视频模型中。由于视频在时空维度上冗余度更高，JiT 的加速效果在视频生成中展现出极高的工业应用价值。

总结与局限

JiT 标志着 Diffusion Transformer 进入了“按需推断”时代。虽然它目前仍依赖预设的阶段比例（Sparsity Schedule），但这种无需训练、即插即用的特性，使其成为当前移动端和实时生成任务的首选方案。未来的挑战在于如何完全自动化地识别每个 Prompt 所需的最优稀疏度。

本文由资深学术技术主编解读。