Thompson Sampling via Fine-Tuning of LLMs

AI summary

Overview

Research area: Bayesian optimization and Thompson sampling applied to large unstructured discrete search spaces, using fine-tuned large language models as the sampling policy (machine learning / sequential decision making).

Technical level: Advanced. The paper combines variational Bayesian optimization theory, regret analysis, and LLM fine-tuning with a reinforcement-learning-style objective.

Scope: The paper introduces ToSFiT, a method that casts Thompson sampling as fine-tuning of a prompt-conditioned LLM toward the posterior probability of maximality, with a new regret bound and experiments on FAQ response refinement, protein search, and quantum circuit design.

What This Paper Is About

Bayesian optimization aims to find the highest-reward point in a space using as few expensive evaluations as possible, but its standard recipe requires maximizing an "acquisition function" — and in large discrete spaces such as amino acid sequences or quantum circuits, there are no gradients and too many points to enumerate. The authors avoid this maximization entirely by teaching an LLM to directly generate candidates with probability equal to the probability that they are the best option, updating the model gradually as new observations arrive.

Key Contributions

1. Tighter regret bound for exact VBOS. The authors improve the cumulative regret bound of exact Variational Bayesian Optimistic Sampling from Õ(√(T|X|)) to Õ(√(T γ^T)), where γ^T is the maximum information gain, thereby accounting for reward correlation across the search space rather than assuming independent arms. They also derive the first regret bound for approximate (e.g., gradient-based) VBOS, which depends on the Bregman divergence between the sampling policy and the exact VBOS maximizer.

2. A principled algorithm design (ToSFiT). The bound motivates initializing the policy from a pre-trained, prompt-conditioned LLM and then adapting it cautiously toward the posterior probability of maximality, yielding the ToSFiT algorithm (Thompson Sampling via Fine-Tuning).

3. Theory of the VBOS optimization landscape. The paper shows the VBOS objective is concave (strictly, when σ_x > 0 for all x), derives its explicit gradients, and interprets VBOS as an energy-based model.

4. Empirical validation across three domains. ToSFiT is evaluated on FAQ response refinement, protein optimization, and quantum circuit design against seven baselines spanning Bayesian optimization, reinforcement learning, and evolutionary search, with state-of-the-art sample and computational efficiency reported.

Main Findings

• State-of-the-art sample efficiency across three tasks. In FAQ Response Refinement, Protein Search, and Quantum Circuit Design, ToSFiT achieves the best best-seen reward among the compared methods. Results are reported as mean and standard error over 25 random seeds, using models between 0.6B and 8B parameters.

• Only LLM-based method combining all four properties. Table 1 evaluates methods on exploiting promising regions, exploring using entropy, exploring using optimism, conditioning model generation, and being instruction-following agnostic. ToSFiT is the only method marked "yes" on all five; FIBO is marked "no" on instruction-following agnostic.

• Optimism is what distinguishes ToSFiT from actor-critic methods. Actor Critic and Soft Actor Critic are described as conducting undirected exploration without optimism in the face of uncertainty, which ToSFiT's optimism counteracts, producing a more stable exploration-exploitation tradeoff.

• Fine-tuning beats in-context Bayesian optimization. FIBO is described as closest in spirit to ToSFiT but relies on in-context learning rather than parameter updates, leading to significantly lower performance. Because FIBO keeps generated candidates as context, its compute scales quadratically with rounds and it runs out of memory; the authors compare against a top-10 truncated version in Figure 4(d), and at 10'000 rounds FIBO runs out of memory in its standard implementation again.

• Batching works. Figure 6 studies batched ToSFiT on Protein Search with up to 16 proteins synthesized and tested in parallel. Batching reduces sample efficiency (more evaluations to reach a given reward) but improves iteration efficiency (target performance in fewer rounds).

• Computational efficiency is state of the art. On fully batched (B = b = 16) Protein Search, ToSFiT's strong sample efficiency compensates for the cost of model fine-tuning. Actor Critics and Post-Generation TS are not plotted because their cost matches ToSFiT or Unguided Generation with strictly worse reward. FIBO is tested with Qwen3-0.6B (fast) and Qwen3-8B (slow).

• Scalable GP inference. Using a linear kernel in a reproducing kernel Hilbert space with feature map φ, both computational and memory complexity scale as Θ(dim(H)²), independently of the number of observations, with negligible overhead.

• Strong priors matter, and so does caution. Masking the number of qubits in the prompt ("ToSFiT weak context") yields significantly worse reward than full context; uniform initialization fails to produce valid code. A large learning rate may initially improve performance by closely following VBOS but eventually causes forgetting of the prior and stagnation.

• Model size helps unevenly. Increasing model size is highly beneficial for FAQ Response Refinement, a natural language task, but Protein Search and Quantum Circuit Design profit far less — matching observations by Romera-Paredes et al. (2024) for code discovery, since novel discovery often requires leaving the training data manifold.

• Compute can be traded for samples. Increasing the number of gradient ascent steps per round (c) improves ToSFiT's sample efficiency on Protein Search.

• Theoretical constants and structure. The exact VBOS bound given is √(2|X|T ln|X|(1+ln T)); the new bound is √(C_{σ_n} H T γ^T) + E Σ_t D_{σ^t}(π^t, π̃^t), with C_{σ_n} = 4/ln(1+σ_n^{-2}), H the expected average entropy of the policy (upper bounded by ln|X|), and γ^T the maximum information gain. The theorem requires K_{x,x} ≤ 1 for all x and additive observation noise, and also holds for heteroscedastic additive Gaussian noise by replacing σ_n with max_x σ_n(x).

• Gradient estimation details. The score-function estimator uses the Reinforce Leave-One-Out (RLOO) baseline, with each baseline set to the average of the other surrogates in the batch, and the advantage is normalized by its empirical standard deviation — a variance-adaptive learning rate. Proposition 3 in Appendix C states that standardized RLOO is mathematically equivalent to the advantage function used in Group Relative Policy Optimization.

Methodology in Plain English

Thompson sampling normally works by imagining many possible reward functions, picking one at random, and choosing the point that would be best under it. In large discrete spaces you cannot search for that best point, so the authors flip the problem: instead of sampling a reward function and maximizing it, they train a language model so that the probability it generates a candidate equals the probability that the candidate is the maximum. That trained model is the Thompson sample generator.

The model starts from a pre-trained LLM conditioned on a task prompt, which supplies useful prior knowledge about plausible candidates. As evaluations come in, a Gaussian process (with a linear kernel over embeddings, making inference closed-form) provides the mean and uncertainty of the reward at any candidate. The model is then nudged with gradient ascent on the VBOS objective, which rewards generating candidates whose predicted reward exceeds a baseline — an optimistic, entropy-aware signal. The authors use a leave-one-out baseline and normalize the advantage to keep gradients stable, and they use small learning rates so the model does not drift away from its useful prior.

The algorithm runs with a burn-in period (m) to fit GP hyperparameters by marginal likelihood maximization, c gradient steps per optimization round, a generation batch size B for stable gradients, and a Bayesian optimization batch size b for parallel evaluations. A multiplicative exploration bonus on the marginal-likelihood-fitted prior amplitude is used to avoid overconfident posteriors. The three experimental settings are: FAQ refinement with Qwen3-1.7B/Qwen3-8B generating responses scored by alignment (judged by Qwen3-Embedding-0.6B, using the first 256 embedding entries as features); protein search with ProtGPT2 (0.738B parameters) generating sequences scored by negative thermal instability index; and quantum circuit design with Qwen2.5-Coder-1.5B/Qwen2.5-Coder-7B using fill-in-the-middle prompting to prepare low-energy states of six qubits, scored by negative energy under an unknown Hamiltonian with strong interaction terms, with a feature map of a validity bit plus all two-qubit Pauli observables.

Why This Matters

Impact on research. The paper connects two previously separate lines of work — principled Bayesian optimization with regret guarantees, and large language models as generative search priors — and shows that acquisition function maximization, long considered essential, can be sidestepped entirely in unstructured discrete domains. The improved regret bound replaces a dependence on the size of the search space |X| (which is vacuous for combinatorially large spaces) with a dependence on the maximum information gain γ^T, which for a linear kernel in d dimensions grows as O(d log T). The extension to approximate VBOS provides the first regret bound for gradient-based policy updates, identifying Bregman divergence from the exact VBOS solution as the quantity that must be controlled.

Real-world applications mentioned or implied by the tasks:

• Drug discovery and industrial biotechnology, via design of thermally stable proteins with improved robustness and shelf life.
• Automated program synthesis and code generation, via the quantum circuit design task using Qiskit circuits.
• Natural language content optimization, via the FAQ response refinement task on semantic alignment.
• High-throughput experimental design where candidate evaluation is slow or parallelized, such as synthesizing and testing up to 16 proteins in parallel.

Industry relevance. Both author affiliations are industrial and academic (IBM Research – Zurich and ETH Zürich), and code is released at a public GitHub repository under a CC BY 4.0 license. The method is designed to reuse existing large-scale LLM pre-training rather than requiring task-specific pre-training of a model or autoencoder, which lowers the barrier to deploying Bayesian optimization on real discrete design problems.

Future Directions

• Learn task-adaptive embeddings jointly with the Gaussian process. The authors note they deliberately used fixed feature maps (pre-trained embeddings or hand-designed ones) for controlled comparisons, and point to Ranković and Schwaller (2025) as a complementary method that could be integrated.

• Replace the Gaussian process with more expressive reward models. The limitations section proposes alternatives such as Bayesian neural networks (the provided text is truncated mid-sentence at this point).

• Extend the theory to prove ToSFiT's own regret guarantees. The paper derives a regret bound for approximate VBOS in terms of Bregman divergence, but does not report a separate end-to-end bound specifically for the fine-tuning procedure.

• Explore how far the pre-training prior can be biased before fine-tuning becomes insufficient. Theorem 1 holds even if the pre-training prior π⁰ is biased toward suboptimal regions, but in that case more gradient steps are required to keep the divergence vanishing — leaving open how much compute this costs in practice.

• Investigate why larger models help unevenly. The paper reports that scaling from small to large models benefits FAQ refinement strongly but protein search and quantum circuit design far less, and frames this as a question of leaving the training data manifold.

Target Audience

Researchers and practitioners working on Bayesian optimization, bandits, and automated scientific discovery, particularly those facing large discrete search spaces where gradient-based acquisition maximization fails. It is also relevant to machine learning engineers applying LLM fine-tuning with reinforcement-learning-style objectives, and to scientists in protein engineering, quantum computing, and natural language processing who need sample-efficient and compute-efficient black-box optimization. A working familiarity with Gaussian processes, Thompson sampling, and policy gradient methods is assumed given the level of the theory and notation.

Authors’ 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. Within a collection of methods covering in-context Bayesian optimization, reinforcement learning, and evolutionary search, ToSFiT exhibits both state-of-the-art sample efficiency and computational efficiency.