LLM Agent：PaSa

阅读原文

LLM Agent：PaSa

以 PaSa（Paper Search）为例，接下来将介绍由 LLM 驱动的先进的论文搜索智能体。PaSa 能够自主做出一系列决策，包括调用搜索工具、阅读论文以及选择相关参考文献，最终为复杂的学术查询获取全面且准确的结果。

学术论文搜索是一项极具挑战性的信息检索任务，它需要涉及长尾的专业知识，达到全面的综述级别的覆盖范围，并且具备处理细粒度查询的能力。例如，考虑这样一个查询：“哪些研究专注于使用基于价值的方法，特别是基于上置信界（UCB）算法的非平稳强化学习？” 虽然像谷歌学术（Google Scholar）这种搜索系统对于一般性的查询很有效，但在处理这些复杂查询时往往力不从心。因此，研究人员常常要花费大量时间来进行文献综述。

大语言模型（LLMs）的发展激发了众多利用大语言模型来增强信息检索的研究，特别是通过优化或重新组织搜索查询来提高检索质量。然而，在学术搜索中，这个过程远不止简单的检索。人类研究人员不仅会使用搜索工具，还会进行更深入的活动，比如阅读相关论文以及核查参考文献，以便进行全面且准确的文献综述。

PaSa，旨在模仿人类行为以实现全面而准确的学术论文搜索。如下图所示，PaSa 由两个大型语言模型（LLM）智能体组成：Crawler（爬虫）和 Selector（选择器）。

Crawler 负责处理用户的查询请求，并且可以访问论文队列中的论文。它能够自主调用搜索工具、拓展参考文献，或者停止对当前论文的处理。Crawler 收集到的所有论文都会被添加到论文队列中。Selector 则会阅读论文队列中的每一篇论文，以判断该论文是否符合用户查询中所指定的标准

对于给定的用户查询，Crawler 可以通过利用搜索工具或从当前论文中提取引用来自主收集相关论文，并将这些论文添加到一个不断增长的论文队列中。Crawler 会迭代处理论文队列中的每一篇论文，通过导航参考文献网络来发现更多相关论文。Selector 则仔细阅读论文队列中的每一篇论文，以确定其是否满足用户查询的要求。PaSa 在 AGILE 框架（一种为 LLM 智能体设计的强化学习框架）中进行了优化。

有效的训练需要高质量的学术搜索数据。人类科学家已经创建了大量高质量的学术论文，其中包含对广泛研究主题的全面综述。PaSa 构建了一个合成但高质量的学术搜索数据集 AutoScholarQuery。

尽管 AutoScholarQuery 仅提供了查询和论文答案，而没有展示科学家收集论文的具体路径，但仍可以利用它进行强化学习训练以改进 PaSa。此外，PaSa 设计了一种新的会话级 PPO（近端策略优化）训练方法，以应对论文搜索任务的独特挑战：

1. 奖励稀疏：AutoScholarQuery 中的论文是通过参考文献收集的，因此它只是实际合格论文集的一个较小子集
2. 轨迹过长：Crawler 的完整轨迹可能涉及数百篇论文，这对于直接输入到 LLM 上下文来说过长

为了评估 PaSa，除了 AutoScholarQuery 的测试集外，还开发了基准数据集 RealScholarQuery。其包含 50个 真实世界的学术查询及其标注的相关论文，用于评估 PaSa 在真实场景中的表现。

该研究的主要贡献总结如下：

• PaSa：一个全面且准确的论文搜索智能体，它能够自主使用在线搜索工具、阅读整篇论文，以及在参考文献网络中进行导航
• 两个用于复杂学术搜索的高质量数据集：AutoScholarQuery、RealScholarQuery
• 尽管 PaSa 仅在合成数据上进行训练，但它在现实世界中取得了显著的性能表现。基于 70 亿参数大语言模型构建的 PaSa 显著优于所有基线模型

8.1、数据集

8.1.1、AutoScholarQuery

为了构建 AutoScholarQuery，收集了 ICLR 2023、ICML 2023、NeurIPS 2023、ACL 2024 以及 CVPR 2024 上发表的所有论文。对于每篇论文的 “相关工作” 部分，提示 GPT-4o 生成学术查询，这些查询的答案对应于 “相关工作” 部分中引用的参考文献。所使用的提示内容如下。

对于每个查询，仅保留可以在 arXiv 网站上检索到的论文，并使用 arXiv 编号作为数据集中文章的唯一标识符。AutoScholarQuery 采用源论文的发表日期作为查询日期。在训练和测试过程中，只考虑在查询日期之前发表的论文。

最终的 AutoScholarQuery 数据集在训练集、开发集和测试集中分别包含 33,551、1,000 和 1,000 个实例。每个实例由一个查询、相关的论文集以及查询日期组成，每个数据集中的查询均来自不同的源论文。下图展示了 AutoScholarQuery 中的一些示例。

为了评估 AutoScholarQuery 数据集的质量，抽取了 100 个 查询-论文（query-paper）对，并评估了每个查询以及相应论文的合理性和相关性。一个合格的查询应该是有意义且明确无误的，一篇合格的论文应该符合学术查询的要求。作者手动检查了每一对数据，确定 94.0% 的查询是合格的。在这些合格的查询中，93.7% 的查询所对应的论文被认为是相关且合适的。

8.1.2、RealScholarQuery

为了在更接近现实的场景中评估 PaSa 构建了真实学术查询数据集，这是一个由 50 个真实世界研究查询组成的测试数据集。在推出 PaSa 的演示版本后，邀请了几位人工智能研究人员使用该系统。从他们提供的查询中随机抽取了一部分查询，并手动过滤掉了过于宽泛的主题（例如，“多模态大语言模型”、“视频生成”）。最终，收集到了 50 个细粒度且贴近现实的查询。

对于每个查询，首先手动收集相关的论文。随后，使用多种方法来检索更多的论文，这些方法包括 PaSa、谷歌搜索、谷歌学术、ChatGPT（启用搜索功能的 GPT-4o），以及搭配 GPT-4o 进行查询改写的谷歌搜索。这些方法得到的结果被汇总到一个候选论文库中。最后，专业的注释人员会对每个查询的所有候选论文进行审查，挑选出那些符合该查询特定要求的论文，从而形成最终的相关论文集。

8.2、方法

8.2.1、概述

PaSa 系统由两个 LLM 智能体组成：Crawler 和 Selector。Crawler 读取用户的查询内容，生成多个搜索查询，并检索相关的论文。检索到的论文会被添加到一个论文队列中。Crawler 进一步处理论文队列中的每一篇论文，以识别值得进一步探索的关键参考文献，并将任何新的相关论文添加到论文列表中。Selector 对论文列表中的每一篇论文进行全面审查，以评估其是否满足用户的查询要求。

总之，Crawler 旨在最大限度地提高相关论文的召回率，而 Selector 则注重在识别符合用户需求的论文时的精确率。

8.2.2、Crawler

用强化学习（RL）的术语来说，Crawler 执行的是基于词元（token）级别的马尔可夫决策过程（MDP）。行动空间 $A$ 对应于大语言模型的词表，其中每个词元都代表一个行动。LLM 充当策略模型。智能体的状态由当前 LLM 的上下文和论文队列来定义。Crawler 通过三个已注册的函数来运行，具体内容见下表。当一个行动与某个函数名称匹配时，相应的函数就会被执行，进而修改智能体的状态。

如下图所示，智能体首先接收用户的查询，将其纳入自身的上下文，并开始采取行动。如果生成的词元是 [search]，大语言模型会继续生成一个搜索查询，然后智能体调用一个搜索工具来检索论文，这些论文随后会被添加到论文列表中。如果词元是 [Expand]，LLM 会继续从当前上下文中的论文提取一个子章节名称。接下来，智能体使用解析工具提取该子章节中的所有引用论文，并将它们添加到论文列表中。如果标记是 [Stop]，智能体就会将其上下文重置为用户查询以及论文队列中下一篇论文的信息。这些信息包括标题、摘要以及所有章节和子章节的提纲。

Crawler 使用多样且互补的查询执行多次 [search] 操作。此外，Crawler 能够评估其行动的长期价值。值得注意的是，当它在参考文献网络中进行更深入的探索时，即使探索路径上的中间论文与用户查询不相符，也能发现许多相关的论文

Crawler 的训练过程包括两个阶段。在第一阶段，为训练数据的一小部分生成轨迹，然后进行模仿学习。在第二阶段，应用强化学习。

模仿学习数据生成

PaSa 以会话为单位生成用于模仿学习的训练数据。会话有两种类型：搜索会话（从状态 $S_q$ 开始）和拓展会话（从状态 $S_{q+p}$ 开始）。

对于从 $S_q$ 开始的搜索会话，从 AutoScholarQuery 的训练集中采样用户查询，并提示 GPT-4o 生成相应的搜索查询。提示模板见下表。

会话轨迹的构建方式是，在每个查询前添加一个 [Search] 词元，将这些查询连接起来，并在末尾添加一个 [Stop] 词元，如下表所示。总共生成了 3011 条搜索会话轨迹。

对于从 $S_{q+p}$ 开始的拓展会话，使用谷歌搜索生成的查询。然后，从搜索结果中采样论文，并获得初始状态，该状态既包括查询也包括一篇论文。为了构建会话轨迹，检查论文的每个子部分。如果该子部分引用了自动学术查询数据集训练集中与该查询对应的至少一篇论文，则选择该子部分。否则，以 10% 的概率选择该子部分，以增加轨迹的多样性。将选定的部分填入上表中的模板，完成会话轨迹的构建。总共构建了 9978 条拓展会话轨迹。

下面将描述强化学习训练的具体实现细节。

a、奖励设计

PaSa 在 AutoScholarQuery 的训练集上进行强化学习训练，其中每个实例都由一个查询 $q$ 和相应的论文集 $P$ 组成。从一个查询 $q$ 开始，Crawler 生成一个轨迹 $\tau =(s_1,a_1,\cdots ,s_T,a_T)$ 。在每个时间步 $t$ ，将当前的论文队列记为 $Q_t$ 。在采取行动 $a_t$ 后，Crawler 将一组新的论文 $(p_1,p_2,\cdots ,p_{n_t})$ 添加到论文队列中。如果 $a_t=[Stop]$ ，则这组论文为空。

在状态 $s_t$ 下执行行动 $a_t$ 的奖励定义为：

$$r(s_t,a_t)=\alpha \times \sum _{i=1}^{n_t}\mathbb{I}(q,p_i,t)-c(a_t)$$

其中，如果 $p_i$ 与查询 $q$ 匹配且不在 $Q_t$ 中，则 $\mathbb{I}(q,p_i,t)=1$ ，否则 $\mathbb{I}(q,p_i,t)=0$ 。这里， $\alpha$ 是奖励系数， $c(a_t)$ 是行动 $a_t$ 的成本。

指示函数 $\mathbb{I}(q,p_i,t)$ 可以通过检查 $p_i$ 是否属于 $P-Q_t$ 来确定。然而，需要注意的是，AutoScholarQuery 数据集可能只包含真实论文集的一个子集，因为参考文献通常只强调有限数量的关键引用。如果 Crawler 仅仅根据与 AutoScholarQuery 数据集中论文的匹配情况来获得奖励，这可能会导致训练过程中的奖励稀疏问题。为了缓解这个问题，将 Selector 作为 Crawler 的辅助奖励模型。修正后的 $\mathbb{I}(q,p_i,t)$ 定义为：

$$\mathbb{I}(q,p_i,t)=\{\begin{array}{ll}1,& \text{if（Selector(q,pi)=1 or pi∈P ）and pi∈/Qt }\\ 0,& \text{otherwise}\end{array}$$

这里，如果 Selector 判定论文 $p_i$ 符合查询 $q$ ，则 $\text{Selector}(q,p_i)=1$ ，否则 $\text{Selector}(q,p_i)=0$ 。

b、强化学习训练

使用强化学习训练 Crawler 的一个关键挑战是，对于给定的查询，采样一个完整的轨迹需要花费大量时间。这是因为每个 [Search] 或 [Expand] 行动都会向论文列表中添加多篇论文，导致最终的论文队列中可能有数百甚至数千篇论文。

为了解决这个问题，将一个会话（session）定义为一个子轨迹，从一个会话的初始状态开始，以 [Stop] 行动结束。如此确定了两种类型的会话初始状态： $S_q$ ，只包含一个查询；以及 $S_{q+p}$ ，由一个查询和一篇论文组成。

形式上，将 Crawler 建模为一个策略 ${\pi }_{\theta }(a_t\mid s_t)$ 。将整个轨迹 $\tau$ 划分为一系列的会话： $({\tau }_{t_1:t_2-1},{\tau }_{t_2:t_3-1},\cdots )$ 。

每个会话是 ${\tau }_{t_i:t_{i+1}-1}=(s_{t_i},a_{t_i},\cdots ,s_{t_{i+1}-1},a_{t_{i+1}-1})$ ，其中初始状态 $s_{t_i}$ 要么属于 $S_q$ 类型，要么属于 $S_{q+p}$ 类型，并且最终行动 $a_{t_{i+1}-1}$ 是 [STOP]。

从这些会话初始状态中采样这样一个子轨迹在计算上是高效的。在 PPO 训练期间，在时间步 $t\in [t_i,t_{i+1})$ ，使用蒙特卡罗采样来估计会话中的回报：

$${\hat{R}}_t=\sum _{k=0}^{t_{i+1}-1-t}{\gamma }_0^k\left[r(s_{t+k},a_{t+k})+{\gamma }_1\sum _{j=1}^{n_{t+k}}{\hat{V}}_{\varphi }(S_{q+p_j})-\beta \cdot \log \frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{sft}(a_t|s_t)}\right]$$

这里， ${\gamma }_0$ 是会话内的折扣因子， ${\gamma }_1$ 是跨会话的折扣因子。 $${\hat{V}}_{\varphi }(\cdot )$$ 是用于近似状态值的价值函数模型。在执行 $a_{t+k}$ 后，论文队列会更新，以包含新找到的论文 $(p_1,p_2,\cdots ,p_{n_{t+k}})$ 。由于 Crawler 随后会启动新的会话来处理这些额外的论文，它们相关的未来奖励应该纳入到回报估计中。

此外，在每个词元上引入一个从学习到的策略 ${\pi }_{\theta }$ 到通过模仿学习得到的初始策略 ${\pi }_{sft}$ 的每个词元的 KL 散度惩罚项，以减轻过度优化问题。这个项由系数 $\beta$ 进行缩放。

然后，优势函数可以近似为：

$$\hat{A}(s_t,a_t)={\hat{R}}_t-{\hat{V}}_{\varphi }(s_t)$$

最后，策略目标和价值目标可以分别表示为：

$${\mathcal{L}}_{policy}(\theta )={\mathbb{E}}_{{\tau }^{\prime }\sim {\pi }_{\theta }^{old}}\left[\min \left(\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{\theta }^{old}(a_t|s_t)}\hat{A}(s_t,a_t),\text{clip}\left(\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{\theta }^{old}(a_t|s_t)},1-ϵ,1+ϵ\right)\hat{A}(s_t,a_t)\right)\right]$$

和

$${\mathcal{L}}_{value}(\varphi )={\mathbb{E}}_{{\tau }^{\prime }\sim {\pi }_{\theta }^{old}}\left[\max \left(({\hat{R}}_t-{\hat{V}}_{\varphi }(s_t){)}^2,({\hat{R}}_t-{\hat{V}}_{\varphi }^{clip}(s_t){)}^2\right)\right]$$

其中，

$${\hat{V}}_{\varphi }^{clip}(s_t)=\text{clip}\left({\hat{V}}_{\varphi }(s_t),V_{\varphi }^{old}(s_t)-ϵ,V_{\varphi }^{old}(s_t)+ϵ\right)$$

这里， ${\pi }_{\theta }^{old}$ 和 $V_{\varphi }^{old}$ 用于采样， ${\tau }^{\prime }$ 是会话轨迹。然后将这些组合成统一的强化学习损失：

$${\mathcal{L}}_{RL}(\theta ,\varphi )={\mathcal{L}}_{policy}(\theta )+\eta \cdot {\mathcal{L}}_{value}(\varphi )$$

其中 $\eta$ 是价值目标的系数。

8.2.3、Selector

Selector 接受两个输入：一个学术查询和一篇研究论文（包括其标题和摘要），生成两个输出：

1. 一个单一的决策词元 $d$，取值为 “True” 或 “False”，表示该论文是否满足查询要求
2. 一个理由 $r=(r_1,r_2,\cdots ,r_m)$，包含 $m$ 个词元，用于支持这一决策。这个理由有两个作用：
   • 通过联合训练模型来生成决策和解释，从而提高决策的准确性
   • 在 PaSa 的应用中提供推理依据，进而增强用户的信任度

为了提高 Crawler 的训练效率，决策词元会在理由之前给出，这使得 Selector 在 Crawler 的训练过程中可以充当一个单词元奖励模型。此外，决策词元的词元概率可用于对搜索结果进行排序。决策和理由的顺序并不影响选择智能体的性能。

PaSa 通过进行模仿学习来优化 Selector。训练数据的收集和训练细节如下。

选择智能体的实现细节

首先从 AutoScholarQuery 的训练集中采样用户查询。对于 AutoScholarQuery 训练集中的每个用户查询及其对应的一篇论文，提示 GPT-4o 生成一个决策词元以及理由（提示内容见下表）。

剔除那些决策词元为 “False” 的数据，因为这与 AutoScholarQuery 的标签相矛盾。剩余的数据将作为正样本 <用户查询，论文> 对保留下来。

接下来，使用基于 GPT-4o 的 PaSa-GPT-4o 模拟部分论文搜索过程。在这个模拟中，每篇论文都有 50% 的概率被添加到论文队列中。如果某篇论文没有被 GPT-4o 选中，并且也不在 AutoScholarQuery 的训练集中，那么对应的 <用户查询，论文> 对将被标记为负样本。

最终的训练数据集由 19,812 个 <用户查询，论文> 对组成，每个对都带有由 GPT-4o 生成的决策词元以及理由，这些数据对均取自 AutoScholarQuery 训练集中的 9000 个实例。

论文选择提示

下表展示了用于 PaSa 的 Selector 和 GPT-4o 判断一篇论文是否符合用户查询要求的提示内容。

下表呈现了用于与 GPT-4o 配合自动生成 AutoScholarQuery 的提示模板。对于每篇论文，提取其 “相关工作” 部分，将其输入到 GPT-4o 中，并使用该提示从 “相关工作” 部分中提取学术查询及其相应的论文答案。

参考文献

1. PaSa: An LLM Agent for Comprehensive Academic Paper Search