Diffusion Transformers have established a new state-of-the-art in image synthesis, but the high computational cost of iterative sampling severely hampers their practical deployment. While existing acceleration methods often focus on the temporal domain, they overlook the substantial spatial redundancy inherent in the generative process, where global structures emerge long before fine-grained details are formed. The uniform computational treatment of all spatial regions represents a critical inefficiency. In this paper, we introduce Just-in-Time (JiT), a novel training-free framework that addresses this challenge by acceleration in the spatial domain. JiT formulates a spatially approximated generative ordinary differential equation (ODE) that drives the full latent state evolution based on computations from a dynamically selected, sparse subset of anchor tokens. To ensure seamless transitions as new tokens are incorporated to expand the dimensions of the latent state, we propose a deterministic micro-flow, a simple and effective finite-time ODE that maintains both structural coherence and statistical correctness. Extensive experiments on the state-of-the-art FLUX.1-dev model demonstrate that JiT achieves up to a 7x speedup with nearly lossless performance, significantly outperforming existing acceleration methods and establishing a new and superior trade-off between inference speed and generation fidelity.
Figures & tables
Figure 2: An overview of our JiT framework, illustrating its core mechanisms and underlying philosophy. (a) The SAG-ODE evolves the latent state by extrapolating a velocity field computed on a sparse subset of tokens. (b) For stage transitions, the DMF evolves newly incorporated tokens to a structurally coherent target with the correct noise level to prevent artifacts. (c) The visualized evolution of the predicted clean image reveals a coarse-to-fine process (global structures first), motivating our strategy to defer computation on detailed regions. (d) The sampling trajectory visualizes our dynamic resource allocation, where the set of active tokens (red flow) starts as a narrow subset and expands over time, reserving full computation for the final detail-refining stages.
Method
NFE
Latency (s) ↓
TFLOPs ↓
Speed ↑
CLIP-IQA ↑
Image Reward ↑
HPSv2.1 ↑
Gen Eval ↑
T2I-Comp. ↑
FLUX.1-dev
50
25.25
2990.96
1.00 ×
0.6139
1.004
30.39
0.6565
0.4836
FLUX.1-dev
12
6.21
729.07
4.10 ×
0.5341
0.9435
29.43
0.6426
0.4843
Bottleneck
14
6.37
709.28
4.22 ×
0.4056
0.9216
29.01
0.6499
0.4779
RALU
18
6.23
723.69
4.13 ×
0.5493
0.8560
29.83
0.6328
0.4620
Teacache
28
6.98
729.32
4.10 ×
0.6003
0.9638
29.68
0.6493
0.4849
JiT (Ours)
18
6.02
706.17
4.24 ×
0.6166
1.017
29.77
0.6540
0.4991
Table 1: Quantitative comparison with other methods. The optimal result is represented in bold, and the sub-optimal result is represented in underlining. FLOPs are measured via torch.profiler and the Speed is calculated based on the FLOPs reduction relative to the base model FLUX.1-dev(50).
Figure 3: Qualitative comparison demonstrating the superior performance of our JiT framework. While competing methods suffer from common acceleration artifacts, including semantic errors, loss of fine detail, and structural inconsistencies, our approach maintains high fidelity across different prompts and acceleration levels.
Opponent Method
Preference for JiT (%) ↑
JiT vs. ∼ 4 × Baselines
FLUX.1-dev (12)
85.6
Bottleneck (14)
90.3
RALU (18)
73.7
Teacache (28)
69.1
JiT vs. ∼ 7 × Baselines
Table 2: Human preference rates for our JiT method in blind pairwise comparisons.
Method
HPSv2.1 ↑
T2I-Comp. ↑
Ours
26.90
0.3727
w/o SAG-ODE
24.18
0.3414
w/o ITA
26.51
0.3670
w/o DMF target
26.04
0.3602
Table 3: Ablation study of each component with JiT.
Figure 4: Visual ablation study of each component within JiT.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: An illustration of the construction process for the initial selector matrix SK .
Parameter
JiT ( ∼4× )
JiT ( ∼7× )
Total NFEs
18
11
Number of Stages
3
3
Stage 2 (Coarse)
Steps
7
4
Token Sparsity
35%
32%
Stage 1 (Medium)
Appendix
Table 4: Detailed configuration of JiT schedules for different acceleration factors. The sparsity ratio denotes the fraction of tokens activated at each stage relative to the full sequence length N .
Figure 6: Visual ablation on the number of stages for a fixed total NFE of 18. (a) A 2-stage schedule offers limited acceleration. (b) Our default 3-stage schedule yields the best trade-off. (c) A 4-stage schedule introduces persistent noise due to a late transition to full resolution. This validates that a 3-stage approach offers a superior balance between speed and quality.
Figure 7: Visual ablation on token allocation for a fixed 3-stage, 18-NFE schedule. (a) An aggressive schedule (20% → 50%) compromises semantic integrity. (b) Our chosen schedule (35% → 62%) provides the best balance. (c) A conservative schedule (50% → 75%) offers similar quality but with a reduced acceleration advantage.
Figure 8: Qualitative results of our JiT framework applied to the Qwen-image model. The images were generated with ∼ 4 × acceleration compared to the standard pipeline.
Figure 9: Qualitative results of our JiT framework applied to the HunyuanVideo-1.5 backbone. Our JiT framework demonstrates robust generalizability in the spatiotemporal domain, maintaining semantic consistency and temporal coherence under both ∼ 4 × and ∼ 7 × acceleration settings.
Figure 10: Additional qualitative comparisons of our JiT framework against baseline methods at ∼ 4 × and ∼ 7 × acceleration. These examples further showcase the superior performance of JiT across a variety of challenging prompts.
College of Electronics and Information Engineering, Shenzhen University · School of Artificial Intelligence, Shenzhen University · College of Computer Science and Software Engineering, Shenzhen University