PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wenhao Sun, Ji Li, Zhaoqiang Liu

arXiv:2603.10744v2cs.CV

TL;DR

Diffusion Transformers are expensive because iterative denoising and quadratic attention process all spatial regions uniformly despite coarse-to-fine generation. JiT introduces training-free spatial acceleration using sparse anchor tokens, a spatially approximated ODE, and deterministic micro-flows for state expansion. On FLUX.1-dev, it achieves up to 7× acceleration with nearly lossless quality and a superior speed-fidelity trade-off.

  • Problem

    Diffusion Transformers incur substantial inference costs from quadratic self-attention, iterative sampling, and uniform computation across spatial regions despite spatial redundancy.

  • Method

    JiT uses a spatially approximated generative ODE driven by dynamically selected sparse anchor tokens and a deterministic micro-flow for coherent state transitions.

  • Results

    Up to 7× acceleration with nearly lossless quality on FLUX.1-dev, establishing a superior speed-fidelity trade-off over existing acceleration methods.

  • Takeaways & Limitations

    JiT provides training-free spatial acceleration that preserves high-quality, artifact-free generation while shifting computation toward dynamically important regions.

Abstract

from arXiv · show

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.

1. Introduction

Diffusion Transformers offer strong generation quality but incur substantial computational costs from quadratic self-attention and iterative denoising. JiT addresses overlooked spatial redundancy by dynamically allocating computation to salient regions without retraining.

  • Computational challenge: Quadratic self-attention and iterative denoising create prohibitive inference latency for high-resolution images and long-duration videos.These costs constrain real-time systems, consumer devices, and large-scale services.
  • Existing acceleration: Existing acceleration methods primarily reduce temporal steps or distill models, but may compromise fidelity or require retraining.These approaches do not directly address spatial computation.
  • Spatial motivation: Uniform computation overlooks spatial redundancy because diffusion models form global structures before refining high-frequency details.This motivates deferring computation on less critical spatial regions.
  • JiT framework: JiT is a training-free framework that approximates the generative ODE using sparse anchor-token computations and a deterministic micro-flow.Its two components are SAG-ODE and DMF.
  • Reported outcome: JiT reaches up to 7× inference speedup with nearly lossless performance on FLUX.1-dev.The framework targets spatial acceleration without model retraining.

2. Related works

Diffusion models have evolved toward deterministic ODE and flow-matching formulations, but iterative sampling remains computationally expensive. Spatial acceleration reduces this burden, while JiT avoids resizing-based transition procedures that can introduce artifacts.

  • Diffusion-model acceleration: Flow-matching models learn an ODE vector field, but iterative sampling still requires many neural function evaluations.JiT contributes a training-free spatial acceleration technique for this ODE-based setting.
  • Spatial acceleration: Early spatial methods reduce latent resolution during noise-dominated stages and progressively upscale it.They often combine resizing with distribution correction at each new scale.
  • Limitations of prior methods: Upsampling can cause information loss or aliasing, while post-hoc correction may fail to ensure visual consistency.These limitations motivate alternatives that avoid explicit resizing.
  • JiT distinction: JiT directly manipulates token subspaces to achieve seamless, artifact-free dimensionality transitions without upsampling or correction procedures.The approach is described as training-free.

3. Methodology

JiT accelerates flow-matching generation by evolving full token states from sparse, dynamically selected anchors and using deterministic micro-flows when expanding the active subspace. Its interpolation and transition mechanisms preserve structure, noise statistics, and continuity while allocating more computation to dynamic regions.

  • 3.1. Preliminaries: Flow matching transports compressed latent representations deterministically from noise distribution p0 to data distribution p1 through an ODE.DiT tokenizes the latent tensor into N tokens before modeling its velocity field.
  • 3.1. Preliminaries: Processing all N tokens at every timestep incurs O(N^2) cost, motivating dynamically changing subsets of activated anchor tokens.The anchor sets form a nested coarse-to-fine hierarchy.
  • 3.2. The spatially approximated generative ODE: SAG-ODE computes the Transformer velocity on anchor tokens and uses an augmented lifter to construct a full-space velocity field.The lifter embeds exact anchor velocities and interpolates velocities for inactive tokens.
  • 3.2. The spatially approximated generative ODE: SAG-ODE preserves exact Transformer dynamics on anchor tokens while interpolation supplies structurally aware motion for inactive regions.The approximation introduces zero error on the anchor-token subspace.
  • 3.2. The spatially approximated generative ODE: JiT selects new tokens by ranking inactive regions using local velocity variance, allocating computation to the most dynamic spatial areas.This content-aware selection supports fine-grained detail synthesis.
  • 3.3. Deterministic micro-flow: DMF constructs newly activated tokens from interpolated structural information and correctly scaled noise, then evolves them toward a coherent target.The target uses a predicted clean latent derived with Tweedie’s formula.
  • 3.3. Deterministic micro-flow: DMF drives new-token states to the target over a short interval, freezes existing anchor-token velocities, and preserves a continuous latent trajectory.The state converges to the target at t = T_k, preventing artifacts and maintaining stability.

4. Experiments

Experiments evaluate JiT against spatial, caching-based, and vanilla FLUX.1-dev baselines using efficiency, generation-quality, alignment, qualitative, human-preference, and ablation studies. JiT maintains nearly lossless quality at approximately 4× and 7× acceleration, while qualitative and ablation results support its design choices.

  • Experimental setup: JiT is evaluated on FLUX.1-dev against RALU, Bottleneck Sampling, TaylorSeer, TeaCache, and vanilla pipelines with varying NFEs.Latency and FLOPs are measured on a single A800 GPU; quality and alignment use CLIP-IQA, ImageReward, HPSv2.1, GenEval, and T2I-CompBench.
  • Quantitative results: ~4× and ~7× acceleration tiers achieve SOTA performance with nearly lossless quality relative to the high-fidelity 50-NFE baseline.JiT avoids upsampling and uses its intrinsic multi-scale knowledge together with seamless DMF state transitions.
  • Qualitative results: JiT produces high-fidelity, artifact-free images at 7× speedup, including coherent rendered text where competing methods produce garbled or incomplete results.Competing methods show semantic errors, blurred textures, and structural inconsistencies associated with upsampling or stale cached features.
  • User study: Human evaluation used 1000 blind pairwise preference votes from 20 participants across 50 prompts, with JiT preferred against all baselines by a significant majority.The comparisons assessed overall visual quality and prompt fidelity.
  • Ablation study: Removing spatial approximation causes inactive regions to degenerate into meaningless noise, showing that structurally aware velocity extrapolation is essential.The ablation was evaluated on the complex compositions subset of T2I-CompBench, with quantitative results in Table 3 and visual results in Figure 4.
  • Ablation study: Static token selection misallocates computation and loses sharpness in high-frequency regions, while naive DMF targets reduce perceptual quality through noise mismatch.These ablations support dynamic importance-guided activation and statistically correct target-state construction.

5. Conclusions

The conclusion presents JiT as a training-free approach that exploits spatial redundancy in DiT to accelerate inference. On FLUX.1-dev, it achieves up to 7× acceleration with negligible quality degradation and a superior speed-fidelity trade-off.

  • 5. Conclusions: JiT exploits spatial redundancy in DiT through a SAG-ODE driven by dynamic token subsets and seamless state transitions ensured by a DMF.The framework moves from uniform computation toward a dynamic, on-demand strategy.
  • 5. Conclusions: Up to 7× acceleration on FLUX.1-dev is achieved with negligible quality degradation while outperforming existing acceleration methods.The paper characterizes this as a superior speed-fidelity trade-off for generative ODE acceleration.

Supplementary Material

The SAG-ODE derives a full-space velocity field from computations on a sparse set of active anchor tokens. It combines projected anchor-token dynamics with an approximation for inactive-token dynamics.

  • The ideal full-dimensional velocity field is decomposed into active and inactive components using the anchor-token subspace and its orthogonal complement.The anchor subspace is defined by the selector matrix S_k and projector P_k.
  • The active component is approximated by evaluating the Transformer only on projected anchor tokens rather than the full token sequence.The resulting sparse velocity is embedded back into the full space through S_k.
  • The inactive component is approximated by inferring its dynamics from active anchors through a predefined interpolation operator I_k.This avoids explicit neural function evaluations for inactive tokens.
  • The augmented lifter Π_k combines the sparse anchor velocity and interpolated inactive velocity to produce a full-space velocity field.Substituting the Transformer output on anchor tokens yields the compact SAG-ODE form.

A.4. Proof of consistency

The consistency analysis shows that SAG-ODE preserves the sparse dynamics on active anchor tokens while interpolation contributes only outside the anchor subspace.

  • The derived ODE preserves exact dynamics on the active anchor tokens relative to the sparse computation.
  • The resulting velocity combines the embedded sparse anchor velocity with the inactive-region interpolation term.
  • The interpolation operator I_k has no effect on the anchor subspace by design.Formally, its action on the anchor subspace is zero.

B. Implementation details

The implementation details specify selector initialization, interpolation, schedules, and timestep allocation for JiT. The framework uses sparse anchors early and expands computation across later stages.

  • Selector initialization: JiT constructs the initial selector matrix S_K with a deterministic strided-grid strategy and explicit boundary constraints.The initialization targets global semantic structure while respecting the sparsity budget.
  • Interpolation: The interpolation pipeline applies nearest-neighbor interpolation, controlled Gaussian blur, and masked composition while preserving exact anchor values.The procedure produces a full tensor from active features and keeps anchor positions unchanged.
  • Interpolation: Subspace extrapolation I_k acts on the velocity field, whereas the structural-prior operator Φ_k acts on data latent tokens.
  • Schedules: JiT schedules specify total NFEs, stages, and token sparsity ratios for different acceleration factors.The reported configurations follow a principled lightweight recipe rather than exhaustive brute-force search.
  • Timestep schedule: The timestep schedule prioritizes early denoising by using a Beta-distribution-based non-uniform sequence with α = 1.4 and β = 0.42.This allocates denser sampling to the high-noise phase where global structure is established.

B.5. Evaluation dataset

The evaluation dataset combines official prompts from GenEval and T2I-CompBench, using multiple random seeds for GenEval and one image per T2I-CompBench prompt.

  • 4,612 images comprise the evaluation dataset: 2,212 from GenEval and 2,400 from T2I-CompBench.GenEval uses 553 prompts with four images per prompt, while T2I-CompBench uses 2,400 prompts with one image each.
  • The GenEval subset contains four generated images per prompt, each using a different random seed.
  • The T2I-CompBench subset contains one generated image for each prompt.

C. Additional ablation study on JiT scheduling

The study examines how JiT scheduling hyperparameters affect the trade-off between computational cost and generation quality. These qualitative ablations justify the configurations used in the main experiments.

  • JiT scheduling hyperparameters are evaluated for their effects on computational cost and generation quality.
  • The ablation studies are designed to clarify how different scheduling choices shape the speed–quality trade-off.
  • The results provide justification for the scheduling configurations adopted in the main experiments.

C.1. Impact of the number of stages

JiT’s scheduling ablations show that a three-stage coarse-to-fine process offers the best balance between acceleration and image quality at fixed computational budgets. Token allocation similarly requires enough early tokens to preserve global structure without eroding the speed advantage.

  • Impact of the number of stages: At a fixed total of 18 NFEs, the default 3-stage schedule provides the best balance between acceleration and generation quality.The 2-stage schedule offers limited acceleration, while the 4-stage schedule introduces persistent noise from a late transition to full resolution.
  • Impact of the number of stages: The 3-stage schedule avoids the 2-stage schedule’s blurred output by allowing early acceleration and later detail refinement.
  • Impact of the number of stages: The 4-stage schedule is computationally cheaper but degrades quality because full resolution is introduced too late for the model to remove resulting noise and artifacts.
  • Impact of token allocation: The chosen token allocation, 35% →62% →100%, preserves quality while retaining substantial acceleration.The aggressive 20% →50% →100% schedule compromises semantic integrity, whereas the conservative 50% →75% →100% schedule reduces the acceleration advantage.
  • Impact of token allocation: Token allocation directly controls computational cost and the quality of the structural prior in a fixed 3-stage schedule.
  • JiT generalizes beyond FLUX.1-dev, achieving approximately 4× acceleration on Qwen-image while maintaining high visual fidelity.Inference time decreases from 26.95s to 6.51s.

E. Additional qualitative comparisons

Additional qualitative comparisons show JiT maintaining high-quality image and video generation under substantial acceleration. The comparisons cover challenging prompts, another image backbone, and a spatiotemporal video backbone.

  • JiT consistently outperforms competing methods at similar computational budgets on challenging prompts.The comparisons particularly highlight preservation of complex compositions and fine-grained details.
  • On Qwen-image, JiT produces qualitative results at approximately 4× acceleration while preserving visual fidelity across complex scenes, textures, and semantic concepts.
  • On HunyuanVideo-1.5, JiT maintains semantic consistency and temporal coherence under approximately 4× and 7× acceleration.
Loading 2603.10744v2…