🤖 AI Summary
Existing diffusion models face a fundamental trade-off between degraded generation quality and vanishing gradient signals when incorporating reward guidance. This work proposes Noise-Tilted Reverse Kernels (NTRK), which, for the first time, enable safe and effective single-sample-per-step sampling by injecting reward gradients—processed through a whitening operator—into the noise term while preserving the pretrained reverse mean. Crucially, NTRK requires no modification to the reverse kernel architecture and substantially improves reward alignment across diverse tasks, outperforming state-of-the-art methods. Notably, in aesthetic image generation, NTRK achieves performance comparable to baseline approaches using 500 function evaluations with only 25 evaluations, reducing computational cost by a factor of 20.
📝 Abstract
We introduce the Noise-Tilted Reverse Kernel (NTRK), a reward-guided diffusion sampler that injects reward gradients through the noise term, leaving the pretrained reverse kernel unchanged and requiring only a single sample per step. Reward-guided sampling at inference time has greatly expanded the versatility of pretrained diffusion models. Yet existing methods face a trade-off. Gradient-based guidance shifts the reverse mean, steering generation but pushing intermediate states outside the region that the model was trained on and degrading quality. Search-based methods preserve quality but gain no gradient signal. No prior method achieves both. NTRK resolves this by keeping the reverse mean fixed and biasing the noise term toward high reward. We introduce a whitening operator, the central mechanism behind NTRK, that makes the reward gradient safe to inject as noise without losing its guiding signal. Across various reward alignment tasks, NTRK outperforms recent state-of-the-art baselines without losing sample quality. Remarkably, on aesthetic generation, NTRK surpasses the reward of the best baseline at 500 NFEs using only 25 NFEs, a 20$\times$ reduction in compute.