Technical Summary: x-Prediction Is All You Need

Problem Statement

Diffusion and flow matching models have established state-of-the-art performance in generative tasks but suffer from high inference costs, typically requiring tens to hundreds of neural function evaluations (NFEs) per sample. Existing acceleration strategies—such as trajectory rectification (Rectified Flow), distillation (progressive, consistency-based, bootstrapping), and advanced ODE solvers—generally require modifications to the training recipe, model architecture, or the trajectory itself. Consequently, these methods are often incompatible with the vast library of existing pretrained checkpoints (e.g., SDXL, SD3.5M, FLUX) without incurring retraining costs. The field lacks a training-free acceleration strategy that works on models "as they are."

Methodology: Truncated Jump Sampling (TJS)

The paper proposes Truncated Jump Sampling (TJS), a training-free inference strategy based on the concept of endpoint decodability.

Core Insight: Endpoint Decodability

The authors formalize that for any non-degenerate affine probability path (the standard structure for diffusion and flow matching models defined as $x_t = \alpha_t x_0 + \sigma_t \epsilon$), the clean sample $x_0$ can be uniquely recovered from any intermediate state $x_t$ and its path velocity $u_t$ via a closed-form decoder.

• Algebraic Recovery: Given the coupled linear system of position and velocity, $x_0$ is recovered as $x_0 = (\sigma_t u_t - \dot{\sigma}_t x_t) / \Delta_t$, where $\Delta_t$ is the path determinant.
• Optimality: Under standard $\ell_2$ training objectives, this algebraic decoder corresponds to the Minimum Mean Square Error (MMSE) estimator $E[x_0 | x_t]$. This implies that standard pretrained models implicitly learn to predict the endpoint $x_0$ at every timestep, even if they are trained to predict noise or velocity.
• Curvature Independence: Crucially, the error of this endpoint prediction depends only on the model's estimation accuracy and the irreducible uncertainty of the state $x_t$, not on the curvature of the trajectory ($\|\ddot{x}_t\|$). This challenges the prevailing assumption that straight trajectories (as in Rectified Flow) are necessary for accurate few-step generation.

The Algorithm

TJS operationalizes this insight by stopping the ODE integration early at a time $t^*$ and immediately outputting the decoded $x_0$ rather than continuing to $t=0$.

1. Partial Integration: The model integrates the ODE for a fraction $\gamma$ of the total steps ($k^* = \lceil \gamma K \rceil$).
2. Early Exit: At step $k^*$, the algorithm computes the endpoint estimate $\hat{x}_0$ using the learned velocity field and the closed-form decoder.
3. Termination: The process stops, returning $\hat{x}_0$ as the final sample.

This approach requires zero retraining, zero distillation, and zero architecture changes. It is compatible with any ODE solver, noise schedule, classifier-free guidance (CFG) scale, or existing checkpoint.

Key Contributions

1. Theoretical Formalization: The paper establishes Endpoint Decodability (Theorem 5) as a universal property of affine paths. It proves that the induced decoder is the MMSE-optimal estimator (Theorem 6) and that the error decomposition is independent of trajectory curvature (Theorem 9).
2. Challenge to Trajectory Straightening: The authors prove that straight trajectories are sufficient but not necessary for accurate endpoint prediction (Proposition 12). This decouples inference acceleration from the need for trajectory rectification, offering a theoretical basis for why TJS works on curved paths where Euler integration would fail.
3. Training-Free Algorithm (TJS): A practical sampler that reduces NFEs by 20–70% across diverse model families without modifying the model or training data.
4. Comprehensive Empirical Validation: Extensive experiments on six model families (SDXL, SD3.5M, Z-Image-Turbo, ImageNet-256, CIFAR-10, MNIST) demonstrate that TJS achieves near-matched or superior quality compared to full ODE trajectories.

Experimental Results

The authors evaluated TJS across class-conditional and text-to-image benchmarks:

• NFE Reduction: TJS reduces NFEs by 20–70% while maintaining near-matched quality.
  • SDXL/SD3.5M: Achieves 95% of full ODE quality (measured by ImageReward) at roughly 37–40% NFE savings (stopping at $k^* \approx 18$ of 30 steps).
  • Z-Image-Turbo: On a distilled 10-step model, TJS achieves 95% quality retention with 70% additional savings (stopping at $k^*=2$, i.e., 3 NFEs total), demonstrating that endpoint decodability persists even in distilled trajectories.
  • Class-Conditional: On MNIST and CIFAR-10, TJS sometimes outperforms the full ODE (the "overshoot" phenomenon) because it avoids the discretization errors accumulated in the final steps of the Euler integration.
• Metric Hierarchy: The experiments reveal a convergence hierarchy where semantic content (CLIP, PickScore) resolves early (often by $k^* \approx 12$), while fine-grained aesthetic details (ImageReward, AES) require deeper integration.
• Robustness: TJS performance is invariant to the choice of ODE solver (DDIM, DPM++, UniPC, etc.), noise schedule (Beta, Karras, etc.), and CFG scale, confirming its orthogonality to existing inference components.

Significance and Claims

The paper claims that moderate few-step acceleration does not require trajectory straightening, distillation, or training redesign. The significance lies in recognizing endpoint decodability as an inherent structural property of generative models on affine paths.

• Practical Impact: TJS provides an immediate, zero-cost acceleration path for the community's existing library of fine-tuned checkpoints, which are often infeasible to retrain or distill.
• Theoretical Shift: By proving that straight trajectories are not necessary for endpoint prediction, the work challenges the foundational motivation of Rectified Flow and Consistency Models regarding inference acceleration.
• Complementary Nature: The authors position TJS as orthogonal to distillation. While distillation compresses the total trajectory length, TJS eliminates the redundant tail of whatever trajectory remains, allowing for additive acceleration (e.g., applying TJS to a distilled model like Z-Image-Turbo).

The paper concludes that recognizing this property "costs nothing" and that ignoring it leaves performance on the table, offering a principled way to trade off integration depth for inference speed without sacrificing the quality of the generated sample.