Truncated Jump Sampling (TJS) in Diffusion Models
- Truncated Jump Sampling (TJS) is a training-free inference acceleration method that decodes the clean sample directly from an intermediate ODE state using endpoint decodability.
- It leverages standard affine probability paths and MMSE estimators to achieve reductions of neural function evaluations by 20–70% while maintaining near-matched sample quality.
- TJS avoids retraining and distillation by truncating ODE integration at an early exit, making it robust across various diffusion and flow matching models.
Searching arXiv for papers on Truncated Jump Sampling and closely related work. Truncated Jump Sampling (TJS) is a training-free accelerated generation method for diffusion and flow matching models in which the ODE trajectory is stopped at an early-exit time and the clean sample is decoded directly from the intermediate state using endpoint decodability (Peng et al., 7 Jul 2026). In this sense, TJS is an inference-time procedure rather than a retrained sampler: it requires no retraining, distillation, or architecture change, and is formulated for standard affine probability paths already used by existing checkpoints. Reported results across SDXL, SD3.5M, Z-Image-Turbo, and class-conditional benchmarks show reductions of neural function evaluations (NFEs) by $20$– with near-matched quality (Peng et al., 7 Jul 2026).
1. Concept and scope
TJS was introduced to address a practical limitation of diffusion and flow matching models: their ODE samplers often require tens to hundreds of NFEs, which remains burdensome when only released checkpoints are available and retraining-based accelerators are undesirable (Peng et al., 7 Jul 2026). Its central claim is that standard affine probability paths already contain sufficient information for endpoint prediction before the trajectory reaches the nominal data endpoint. Rather than integrating the sampler to , TJS truncates the path at an intermediate time and “jumps” to a decoded estimate of the clean sample.
The method is explicitly positioned against acceleration strategies that depend on retraining, distillation, or trajectory redesign. Its premise is not that the path should be straighter, but that an intermediate state together with the model’s learned path velocity already determines a principled estimate of . This places TJS within the broader literature on training-free inference acceleration while giving it a narrower technical meaning than generic early stopping: the early exit is coupled to an algebraic endpoint decoder and justified as an MMSE estimator under the usual training objective (Peng et al., 7 Jul 2026).
2. Affine probability paths and endpoint decodability
The formal setting is the affine path family
with as pure noise and as pure data, and with schedules such as linear, VP, VE, and EDM included in this formulation (Peng et al., 7 Jul 2026). The associated conditional velocity is
0
Stacking 1 and 2 yields a 3 linear system in 4: 5
The determinant
6
is the path determinant. When 7, the map from 8 to 9 is invertible, a property termed endpoint decodability. The endpoint decoder is then
$20$0
In practice $20$1 is not observed exactly; the model supplies a predictor $20$2, giving
$20$3
A second theorem identifies this decoder with the minimum mean-square error estimator. If the optimal predictor is $20$4, then plugging it into the decoder yields
$20$5
Accordingly, standard diffusion and flow matching models trained with the usual $20$6 objective are implicitly trained to support endpoint prediction, even when they were not explicitly parameterized as $20$7-predictors (Peng et al., 7 Jul 2026).
3. Algorithmic form and unified decoders
Operationally, TJS is simple. Given a pretrained model and a nominal step budget $20$8, one chooses an early-exit fraction $20$9, integrates the ODE only to 0, evaluates the network once more at 1, decodes 2, and returns that decoded sample. In the notation of the paper, “TJS-3” denotes this early-exit policy, and the NFE count is reduced from 4 to 5 (Peng et al., 7 Jul 2026).
Because endpoint decodability is tied to the affine path rather than to a single network parameterization, several standard prediction forms admit explicit endpoint decoders:
| Parameterization | Endpoint decoder |
|---|---|
| Velocity prediction | 6 |
| Noise prediction | 7 |
| Score prediction | 8 |
| Direct 9-prediction | 0 |
The paper states that direct 1-prediction is theoretically optimal for TJS, especially for extreme early exits, because it avoids divisions by 2 or 3 that can amplify errors at low SNR. Other parameterizations remain compatible through the algebraic conversions above, so the method is not restricted to a single training convention (Peng et al., 7 Jul 2026).
4. Error decomposition and theoretical position
The principal error decomposition for TJS writes the decoded output as
4
where 5 is the MMSE estimator and 6 is model error. The expected output error at exit time 7 is
8
with
9
The first term is predictor error; the second is irreducible endpoint uncertainty. A notable claim is that no term in this expression depends on trajectory curvature such as 0, 1, or 2, unlike Euler-type ODE integration error (Peng et al., 7 Jul 2026).
This yields a direct comparison with trajectory-straightening approaches. TJS does not require straight trajectories or trajectory redesign, including rectified-flow style arguments. For equal NFE budgets, the paper compares TJS to coarse Euler integration and proves
3
where
4
In this comparison, TJS is strictly superior when endpoint uncertainty is smaller than the Euler curvature penalty (Peng et al., 7 Jul 2026).
The same framework admits an information-theoretic reading. The quantity 5 decreases as 6, and the paper relates it to the I-MMSE theorem through
7
This suggests that the utility of early exit depends on how rapidly 8 concentrates information about 9: if 0 drops quickly at moderate 1, further integration has diminishing benefit relative to its NFE cost (Peng et al., 7 Jul 2026).
5. Empirical behavior and operational regimes
Across reported experiments, TJS shows a monotonic quality trend with respect to exit depth: FID, ImageReward, HPSv2, PickScore, and CLIP score all improve as 2 increases toward the full trajectory (Peng et al., 7 Jul 2026). The practical interest lies in the intermediate regime, where many models exhibit a “sweet spot” around 3–4, achieving at least 5–6 of full-trajectory quality with substantially fewer NFEs. In several cases, moderate early exit even exceeds the quality of the original full ODE sample, which the paper attributes to avoidance of integration-induced discretization errors in late steps.
The method is reported on SDXL, SD3.5M, Z-Image-Turbo, FLUX, DiT, and class-conditional U-Nets, as well as ImageNet-256, CIFAR-10, and MNIST. Over these settings, NFEs are reduced by 7–8 at near-matched quality, and on already-distilled models such as Z-Image-Turbo with 10 steps, an additional 9 NFE reduction at 0 quality is reported (Peng et al., 7 Jul 2026).
The same experiments characterize TJS as sampler-agnostic and schedule-agnostic. It is reported to work with DDIM, UniPC, and DPM++, and with 1, exponential, Karras, Laplace, and EDM schedules, provided 2. It composes with classifier-free guidance, with distillation and fast solvers such as Consistency, Rectified Flow, or Turbo, and is described as robust to whether the checkpoint was trained with vanilla training, distillation, or trajectory straightening (Peng et al., 7 Jul 2026).
6. Relations, misconceptions, and overlapping uses of the term
A common misconception is that TJS is simply the reuse of intermediate 3-predictions already present in DDIM. The distinction drawn in the literature is narrower and more formal: DDIM computes an intermediate 4 at each step, but did not propose using that estimate as the final sample and did not supply the endpoint-decodability or MMSE justification for early stopping. TJS generalizes the argument to all affine paths, including flow matching, and frames the early exit as a principled estimator rather than a heuristic truncation (Peng et al., 7 Jul 2026).
TJS is also distinct from schedule-optimization methods in discrete diffusion. “Jump Your Steps” optimizes the allocation of discrete sampling timesteps by minimizing Compounding Decoding Error through a KL-based schedule search, without extra computational cost and without changing the model architecture or loss (Park et al., 2024). That objective concerns discrete parallel sampling schedules rather than algebraic endpoint decoding on affine continuous paths. The two methods are therefore complementary rather than interchangeable.
The acronym itself is not globally unique. In continuous-time Markov-jump-process inference, “TJS” appears in comparison with endpoint-conditioned path samplers that attempt to sample the random time grid directly from its posterior, independent of the current trajectory, typically using matrix exponentiation together with rejection or importance sampling (Rao et al., 2012). In noisy quantum-circuit simulation, tensor jump methods are described as incorporating or extending TJS-style variance-aware unravelings in an MPS framework (Fröhlich et al., 1 Jul 2026). These usages are technically distinct from the endpoint-decoding method introduced for diffusion and flow matching models, and the modern generative-model sense is best understood as the specific training-free early-exit procedure grounded in endpoint decodability (Peng et al., 7 Jul 2026).