Thompson Sampling via Fine-Tuning of LLMs

本文提出了基于大语言模型微调的汤普森采样方法（ToSFiT），通过直接参数化候选解的最优概率来规避无梯度离散空间中的采集函数最大化难题，在理论上证明了其后悔界并验证了其在 FAQ 优化、蛋白质搜索及量子电路设计等任务中兼具卓越的样本效率与计算效率。

要点

这篇论文介绍了一种名为 TOSFIT 的新方法，它利用大型语言模型（LLM）来解决一个非常棘手的科学难题：如何在巨大的、杂乱无章的“可能性海洋”中，快速找到那个“最完美的宝藏”。

为了让你轻松理解，我们可以把整个过程想象成在一个巨大的、没有地图的迷宫里寻找“最完美的宝藏”。

1. 核心难题：迷宫太大，找不到路

想象一下，你正在寻找一种能抵抗高温的超级蛋白质（就像在寻找一种能防火的超级材料），或者设计一个完美的量子电路。

• 迷宫有多大？ 这里的“可能性”多到不可思议。比如，如果蛋白质由 100 个氨基酸组成，每个位置有 20 种选择，那么可能的组合数量比宇宙中所有原子的总数还要多。
• 传统方法的困境： 以前，科学家试图用“贝叶斯优化”来寻找宝藏。这就像派一个侦探去迷宫里，先画一张地图（概率模型），然后试图计算哪条路最好。但在这么大的迷宫里，没有“路标”（梯度），侦探根本没法算出哪条路最好，因为要检查所有路是不可能的。这就像试图在茫茫大海里用肉眼找到一根特定的针。

2. TOSFIT 的绝招：让“探险家”直接学会直觉

TOSFIT 的核心思想非常巧妙：既然算不出哪条路最好，那就让“探险家”（AI 模型）直接学会“直觉”，让它觉得哪条路最好，就走去哪条路。

• 以前的做法（像笨侦探）： 侦探手里拿着一张不完整的地图，每走一步都要停下来，花大量时间计算：“如果我去左边，概率是多少？如果我去右边，概率是多少？”然后还要在成千上万个路口里挑一个最好的。这太慢了，而且容易算错。
• TOSFIT 的做法（像直觉大师）：
  1. 起点（预训练）： 我们派出的“探险家”是一个已经读过万卷书的超级 AI（预训练的大语言模型）。它虽然还没见过这个特定的迷宫，但它已经具备了很强的常识和语言理解能力（比如知道蛋白质怎么折叠，代码怎么写）。
  2. 行动（采样）： 探险家不需要计算所有路。它直接根据直觉说：“我觉得走这条路（生成一个候选方案）感觉不错！”然后我们就去测试这条路。
  3. 反馈与微调（Fine-tuning）： 测试完发现，这条路确实不错（或者不好）。这时，我们不是重新画地图，而是直接修改探险家的大脑（微调模型参数）。我们告诉它：“刚才那个直觉是对的，下次多往那个方向走；如果错了，下次少往那个方向走。”
  4. 循环： 这个过程不断重复。探险家越来越聪明，越来越接近那个“完美的宝藏”。

3. 为什么要叫“汤普森采样”？

论文里提到的“汤普森采样”（Thompson Sampling），可以想象成一种**“带着不确定性去冒险”**的策略。

• 普通的探险家可能只敢走它认为 100% 安全的路（太保守，找不到新大陆）。
• 或者只敢乱走（太冒险，浪费时间）。
• TOSFIT 的探险家则不同：它知道哪里“可能”有宝藏，但也知道哪里“不确定”。它会故意去一些它觉得“可能很好，但还没完全确认”的地方。这种**“既贪婪又好奇”**的平衡，让它能最快找到最优解。

4. 这个方法的厉害之处

• 不用算死题： 它避开了那个最难的“计算哪条路最好”的数学题，直接让 AI 生成答案。
• 越用越聪明： 它利用了 AI 原本就有的知识（预训练），然后像教学生一样，通过不断的“考试 - 反馈 - 复习”（微调），让它专门适应当前的任务。
• 省时间又省钱： 在三个实际任务中（优化 FAQ 回答、寻找耐热蛋白质、设计量子电路），TOSFIT 都比其他方法（如传统的进化算法、强化学习）找得更快、更好，而且计算成本更低。

5. 生活中的类比

想象你在教一个刚毕业的大学生（LLM）如何写一份完美的“故障排除指南”：

• 传统方法： 你让他先列出所有可能的故障，然后对每一种故障，你都要花一小时去计算哪种解释最好。这太慢了。
• TOSFIT 方法：
  1. 你让他先凭自己的知识写一份指南（利用预训练知识）。
  2. 你读一下，发现某句话写得太啰嗦，或者某个步骤漏了。
  3. 你直接给他“补课”（微调），告诉他：“下次写这种指南时，要更简洁，步骤要更清晰。”
  4. 他再写一份，你发现进步很大。
  5. 重复几次，他就能写出完美的指南，而且你不需要教他从头学起，只需要修正他的“直觉”。

总结

这篇论文就像是在说：面对那些复杂到让人头秃的搜索问题，不要试图用笨办法去“算”出答案，而是利用大语言模型强大的“直觉”和“常识”，通过不断的“微调”和“反馈”，让模型自己学会如何找到那个完美的答案。

这种方法不仅让 AI 在科学发现（如新药研发、量子计算）中变得更快，也让它变得更像一个聪明的、会学习的助手，而不是一个只会死算的计算器。

技术摘要

1. 研究背景与问题 (Problem)

核心挑战：
在大规模非结构化离散空间（Unstructured Discrete Spaces）中进行贝叶斯优化（Bayesian Optimization, BO）面临巨大的计算瓶颈。

• 传统方法的局限： 传统的贝叶斯优化通常依赖于最大化“采集函数”（Acquisition Function）来选择下一个候选点。在连续空间中，这可以通过梯度上升高效完成；但在非结构化离散空间（如蛋白质序列、量子电路代码、自然语言文本）中，由于缺乏梯度且搜索空间呈组合爆炸式增长（例如，100 个氨基酸长度的序列空间远超宇宙原子总数），最大化采集函数变得不可行（Intractable）。
• 现有替代方案的不足： 现有的方法（如基于强化学习、进化搜索或上下文内贝叶斯优化）往往在样本效率（Sample Efficiency）或计算效率（Computational Efficiency）上存在权衡，或者无法有效利用大语言模型（LLM）的先验知识。

目标：
提出一种可扩展的贝叶斯优化方法，能够直接处理大规模非结构化离散空间，无需显式地最大化采集函数，同时保持高样本效率和计算效率。

---

2. 方法论 (Methodology)

论文提出了 TOSFIT (Thompson Sampling via Fine-Tuning)，一种基于大语言模型微调的汤普森采样（Thompson Sampling）变体。

核心思想

TOSFIT 将候选点的生成过程直接建模为从**最大概率（Probability of Maximality, PoM）**分布中采样，而不是通过优化采集函数来寻找最大值。

• 原理： 汤普森采样通过从奖励后验分布中采样一个函数，并选择该函数上的最大值点来进行决策。TOSFIT 利用预训练的 LLM 作为策略（Policy），将其参数化为 PoM 分布。
• 流程：
  1. 初始化： 使用带有提示（Prompt）的预训练 LLM 作为初始策略 $\pi_\theta$，利用其强大的先验知识。
  2. 生成与评估： 从当前策略生成一批候选点，评估其奖励（Reward）。
  3. 后验更新： 使用高斯过程（Gaussian Process, GP）更新奖励模型的后验分布。
  4. 策略微调（关键步骤）： 不直接最大化采集函数，而是通过**变分贝叶斯乐观采样（Variational Bayesian Optimistic Sampling, VBOS）**的目标函数，利用梯度上升微调 LLM 的参数，使其分布逐渐逼近后验 PoM 分布。

关键技术细节

• VBOS 目标函数： 论文基于 O'Donoghue & Lattimore (2021) 的工作，定义了一个变分目标 $V(\pi)$，该目标结合了奖励均值 $\mu_x$ 和基于熵的探索奖励（$\sqrt{-2 \ln \pi_x} \cdot \sigma_x$）。
• 梯度估计与稳定化：
  • 推导了 VBOS 目标关于 LLM 参数的显式梯度。
  • 引入了 RLOO (Reinforce Leave-One-Out) 基线来降低梯度估计的方差。
  • 对优势函数（Advantage Function）进行标准化，实现了类似 GRPO (Group Relative Policy Optimization) 的方差自适应学习率，防止训练不稳定。
• 可扩展的高斯过程： 为了处理离散空间，论文使用了线性核的高斯过程，结合预训练的嵌入（Embeddings）。通过 Sherman-Morrison-Woodbury 公式，实现了条件概率、后验均值/方差及边际似然最大化的常数时间复杂度（相对于观测数量），仅依赖于嵌入维度。

---

3. 主要贡献 (Key Contributions)

1. 理论突破：改进的遗憾界（Regret Bound）

   • 论文为精确的 VBOS 推导了新的累积遗憾上界，从 $\tilde{O}(\sqrt{T|X|})$ 改进为 $\tilde{O}(\sqrt{T\gamma_T})$。
   • 其中 $\gamma_T$ 是最大信息增益（Maximum Information Gain）。这一改进使得遗憾界在组合爆炸的离散空间中依然有效（不再随空间大小 $|X|$ 线性增长），并匹配了标准汤普森采样和 GP-UCB 的强保证。
   • 进一步推导了近似 VBOS（即使用梯度下降微调 LLM）的遗憾界，证明了策略与最优解之间的 Bregman 散度直接决定了遗憾的上界。

2. 算法创新：TOSFIT 算法

   • 提出了利用预训练 LLM 的先验知识初始化策略，并谨慎地微调至后验 PoM 分布的算法。
   • 理论分析表明，“谨慎的微调”（Careful Adaptation）至关重要：既要利用先验知识避免在巨大的策略单纯形中迷失，又要通过微调适应新的观测数据。

3. 实证验证：多领域 SOTA 性能

   • 在三个截然不同的任务中验证了方法：
     • FAQ 回复优化： 自然语言生成与语义对齐。
     • 热稳定蛋白质搜索： 设计具有高热稳定性的氨基酸序列（搜索空间极大）。
     • 量子电路设计： 生成低能量量子态的 Qiskit 电路。
   • 在与贝叶斯优化、强化学习（Actor-Critic, SAC）、进化搜索（Evolutionary Search）及上下文内采样（FIBO）的对比中，TOSFIT 在样本效率和计算效率上均达到了最先进（State-of-the-Art）水平。

---

4. 实验结果 (Results)

• 样本效率： 在所有三个任务中，TOSFIT 都能比基线方法（包括未引导生成、后生成汤普森采样、进化搜索等）更快地找到高奖励解。特别是在蛋白质搜索和量子电路设计中，TOSFIT 能够探索到训练数据流形之外的新解。
• 计算效率： 尽管涉及微调，但由于 TOSFIT 避免了昂贵的采集函数最大化（通常需要在离散空间进行暴力搜索或复杂的启发式搜索），其整体计算效率（达到特定奖励所需的总时间）优于基线。
• 批处理（Batched）能力： TOSFIT 天然支持批处理贝叶斯优化，能够在并行评估多个候选点时保持高效，显著减少了墙钟时间（Wall-clock time）。
• 消融实验：
  • 先验的重要性： 使用强先验（完整 Prompt 上下文）比弱先验或均匀初始化效果好得多。
  • 微调的谨慎性： 过大的学习率会导致模型“遗忘”先验知识并陷入局部最优，而较小的学习率配合多步梯度上升能平衡探索与利用。

---

5. 意义与影响 (Significance)

• 解决离散优化瓶颈： TOSFIT 为在大规模非结构化离散空间（如生物分子设计、代码生成、复杂文本优化）中进行高效贝叶斯优化提供了一条新途径，解决了传统方法无法处理梯度缺失和空间爆炸的问题。
• LLM 与贝叶斯优化的深度融合： 该工作展示了如何将基础模型（Foundation Models）的生成能力与 principled 的贝叶斯优化理论（如汤普森采样、遗憾界）相结合。它不仅仅是利用 LLM 作为黑盒生成器，而是将其作为可微分的概率分布参数化器，通过理论指导进行微调。
• 理论指导实践： 论文不仅提出了算法，还通过严格的理论分析（改进的遗憾界）解释了为什么“预训练初始化 + 谨慎微调”是有效的，为未来在离散空间使用 LLM 进行优化提供了理论依据。
• 实际应用价值： 在药物发现（蛋白质设计）、量子计算（电路设计）和自动化系统（FAQ 优化）等领域具有直接的应用潜力，能够加速科学发现和技术创新的过程。

总结： TOSFIT 通过巧妙地将汤普森采样转化为 LLM 的微调问题，成功克服了大规模离散空间贝叶斯优化的计算障碍，在理论和实践上均取得了显著突破，是基础模型与优化算法结合的一个典范。