[2025] Just-in-Time: Breaking the Spatial Efficiency Barrier in Diffusion Transformers

TL;DR

Diffusion Transformers (DiTs) like FLUX.1 have revolutionized image synthesis but are notoriously slow due to the quadratic cost of self-attention. Just-in-Time (JiT) is a training-free framework that accelerates DiT inference by up to 7×. It works by identifying that global image structures don't need full-resolution computation in the early denoising steps. By calculating "anchor tokens" and interpolating the rest through a novel SAG-ODE and Deterministic Micro-Flow, JiT achieves blazing speed without the typical artifacts of upsampling.

The Motivation: Why Waste Pixels on Noise?

Standard DiT inference treats every patch of an image with equal importance from the very first step. However, diffusion models naturally follow a coarse-to-fine trajectory: they first establish the "gist" (low-frequency layout) and only later refine "details" (high-frequency textures).

Current SOTA acceleration methods like TeaCache (temporal caching) or RALU (spatial upsampling) often hit a "quality ceiling" or introduce blurring. The authors of JiT asked: Can we dynamically focus computation only on salient regions and expand the "active" token set as the image matures?

Methodology: The Core Mechanics

1. SAG-ODE: The Augmented Lifter

The Spatially Approximated Generative ODE (SAG-ODE) allows the model to run a forward pass on only a tiny subset of tokens (e.g., 35%). An "augmented lifter" then takes the computed velocities and extrapolates them to the rest of the latent space. Crucially, the math ensures that the "anchor tokens" remain statistically exact, preserving the learned dynamics where it matters most.

2. Deterministic Micro-Flow (DMF)

Transitioning from a sparse grid to a dense grid is usually where artifacts (like blockiness) appear. JiT solves this with DMF, a short, finite-time ODE that "hits" a target state for newly activated tokens. Instead of just injecting noise, it blends structural priors from existing anchors with the correct noise level, ensuring a seamless dive into the next stage of resolution.

Figure 1: The JiT framework overview. (a) SAG-ODE for sparse evolution, (b) DMF for artifact-free transitions, and (c) the Importance-guided strategy.

3. ITA: Importance-guided Token Activation

Rather than using a fixed grid, JiT uses ITA to find where the "action" is. It measures the local variance of the velocity field—regions with high variation (like edges and complex textures) are prioritized for activation in the next stage.

Performance: SOTA Results

JiT was tested against the heavyweights of the field, including FLUX.1-dev.

• Speed: Achieves 4.24× and 7.07× speedups.
• Quality: Unlike other methods that fail at rendering text under high acceleration, JiT maintains legible typography and sharp edges.
• User Preference: In blind tests, JiT-generated images were preferred over baselines in over 70% of votes.

Figure 2: Qualitative comparison showing JiT's ability to maintain structural integrity and text legibility at 7× acceleration.

Deep Insight: Beyond 2D Images

What makes JiT particularly exciting is its extensibility. Because it treats the latent space as a sequence of tokens, it is backbone-agnostic. The authors successfully applied it to:

1. Qwen-image: Reaching ~4× speedup.
2. HunyuanVideo-1.5: Proving that spatial-temporal redundancy in video can be exploited using the same token-subset logic.

Conclusion & Limitations

JiT represents a shift from "brute-force" inference to intelligent, content-aware sampling. While it works remarkably well training-free, future work might involve fine-tuning the lifter operators to be even more "semantically aware."

For developers and researchers, JiT offers a practical way to deploy massive models like FLUX on consumer-grade hardware without the typical ~20-second wait times, bringing us one step closer to real-time high-fidelity generation.