raink：使用LLM进行文档排序

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

背景

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

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

排名TLD的挑战

当简单"向AI抛出问题"时，AI可能显得很神奇——即使没有完全定义问题约束，也能获得有意义的结果。例如：哪个顶级域名（TLD）最具数学相关性？

当处理小规模数据时，正常的交互式ChatGPT会话效果很好。但当我们尝试处理IANA的所有1445个可用TLD时，会出现几个问题：

1. ChatGPT承认几个结果实际上不在原始列表中
2. 只得到16个结果（模型以"这将是一个非常广泛的过程"为由拒绝提供所有结果）
3. 可能提供超出上下文窗口容量的数据

文档排名的历史背景

面对这些挑战，值得庆幸的是我们远非第一个解决文档排名问题的人：

• PageRank：谷歌开发的最早排名算法之一，通过将互联网视为大型图来确定网页的"重要性"
• 学习排序（LTR）：在PageRank等思想基础上添加机器学习，使用训练数据教ML模型如何基于相关性或上下文等特征对项目排序
  • Pointwise：单独处理每个项目，预测每个项目的单个"相关性分数"
  • Pairwise：一次比较两个项目（“A比B好吗？”）并将这些结果组合成完整的排序列表
  • Listwise：一次查看一组项目，试图直接一次性改进最终排序
• Pairwise Ranking Prompting（PRP）：2024年谷歌论文引入，使用简单提示询问LLM"这两个中哪个更好？"

PRP论文专门解决了我们在TLD排名问题中遇到的挑战：

• Pointwise相关性预测需要模型输出校准的点预测，这在不同提示间很难实现
• 由于LLM列表排序任务的难度，经常出现预测失败：
  • 缺失：LLM只输出输入文档的部分列表
  • 拒绝：LLM拒绝执行排序任务并产生不相关输出
  • 重复：LLM多次输出相同文档
  • 不一致：相同的文档列表在不同顺序或上下文下有不同的输出排名

列表排序解决方案

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

• 批处理：将原始列表分成相对较小的子集，以适应上下文窗口且不会压垮模型
• 验证：检查LLM调用的输出并根据需要实施重试
• 重复：在洗牌输入上多次运行过程，使每个项目与许多其他项目充分比较

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

1. 初始批处理和排名

   • 洗牌所有项目
   • 将它们分成小批次（例如，10个一组）
   • 单独排名每个批次以获得本地排序，同时验证LLM调用返回了我们放入的所有项目
   • 保存每个项目在其批次中的相对位置作为数值分数

2. 重复传递

   • 多次运行步骤1（洗牌-批处理-排名）
   • 平均每个项目每次传递的相对位置以形成初始排名

3. 细化

   • 基于当前排名选择顶部部分（例如上半部分）
   • 在此上部子集上重复步骤1和2（多遍洗牌-批处理-排名）
   • 继续递归细化，直到分离出顶部项目

4. 重建完整列表

   • 将细化排序与其余项目合并，生成最终排序列表

算法的三个主要参数：

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

测试结果

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

raink命令：

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".'

前5%最具数学相关性的TLD：

1. edu
2. university
3. academy
4. education
5. school
6. institute
7. mit
8. courses
9. phd
10. engineering

漏洞识别应用

我们已经显示raink可用于根据与给定主题的紧密程度对对象列表进行排序。考虑我们可以输入软件补丁中更改的函数，并尝试找到哪些函数与给定安全公告最密切相关。

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

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

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

我们看到在进攻性安全工程师以这种方式使用raink方面有很大的效率提升潜力。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）
公众号二维码