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

Accurate sampling from continuous-time generative models can require many model evaluations. Low-step sampling reduces this cost, but coarse updates introduce endpoint errors that can accumulate along the sampling trajectory and degrade generation quality. We introduce TraceSampler, a post-training method that corrects these errors without modifying the pretrained model. Rather than directly optimising the low-step update itself, TraceSampler models its 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 the high-dimensional error predictable from the current sampling trace. This representation combines shared residual structure, directions derived from the trace, and a low-rank remainder, so shallow per-step MLP heads need only predict sample-dependent coefficients. These coefficients drive a two-stage learned residual correction: the first stage predicts the full endpoint residual, while a backbone evaluation at the intermediate state supplies new trace information for predicting the residual that remains; 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. Code is provided in the supplementary file.