arXiv cs.AI· Nicolas Menet, Andreas Krause, Abbas Rahimi·· 12 小时前AI 评分32
POETS:用高效策略集成实现不确定性感知的 LLM 优化
POETS: Uncertainty-Aware LLM Optimization via Compute-Efficient Policy Ensembles
AI 导读
POETS 框架通过共享预训练主干、独立 LoRA 分支的策略集成来捕捉 LLM 的认知不确定性,无需单独训练奖励模型即可近似 KL 正则化 Thompson 采样。理论上其累积遗憾界为 O(√(T γ_T)),在蛋白质搜索、治疗性分子设计和量子电路设计等任务上取得 SOTA 样本效率。该工作已被 NeurIPS 2026 接收。
正文 · 原文
Abstract:Balancing exploration with exploitation is a core challenge in black-box optimization and reinforcement learning. We introduce POETS ($\textbf{Po}$licy $\textbf{E}$nsembles for $\textbf{T}$hompson $\textbf{S}$ampling), 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(\sqrt{T \gamma_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.
| Comments: | Accepted at NeurIPS 2026 |
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Machine Learning (stat.ML) |
| Cite as: | arXiv:2605.07775 [cs.LG] |
| (or arXiv:2605.07775v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2605.07775 arXiv-issued DOI via DataCite |
Submission history
From: Nicolas Menet [view email]
[v1]
Fri, 8 May 2026 14:16:32 UTC (1,207 KB)
[v2]
Fri, 23 Oct 2026 11:46:17 UTC (1,029 KB)
来源:arXiv cs.AI · arxiv.org