💻 computer science

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

This paper introduces Just-in-Time (JiT), a training-free framework that accelerates Diffusion Transformers by leveraging spatial redundancy to compute updates on a sparse subset of anchor tokens, achieving up to a 7x speedup with nearly lossless image quality on the FLUX.1-dev model.

Original authors: Wenhao Sun, Ji Li, Zhaoqiang Liu

Published 2026-03-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Wenhao Sun, Ji Li, Zhaoqiang Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are an artist tasked with painting a massive, hyper-realistic mural of a bustling city.

The Old Way (The Problem):
Currently, the most advanced AI artists (called Diffusion Transformers) work like a perfectionist who paints the entire mural one tiny square inch at a time, from the very first stroke to the last. They start with a blank canvas covered in static (noise) and slowly refine every single pixel simultaneously.

The problem? This is incredibly slow. Even for a simple image, the AI has to "think" about every single pixel thousands of times before the picture is done. It's like trying to fill a swimming pool by pouring water in one drop at a time, checking the level after every drop. It's accurate, but it takes forever and requires a massive amount of energy.

The New Idea (JiT - Just-in-Time):
The paper introduces a new method called JiT (Just-in-Time). Instead of painting the whole wall at once, JiT acts like a smart, efficient foreman who knows exactly when and where to send the painters.

Here is how JiT works, broken down into three simple concepts:

1. The "Big Picture First" Strategy (Spatial Redundancy)

When you look at a blurry photo coming into focus, what do you see first? You see the big shapes: the sky, the horizon, the outline of a building. You don't see the texture of the bricks or the individual leaves on a tree until the very end.

The Analogy: Imagine building a house. You wouldn't start by installing the doorknobs and painting the trim before you've even poured the foundation.
How JiT uses this: In the early stages of generating an image, JiT realizes that most of the "pixels" (called tokens) are just filling in empty space or repeating the same background information. So, it tells the AI: "Hey, don't waste brainpower on the empty sky or the blurry background yet. Just focus on the few key spots that define the main shape."
The Result: The AI only does the heavy math on a small, sparse set of "anchor" tokens (the important bits) and guesses the rest. This saves a massive amount of time.

2. The "Smart Guess" (SAG-ODE)

You might ask, "If it's only looking at a few spots, won't the rest of the image be a mess?"

The Analogy: Think of a jigsaw puzzle. If you have the corner pieces and a few key pieces in the middle, you can easily guess what the picture in between looks like without actually looking at every single piece.
How JiT uses this: JiT uses a clever mathematical trick (called a Lifter) to take the calculations from those few important "anchor" spots and smoothly "stretch" them to fill the whole image. It's like taking a low-resolution sketch and using a smart algorithm to fill in the gaps so it looks like a high-res photo, without actually doing the work for every single pixel.

3. The "Seamless Handoff" (Deterministic Micro-Flow)

As the image gets clearer, the AI needs to start paying attention to more details (like the texture of the bricks or the eyes of a person). JiT gradually "wakes up" more parts of the image.

The Analogy: Imagine a construction crew. At first, only the foundation crew is working. Then, the framing crew arrives. If the framing crew just jumps in and starts hammering randomly, the house might collapse. They need a smooth transition where the new workers fit perfectly with the old ones.
How JiT uses this: When JiT decides to "wake up" a new section of the image, it uses a special Micro-Flow process. It gently guides the new, inactive parts of the image to match the style and noise level of the active parts. This prevents "glitches" or weird artifacts (like a blurry face suddenly appearing on a sharp body). It ensures the transition is invisible to the human eye.

The "Just-in-Time" Magic

The name Just-in-Time comes from manufacturing (like Toyota's car factories), where parts are delivered exactly when they are needed, not before.

Old Way: The AI tries to process the whole image (all parts) at every single step, even when it doesn't need to.
JiT Way: JiT delivers the "computational power" (the parts) Just-in-Time.
- Early stage: Only the "skeleton" of the image is processed.
- Middle stage: The "muscles" and "skin" are added.
- Final stage: The "makeup" and "jewelry" (fine details) are applied.

The Results

The paper tested this on FLUX.1-dev, one of the best image generators available today.

Speed: It made the AI 4 to 7 times faster.
Quality: The images looked almost identical to the slow, perfect version. There were no blurry patches or weird glitches.
No Training Needed: The best part? They didn't have to retrain the AI model. They just changed how the AI runs, like giving a supercomputer a better set of instructions rather than building a new computer.

In Summary:
JiT is like hiring a team of painters who stop trying to paint the whole wall at once. Instead, they focus on the big shapes first, guess the empty spaces, and only bring in the detailed painters when the picture is almost done. The result? A masterpiece created in a fraction of the time.

1. Problem Statement

Diffusion Transformers (DiTs), such as FLUX.1-dev, have achieved state-of-the-art (SOTA) performance in image synthesis. However, their practical deployment is hindered by high computational costs stemming from two factors:

Quadratic Complexity: The self-attention mechanism scales quadratically ( $O(N^2)$ ) with the number of input tokens ( $N$ ), making high-resolution generation expensive.
Iterative Sampling: The denoising process requires hundreds of neural function evaluations (NFEs).

Existing acceleration methods primarily focus on the temporal domain (reducing the number of sampling steps via distillation or advanced solvers) or parameter efficiency (quantization, feature caching). These approaches often overlook spatial redundancy. Diffusion models naturally follow a "coarse-to-fine" generation process: global low-frequency structures emerge early, while high-frequency details are refined later. Current methods treat all spatial regions uniformly throughout the process, wasting computation on regions that do not yet require fine-grained processing.

2. Methodology: Just-in-Time (JiT)

The authors propose JiT, a training-free framework that accelerates generation by dynamically reducing computational load in the spatial domain. Instead of processing all tokens at every step, JiT computes the full velocity field based on a dynamically selected, sparse subset of "anchor tokens," approximating the rest.

The framework consists of two core components:

A. Spatially Approximated Generative ODE (SAG-ODE)

This component governs the evolution of the latent state during a specific stage where only a subset of tokens is active.

Mechanism: It formulates a generative Ordinary Differential Equation (ODE) where the velocity field is computed only on a sparse set of anchor tokens ( $\Omega_k$ ).
Augmented Lifter Operator ( $\Pi_k$ ): To drive the evolution of the entire latent state (including inactive tokens) using only anchor computations, JiT introduces a lifter operator.
- Exact Part: The computed velocity for anchor tokens is placed back into their correct positions.
- Approximation Part: The velocity for inactive tokens is extrapolated using a smooth spatial interpolation operator ( $I_k$ ) based on the anchor velocities.
Consistency: The design ensures that the dynamics of the anchor tokens remain exact (zero approximation error on active tokens), while inactive tokens are guided by a structurally aware approximation.

B. Deterministic Micro-Flow (DMF)

When the generation process transitions from a sparse stage to a denser stage (activating new tokens), a naive injection of state can cause artifacts or statistical mismatches.

Mechanism: The DMF is a short, finite-time ODE that deterministically evolves newly activated tokens from their interpolated state to a statistically correct target.
Target Construction: The target state combines structural information from existing anchors (interpolated to the new dimension) with the correct noise level required by the flow matching trajectory.
Function: This ensures seamless transitions, maintaining latent space continuity and preventing visual artifacts (e.g., blurring or structural inconsistencies) during dimension expansion.

C. Importance-Guided Token Activation (ITA)

Instead of using a static grid to select tokens, JiT employs a dynamic strategy:

Metric: It calculates a spatial importance map based on the local variance of the predicted velocity field ( $u_\theta$ ).
Strategy: Regions with high velocity variance (indicating active generation of details like edges or textures) are prioritized. New tokens are activated in these high-information regions first, ensuring computational resources are allocated where they are most needed.

3. Key Contributions

Novel Training-Free Framework: JiT is the first training-free method to achieve spatial acceleration in DiTs by completely obviating the need for error-prone upsampling/correction steps found in hierarchical methods.
SAG-ODE & Lifter Operator: Introduces a mathematically grounded ODE that drives full-state evolution using sparse computations, preserving exact dynamics on critical tokens.
Deterministic Micro-Flow (DMF): A novel mechanism to handle stage transitions, ensuring statistical correctness and structural coherence when expanding the token set.
Dynamic Resource Allocation: The ITA strategy dynamically focuses computation on salient spatial regions, aligning with the intrinsic coarse-to-fine nature of diffusion.

4. Experimental Results

The method was evaluated on the FLUX.1-dev model (a leading DiT) and extended to Qwen-image and HunyuanVideo-1.5.

Speedup: JiT achieves up to 7× speedup (reducing inference time from ~25s to ~3.6s) and 4× speedup with negligible quality loss.
Quality Metrics:
- CLIP-IQA / ImageReward / HPSv2.1: JiT outperforms SOTA baselines (including RALU, Bottleneck Sampling, TaylorSeer, and TeaCache) in both automated metrics and human preference.
- Text-Image Alignment: Maintains high scores on GenEval and T2I-CompBench, preserving semantic coherence even at 7× acceleration.
Visual Quality: Unlike competing methods that suffer from blurring, semantic errors, or structural artifacts, JiT produces high-fidelity images with legible text and fine details even at extreme acceleration.
Ablation Studies: Removing the SAG-ODE approximation or the DMF target leads to catastrophic quality drops, confirming the necessity of both components.

5. Significance

Paradigm Shift: Moves the field from uniform temporal acceleration to dynamic spatial acceleration, exploiting the inherent redundancy in generative processes.
Practical Deployment: Offers a "plug-and-play" solution (training-free) that significantly lowers hardware requirements for high-fidelity image and video generation, making DiTs viable for real-time and consumer-grade applications.
Generalizability: Demonstrates effectiveness across different model architectures (FLUX, Qwen) and modalities (Image, Video), proving that the coarse-to-fine spatial redundancy is a fundamental property of diffusion probability flows.

In summary, JiT redefines the speed-fidelity trade-off in Diffusion Transformers by intelligently deferring computation on less critical spatial regions until they are needed, achieving massive acceleration without retraining or sacrificing generation quality.