即时计算：面向扩散Transformer的免训练空间加速框架

通过空间稀疏化和动态token激活实现DiT免训练7倍加速

前置知识

扩散Transformer (DiT)

扩散Transformer是一种将Transformer架构应用于扩散模型的生成式架构。与传统的U-Net扩散模型不同，DiT使用自注意力机制来建模图像或视频token之间的长程依赖关系。在去噪过程中，模型将连续潜在表示通过空间patch化转换为N个token序列，然后使用Transformer网络 $u_\theta$ 来建模底层速度场。这种架构在FLUX.1-dev等模型中展现了卓越的可扩展性和生成质量，但自注意力的 $O(N^2)$ 复杂度使其在高分辨率生成时面临巨大计算开销。

理解DiT的token化表示和自注意力瓶颈是理解本文空间加速策略的基础，本文正是针对这一架构的token维度进行优化。

Flow Matching (流匹配)

流匹配是一种将生成建模表述为求解常微分方程(ODE)的框架，它在噪声分布 $p_0$ 和数据分布 $p_1$ 之间建立确定性传输。具体来说，从纯高斯噪声 $t=0$ 出发，沿着学习到的速度场 $u_\theta(y(t), t)$ 积分到 $t=1$ 来生成数据。该框架直接学习ODE向量场，是FLUX.1-dev等现代DiT模型采用的训练范式，相较于传统的DDPM/DDIM具有更直观的数学表述。

本文提出的SAG-ODE正是在流匹配框架下对ODE进行空间近似，理解流匹配的基本原理是理解本文方法数学推导的前提。

空间冗余

扩散模型在生成过程中存在固有的空间冗余特性：模型首先合成低频全局结构，然后逐步细化高频细节。这意味着在生成的早期阶段，所有空间区域并不需要同等的计算资源。对所有token进行均匀计算处理代表了一种关键的效率瓶颈，尤其是在早期噪声主导的阶段，大部分空间区域的计算是冗余的。

空间冗余是本文的核心观察和切入点，JiT框架正是通过利用这种冗余来实现空间维度的加速，而非传统的时间维度加速。

正交投影与子空间分解

论文中使用选择矩阵 $S_k \in \{0,1\}^{Nd \times m_k d}$ 来提取锚点token，并定义正交投影矩阵 $P_k = S_k S_k^\top$ 将全空间投影到锚点token张成的子空间。新激活token的投影算子定义为 $Q_k = P_{k-1} - P_k$，精确隔离下一阶段要激活的token集合。这种线性代数工具为构建空间近似ODE和确保锚点token的动力学精确性提供了严格的数学基础。

投影矩阵和子空间分解是SAG-ODE一致性和DMF无缝过渡的数学基石，不理解这些工具就无法理解方法的理论保证。

研究动机

扩散Transformer（DiT）在图像合成领域建立了新的state-of-the-art，但其卓越性能伴随着巨大的计算开销。核心瓶颈在于自注意力机制的复杂度随输入token数量呈二次方增长 $O(N^2)$，在生成高分辨率图像或长视频时导致计算量雪崩式增长。这种问题被扩散去噪过程的固有迭代性进一步放大，导致推理延迟极高且对硬件要求严苛，严重制约了DiT在实时交互系统、消费级设备和大规模商业服务中的应用。现有加速方法主要沿两个方向推进：时间域加速（通过减少采样步数）和模型效率优化（如特征缓存、低比特量化）。然而，这些方法忽略了生成过程中固有的空间冗余——即所有空间区域在整个生成过程中被均匀对待，而实际上全局结构在精细细节之前就已经形成。

本文的目标是本文的目标是提出一种名为Just-in-Time（JiT）的全新免训练框架，通过在空间域实现加速来填补这一空白。具体目标包括：（1）构建一个空间近似的生成ODE（SAG-ODE），仅基于稀疏锚点token的计算来驱动完整潜在状态的演化；（2）设计一个确定性微流（DMF）机制，确保新token激活时的无缝过渡；（3）在state-of-the-art的FLUX.1-dev模型上实现显著的推理加速，同时保持接近无损的生成质量，目标加速比达到4倍甚至7倍。

与已有工作不同的是，本文的独特切入角度在于将加速的视角从传统的时间域（减少步数）转向空间域（减少每步处理的token数）。与现有的分层/金字塔方法依赖显式上采样和分布校正不同，JiT完全避免了这些容易出错的操作，直接操纵token子空间实现无缝的维度转换。其核心洞察是：扩散模型天然具有从粗到细的生成特性（先全局结构后细节），因此可以在早期阶段使用少量token，在后期逐步扩展，实现动态的空间资源分配。这是第一个完全免训练、免上采样的空间加速方法。

核心方法

JiT框架的整体思路可以概括为：在扩散Transformer的采样过程中，不是对所有N个token在每个时间步都进行均匀计算，而是动态地选择一个稀疏的token子集（锚点token）进行实际的神经网络计算，然后通过空间插值外推到完整空间。直觉上，这类似于一位画家先用少量笔触勾勒全局构图，再逐步添加细节——早期阶段只需要粗粒度的信息来建立结构，后期才需要完整的像素级精度。技术路线包含两个协同组件：SAG-ODE负责在每个阶段内用稀疏token驱动状态演化，DMF负责在阶段切换时确保新激活token的平滑过渡。整个过程遵循一个嵌套的token子集链 $\Omega_K \subset \Omega_{K-1} \subset \cdots \subset \Omega_1 \subset \Omega_0$，从最稀疏逐步扩展到完整token集合。

JiT的核心创新在于提出了一种空间近似的生成ODE（SAG-ODE），通过增广提升算子（augmented lifter operator）$\Pi_k$ 将仅在 $m_k$ 个锚点token上计算的速度场 $u_\theta(S_k^\top y(t), t)$ 外推到完整的 $N$ 维token空间。与现有方法的本质区别体现在三个方面：首先，与时间域加速方法（如高阶求解器、蒸馏）不同，JiT在空间维度工作，可以与时间域方法正交组合；其次，与现有的空间加速方法（如RALU、Bottleneck Sampling）依赖显式上采样和分布校正不同，JiT完全免除了这些容易引入伪影的操作；最后，与缓存方法（如TeaCache、TaylorSeer）不同，JiT不是在时间维度复用特征，而是在空间维度选择性计算。提升算子的设计保证了一致性性质：$S_k^\top(\Pi_k u_\theta) = u_\theta$，即锚点token的动力学完全由Transformer输出精确驱动，不会因近似而损失。

方法步骤详情

JiT采样算法包含以下关键步骤：（1）初始化全维状态向量 $y_{t_0} \sim \mathcal{N}(0, I)$ 和最小token集合 $\Omega_K$（通过步长为2的网格采样加边界约束构建）；（2）对于每个去噪步骤 $i=0,\ldots,N-1$，检查是否到达阶段转换时间 $T_k$：若是，则计算重要性图 $I(T_k)$（基于速度场的局部方差），选择top候选token激活，构建目标状态 $y_k^\star = Q_k T_k \Phi_k(S_k^\top \hat{y}(1)) + (1-T_k)\epsilon$，并通过DMF的有限时间ODE将新token演化到目标状态；（3）计算完整维度速度 $v_{t_i} = \Pi_k u_\theta(S_k^\top y_{t_i}, t_i)$，其中 $\Pi_k = S_k(\cdot) + I_k(\cdot)$ 将稀疏速度嵌入全空间并对非活跃token进行最近邻插值加高斯模糊；（4）更新状态 $y_{t_{i+1}} = y_{t_i} + v_{t_i} \cdot \Delta t$；（5）最终通过VAE解码器 $D_\phi$ 生成图像。整个过程采用Beta分布时间步调度（$\alpha=1.4, \beta=0.42$），将更多NFE分配给早期阶段以确保全局结构稳固建立。

技术新颖性

JiT的技术新颖性体现在多个层面。首先，增广提升算子 $\Pi_k$ 的设计是原创的：它将精确嵌入项 $S_k u_\theta$（保证锚点token精确）与插值项 $I_k(u_\theta)$（外推非活跃token）结合，实现了一个结构感知的全空间速度场，仅需在稀疏子集上进行神经网络评估。一致性性质 $S_k^\top(\Pi_k u_\theta) = u_\theta$ 确保了近似不会损害关键token的动力学。其次，确定性微流（DMF）是处理子空间维度变化的全新机制，其有限时间hitting ODE $Q_k \dot{y}(t) = \frac{y_k^\star - Q_k y(t)}{T_k - t}$ 确保新token从插值状态平滑收敛到统计正确的目标，避免了传统方法中上采样引入的伪影。第三，重要性引导的token激活（ITA）通过速度场局部方差动态选择最具信息密度的区域，而非使用静态预定义模式。这些创新的组合使得JiT成为首个完全免训练、免上采样的空间加速框架。

**Figure 2**: An overview of our JiT framework, illustrating its core mechanisms and underlying philosophy

**Figure 4**: Visual ablation study of each component within JiT

**Figure 5**: An illustration of the construction process for the initial selector matrix $S_K$

**Figure 6**: Visual ablation on the number of stages for a fixed total NFE of 18

**Figure 7**: Visual ablation on token allocation for a fixed 3-stage, 18-NFE schedule

实验结果

论文在FLUX.1-dev模型上进行了全面实验，结果表明JiT在多个加速级别均达到了state-of-the-art性能。在约4倍加速设置下（总NFE=18，延迟6.02秒，706.17 TFLOPs），JiT的CLIP-IQA达到0.6166，甚至略优于50步基线的0.6139，ImageReward达到1.017同样超过基线的1.004，实现了真正的近无损加速。相比之下，同等计算预算的TeaCache（28步，6.98秒）CLIP-IQA仅为0.6003，Bottleneck（14步）更是低至0.4056。在更激进的约7倍加速设置下（总NFE=11，延迟3.67秒，423.26 TFLOPs），JiT的CLIP-IQA为0.5397，显著优于7步FLUX.1-dev基线的0.4134和TaylorSeer的0.4164，ImageReward达到0.9746远超基线的0.7474。用户研究进一步证实了这些结果：在与所有基线的盲评对比中，JiT被显著偏好，例如对比FLUX.1-dev(12)获得85.6%偏好率，对比FLUX.1-dev(7)获得93.1%偏好率。消融实验证实了各组件的必要性：移除SAG-ODE的空间近似项导致T2I-CompBench从0.3727骤降至0.3414，移除ITA或DMF目标也均导致性能下降。此外，论文还展示了JiT在Qwen-image（约4倍加速，从26.95秒降至6.51秒）和HunyuanVideo-1.5视频模型（约4倍和7倍加速）上的成功扩展，证明了方法的通用性。

**Table 1**: Quantitative comparison with other methods

**Table 2**: Human preference rates for our JiT method in blind pairwise comparisons

**Table 3**: Ablation study of each component with JiT

**Table 4**: Detailed configuration of JiT schedules for different acceleration factors

**Figure 1**: Visual showcases of our JiT framework applied to the FLUX.1-dev model

**Figure 3**: Qualitative comparison demonstrating the superior performance of our JiT framework

**Figure 8**: Qualitative results of our JiT framework applied to the Qwen-image model

**Figure 9**: Qualitative results of our JiT framework applied to the HunyuanVideo-1.5 backbone

查看结构化数据

任务	指标	本文	基线	提升
图像生成质量（~4×加速）	CLIP-IQA ↑	0.6166	FLUX.1-dev(50步) 0.6139	略优，实现近无损4倍加速
图像生成质量（~7×加速）	CLIP-IQA ↑	0.5397	FLUX.1-dev(7步) 0.4134	+30.5%相对提升
人类偏好（~4×加速）	ImageReward ↑	1.017	FLUX.1-dev(50步) 1.004	+1.3%超越50步基线
文本-图像对齐（~4×加速）	T2I-CompBench ↑	0.4991	FLUX.1-dev(50步) 0.4836	+3.2%
用户研究偏好率	偏好率 ↑	85.6%-93.1%	各基线方法	压倒性优势

局限与改进

论文存在若干值得讨论的局限性。首先，实验主要在FLUX.1-dev这一单一文本到图像模型上进行详尽验证，虽然补充了Qwen-image和HunyuanVideo-1.5的初步实验，但这些扩展实验的消融分析不够深入，方法在不同架构（如非流匹配框架、不同tokenizer设计）上的适用性仍需进一步验证。其次，JiT的超参数配置（阶段数、每阶段token稀疏度、时间步调度）目前依赖基于原则的轻量级调参，并非完全自动化，不同模型和任务可能需要不同的调度策略。第三，当前的重要性引导token激活（ITA）基于速度场方差这一单一指标，可能无法完全捕捉所有类型的语义重要性，例如在生成包含多个同等重要对象的场景时。第四，虽然论文展示了7倍加速的效果，但在更极端的加速比（如10倍以上）下，初始阶段的过稀疏可能导致全局结构信息不足，这在消融实验的激进调度（20%→50%）结果中已有体现。最后，论文未充分讨论JiT在视频生成中的时间一致性问题，虽然展示了HunyuanVideo-1.5的结果，但缺乏定量的时序一致性评估。

独立分析的弱点

独立分析发现以下弱点需要改进：（1）自适应阶段数：当前固定3阶段设计可能不是最优的，不同图像内容（简单vs复杂构图）可能需要不同的阶段划分，一个自适应机制可以根据生成内容的复杂度动态调整阶段数和过渡时机；（2）ITA指标的局限性：基于速度场局部方差的重要性度量可能在某些场景下失效，例如均匀纹理区域可能具有低方差但对视觉质量很重要，可以考虑结合语义分割信息或注意力图来改进重要性评估；（3）插值策略：当前的最近邻+高斯模糊方案虽然避免了过度平滑，但在高度不规则的token分布下可能产生次优结果，可探索基于图神经网络或扩散先验的更高级插值方法；（4）与时间域加速的联合优化：论文声称JiT可与时间域方法正交组合，但缺乏实际的联合实验验证，且最优的时空联合调度策略未被探索。

未来方向

基于JiT的成果，未来研究可以从多个方向延伸。作者提出的直接方向是将框架扩展到更多backbone模型和生成任务。基于已有成果可延伸的方向包括：（1）时空联合加速——将空间稀疏化与时间域加速（如高阶求解器、步数蒸馏）进行联合优化，设计最优的时空资源分配策略，可能实现超过10倍的综合加速；（2）自适应空间调度——利用生成过程中的中间信号（如注意力分布、梯度信息）动态决定何时进行阶段转换以及激活多少token，而非使用预设的时间表；（3）视频生成的时序感知加速——针对视频扩散模型，设计同时考虑空间和时间冗余的token选择策略，例如相邻帧中相似区域的跨帧复用；（4）与模型压缩技术的结合——将JiT的空间稀疏化与量化、剪枝等技术结合，进一步降低单token的计算成本；（5）理论分析——严格分析空间近似误差的传播和累积，为加速比的理论上界提供保证。

复现评估

从复现角度来看，JiT具有较好的可复现性。论文在FLUX.1-dev这一公开可用的模型上进行实验，该模型权重可从公开渠道获取。方法本身是完全免训练的，不需要任何额外的模型微调或数据收集，仅需要定义多阶段调度和对应的SAG-ODE近似。论文提供了详细的实现细节，包括插值算子（Algorithm 2）、完整采样算法（Algorithm 1）以及超参数配置（Table 4），Beta分布时间步调度的参数也已明确给出（$\alpha=1.4, \beta=0.42$）。计算资源方面，所有延迟和FLOPs在单张A800 GPU上测量，这是一张常见的数据中心GPU，复现门槛适中。评估使用了标准benchmark（GenEval的553个prompt和T2I-CompBench的2400个prompt），共4612张图像。不过，论文未提供完整的代码仓库链接，这可能会影响精确复现，但核心算法的描述足够详细，有经验的研究者应该能够在数天内完成复现。

即时计算：面向扩散Transformer的免训练空间加速框架 Just-in-Time: Training-Free Spatial Acceleration for Diffusion Transformers