Technical Summary: POETS (Policy Ensembles for Thompson Sampling)

Problem Statement

Balancing exploration and exploitation remains a core challenge in sequential decision-making and black-box optimization, particularly when applied to Large Language Models (LLMs). Effective exploration requires principled uncertainty quantification. While ensembling independently trained models is a proven, model-agnostic method for estimating epistemic uncertainty, directly applying naive ensembling to LLMs is computationally prohibitive due to memory and compute constraints. Furthermore, existing exploration methods for LLMs often rely on explicit reward models (which introduce nested optimization complexities) or depend heavily on the model's instruction-following capabilities.

Methodology

The paper introduces POETS (Policy Ensembles for Thompson Sampling), a framework that bridges uncertainty quantification and policy optimization without requiring an explicit probabilistic reward model.

Core Insight: Implicit Reward Encoding

POETS is grounded in the duality between reward functions and optimal policies under regularized objectives. Specifically, policies trained with Kullback-Leibler (KL) regularization implicitly encode an underlying reward function. Instead of learning a reward model and then fitting a policy to it (a two-step process), POETS directly matches the implicit reward functions encoded by a policy ensemble to observed online data.

The Algorithm

1. Policy Ensemble: POETS maintains an ensemble of $n$ independent policies, $\Pi = \{\pi_1, \dots, \pi_n\}$.
2. Poisson Bootstrapping: To capture epistemic uncertainty, the ensemble members are diversified via Poisson bootstrapping. For each incoming batch of data, every ensemble member is assigned an independent integer weight drawn from a Poisson(1) distribution. This allows the ensemble to learn an approximate Bayesian posterior over reward functions without requiring a static dataset.
3. Loss Function: The method optimizes a loss function that aligns the implicit reward of a single policy with the true reward. Theoretical derivation shows that optimizing this loss is mathematically equivalent to a policy gradient update (specifically, a variant of GRPO) scaled by the Poisson weights.
4. Thompson Sampling: By fitting an implicit reward ensemble to true rewards using independent bootstraps, POETS learns a posterior over reward functions. Executing KL-regularized Thompson sampling is achieved simply by uniformly sampling a policy from the ensemble $\Pi$ and acting according to it.

Efficient Architecture: "Trunk & Branch"

To overcome the prohibitive cost of maintaining $n$ full LLM copies, POETS utilizes a parameter-efficient architecture:

• Shared Trunk: A pre-trained backbone (the majority of weights) is shared across all ensemble members and evaluated only once per forward/backward pass.
• Independent Branches: Ensemble diversity is encapsulated entirely in independent, low-rank adaptation (LoRA) modules acting as the final readout layers.
• Efficiency: This design ensures that generating responses and accumulating gradients for $n$ policies takes virtually the same wall-clock time as a single policy, with memory overhead scaling only with the size of the LoRA branches.

Key Contributions

1. Novel Framework: Introduction of POETS, which captures epistemic uncertainty via policy ensembles and implicitly conducts KL-regularized Thompson sampling without an explicit reward model.
2. Theoretical Guarantees: Proof that POETS inherits the strong cumulative soft regret bounds of $O(\sqrt{T\gamma_T})$ associated with KL-regularized Thompson sampling, where $\gamma_T$ is the maximal information gain.
3. Computational Efficiency: Demonstration that a "Trunk & Branch" LoRA architecture allows for deep ensembling with minimal overhead (e.g., ~7.7% runtime increase for $n=16$ on an H200 GPU).
4. Empirical Performance: State-of-the-art sample efficiency across diverse combinatorial black-box optimization domains (FAQ refinement, protein search, quantum circuit design) and improved robustness in reinforcement learning settings, particularly in off-policy scenarios with experience replay.

Experimental Results

The authors evaluated POETS against baselines including GRPO, TOSFIT (which uses an explicit reward model), FIBO (in-context Thompson sampling), and EVOSEARCH (evolutionary search).

• Black-Box Optimization: POETS achieved state-of-the-art sample efficiency in protein search, quantum circuit design, and FAQ refinement. Notably, it outperformed baselines even without a replay buffer, whereas standard GRPO failed to maintain diversity and overfit to early candidates.
• Reinforcement Learning (RL): In the AIME dataset setting (contextual bandit with verifiable rewards), POETS improved optimization trajectories for both on-policy and off-policy (experience replay) settings. It successfully mitigated the overfitting issues observed in GRPO when exposed to repeated contexts.
• Diversity: The ensemble maintained generation diversity (measured by Jensen-Shannon divergence) throughout training, preventing the entropy collapse common in standard policy gradient methods.
• Ablations: Experiments confirmed that Poisson bootstrapping is essential for diversity, that small ensembles ($n=8$) suffice for high performance, and that the "Trunk & Branch" architecture effectively captures uncertainty with minimal computational cost.

Significance and Claims

The paper claims that POETS represents a significant advancement by enabling principled, uncertainty-aware exploration for LLMs in scientific discovery and RL without the computational burden of explicit reward modeling or the fragility of in-context instruction following.

• Theoretical: It establishes that policy ensembles with bootstrapping are sufficient to implement KL-regularized Thompson sampling, inheriting optimal regret bounds.
• Practical: It provides a scalable solution that makes deep ensembling feasible for LLMs, allowing for effective use of experience replay and robust exploration in high-uncertainty regimes.
• Scope: The work is positioned as a foundational algorithmic contribution for LLM post-training and scientific discovery, specifically targeting combinatorial optimization problems where evaluations are costly.

The authors note that while POETS improves sample efficiency and robustness, it does not address reward misspecification, fairness, or alignment, which remain responsibilities of system designers. Future work is suggested to extend the framework from contextual bandits to multi-step settings with fine-grained credit assignment.