TraceSampler: Making Sampling Errors Predictable and Correctable for Accelerated Generative Sampling

George Matthew; Guoqiang Zhang; Kenta Niwa

cs.CVOct 9, 2026

TraceSampler: Making Sampling Errors Predictable and Correctable for Accelerated Generative Sampling

Authors: George Matthew, Guoqiang Zhang, Kenta Niwa

Organizations: University of Exeter, United Kingdom · NTT Communication Science Laboratories, Japan

Abstract

Accurate sampling from continuous-time generative models can require many model evaluations. Low-step sampling reduces cost, but coarse updates introduce endpoint errors that can accumulate and degrade generation quality. We introduce TraceSampler, a post-training method that corrects these errors without modifying the pretrained model. Rather than optimising low-step updates directly, TraceSampler models endpoint error as a sample-specific prediction problem. A higher-step reference integration over the same interval provides the supervised target, while a compact representation makes high-dimensional error predictable from the sampling trace. This representation combines shared residual structure, trace-derived directions, and a low-rank remainder, so shallow per-step MLP heads predict sample-dependent coefficients. These coefficients drive a two-stage learned residual correction: the first stage predicts the full endpoint residual, while an intermediate backbone evaluation provides new trace information to predict the remaining residual; this evaluation is reused by the following step. TraceSampler is evaluated with Rectified Flow on CIFAR-10, SiT-XL/2 on ImageNet-256, and SD3.5 Medium for text-to-image generation. At 10 steps, it reduces trajectory MSE relative to matched-step Euler by 90.5% on CIFAR-10 and 85.6% on ImageNet-256, while achieving FID scores of 2.692 and 2.189, close to 120-step Heun references of 2.601 and 2.056 with 24 $\times$ fewer backbone evaluations. On SD3.5 Medium, TraceSampler improves image quality and prompt following; at 10 steps, ImageReward rises from 0.752 to 0.906 while trajectory MSE falls by 70.7%. These gains preserve Euler's backbone-evaluation budget, add less than 0.4% inference FLOPs, and incur negligible wall-clock overhead, showing that coarse-step sampling error can be made predictable from the sampling trace and efficiently corrected.

Figures & tables

Figure 1: Two-stage learned residual correction for one sampling step.

Figure 2: Offline fitting and inference workflow of TraceSampler.

Model/config.	Offline reference	PCA rank	MLP width	Schedule
RF CIFAR-10	120-step Heun	$K=32$	$h=12288$	Optimised
SiT-XL/2	120-step Heun	$K=32$	$h=12288$	Optimised
SD3.5 Medium (8-step)	64-step Euler	$K=32$	$h=24576$	Shift $=3$
SD3.5 Medium (10-step)	60-step Euler	$K=32$	$h=24576$	Shift $=3$

Table 1: Main TraceSampler fitting and sampling configurations.

Method	Schedule	$K$	FID $\downarrow$	IS $\uparrow$	Trajectory MSE reduction $\uparrow$
Base Euler	Uniform	0	15.729	6.321	0.0%
Predictor directions only	Uniform	0	9.705	7.757	34.7%
Predictor + corrector directions	Uniform	0	4.101	8.612	80.5%
Full TraceSampler	Uniform	32	3.666	8.626	87.5%
Full TraceSampler + schedule	Optimised	32	2.798	8.552	90.8%

Table 2: Component build-up at 8 steps.

Figure 3: Per-step trajectory MSE and the fixed-to-sample-dependent error ratio.

Coefficient prediction	FID $\downarrow$	IS $\uparrow$	Trajectory MSE reduction $\uparrow$
Fixed mean coefficients	3.240	8.563	81.5%
Sample-dependent heads	2.798	8.552	90.8%

Table 3: Coefficient-prediction ablation at 8 steps.

Budget	Euler		IIA–Euler		UniPC		TraceSampler		Traj. MSE red. $\uparrow$
Budget	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	Traj. MSE red. $\uparrow$
5 steps	33.149	4.566	22.253	5.783	12.817	7.238	4.373	8.239	89.9%
6 steps	24.143	5.327	14.940	6.510	9.333	7.730	3.560	8.344	88.8%
8 steps	15.729	6.321	8.846	7.357	7.751	8.011	2.798	8.552	90.8%
10 steps	11.782	6.886	6.668	7.670	6.429	8.169	2.692	8.556	90.5%
12 steps	9.632	7.205	5.671	7.850	5.778	8.256	2.641	8.517	89.5%

Table 4: CIFAR-10 matched-budget comparison.

Budget	Euler		IIA–Euler		UniPC		TraceSampler		Traj. MSE red. $\uparrow$
Budget	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	Traj. MSE red. $\uparrow$
6 steps	19.786	114.66	13.233	151.307	4.630	206.75	3.624	209.74	79.8%
8 steps	10.675	161.18	7.405	185.997	3.105	227.55	2.418	235.21	83.2%
10 steps	7.189	186.80	5.204	204.654	2.544	238.45	2.189	245.15	85.6%
Heun 120-step reference	FID $2.056$ , IS $254.16$

Table 5: ImageNet-256 matched-budget comparison.

Sampler	IR $\uparrow$	FID $\downarrow$	IS $\uparrow$	CLIP $\uparrow$	GenEval $\uparrow$	HPS $\uparrow$	Traj. MSE red. $\uparrow$
Euler 8	0.611	17.320	39.23	0.2698	0.60997	0.2572	–
Euler 10	0.752	17.368	40.48	0.2697	0.64152	0.2702	–
TraceSampler 8	0.862	14.616	43.28	0.2661	0.64554	0.2807	67.9%
TraceSampler 10	0.906	15.478	43.05	0.2659	0.66811	0.2878	70.7%
Euler 64 eval. ref.	0.973	18.059	43.95	0.2651	0.68275	0.2968	–

Table 6: SD3.5 Medium results.

Figure 4: SD3.5 Medium qualitative comparison at matched noise.

Model	GFLOPs/step			Throughput (images/s)		Peak VRAM (GiB)
	Euler	Trace	Increase	Euler	Trace	Euler	Trace
Rectified Flow	21.6931	21.7714	0.3610%	321.45	316.05	6.77	7.52
SiT-XL/2	474.4857	475.0351	0.1158%	31.07	30.97	7.92	8.51
SD3.5 Medium	4275.6574	4278.4379	0.0650%	4.144	4.126	11.60	14.22

Table 7: Inference cost at 8 steps; FLOPs are per step.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Targeted blocks	Compression	Input dim.
Rectified Flow CIFAR-10	Zero-indexed NCSN++ residual blocks $4$ , $19$ , $26$ , and $39$	Spatial mean and standard deviation for each activation channel; no learned activation projection	$1639/1641$
SiT-XL/2	Conditional halves of block.0 , block.26 , and block.27	Fixed block-specific PCA projection to 256 dimensions, followed by token-wise mean and standard deviation	$1928/1930$
SD3.5 Medium	Conditional image streams of block.7 , block.14 , and block.20	Fixed block-specific PCA projection to 256 dimensions, followed by token-wise mean and standard deviation	$3092/3094$

Appendix

Table 8: Model-specific hidden-activation features. Input dimensions are shown as predictor/corrector (P/C).

Setting	Configuration
Architecture	Independent one-hidden-layer ReLU MLP for each predictor step and each non-final corrector step; hidden width $12288$ for Rectified Flow and SiT-XL/2 and $24576$ for SD3.5 Medium
Optimisation	AdamW; initial learning rate $10^{-3}$ , weight decay $10^{-3}$ , dropout $0.1$ , batch size $2048$ , and 20-epoch linear warm-up
Learning-rate schedule	ReduceLROnPlateau with factor $0.5$ and patience 50
Stopping and selection	Up to $20{,}000$ epochs; early stopping with patience 400 on the $4{,}915$ -trajectory validation set; best validation checkpoint retained
Numerical precision	Backbone evaluations in bfloat16 autocast; residual decomposition, coefficient targets, feature normalisation, and MLP fitting in float32

Appendix

Table 9: Prediction-head architecture and training settings.

Figure 5: Overview of the optional timestep-schedule optimisation procedure.

	Rectified Flow CIFAR-10		SiT-XL/2
$k$	$t_{k}$	$\Delta t_{k}$	$t_{k}$	$\Delta t_{k}$
0	0.001000	0.092182	0.001000	0.103061
1	0.093182	0.104293	0.104061	0.104207
2	0.197475	0.115650	0.208269	0.092542
3	0.313125	0.126756	0.300811	0.095130
4	0.439881	0.158832	0.395941	0.105806

Appendix

Table 10: Selected sampling times and step sizes for the representative 8-step schedules.

Setting	Value
Probe trajectories	$4096$ for Rectified Flow CIFAR-10; $2048$ for SiT-XL/2
Sampling-time range	$[10^{-3},1-10^{-3}]$
Minimum step size	$0.001$
Error-scaling exponent	$p=2$
Optimiser	L-BFGS-B
Maximum optimisation passes	$10$

Appendix

Table 11: Timestep-schedule search settings. All settings are shared across the optimised configurations except for the number of probe trajectories.

Ablation	Setting	FID $\downarrow$	IS $\uparrow$	Trajectory MSE $\downarrow$
Reference accuracy	5 Heun substeps	2.977	8.526	$1.730\times 10^{-4}$
	10 Heun substeps	2.925	8.508	$1.900\times 10^{-4}$
	15 Heun substeps	2.798	8.552	$1.647\times 10^{-4}$
	20 Heun substeps	2.845	8.581	$1.804\times 10^{-4}$
PCA rank	$K=16$	2.883	8.530	$2.001\times 10^{-4}$
	$K=32$	2.798	8.552	$1.647\times 10^{-4}$

Appendix

Table 12: Key Rectified Flow CIFAR-10 design-space ablation results.

Figure 7: Qualitative comparison between Euler and TraceSampler on SiT-XL/2 ImageNet-256.

Label	Prompt
P1	A fluffy white cat sitting on a mossy mountain ledge under a huge full moon, soft starlight, crisp detailed fur, cinematic night landscape.
P2	A close-up portrait of a cyberpunk geisha on a rainy Tokyo rooftop, detailed expressive face, neon skyline, holographic lanterns, wet silk kimono, cinematic lighting.
P3	A masked knight in ornate silver armour standing inside a ruined cathedral, sword point resting on the stone floor in front of him, both hands on the hilt, glowing stained-glass light, misty air, cinematic atmosphere.
P4	A close-up portrait of an old lighthouse keeper in a storm, detailed weathered face, thick wool coat, lantern glow, rain on glass, cinematic lighting.

Appendix

Table 13: Prompts used in the SD3.5 Medium qualitative comparison.

Model	Schedule optimisation	Artefact fitting	Total
RF CIFAR-10	0:26:57	2:32:42	2:59:39
SiT-XL/2	1:07:04	9:02:39	10:09:42
SD3.5 Medium	–	8:35:00	8:35:00

Appendix

Table 14: Measured offline preparation time for 8-step TraceSampler configurations.

Model	Predictor (M)	Corrector (M)
Rectified Flow	20.570	20.619
SiT-XL/2	25.30	25.35
SD3.5 Medium	88.646	88.744

Appendix

Table 15: Parameters per TraceSampler coefficient head.

Baseline	Configuration
IIA–Euler	Basic $r=1$ IIA adapted to the flow-matching update. Two global coefficients combine the current and previous Euler increments. Calibration uses 512 trajectories, three fine-Euler substeps per coarse step, and ridge regularisation of $10^{-10}$ .
UniPC	Multistep predictor with order 3, the bh2 variant, log-SNR spacing, data_prediction , lower-order final steps, and one model evaluation per deployed step.

Appendix

Table 16: Configurations used for the sampler baselines.

Evaluation	Configuration
CIFAR-10	$50000$ generated samples. FID uses 2048-dimensional Inception-V3 features against CIFAR-10 training-set statistics; IS uses ten splits.
ImageNet-256	$50000$ class-conditional samples. FID and IS use the VIRTUAL_imagenet256_labeled.npz reference statistics and ten IS splits.
COCO captions	One image for each of $30000$ captions. The evaluation reports ImageReward-v1.0, Inception FID and IS, raw CLIP-L/14 image–text cosine similarity, reported as CLIP similarity, and trajectory MSE against the matched higher-step reference.
GenEval	Standard 553-prompt metadata with four images per prompt, CFG $4.5$ , and the official MMDetection-2 scoring pipeline.
HPSv2.1	Official anime, concept-art, paintings, and photo prompt subsets, with one image per retained prompt and the mean of the four subset scores. Prompts exceeding 77 tokens for any SD3 text encoder are excluded.

Appendix

Table 17: Evaluation sets and scoring configurations.

Few-Step Diffusion Sampling

Few-step diffusion sampling develops ways to generate data with far fewer denoising evaluations while retaining quality and diversity.

See all 173 papers

May 8, 2026

Normalizing Trajectory Models

Diffusion-based models decompose sampling into many small Gaussian denoising steps -- an assumption that breaks down when generation is compressed to a few coarse transitions. Existing few-step methods address this through distillation, consistency training, or adversarial objectives, but sacrifice the likelihood framework in the process. We introduce Normalizing Trajectory Models (NTM), which models each reverse step as an expressive conditional normalizing flow with exact likelihood training. Architecturally, NTM combines shallow invertible blocks within each step with a deep parallel predictor across the trajectory, forming an end-to-end network trainable from scratch or initializable from pretrained flow-matching models. Its exact trajectory likelihood further enables self-distillation: a lightweight denoiser trained on the model's own score produces high-quality samples in four steps. On text-to-image benchmarks, NTM matches or outperforms strong image generation baselines in just four sampling steps while uniquely retaining exact likelihood over the generative trajectory.

Jiatao Gu, Tianrong Chen, Ying Shen +3

Apple · UIUC

Normalizing Flows Diffusion Model Distillation Few-Step Diffusion Sampling Diffusion-Based Image Generation Flow-Based Generative Modeling Diffusion Models

Ask PDF

Jul 7, 2026

x-Prediction Is All You Need:Training-Free Accelerated Generation via Endpoint Decodability

Diffusion and flow matching models generate high-quality samples, but their ODE samplers often need tens to hundreds of neural function evaluations (NFEs). This remains a practical challenge for released checkpoints, since many accelerators require additional design choices and training cost through retraining, distillation, or trajectory redesign. We investigate a different route based on

x

-prediction. During sampling, standard affine probability paths already expose

x_0

information: an intermediate state and its path velocity determine a principled estimate of the clean sample. We formalize this property as endpoint decodability and show that the decoder is the minimum-MSE estimator

\mathbb{E}[x_0\mid x_t]

under the usual

\ell_2

objective. This yields Truncated Jump Sampling (TJS): stop the ODE at an early-exit time

t^*

and return the decoded

x_0

. TJS requires no retraining, distillation, or architecture change. Across SDXL, SD3.5M, Z-Image-Turbo, and three class-conditional benchmarks, it reduces NFEs by 20--70% with near-matched quality. The analysis also shows why endpoint prediction can work without straightening the trajectory, providing inference acceleration without trajectory redesign.

Xin Peng, Ang Gao

School of Physical Science and Technology, Beijing University of Posts and Telecommunications, Beijing, China

Diffusion Sampling Few-Step Diffusion Sampling Flow-Based Generative Modeling Diffusion Model Inference Acceleration

Ask PDF

May 12, 2026

Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching

Flow matching models generate samples by numerically integrating a learned velocity field, with each integration step requiring a neural network evaluation. Fast generation therefore requires using a small fixed evaluation budget effectively: the key question is not only how to integrate the flow, but where the sampler should spend its steps. We propose SharpEuler, a training-free sampler that profiles a pretrained model offline by estimating where the learned velocity field changes most rapidly along calibration trajectories. This finite-difference estimate defines a solver-aware sharpness profile, which is smoothed and converted by a quantile transform into a timestep grid for any desired inference budget. At test time, sampling remains ordinary Euler integration with the same number of model evaluations as a uniform schedule. We justify SharpEuler using three principles: a numerical principle identifying trajectory acceleration as the leading source of Euler discretization error, a variational principle deriving sharpness-based power-law timestep densities, and a statistical guarantee showing that the finite-sample calibrated sampler is stable at the terminal distribution level. Our experiments show that SharpEuler improves sample quality at fixed budgets, reducing inter-mode leakage and increasing mode coverage.

Aditi Gupta, Soon Hoe Lim, Annan Yu +1

Lawrence Berkeley National Laboratory · International Computer Science Institute · 1Lawrence Berkeley National Laboratory, 2International Computer Science Institute +4

Flow Matching Flow-Based Generative Modeling

Ask PDF

Abstract

Few-Step Diffusion Sampling

Explore similar work

Normalizing Trajectory Models

x-Prediction Is All You Need:Training-Free Accelerated Generation via Endpoint Decodability

Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching