raink：使用LLM进行文档排序
TL;DR：Bishop Fox发布了raink，这是一个基于新型LLM列表排序算法的命令行工具。该工具最初在RVASec 2024上展示，能够解决复杂排序问题，包括将代码差异与安全公告关联。
背景与挑战
在RVASec 2024上，Bishop Fox展示了如何利用基于LLM的新型算法，在N-day漏洞分析中将软件补丁中的代码差异与相应安全公告关联。现在，我们开源了raink工具，将列表排序算法实现为命令行工具。
LLM在处理大规模排序任务时面临以下挑战：
- 缺失问题：LLM仅输出部分输入文档
- 拒绝问题：LLM拒绝执行排序任务
- 重复问题：LLM重复输出相同文档
- 不一致问题：相同输入在不同上下文下产生不同排序结果
排序算法演进
传统方法
- PageRank：基于链接分析的网页重要性排序算法
- 学习排序(LTR)：使用机器学习模型进行排序，包含三种方法：
- Pointwise：为每个项目单独预测相关性分数
- Pairwise：两两比较项目优劣
- Listwise：直接优化整个列表的排序
最新进展
2024年Google提出的Pairwise Ranking Prompting (PRP)方法：
- PRP-Allpair：枚举所有文档对，复杂度O(N²)
- PRP-Sorting：使用排序算法，复杂度O(N log N)
raink算法原理
raink采用改进的列表排序方法，通过以下技术解决LLM排序问题：
核心机制
- 批处理：将输入列表分成小批次处理
- 验证：检查LLM输出并实现重试机制
- 重复处理：对打乱顺序的输入进行多次处理
算法步骤
-
初始批处理和排序
- 打乱所有项目
- 分成小批次（如每组10个）
- 对每个批次单独排序
- 保存每个项目的相对位置作为分数
-
多次处理
- 重复步骤1多次（如10次）
- 计算每个项目的平均位置作为初始排名
-
精细化处理
- 选择排名靠前的部分项目（如前50%）
- 对子集重复步骤1和2
- 递归细化直到确定顶部项目
关键参数
- 批次大小：每批处理的项目数量
- 处理次数：重复排序的次数
- 细化比例：递归处理的顶部项目比例
实际应用测试
TLD排序测试
使用GPT-4o mini，在2分钟内完成1445个TLD的数学相关性排序：
|
|
前5%数学相关TLD结果：
- edu
- university
- academy
- education
- school …（共40个结果）
漏洞识别应用
在补丁差异分析中，raink可识别与安全公告最相关的代码变更：
|
|
测试结果：
- 60%的情况下正确识别修复函数（3/5次）
- 平均排名在前7%以内
技术优势
- 高效性：相比PRP-Allpair的O(N²)复杂度，raink通过批处理和递归优化显著减少LLM调用次数
- 准确性：通过多次处理和验证机制提高排序准确性
- 灵活性：支持自定义提示词和参数调整
- 并行处理：支持批量LLM调用并行处理
应用场景
- 漏洞识别：关联代码差异与安全公告
- 结果排序：对安全扫描结果进行优先级排序
- 数据分类：基于语义相似性的文档排序
raink已在GitHub开源，为 offensive security 工程师提供了高效的LLM排序解决方案。