Technical Summary: NoiseTilt: Noise-Tilted Reverse Kernels for Diffusion Reward Alignment

Problem Statement

Inference-time scaling has emerged as a powerful mechanism for aligning pretrained diffusion models with user preferences without costly retraining. When preferences are expressed as reward functions, the iterative denoising process of diffusion models offers a model-agnostic way to steer sampling toward high-reward regions. However, existing methods face a fundamental trade-off between gradient guidance and noise compatibility:

1. Mean-Shifted Methods (Gradient Guidance): Approaches like DPS [10] inject reward gradients directly into the mean of the reverse Gaussian kernel. While effective at steering generation, this shifts the sampled states outside the "noise-compatible regime"—the narrow region of the latent space where the pretrained model was trained to operate. This displacement causes intermediate states to become out-of-distribution, degrading generation quality and making the system prone to reward hacking.
2. Search-Based Methods: Approaches like SVDD [38] preserve the base reverse mean by drawing $K$ candidates from the stochastic kernel and selecting the best one. While this maintains noise compatibility and sample quality, it forgoes direct gradient guidance, relying on random sampling to find reward-improving directions. This is computationally inefficient, requiring a prohibitively large number of samples to achieve strong alignment.

No prior method simultaneously achieves gradient guidance and noise compatibility.

Methodology: Noise-Tilted Reverse Kernel (NTRK)

The authors propose Noise-Tilted Reverse Kernels (NTRK), a sampler that resolves this trade-off by keeping the reverse mean fixed and routing reward information exclusively through the noise term.

Core Mechanism

The standard reverse step is defined as $\bm{x}_{t-1} = \bm{\mu}_\theta(\bm{x}_t, t) + \sigma_t \bm{\epsilon}_t$, where $\bm{\epsilon}_t \sim \mathcal{N}(\bm{0}, \bm{I})$.
NTRK modifies the noise term to:
$$\bm{x}_{t-1} = \bm{\mu}_\theta(\bm{x}_t, t) + \sigma_t \tilde{\bm{\epsilon}}_t$$
where the guided noise $\tilde{\bm{\epsilon}}_t$ is a mixture of a whitened reward direction and unbiased stochasticity:
$$\tilde{\bm{\epsilon}}_t = \sqrt{\rho_t} \, \mathcal{W}(\nabla_{\bm{x}_t} r(\hat{\bm{x}}_{0|t})) + \sqrt{1-\rho_t} \, \bm{\epsilon}_t$$
Here, $\rho_t \in [0, 1]$ controls the guidance strength, and $\hat{\bm{x}}_{0|t}$ is the Tweedie estimate of the clean image.

The Whitening Operator ($\mathcal{W}$)

The central challenge is that a raw reward gradient is structured and deterministic, which violates the statistical properties of typical Gaussian noise required by the pretrained model. Injecting it directly would create "atypical noise," leading to artifacts.

To solve this, the authors introduce a whitening operator $\mathcal{W}$ that transforms the reward gradient into a "typical" noise vector before injection. This operator ensures the perturbation remains within the noise-compatible regime while preserving the directional signal of the gradient. The operator is constructed via a sequence of projections onto high-confidence sets derived from standard Gaussian statistics:

1. Two-Level Order-Statistic (2OS) Projection: The input vector is reshaped into tiles. Values are sorted within each tile and then across tiles. The resulting doubly-sorted statistics are clipped to 99.99% confidence intervals derived from the Beta distribution. This suppresses localized structured artifacts while maintaining the global histogram.
2. Tile-Wise Statistics: The operator further constrains the mean and centered energy of each tile to match the expected distributions of standard Gaussian noise.
3. Orthogonal Transform Domains: The same confidence-interval projections are applied in transformed domains (e.g., compact Fourier and Hadamard-style mixing) to suppress residual global structures that are not visible in the spatial domain.

Crucially, for inputs that are already typical Gaussian noise, $\mathcal{W}$ leaves them nearly unchanged, whereas structured inputs are substantially whitened.

Key Contributions

1. Resolution of the Trade-off: NTRK is the first method to simultaneously achieve gradient-based guidance and noise compatibility. It preserves the pretrained reverse mean (unlike mean-shifted methods) while using a single sample per step (unlike search-based methods).
2. The Whitening Operator: The paper identifies that the quality of reward alignment is directly tied to the quality of the whitening operator. The proposed 2OS-based projection is shown to be superior to simple norm projection or hard equality constraints (like WGNC [24]), as it balances structure suppression with the preservation of natural noise variability.
3. Efficiency: By routing gradients through the noise term, NTRK achieves state-of-the-art performance with significantly fewer function evaluations (NFE).

Experimental Results

The authors evaluated NTRK across three primary tasks using FLUX [35] for images and Wan2.1 [70] for video:

• Aesthetic Image Generation: NTRK achieved the highest target reward (Aesthetic Score) among all baselines. Notably, NTRK with only 25 NFEs surpassed the reward of the best baseline (SVDD) which used 500 NFEs, representing a 20$\times$ reduction in compute. It also maintained or improved held-out metrics like ImageReward and HPSv2.
• Text-Aligned Image Generation: On complex compositional prompts (T2I-CompBench++), NTRK outperformed baselines in PickScore (target reward) and held-out metrics like VQA Score and HPSv2.
• Preference-Aligned Video Generation: Using VideoReward as the target, NTRK generated videos with better text alignment, motion quality, and visual quality compared to DPS and FreeDoM.
• Additional Applications: The method was successfully applied to quantity-aware generation (counting objects) and VLM-based reward alignment (horizon and depth positioning), demonstrating broad applicability.
• Integration with Fine-Tuning: NTRK was shown to be orthogonal to fine-tuning (e.g., MixGRPO [37]), providing additional gains when applied on top of fine-tuned models.

Significance and Claims

The paper claims that NTRK offers a simple, broadly applicable mechanism for inference-time reward alignment that leaves the pretrained sampling dynamics intact. By treating the reward gradient as a perturbation to be whitened rather than a shift to the mean, NTRK avoids the out-of-distribution issues that plague gradient-based guidance.

The authors emphasize that the effectiveness of the method relies heavily on the whitening operator. Their ablation studies show that weaker operators fail to remove harmful structure, while overly rigid operators distort genuine noise. The proposed 2OS-based approach successfully navigates this trade-off, enabling high-reward generation without sacrificing the quality inherent to the pretrained model.

The work suggests that future inference-time scaling can move beyond brute-force search or destructive mean shifts, instead leveraging statistical properties of noise to guide generation efficiently.