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

arXiv: https://arxiv.org/abs/2610.14332 · HTML: https://mlsanity.com/paper/2610.14332/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 20 references that are in the map.
- First posted on arXiv 9 Oct 2026.

## Builds on (most cited first, 17 of 17)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment](https://mlsanity.com/paper/2310.11513/) (2023)
- [DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models](https://mlsanity.com/paper/2211.01095/) (2022)
- [Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis](https://mlsanity.com/paper/2306.09341/) (2023)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models](https://mlsanity.com/paper/2302.04867/) (2023)
- [DC-Solver: Improving Predictor-Corrector Diffusion Sampler via Dynamic Compensation](https://mlsanity.com/paper/2409.03755/) (2024)
- [On Accelerating Diffusion-Based Sampling Process via Improved Integration Approximation](https://mlsanity.com/paper/2304.11328/) (2023)
- [Differentiable Solver Search for Fast Diffusion Sampling](https://mlsanity.com/paper/2505.21114/) (2025)
- [Diffusion Sampling Correction via Approximately 10 Parameters](https://mlsanity.com/paper/2411.06503/) (2024)

## Abstract (arXiv)

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× 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.
