Gradient-Guided Reward Optimization for Inference-time Alignment

Hankun Lin, Ruqi Zhang
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:3867-3887, 2026.

Abstract

Ensuring the reliability of Large Language Models ({LLMs}) under distribution drift requires inference-time adaptation. While inference-time alignment methods such as Best-of-$N$ and rejection sampling are widely used, they frame the task as a sampling-intensive, reward-guided search, leading to two key limitations: their performance is bounded by the base model’s generation quality, and their reliance on imperfect reward models makes them vulnerable to reward hacking. To address these challenges, we introduce *Gradient-Guided Reward Optimization* (GGRO), a lightweight inference-time method that performs targeted, minimal intervention during decoding via gradient guidance. Specifically, GGRO monitors token-level entropy to identify high-uncertainty regions indicative of drift or misalignment. Upon detection, it responds by injecting nudging tokens, generated using gradient signals from an off-the-shelf reward model, to steer the generation trajectory rather than merely re-ranking samples. Experiments show that GGRO consistently improves inference-time alignment across safety, helpfulness, and reasoning benchmarks. It also increases coverage of high-quality responses and robustness to reward hacking, with minimal computational overhead. Code is available at https://github.com/lhk2004/GGRO.

Cite this Paper

BibTeX
@InProceedings{pmlr-v337-lin26a, title = {Gradient-Guided Reward Optimization for Inference-time Alignment}, author = {Lin, Hankun and Zhang, Ruqi}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {3867--3887}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/lin26a/lin26a.pdf}, url = {https://proceedings.mlr.press/v337/lin26a.html}, abstract = {Ensuring the reliability of Large Language Models ({LLMs}) under distribution drift requires inference-time adaptation. While inference-time alignment methods such as Best-of-$N$ and rejection sampling are widely used, they frame the task as a sampling-intensive, reward-guided search, leading to two key limitations: their performance is bounded by the base model’s generation quality, and their reliance on imperfect reward models makes them vulnerable to reward hacking. To address these challenges, we introduce *Gradient-Guided Reward Optimization* (GGRO), a lightweight inference-time method that performs targeted, minimal intervention during decoding via gradient guidance. Specifically, GGRO monitors token-level entropy to identify high-uncertainty regions indicative of drift or misalignment. Upon detection, it responds by injecting nudging tokens, generated using gradient signals from an off-the-shelf reward model, to steer the generation trajectory rather than merely re-ranking samples. Experiments show that GGRO consistently improves inference-time alignment across safety, helpfulness, and reasoning benchmarks. It also increases coverage of high-quality responses and robustness to reward hacking, with minimal computational overhead. Code is available at https://github.com/lhk2004/GGRO.} }
Endnote
%0 Conference Paper %T Gradient-Guided Reward Optimization for Inference-time Alignment %A Hankun Lin %A Ruqi Zhang %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-lin26a %I PMLR %P 3867--3887 %U https://proceedings.mlr.press/v337/lin26a.html %V 337 %X Ensuring the reliability of Large Language Models ({LLMs}) under distribution drift requires inference-time adaptation. While inference-time alignment methods such as Best-of-$N$ and rejection sampling are widely used, they frame the task as a sampling-intensive, reward-guided search, leading to two key limitations: their performance is bounded by the base model’s generation quality, and their reliance on imperfect reward models makes them vulnerable to reward hacking. To address these challenges, we introduce *Gradient-Guided Reward Optimization* (GGRO), a lightweight inference-time method that performs targeted, minimal intervention during decoding via gradient guidance. Specifically, GGRO monitors token-level entropy to identify high-uncertainty regions indicative of drift or misalignment. Upon detection, it responds by injecting nudging tokens, generated using gradient signals from an off-the-shelf reward model, to steer the generation trajectory rather than merely re-ranking samples. Experiments show that GGRO consistently improves inference-time alignment across safety, helpfulness, and reasoning benchmarks. It also increases coverage of high-quality responses and robustness to reward hacking, with minimal computational overhead. Code is available at https://github.com/lhk2004/GGRO.
APA
Lin, H. & Zhang, R.. (2026). Gradient-Guided Reward Optimization for Inference-time Alignment. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:3867-3887 Available from https://proceedings.mlr.press/v337/lin26a.html.

Related Material