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次重复通过，同时递归细化列表上半部分，我们在使用GPT-4o mini下2分钟内获得排名列表。

1

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 id 3 academy 13 prime 23 ma 33 accountant 4 education 14 cal 24 zero 34 guru 5 school 15 mm 25 mu 35 py 6 institute 16 mt 26 scholarships 36 science 7 mit 17 college 27 financial 37 plus 8 courses 18 solutions 28 training 38 technology 9 phd 19 study 29 ieee 39 expert 10 engineering 20 data 30 engineer 40 foundation

很难争论“edu”作为数学相关TLD。您可以注意到此模型的偏见：一些顶部结果 heavily lean on math as an academic concept，应用概念（accountants等）在列表中稍低。如果需要，我们可以轻松调整提示，为LLM比较器提供更多方向，但对于第一次通过，它 clearly works as intended。

用于漏洞识别

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

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

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

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

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