POETS: Uncertainty-Aware LLM Optimization via Compute-Efficient Policy Ensembles

Abstract

Balancing exploration with exploitation is a core challenge in black-box optimization and reinforcement learning. We introduce POETS (Policy Ensembles for Thompson Sampling), a novel framework that bridges uncertainty quantification and policy optimization. Our approach is grounded in the insight that policies trained with Kullback-Leibler (KL) regularization implicitly encode an underlying reward function. Building on this, POETS bypasses the complex, nested process of training an uncertainty-aware reward model and separately fitting a policy to this model. Instead, we directly train a policy ensemble to capture epistemic uncertainty by matching implicitly encoded reward functions to online, bootstrapped data. To overcome the prohibitive compute and memory constraints of ensembling Large Language Models (LLMs), POETS utilizes an efficient architecture: the ensemble shares a pre-trained backbone while maintaining diversity through independent Low-Rank Adaptation (LoRA) branches. Theoretically, we show that POETS implicitly approximates KL-regularized Thompson sampling, a principled optimization strategy with optimal cumulative regret bounds of O(√(T γ_T)). Empirically, we demonstrate that POETS achieves state-of-the-art sample efficiency across diverse scientific discovery domains, including protein search, therapeutic molecular design, and quantum circuit design. POETS also improves the optimization trajectory and out-of-distribution transfer of reinforcement learning across mathematical reasoning and code generation tasks, proving particularly robust in off-policy settings with experience replay or in small dataset regimes.

ReAPO: Retry-Aware Policy Optimization for Efficient Exploration in LLM Reasoning
Nghi Dao, Calder Morton-Ferguson, Noah Golowich +1
arXiv20 Oct 2026PreprintCC BY 4.0Full textStatisticsarXiv
A Path-Dependent Agent-Based Microsimulation of Crime and Violence Reduction Policy Portfolios in Bolivia: Survey Calibration, Adaptive Emulator Ensembles, Global Sensitivity, Tempered MCMC, and Multi-Objective Policy Analysis
Ricardo Alonzo Fernández-Salguero
arXiv2 Oct 2026PreprintCC BY 4.0Full textStatisticsarXiv
Diffusion Distillation for Efficient Weather Ensembles
Yiming Yang, Valentin Brekke, James Briant +1
arXiv27 Aug 2026PreprintCC BY 4.0Full textStatisticsarXiv
Calibrated for Which Policy? Conformal Uncertainty for Fixed World Models Under Intervention
Ayush Mittal, Dhruv Gupta
arXiv1 Oct 2026PreprintCC BY 4.0Full textStatisticsarXiv
Adjustable Robust Optimization with Discrete Uncertainty
Henri Lefebvre, Enrico Malaguti, Michele Monaci
INFORMS Journal on Computing14 Sep 2026Peer-reviewedLicence not statedFull text6 citedOptimisation & controlHAL

Cite

BibTeX RIS

BibTeX

@misc{menet2026poets,
  title = {{POETS: Uncertainty-Aware LLM Optimization via Compute-Efficient Policy Ensembles}},
  author = {Nicolas Menet and Andreas Krause and Abbas Rahimi},
  howpublished = {arXiv},
  year = {2026},
  month = oct,
  doi = {10.48550/arxiv.2605.07775},
  eprint = {2605.07775},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2605.07775},
  note = {CC BY 4.0},
}

RIS

TY  - UNPB
TI  - POETS: Uncertainty-Aware LLM Optimization via Compute-Efficient Policy Ensembles
AU  - Nicolas Menet
AU  - Andreas Krause
AU  - Abbas Rahimi
PY  - 2026
DA  - 2026/10/23
PB  - arXiv
DO  - 10.48550/arxiv.2605.07775
UR  - https://arxiv.org/abs/2605.07775
AB  - Balancing exploration with exploitation is a core challenge in black-box optimization and reinforcement learning. We introduce POETS (Policy Ensembles for Thompson Sampling), a novel framework that bridges uncertainty quantification and policy optimization. Our approach is grounded in the insight that policies trained with Kullback-Leibler (KL) regularization implicitly encode an underlying reward function. Building on this, POETS bypasses the complex, nested process of training an uncertainty-aware reward model and separately fitting a policy to this model. Instead, we directly train a policy ensemble to capture epistemic uncertainty by matching implicitly encoded reward functions to online, bootstrapped data. To overcome the prohibitive compute and memory constraints of ensembling Large Language Models (LLMs), POETS utilizes an efficient architecture: the ensemble shares a pre-trained backbone while maintaining diversity through independent Low-Rank Adaptation (LoRA) branches. Theoretically, we show that POETS implicitly approximates KL-regularized Thompson sampling, a principled optimization strategy with optimal cumulative regret bounds of O(√(T γ_T)). Empirically, we demonstrate that POETS achieves state-of-the-art sample efficiency across diverse scientific discovery domains, including protein search, therapeutic molecular design, and quantum circuit design. POETS also improves the optimization trajectory and out-of-distribution transfer of reinforcement learning across mathematical reasoning and code generation tasks, proving particularly robust in off-policy settings with experience replay or in small dataset regimes.
ER  -

Abstract

Related

Cite