raink：使用LLM进行文档排序

TL;DR：Bishop Fox发布了raink，这是一个命令行工具，采用新颖的基于LLM的列表排序算法。最初在RVASec 2024上展示，raink解决了复杂的排序问题，包括将代码差异与安全公告关联。

2024年6月，Bishop Fox在RVASec上展示了“Patch Perfect: Harmonizing with LLMs to Find Security Vulns”，演示了如何使用基于LLM的新算法在N-day漏洞分析中将软件补丁中的代码差异与相应的安全公告关联。现在，我们开源了Bishop Fox的新工具raink：将我们的列表排序算法实现为命令行工具。

本文将展示raink如何解决LLM难以处理的一般排序问题。我们将通过一个相对简单的问题（排名域名）来说明raink的工作原理，然后简要介绍在补丁差异分析场景中用于漏洞识别的建议用法。

排名TLD

AI有时感觉很神奇，当你简单“把问题扔给它”（即使没有完全定义问题约束）时，仍然能得到有意义的结果。例如：哪个顶级域（TLD）最数学？

当我看到这个帖子时，我渴望自动化找到答案。如何确定一个给定术语是否与数学相关？你是严格解释这个术语吗？你把它当作首字母缩写词吗？重要的是，你如何决定一个术语比另一个更数学？这是一个模糊的问题，感觉像是“我看到就知道”。

在小规模上，正常的交互式ChatGPT会话效果很好。以下是我从DNS注册商随机抽取的10个TLD样本。我们可以简单地要求答案，而不解决上述任何细微差别：

结果检查通过。我自己会选择“int”。酷，让我们放入IANA的所有1445个可用TLD，找出最数学的那个！

好吧，有几个问题：

• ChatGPT承认几个结果实际上不在原始列表中。感谢但无帮助。
• 我们只得到了16个结果。（模型拒绝进一步提示所有结果，理由是“这将是一个非常广泛的过程”。）

我们可以尝试更强大的模型或编写更好的提示——但如果你经常使用LLM，你会从经验中知道，无论我们尝试什么，我们都不会完全停止幻觉或得到完整的原始输入列表。似乎LLM“记住”所有你给它的数据的能力是有限的，在某些情况下，你可能会提供比它上下文窗口能容纳的更多的数据：

不要害怕——我们有计算机科学学位。让我们简单地分而治之。

如果我们将大的TLD列表分成多个较小的列表，我们可以对每个列表排序，然后将结果合并成最终排序列表，对吗？这听起来起初很容易，但变得棘手。排序一个小列表工作良好，正如我们最初证明的——但我们如何将两个排序列表合并在一起？也就是说，我们如何比较列表A的顶部项目与列表B的顶部项目，因为它们的排名仅相对于各自单独列表中的同行有意义？

我们可以尝试更清楚地定义问题，并给LLM一个规定的评分算法，为每个列表项生成一个数值“客观”分数，可以确定性地排序。我会为你节省一些时间，并报告我使用这种方法的经验：LLM在分配分数方面表现不佳。它们往往不一致，并且将分数膨胀到你正在使用的任何尺度的上端。

文档排名的TLD-R

面对这些挑战，让我们庆幸我们远非第一个解决文档排名问题的人。简要历史：

• Google开发了PageRank，最早的排名算法之一，用于确定网页的“重要性”。它通过将互联网视为一个巨大图来工作，其中每个链接都是信任投票。拥有更多链接的页面——尤其是来自其他重要页面的页面——获得更高分数。简单且严格数值。
• 后来出现了学习排名（LTR），它建立在像PageRank这样的想法上，但添加了机器学习。LTR使用训练数据教ML模型如何基于相关性或上下文等特征对项目排名。它引入了相对比较项目的概念，并通常使用三种主要方法之一：
  • 点式：单独对待每个项目，预测每个项目的单个“相关性分数”。类似于基本分类问题：“项目X有多相关？”然后你按这些预测分数排序。
  • 对式：一次比较两个项目（“A比B好吗？”）并将这些结果组合成完整排序列表。这将排名任务简化为简单、重复的是/否选择，但在比较每对可能时可能更耗时。
  • 列表式：一次查看一组项目，试图直接一次性改进最终排序。它可以更准确，但更难实现，因为它处理排名完整列表的复杂性，而不是单个项目或对。
• 最近，一篇2024年论文（也由Google发表）引入了对式排名提示（PRP），利用LLM的力量。而不是训练模型或分配分数，PRP使用简单提示问LLM，“这两个中哪个更好？”这种方法简化了排名任务，并使用开源模型实现了最先进的性能，表明即使中等大小的LLM在排名问题上也能胜过像GPT-4这样更大的商业替代品。

PRP论文具体解决了我们在TLD排名问题中遇到的挑战（这些问题应该听起来熟悉！）。一些直接引用为清晰编辑：

• 点式相关性预测要求模型输出校准的点式预测，以便它们可以用于排序比较。这在不同提示之间很难实现。（即，LLM很难为单个项目生成一致的客观分数。）
• 由于LLM列表排名任务的困难，经常出现预测失败，采取以下模式：
  • 缺失：当LLM只输出输入文档的部分列表。
  • 拒绝：LLM拒绝执行排名任务并产生不相关输出。
  • 重复：LLM多次输出相同文档。
  • 不一致：相同的文档列表在以不同顺序或上下文输入时有不同的输出排名。

PRP研究人员选择对式方法因其简单性和鲁棒性（显著降低LLM的任务复杂性并解决校准问题），同时注意到效率问题，因为对式方法通常比点式或列表式计算更复杂。考虑这两种PRP变体，其各自的复杂性O(N2)和O(N log N)都低于列表式的O(N)：

• PRP-Allpair枚举每对可能的文档，并提示LLM决定哪篇文档更相关。每篇文档的整体相关性分数是其在成对匹配中的总“胜利”。虽然这可以实现高准确性，但需要O(N2)次LLM调用用于N篇文档。
• PRP-Sorting使用标准排序算法（如Heapsort或Quicksort），但用基于LLM的比较器替换通常的“比较两个项目”步骤。而不是穷举比较所有对，算法只比较排序过程中需要的对，导致更高效的O(N log N)复杂性。

所以，对式方法效果很好，但需要大量LLM API调用。我们能否减少LLM调用，同时仍然得到“足够好”的结果？

列表式或失败

列表式方法有潜力，但我们必须解决几个问题。让我们引入一些概念来修复PRP论文中概述的问题：

• 批处理：将原始列表分成相对较小的子集，这些子集将适合上下文窗口且不会压垮模型。这有助于解决缺失和拒绝问题。
• 验证：检查LLM调用的输出，并根据需要实施重试。这有助于解决缺失、拒绝和重复问题。
• 重复：对洗牌输入多次运行过程，因此每个项目充分与许多其他项目比较（但不是每个其他项目，如PRP-Allpair）。这有助于解决不一致问题。

以下是Bishop Fox新raink工具的算法工作原理：

1. 初始批处理和排名
   • 洗牌所有项目。
   • 将它们分成小批次（例如，10组）。减少批次大小，直到批次安全适合上下文窗口（这可能取决于标记化输入项目的大小）。
   • 单独排名每个批次以获得本地排序，同时验证LLM调用返回了我们放入的所有项目（否则重试）。
   • 保存每个项目在其批次中的相对位置，用作数值分数。
2. 重复传递
   • 多次运行步骤1（洗牌-批次-排名）（例如，10次传递）。
   • 平均每个项目每次传递的相对位置，形成初始排名。
3. 细化
   • 选择顶部部分（例如，上半部分）基于当前排名。
   • 在此上部子集上重复步骤1和2（多传递洗牌-批次-排名）。
   • 继续递归细化，直到隔离顶部项目。
4. 重建完整列表
   • 将你的细化排序与其余项目合并，产生最终排序列表。

算法的三个主要参数是：

• 批次大小：一个批次中可以容纳多少项目
• 传递次数：重复洗牌-批次-排名的次数
• 细化比率：递归细化的上部部分有多大

细化步骤有助于确保我们很好地花费精力——即，更积极地比较似乎最相关的项目，而不是执行每对项目的完整成对比较。LLM调用也可以为每次洗牌-批次-排名传递并行进行。

测试

使用10个项目批次大小，10次重复传递，同时递归细化列表上半部分，我们在不到2分钟内使用GPT-4o mini获得排名列表。

raink -f tlds-iana.lst -r 10 -s 10 -p ‘Rank each of these top-level domains (descending order, where most relevant is first) according to their relevancy to the concept of “math”.’

以下是raink认为前5%最数学的TLD： 1 edu 11 analytics 21 int 31 accountants 2 university 12 degree 22 iq 32 accountant 3 academy 13 prime 23 ma 33 guru 4 education 14 cal 24 zero 34 py 5 school 15 mm 25 mu 35 science 6 institute 16 mt 26 scholarships 36 plus 7 mit 17 college 27 financial 37 technology 8 courses 18 solutions 28 training 38 expert 9 phd 19 study 29 ieee 39 foundation 10 engineering 20 data 30 engineer 40

很难争论“edu”作为数学相关TLD。你可以注意到这个模型的偏见：一些顶级结果严重倾向于数学作为学术概念，应用概念（会计师等）在列表中稍低。如果需要，我们可以用更多方向轻松调整提示给我们的LLM比较器，但作为第一次通过，它显然按预期工作。

用于漏洞识别

我们已经显示raink可用于排序对象列表，关于它们与某个给定主题的密切程度。而不是与数学概念相关的TLD，考虑我们可以输入软件补丁中更改的函数，并尝试找出哪些与给定安全公告最密切相关。或者也许排名Semgrep结果，关于哪些似乎对分析师调查最有影响。

要将raink应用于补丁差异分析中的漏洞识别：我们可以传递最近安全公告的文本，以及从Ghidriff补丁差异生成的代码更改列表，并要求它排名最可能与公告中描述的问题相关的更改函数：

在上述配置中连续运行五次后，raink能够成功识别固定函数：

• 作为顶部排名项目，60%的时间（3/5次）。
• 平均在顶部7%排名项目内。

我们看到在进攻性安全工程师以这种方式使用raink时有很多效率增益的潜力；我们将为未来的博客文章保存补丁差异分析的详细结果。同时，在GitHub上下载raink，并查看我们的演讲Patch Perfect，了解我们应用LLM到N-day漏洞识别的一些成功。