TraceSampler improves image and text generative model quality fast
TraceSampler: Making Sampling Errors Predictable and Correctable for Accelerated Generative Sampling
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Sampling from complex AI models that generate images or text can take a long time, and speeding it up usually lowers quality. The authors present TraceSampler, a method that fixes quality errors caused by taking fewer steps without changing the original model. It uses extra calculations from longer-step references to learn how to predict and correct the mistakes made in quick sampling. This approach keeps the AI's speed fast but delivers outputs that nearly match the quality of much slower methods.
What this means in practice
- •For image generation engineers: Improve the quality and fidelity of image outputs from generative models while reducing runtime by up to 24 times using low-step corrected sampling.
- •For text-to-image system developers: Enhance prompt accuracy and image quality in text-to-image synthesis applications with minimal added computation by applying learned residual error corrections.
Authors
George Matthew, Guoqiang Zhang, Kenta Niwa
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.