Just-in-Time: Training-Free Spatial Acceleration for Diffusion Transformers¶

会议: CVPR2026 arXiv: 2603.10744 代码: 项目主页领域: 图像生成 / 扩散模型加速 关键词: Diffusion Transformer, 空间加速, 免训练, Flow Matching, token稀疏化, ODE求解

一句话总结¶

提出 Just-in-Time (JiT) 框架，通过在空间域动态选择稀疏 anchor token 驱动生成 ODE 演化，并设计确定性 micro-flow 保证新 token 无缝激活，在 FLUX.1-dev 上实现最高 7× 加速且几乎无损。

背景与动机¶

DiT 计算瓶颈：Diffusion Transformer 的 self-attention 复杂度为 \(\mathcal{O}(N^2)\)，高分辨率图像/视频生成时推理延迟极高，严重制约实时交互和消费级部署
时域加速的局限：现有加速方法主要关注时域（高阶求解器、蒸馏少步模型），但在超低步数时质量显著下降，且蒸馏需要大量重训练资源
缓存方法的天花板：特征缓存（TeaCache、TaylorSeer）复用中间激活来减少计算，但其质量上界受限于对应 NFE 的基线表现，存在特征陈旧问题
空间冗余被忽视：扩散生成过程具有从低频全局结构到高频细节的渐进特性，但现有方法对所有空间区域统一计算——这是不必要的浪费
现有空间方法的缺陷：已有的金字塔/层级式空间加速方法依赖显式上采样和分布校正，容易引入混叠伪影和信息损失
核心洞察：生成早期阶段全局结构已形成，只需在少量关键区域计算即可驱动完整潜在状态演化，细节区域可延迟处理

方法详解¶

整体框架¶

JiT 是一个无训练的空间域加速框架，包含两个核心组件：

SAG-ODE（空间近似生成ODE）：在稀疏 anchor token 上计算速度场，通过增广提升算子外推到全空间
DMF（确定性微流）：在阶段转换时，通过有限时间 ODE 将新激活 token 从插值状态平滑演化到统计正确的目标状态

SAG-ODE 设计¶

构建嵌套 token 子集链 \(\Omega_K \subset \Omega_{K-1} \subset \cdots \subset \Omega_0 = \{1,...,N\}\)，从最小子集逐步扩展。

核心方程：

\[\frac{d\mathbf{y}(t)}{dt} = \mathbf{\Pi}_k \, \boldsymbol{u}_\theta(\mathbf{S}_k^\top \mathbf{y}(t), t)\]

其中增广提升算子 \(\mathbf{\Pi}_k\) 由两部分组成： - 嵌入映射 \(\mathbf{S}_k \boldsymbol{u}_\theta\)：将 anchor token 的精确速度放回全空间对应位置 - 插值算子 \(\mathcal{I}_k(\boldsymbol{u}_\theta)\)：对非活跃 token 进行空间插值approximation

一致性保证：\(\mathbf{S}_k^\top(\mathbf{\Pi}_k \boldsymbol{u}_\theta) = \boldsymbol{u}_\theta\)，即 anchor token 的动力学由 Transformer 精确控制，加速不损害关键区域质量

DMF（确定性微流）¶

阶段转换时为新 token 构建目标状态：

\[\mathbf{y}_k^\star = \mathbf{Q}_k \left( T_k \Phi_k(\mathbf{S}_k^\top \hat{\mathbf{y}}(1)) + (1-T_k)\epsilon \right)\]

通过 Tweedie 公式预测干净数据，结合结构先验插值和正确噪声水平，确保统计一致性。然后用有限时间打靶 ODE 在极短区间内将新 token 精确收敛到目标。

重要性引导的 token 激活（ITA）¶

不使用固定网格模式，而是根据速度场的局部方差计算重要性图：

\[\mathbf{I}(t) = \mathbb{E}_\mathcal{W}[\boldsymbol{u}_\theta \odot \boldsymbol{u}_\theta] - (\mathbb{E}_\mathcal{W}[\boldsymbol{u}_\theta])^{\odot 2}\]

选择方差最大（生成过程最活跃）的区域优先激活，将算力分配到高频细节区域。

实验关键数据¶

主实验（FLUX.1-dev，Tab.1）¶

方法	NFE	延迟(s)	TFLOPs	加速比	CLIP-IQA↑	ImageReward↑	HPSv2.1↑	GenEval↑	T2I-Comp↑
FLUX.1-dev	50	25.25	2991	1.0×	0.6139	1.004	30.39	0.6565	0.4836
TeaCache	28	6.98	729	4.1×	0.6003	0.964	29.68	0.6493	0.4849
JiT (Ours)	18	6.02	706	4.24×	0.6166	1.017	29.77	0.6540	0.4991
TeaCache	28	4.53	432	6.9×	0.5183	0.773	27.86	0.5837	0.4625
JiT (Ours)	11	3.67	423	7.07×	0.5397	0.975	29.02	0.6457	0.4961

4× 加速时：JiT 在 CLIP-IQA、ImageReward、GenEval、T2I-Comp 均为最优，接近 50-NFE 基线
7× 加速时：JiT 大幅超越所有竞品，ImageReward 从 0.773 提升到 0.975

用户研究¶

对比方法	JiT 偏好率
vs FLUX.1-dev (12 NFE)	85.6%
vs Bottleneck (14 NFE)	90.3%
vs FLUX.1-dev (7 NFE)	93.1%
vs TaylorSeer (28 NFE)	89.5%

20位参与者在1000次盲测中显著偏好 JiT 生成结果。

消融实验（T2I-CompBench complex compositions）¶

变体	HPSv2.1↑	T2I-Comp↑
完整 JiT	26.90	0.3727
去除 SAG-ODE 插值	24.18	0.3414
去除 ITA（用固定网格）	26.51	0.3670
去除 DMF 目标构建	26.04	0.3602

去除空间插值导致灾难性下降（非活跃区域退化为噪声），验证了各组件的必要性。

亮点¶

完全免训练：无需重训练或微调，直接应用于预训练 DiT 模型
无上采样设计：摆脱了传统空间加速方法对显式上采样/下采样的依赖，从根源避免伪影
数学优雅：SAG-ODE 具有一致性证明（anchor token 无损），DMF 有严格的打靶 ODE 收敛保证
动态资源分配：ITA 基于速度场方差的 content-aware 策略，比固定模式更高效
极端加速下仍保持质量：7× 加速时仍能正确渲染文字等高频细节，优势在极限场景更突出

局限性 / 可改进方向¶

仅在 FLUX.1-dev 一个模型上验证，未展示对其他 DiT（SD3、PixArt 等）的泛化性
阶段调度（\(\{T_k, m_k\}\)）需要手动设计，缺乏自适应调度机制
插值算子 \(\mathcal{I}_k\) 的设计较简单（空间平滑插值），对纹理丰富区域可能不够精确
仅验证图像生成，未扩展到视频生成场景（token 数更多，空间冗余可能更显著）
与时域加速方法（步数蒸馏）的组合潜力未探索，理论上两者正交可叠加
DMF 中的噪声 \(\epsilon\) 每次阶段转换重新采样，可能引入微小随机性

与相关工作的对比¶

类别	方法	对比
空间加速	RALU、Bottleneck Sampling	依赖显式上采样+分布校正，易引入伪影；JiT 无上采样设计
缓存加速	TeaCache、TaylorSeer	质量上界受低 NFE 基线限制；JiT 不受此约束
子空间扩散	Subspace Diffusion	概念启发但限于低维子空间；JiT 动态操作 token 子集更灵活
金字塔方法	Pyramidal Flow	逐级上采样+校正；JiT 通过 DMF 实现无损维度转换

评分¶

新颖性: ⭐⭐⭐⭐ 空间域稀疏 token 加速 + 微流转换的组合设计新颖，数学框架清晰
实验充分度: ⭐⭐⭐⭐ 多指标定量+定性+用户研究+消融齐全，但只在单一模型验证
写作质量: ⭐⭐⭐⭐⭐ 问题动机清晰，数学推导严谨，图示直观
价值: ⭐⭐⭐⭐ 免训练7×加速实用价值高，但泛化性和视频扩展有待验证