Noise-Tilted Reverse Kernel (NTRK)

• The paper introduces NTRK, which injects reward gradients exclusively through noise tilting while preserving the reverse kernel mean, ensuring samples shift toward higher-reward regions.
• The whitening operator uses a two-level order-statistic projection and tile-wise transformations to enforce Gaussian constraints and suppress non-Gaussian artifacts.
• NTRK achieves superior reward alignment with significantly fewer network function evaluations, demonstrating efficiency across diverse image and video generation tasks.

The Noise-Tilted Reverse Kernel (NTRK) is a reward-guided diffusion sampling method that injects reward gradients exclusively through the noise channel of a pretrained reverse diffusion kernel, preserving the mean while biasing the sampling process toward higher-reward regions. NTRK achieves reward alignment at inference time without sample quality degradation or heavy search overhead, addressing key limitations of previous mean-shifted and search-based diffusion guidance approaches. The method’s central mechanism is a whitening operator enforcing high-probability Gaussian constraints on injected gradients, ensuring compatibility with the pretrained noise distribution and maintaining the original kernel statistics. Experimental evaluations show NTRK outperforms previous state-of-the-art methods across diverse alignment tasks using an order of magnitude fewer network function evaluations (NFEs) (Hwang et al., 16 Jun 2026).

1. Formal Definition and Theoretical Guarantees

Let $$p_\theta(\mathbf{x}_{t-1}\mid \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t), \sigma_t^2 \mathbf{I})$$ denote the pretrained reverse kernel of a diffusion model at time $t$. NTRK sampling holds the pretrained mean, modifying the noise term via a reward-dependent tilt: $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$ where

$$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$

Here, $$\mathbf{w}_t = \mathcal{W}(\nabla_{\mathbf{x}_t} r(\hat{\mathbf{x}}_{0|t}))$$ is the whitened reward gradient using operator $\mathcal{W}$. The scalar $\rho_t \in [0,1]$ controls guidance strength.

This parameterization leaves the mean and variance of the kernel unchanged: convex Gaussianity of noise ensures $$\tilde{\boldsymbol\epsilon}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$. Unlike mean-shifted approaches, NTRK’s transitions remain on the pretrained distributional manifold while still biasing samples toward reward maxima. The gradient-based guidance is thus injected safely and unambiguously through the stochastic channel.

2. Whitening Operator Construction

The whitening operator $$\mathcal{W} : \mathbb{R}^N \rightarrow \mathbb{R}^N$$ is central to NTRK. Its function is to remap an arbitrary reward gradient to a typical instance of Gaussian noise, enforcing rigorous constraints that preserve high-probability structure under $\mathcal{N}(\mathbf{0}, \mathbf{I})$. The process consists of:

• Two-Level Order-Statistic Projection (2OS):

1. Partition $t$0 into $t$1 chunks of size $t$2, reshape to $t$3.
2. Sort each row ($t$4), then each column ($t$5), producing $t$6.
3. For each entry, clip to $t$7 confidence intervals under the reference Beta distribution.
4. Undo the serial sorts.

• Tile-Wise Mean and Energy Constraints:

1. Partition into non-overlapping tiles of size $t$8.
2. For each tile $t$9, calculate $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$0 (mean) and $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$1 (energy).
3. Project these to high-confidence intervals under $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$2 and $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$3 for mean and energy, respectively.
4. Shift and rescale tile to match projected statistics with minimal $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$4 deviation.

• Multi-Domain Projections:

1. Repeat both previous stages in additional orthogonal bases (e.g., real Fourier, Hadamard).
2. Each transform preserves the reference Gaussian statistics due to orthogonality.

Stacking these procedures yields a transformed gradient $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$5 that is highly regular as Gaussian noise, suppressing structured artifacts and preserving stochasticity: genuine noise is preserved (cosine similarity $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$6), while outliers and spurious correlations are actively removed.

3. Mean Preservation and Noise Tilting

NTRK’s distinguishing property is that $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$7 (the reverse mean) is never shifted. This is in contrast to standard gradient-based or mean-shifted samplers, which modify $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$8 by a reward gradient step, risking transition dynamics departing from the pretrained manifold and sample quality degradation. In NTRK, only the noise term is “tilted”: $$\mathbf{x}_{t-1} = \mu_\theta(\mathbf{x}_t, t) + \sigma_t\, \tilde{\boldsymbol\epsilon}_t$$9 As both $$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$0 and $$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$1 are individually standard normal (the former by construction), their convex combination preserves the overall kernel’s zero-mean, unit-covariance statistics. The reward alignment is achieved by biasing the noise in the direction of whitened local reward gradients, without direct displacement of the distribution mean.

This approach resolves the guidance-quality trade-off: gradient signal is utilized while sample transitions remain consistent with pretraining.

4. Practical Sampling Procedure

The NTRK sampling algorithm iteratively updates the state from the final noise draw, applying reward-aligned tilting to the noise at each step. The full pseudocode is:

$$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$7

The hyperparameter $$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$2 governs the trade-off between exploitation (gradient signal) and exploration (stochasticity). Intermediate values ($$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$3) are empirically found to optimize reward alignment and sample fidelity. The number of function evaluations $$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$4 is typically as low as 25, an order of magnitude smaller than search-based diffusion guidance.

5. Empirical Performance and Efficiency

NTRK demonstrates computational efficiency and robust empirical performance across a suite of alignment tasks using both image and video diffusion backbones. Table 1 summarizes key comparative benchmark results:

Task	Configuration	Score (Base / DPS / NTRK)
Aesthetic Image Generation	Aesthetic Reward @ 25 NFE / 500 NFE	6.03 / 6.76 / 7.45
Text-Aligned Image Generation	PickScore @ 25 NFE / 500 NFE	0.205 / 0.214 / 0.222
Preference-Aligned Video Generation	VideoReward @ 25 NFE (Base / DPS / NTRK)	-0.399 / -0.130 / 3.465
Quantity-Aware Image Counting	MAE / % Accuracy	~17, 5% / — / 2.3, 37.5%

Reward and perceptual quality are maintained or improved relative to search-based methods (DPS) and prior baseline approaches, with NTRK using $$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$525 NFEs where DPS typically requires 500. In practice, for a FLUX image generation model, wall-clock time for 25 steps increases by only 7.8% with NTRK’s whitening overhead, yet achieves a 20× reduction in NFEs and a 4.6× wall-time speedup over a 500-step DPS run (Hwang et al., 16 Jun 2026).

NTRK’s guidance generalizes to different backbones (e.g., Z-Image), compound models (fine-tuned with MixGRPO), and varied dataset conditions, supporting broad applicability. The guidance is shown to be orthogonal to model parameter updates, allowing combined use with fine-tuning for further reward maximization.

6. Limitations and Open Research Questions

Key open limits and research challenges remain:

• Hyperparameter Sensitivity: The choice of $$\tilde{\boldsymbol\epsilon}_t = \sqrt{\rho_t}\, \mathbf{w}_t + \sqrt{1-\rho_t}\, \boldsymbol\epsilon_t, \quad \boldsymbol\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$6 strongly affects guidance; excessive values induce potential out-of-distribution drift, while insufficient guidance fails to improve reward. This suggests practical tuning or adaptive schedules may be required for different settings.
• Whitening Configuration: The whitening operator, although efficient in practice, exhibits multiple configuration options (chunk size, tile shape, transform domain). Optimal choices and automation for diverse data modalities remain open.
• Theoretical Guarantees: Current analysis relies on local linear approximations and heuristic correspondence to search-based methods; convergence to the entropy-regularized reward optimum is not formally established.
• Adversarial Robustness: As for all gradient-based inference-time alignment, reward models susceptible to manipulation or adversarial optimization may expose vulnerabilities.
• High-Dimensional Scaling: The reliability and computational feasibility of the whitening procedure for highly structured, large-scale latent spaces (notably in video diffusion) require further empirical and theoretical investigation.

7. Context and Significance

NTRK directly addresses the longstanding trade-off in diffusion reward alignment between sample quality and reward maximization. Unlike mean-shifted and search-based approaches, NTRK decouples reward guidance from the pretrained kernel mean, enabling efficient, gradient-informed, and high-fidelity sampling. Its model-agnostic design and demonstrated empirical success across vision and generative tasks suggest broad relevance for inference-time preference alignment in large, pretrained diffusion models (Hwang et al., 16 Jun 2026). Future work aims to increase automation, improve robustness, and rigorously characterize optimization behavior and generalization properties.