Thompson Sampling via Fine-Tuning of LLMs

TL;DR

This work tackles challenging Bayesian optimization in large, unstructured discrete spaces by reframing Thompson sampling as fine-tuning of a prompt-conditioned LLM to approximate the probability of maximality ($\mathrm{PoM}$). It builds on Variational Bayesian Optimistic Sampling (VBOS) and proves a refined regret bound that scales with the maximal information gain $\gamma^T$, justifying initialization from strong priors and cautious adaptation toward the posterior. The ToSFiT algorithm avoids intractable acquisition-function maximization, achieving sublinear regret while leveraging pre-trained models to improve sample efficiency. Empirically, ToSFiT outperforms baselines across FAQ refinement, protein design, and quantum circuit design, and extends naturally to batched Bayesian optimization with modest computational overhead. These results highlight the potential of uniting foundation-model priors with principled Bayesian objectives for efficient search in complex discrete domains.

Abstract

Bayesian optimization in large unstructured discrete spaces is often hindered by the computational cost of maximizing acquisition functions due to the absence of gradients. We propose a scalable alternative based on Thompson sampling that eliminates the need for acquisition function maximization by directly parameterizing the probability that a candidate yields the maximum reward. Our approach, Thompson Sampling via Fine-Tuning (ToSFiT) leverages the prior knowledge embedded in prompt-conditioned large language models, and incrementally adapts them toward the posterior. Theoretically, we derive a novel regret bound for a variational formulation of Thompson Sampling that matches the strong guarantees of its standard counterpart. Our analysis reveals the critical role of careful adaptation to the posterior probability of maximality--a principle that underpins our ToSFiT algorithm. Empirically, we validate our method on three diverse tasks: FAQ response refinement, thermally stable protein search, and quantum circuit design. We demonstrate that online fine-tuning significantly improves sample efficiency, with negligible impact on computational efficiency.

Figures (7)

• Figure 1: ToSFiT treats candidate generation as Thompson sampling and fine-tunes the LLM toward the posterior PoM to ensure consistency.
• Figure 2: Whereas in continuous domains gradient ascent can be used to maximize the acquisition function, in unstructured discrete domains maximization would require iterating over all points. To scale to combinatorially large discrete domains, we propose to circumvent intractable acquisition function maximization by directly sampling from the parameterized probability of maximality.
• Figure 3: Left: inexact gradient-based VBOS struggles to find the best arm in very high-dimensional probability simplices. Right: ToSFiT starts in the right region of the simplex. Still, careful adaptation towards VBOS is paramount to retain the prior knowledge necessary to find the best arm.
• Figure 4: Unguided Generation quickly plateaus, and classical Bayesian optimization over a fixed LLM-generated candidate set saturates early (Post-Generation TS). In contrast, ToSFiT explores the full solution space and consistently discovers higher-reward candidates. Via optimism in the face of uncertainty, it explores more efficiently than Actor Critic and Soft Actor Critic.
• Figure 5: ToSFiT benefits markedly from a strong initial policy. On Quantum Circuit Design, using the most informative initial policy yields much better performance.

Theorems & Definitions (30)

• Proposition 1
• Proposition 2
• Theorem 1
• Proposition 3
• Proposition 4
• Proposition 1
• proof : Proof of Proposition \ref{['prop:derivatives_of_vbos']}
• Proposition 2
• proof : Proof of Proposition \ref{['prop:bregman_divergence']}
• Theorem 1