Gradient Preconditioning for Efficient and Reliable Reward-Guided Generation

Official implementation for the ICML 2026 paper:

Gradient Preconditioning for Efficient and Reliable Reward-Guided Generation
Jisung Hwang and Minhyuk Sung
[arXiv]

Abstract. We propose a gradient preconditioning method that makes reward-guided generation with one-step generative models both efficient and reliable. Test-time noise optimization can unlock substantially better reward-guided generations from pretrained generative models, but it is prone to reward hacking that degrades quality and is often too slow for practical use. We precondition reward gradients by projecting them onto a carefully designed white Gaussian noise feasible set, a compact spectral set with blockwise norm constraints that captures the statistics and spatial uncorrelatedness of white Gaussian noise. This preconditioning reshapes each gradient update into a noise-aligned direction, driving faster and more effective reward ascent while preventing reward hacking. The projection is closed-form and matches the $\mathcal{O}(N \log N)$ complexity of FFT, adding negligible overhead in practice. In experiments on FLUX with four reward models, the method reaches a comparable Aesthetic Score using only 30% of the wall-clock time required by the state-of-the-art regularization-based method.

Release

• FLUX.1-schnell runner for 45-animal aesthetic optimization.
• Aesthetic Score and PickScore reward support.
• WGNC projection in wgnc.py, separated from the optimizer in methods.py.

Method

This runner performs test-time latent optimization for one-step text-to-image models. Each run is initialized from white Gaussian noise and optimized with Adam to maximize the reward.

The method preconditions the reward gradient by projecting it onto a white Gaussian noise feasible set before applying the update. The feasible set is defined in the compact spectral domain with blockwise $\ell_1$ and $\ell_2$ norm constraints, giving a closed-form projection with $O(N \log N)$ complexity. This reshapes each update into a noise-aligned direction while avoiding the extra regularization coefficient used by soft regularization methods.

$$ J = r(M(x)), \qquad x \leftarrow \mathrm{Adam}\left(x, \mathrm{Proj}_{\mathcal{G}}(\nabla_x J)\right). $$

By default, the runner applies both projections used in the FLUX experiments: the reward gradient is projected before the optimizer step, and the latent vector is projected after each update. To isolate gradient projection alone, pass --no-latent-projection.

With the default WGNC block size $B=16$, the flattened latent dimension $N$ should be divisible by $2B=32$. For an unusual latent size that is not divisible by 32, choose a block size that divides $N/2$; for example, $N=3000$ can use $B=15$.

Setup

From this directory, create or activate a Python environment:

conda create -n gradient-preconditioning python=3.11
conda activate gradient-preconditioning

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

The first 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/
|   `-- .gitignore
|-- datasets/
|   `-- animal_prompts.txt
|-- models.py                 # model registry and differentiable FLUX wrapper
|-- rewards.py                # reward registry and scorer implementations
|-- wgnc.py                   # white Gaussian noise constraint projection
|-- methods.py                # optimization 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.

Run

Run the 45-animal benchmark:

python run.py \
  --model flux \
  --dataset animals \
  --target-reward aesthetic \
  --eval-rewards pickscore \
  --max-prompts 45 \
  --num-iters 200 \
  --save-images \
  --outdir outputs/gradient_preconditioning_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.

Outputs

Each run writes:

• results.csv: prompt, seed, 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

• --num-iters: number of optimization steps.
• --lr: Adam learning rate for initial-noise optimization.
• --grad-clip: RMS clipping value before the projected gradient update.
• --guidance: classifier-free guidance scale used by the generator.
• --no-latent-projection: disable projection of the optimized latent.
• --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.

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{hwang2026gradientpreconditioning,
  title={Gradient Preconditioning for Efficient and Reliable Reward-Guided Generation},
  author={Hwang, Jisung and Sung, Minhyuk},
  year={2026},
  eprint={2602.08646},
  archivePrefix={arXiv},
  primaryClass={cs.LG}
}