---
title: Truncated Jump Sampling (TJS) in Diffusion Models
url: https://www.emergentmind.com/topics/truncated-jump-sampling-tjs
type: topic
---

# Truncated Jump Sampling (TJS) in Diffusion Models

Searching arXiv for recent 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 \(t^*<1\) and the clean sample \(x_0\) is decoded directly from the intermediate state using endpoint decodability [2607.06114]. 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\)–\(70\%\) with near-matched quality [2607.06114].

## 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 [2607.06114]. 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 \(t=1\), 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 \(x_0\). 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 \(\ell_2\) training objective [2607.06114].

## 2. Affine probability paths and endpoint decodability

The formal setting is the affine path family
\[
x_t=\alpha_t x_0+\sigma_t \epsilon,\qquad \epsilon\sim\mathcal N(0,\mathbf I),\qquad t\in[0,1],
\]
with \(t=0\) as pure noise and \(t=1\) as pure data, and with schedules such as linear, VP, VE, and EDM included in this formulation [2607.06114]. The associated conditional velocity is
\[
u_t=\frac{dx_t}{dt}=\dot{\alpha}_t x_0+\dot{\sigma}_t \epsilon.
\]
Stacking \(x_t\) and \(u_t\) yields a \(2\times 2\) linear system in \((x_0,\epsilon)\):
\[
\begin{bmatrix} x_t \\ u_t \end{bmatrix}
=
\begin{bmatrix}
\alpha_t & \sigma_t\\
\dot{\alpha}_t & \dot{\sigma}_t
\end{bmatrix}
\begin{bmatrix} x_0 \\ \epsilon \end{bmatrix}.
\]

The determinant
\[
\Delta_t:=\alpha_t\dot{\sigma}_t-\sigma_t\dot{\alpha}_t
\]
is the path determinant. When \(\Delta_t\neq 0\), the map from \((x_t,u_t)\) to \((x_0,\epsilon)\) is invertible, a property termed **endpoint decodability**. The endpoint decoder is then
\[
x_0=\frac{\sigma_t u_t-\dot{\sigma}_t x_t}{\Delta_t}.
\]
In practice \(u_t\) is not observed exactly; the model supplies a predictor \(v_\theta(x_t,t)\approx u_t\), giving
\[
\hat{x}_0^{\mathrm{vel}}=\frac{\sigma_t v_\theta(x_t,t)-\dot{\sigma}_t x_t}{\Delta_t}.
\]

A second theorem identifies this decoder with the minimum mean-square error estimator. If the optimal predictor is \(v^\star(x,t)=\mathbb E[u_t\mid x_t=x]\), then plugging it into the decoder yields
\[
\mathbb E[x_0\mid x_t]=\frac{\sigma_t v^\star(x_t,t)-\dot{\sigma}_t x_t}{\Delta_t}.
\]
Accordingly, standard diffusion and flow matching models trained with the usual \(\ell_2\) objective are implicitly trained to support endpoint prediction, even when they were not explicitly parameterized as \(x_0\)-predictors [2607.06114].

## 3. Algorithmic form and unified decoders

Operationally, TJS is simple. Given a pretrained model and a nominal step budget \(K\), one chooses an early-exit fraction \(\gamma\in(0,1]\), integrates the ODE only to \(t^*=k^*/K=\gamma\), evaluates the network once more at \(t^*\), decodes \(\hat{x}_0\), and returns that decoded sample. In the notation of the paper, “TJS-\(\gamma\)” denotes this early-exit policy, and the NFE count is reduced from \(K\) to \(\lceil \gamma K\rceil+1\) [2607.06114].

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 | \(\hat{x}_0^{\mathrm{vel}}(x,t)=\dfrac{\sigma_t v_\theta(x,t)-\dot{\sigma}_t x}{\Delta_t}\) |
| Noise prediction | \(\hat{x}_0^{\mathrm{noise}}(x,t)=\dfrac{x-\sigma_t\epsilon_\theta(x,t)}{\alpha_t}\) |
| Score prediction | \(\hat{x}_0^{\mathrm{score}}(x,t)=\dfrac{x+\sigma_t^2 s_\theta(x,t)}{\alpha_t}\) |
| Direct \(x_0\)-prediction | \(\hat{x}_0=x_\theta(x,t)\) |

The paper states that direct \(x_0\)-prediction is theoretically optimal for TJS, especially for extreme early exits, because it avoids divisions by \(\Delta_t\) or \(\alpha_t\) 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 [2607.06114].

## 4. Error decomposition and theoretical position

The principal error decomposition for TJS writes the decoded output as
\[
\hat{x}_0(x_t,t)=m_t(x_t)+e_t(x_t),
\]
where \(m_t(x_t)=\mathbb E[x_0\mid x_t]\) is the MMSE estimator and \(e_t(x_t)\) is model error. The expected output error at exit time \(t^*\) is
\[
\mathbb E[\|x_{\mathrm{out}}-x_0\|^2]
=
\mathbb E[\|e_{t^*}\|^2]+\mathcal U(t^*),
\]
with
\[
\mathcal U(t^*)=\mathbb E_{x_{t^*}}\!\left[\operatorname{Tr}\,\operatorname{Var}(x_0\mid x_{t^*})\right].
\]
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 \(\ddot{x}_t\), \(\ddot{\alpha}_t\), or \(\ddot{\sigma}_t\), unlike Euler-type ODE integration error [2607.06114].

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
\[
\mathrm{MSE}_{\mathrm{TJS}}-\mathrm{MSE}_{\mathrm{Euler}}
\le
\mathcal U(t^*)-\frac{h^2}{2}C_{\alpha,\sigma}\,\mathbb E[\|x_0\|^2+\|\epsilon\|^2]+2\epsilon,
\]
where
\[
C_{\alpha,\sigma}=\sup_{\tau\in[t^*,1]}\left(|\ddot{\alpha}_\tau|^2+|\ddot{\sigma}_\tau|^2\right).
\]
In this comparison, TJS is strictly superior when endpoint uncertainty is smaller than the Euler curvature penalty [2607.06114].

The same framework admits an information-theoretic reading. The quantity \(\mathcal U(t)=\mathrm{MMSE}(x_0\mid x_t)\) decreases as \(t\to 1\), and the paper relates it to the I-MMSE theorem through
\[
\frac{dI(x_0;x_t)}{d\mathrm{SNR}}=\frac{1}{2}\mathcal U(t).
\]
This suggests that the utility of early exit depends on how rapidly \(x_t\) concentrates information about \(x_0\): if \(\mathcal U(t)\) drops quickly at moderate \(t\), further integration has diminishing benefit relative to its NFE cost [2607.06114].

## 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 \(k^*\) increases toward the full trajectory [2607.06114]. The practical interest lies in the intermediate regime, where many models exhibit a “sweet spot” around \(k^*\approx 0.6\)–\(0.7K\), achieving at least \(90\)–\(95\%\) 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 \(20\)–\(70\%\) at near-matched quality, and on already-distilled models such as Z-Image-Turbo with 10 steps, an additional \(70\%\) NFE reduction at \(95\%+\) quality is reported [2607.06114].

The same experiments characterize TJS as sampler-agnostic and schedule-agnostic. It is reported to work with DDIM, UniPC, and DPM++, and with \(\beta\), exponential, Karras, Laplace, and EDM schedules, provided \(\Delta_t\neq 0\). 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 [2607.06114].

## 6. Relations, misconceptions, and overlapping uses of the term

A common misconception is that TJS is simply the reuse of intermediate \(x_0\)-predictions already present in DDIM. The distinction drawn in the literature is narrower and more formal: DDIM computes an intermediate \(\hat{x}_0\) 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 [2607.06114].

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 [2410.07761]. 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 [1208.4818]. In noisy quantum-circuit simulation, tensor jump methods are described as incorporating or extending TJS-style variance-aware unravelings in an MPS framework [2607.01323]. 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 [2607.06114].

Source: https://www.emergentmind.com/topics/truncated-jump-sampling-tjs