NoiseTilt: Noise-Tilted Reverse Kernels for Diffusion Reward Alignment

Official implementation for the ECCV 2026 paper:

NoiseTilt: Noise-Tilted Reverse Kernels for Diffusion Reward Alignment
Jisung Hwang, Yunhong Min, Jaihoon Kim, I-Chao Shen, and Minhyuk Sung
[arXiv]

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, but existing methods face a trade-off: gradient-based guidance shifts the reverse mean and can degrade quality, while search-based methods preserve quality but gain no gradient signal. NTRK resolves this by keeping the reverse mean fixed and biasing the noise term toward high reward. This is enabled by a whitening operator, the central mechanism behind NTRK, which converts reward gradients into noise-compatible perturbations without losing their guiding signal. Across reward alignment tasks, NTRK outperforms recent state-of-the-art baselines without losing sample quality; on aesthetic generation, it surpasses the reward of the best baseline at 500 NFEs using only 25 NFEs.

Release

• NoiseTilt whitening operator in whitening.py, separated from the sampler in methods.py.
• CUDA extension for the whitening bottleneck in whitening_cuda/.
• FLUX.1-schnell runner for 45-animal aesthetic optimization.
• Aesthetic Score and PickScore reward support.

Method

This runner implements the Noise-Tilted Reverse Kernel (NTRK). Instead of shifting the reverse-kernel mean, NTRK preserves the pretrained reverse mean and injects reward information only through the stochastic noise term.

A raw reward gradient is structured and deterministic, so it cannot directly serve as a noise-compatible perturbation. The whitening operator maps the reward gradient to a noise-compatible direction. The guided injected noise then mixes the whitened direction with unbiased stochasticity:

$$\tilde{\epsilon}_t = \sqrt{\rho}\, \mathcal{W}\!\left(\nabla_{x_t} r(\hat{x}_{0|t})\right) + \sqrt{1-\rho}\,\epsilon_t, \qquad \epsilon_t \sim \mathcal{N}(0,I).$$ $$x_{t-\Delta t} = \mu_\theta(x_t,t) + \sigma_t \tilde{\epsilon}_t.$$

Here $\rho$ controls the guidance strength. The reverse-kernel structure keeps the same base mean and nominal noise scale $\sigma_t$, while the whitening operator uses confidence-interval projections to suppress atypical structure and preserve typical standard Gaussian noise. The code exposes this separation directly: methods.py contains the NTRK sampling loop, while whitening.py contains the whitening operator.

Setup

From this directory, create or activate a Python environment:

conda create -n noisetilt python=3.11
conda activate noisetilt

Install dependencies:

pip install torch==2.5.1 torchvision==0.20.1 \
  --index-url https://download.pytorch.org/whl/cu124

pip install \
  diffusers==0.34.0 \
  transformers==4.49.0 \
  accelerate==1.3.0 \
  numpy==2.4.4 \
  scipy==1.15.1 \
  tqdm==4.67.1 \
  pillow==12.2.0 \
  sentencepiece==0.2.0 \
  protobuf==3.20.3

For faster whitening on CUDA, build the extension:

pip install --no-build-isolation ./whitening_cuda

The default run downloads the FLUX and CLIP checkpoints unless they are already in your Hugging Face cache. Reward checkpoints are downloaded on demand unless the expected files are already available in assets/ or the Hugging Face cache. Use --local-files-only when all required checkpoints are already available locally.

Code Structure

.
|-- assets/
|   `-- teaser.png
|-- whitening_cuda/       # CUDA whitening acceleration
|-- datasets/
|   `-- animal_prompts.txt
|-- models.py                 # model registry and default FLUX wrapper
|-- rewards.py                # reward registry and default scorer implementations
|-- whitening.py              # NoiseTilt whitening operator
|-- methods.py                # NoiseTilt sampling loop
`-- run.py                    # experiment CLI

Customization

The runner keeps the model and reward components separate:

• Select the generator with --model, --variant, and --model-id.
• Set the generation resolution with --height and --width.
• Select the target reward with --target-reward.
• Select held-out rewards with --eval-rewards.
• Select the prompt set with --dataset, or pass a text or JSON list with --prompt-file.

The release includes the FLUX wrapper and Aesthetic Score/PickScore scorers used in the paper experiments. Additional generators and rewards can be used by adding entries to the registries in models.py and rewards.py.

Run

Run the 45-animal benchmark:

python run.py \
  --model flux \
  --dataset animals \
  --target-reward aesthetic \
  --eval-rewards pickscore \
  --max-prompts 45 \
  --num-samples 20 \
  --num-steps 25 \
  --rho 0.3 \
  --g-coeff 0.2 \
  --fast-whitening \
  --save-images \
  --outdir outputs/noisetilt_animals

The default prompts are the 45 animal classes in datasets/animal_prompts.txt. The command above uses Aesthetic Score as the target reward and reports PickScore as a held-out reward. Omit --fast-whitening to use the PyTorch whitening implementation.

Outputs

Each run writes:

• results.csv: prompt, seed, Best-of-N index, target reward, held-out rewards, and runtime.
• summary.csv: mean and standard deviation for the target reward, held-out rewards, and runtime.
• config.json: run-level parameters, dataset path, and reward choices.
• 000.png, 001.png, ... when --save-images is enabled.

Useful Options

• --rho: set $\rho$, the interpolation between the whitened reward direction and fresh noise.
• --num-samples: number of independent NoiseTilt trajectories per prompt.
• --num-steps: number of denoising steps.
• --g-coeff, --g-exp: diffusion-coefficient schedule parameters.
• --guidance: classifier-free guidance scale used by the generator.
• --max-prompts: limit the number of prompts in a run.
• --seed: base random seed.
• --save-images: save generated images.
• --outdir: output directory.
• --cache-dir: set a local Hugging Face cache directory.

Whitening Parameters

By default, whitening.py uses a shape-dependent manual setting for the whitening tile sizes. The command-line overrides below are provided for users who want finer control over the whitening operator:

The whitening layout is derived from the actual latent shape produced by the model. For FLUX at $512 \times 512$, the latent shape is treated as $1024 \times 64$ for whitening; for higher or lower resolutions, the same logic is applied to the new latent shape before validating the tile sizes.

python run.py \
  --whitening-band 2x2,8x8 \
  --whitening-mvar 1x1,2x2,8x8 \
  --whitening-hada 1024x64,4x4 \
  --whitening-mix 32,8 \
  --whitening-alpha 1e-4

The example above is the FLUX setting used in the paper, where the latent whitening layout is $1024 \times 64$. Tile overrides use comma-separated HxW pairs. Each mvar tile must divide the whitening layout, each band tile must divide the corresponding reduced layout after the mvar tile, and hada and mix must have the same length. When --fast-whitening is used, each band tile area must be at most 256.

Use --whitening-alpha to set the confidence-interval level and --fast-whitening to use the CUDA extension when available.

Third-Party Weights

Reward weights are not bundled in this repository. When needed, the code downloads the Aesthetic Score weights sac+logos+ava1-l14-linearMSE.pth and the PickScore checkpoints yuvalkirstain/PickScore_v1 and laion/CLIP-ViT-H-14-laion2B-s32B-b79K.

Citation

@misc{hwang2026noisetilt,
  title={NoiseTilt: Noise-Tilted Reverse Kernels for Diffusion Reward Alignment},
  author={Hwang, Jisung and Min, Yunhong and Kim, Jaihoon and Shen, I-Chao and Sung, Minhyuk},
  year={2026},
  eprint={2606.18066},
  archivePrefix={arXiv},
  primaryClass={cs.LG}
}