LLM 路由用于批量指令
引言
大型语言模型 (LLM) 已成为众多 AI 应用不可或缺的一部分。然而,大规模使用 LLM 存在性能和成本之间的关键权衡。没有一个 LLM 对所有任务都是最优的,并且在管理预算限制的同时为每个指令选择最合适的模型是一个根本性挑战。
在本报告中,我们描述了一种预测性路由方法,用于将指令分配给既经济又高性能的 LLM。我们提出了三种分配策略
- 基于阈值的分配:选择满足目标性能阈值的最便宜的 LLM。
- 使用非递减凸包 (NDCH) 的固定预算分配:确定 LLM 的最佳组合,以在成本约束下最大化性能。
- 使用混合整数规划 (MIP) 的优化预算分配:将分配问题表述为受约束的优化问题,以最大化平均精度。
我们将在真实世界数据集上评估这些策略,并提供对其有效性的见解。
问题定义
给定一批指令 ,我们的目标是将每个指令 分配给 LLM ,该 LLM 从可用模型池 ,同时优化以下任一目标:
- 性能阈值:选择满足所需准确率的最便宜的模型。
- 固定预算:在给定预算内最大化性能。
为此,我们使用一个多标签分类器,该分类器经过训练,可以预测每个 LLM 每条指令的成功概率,并使用温度标定和直方图分桶进行校准。
预测 LLM 性能
给定指令数据集和记录的 LLM 性能指标,我们可以将预测 LLM 在给定指令上成功概率的任务构架为多标签分类问题。这类似于为每个 LLM 训练一个单独的二元分类器,但大多数模型参数在分类器之间共享,从而实现高效学习。
我们的分类器建立在嵌入模型上,例如像 BERT 这样的掩码语言模型,它经过专门训练用于句子嵌入。然后,在其之上添加一个分类头,以估计每个 LLM 的成功概率。
为了提高概率估计的可靠性,部分训练集被保留为校准数据集。应用温度缩放以调整模型的置信度分数,确保它们更好地与真实世界性能保持一致。此步骤可减少过度自信并增强可解释性。
校准后,成功概率的范围被划分为离散的桶。对于每个桶,我们计算每个 LLM 的预期成功率。此过程将原始概率分数转换为更真实和可解释的成功估计,从而实现更明智的路由决策。
分配策略
一旦我们有了能够估计给定指令上每个 LLM 成功概率的分类器,我们就可以利用这些概率来指导指令到不同 LLM 的分配,使用以下策略。
1. 基于性能阈值的分配
给定目标性能阈值(例如,0.75 准确率),我们从校准后的分类器中识别出满足目标性能的性能桶。然后,我们选择预测成功率高于目标性能阈值的最便宜的 LLM。
伪代码:
for instruction in batch:
candidate_models = [M_k for M_k in models if P(instruction | M_k) >= threshold]
best_model = min(candidate_models, key=lambda m: cost[m])
assign(instruction, best_model)
2. 使用 NDCH 的固定预算分配
非递减凸包 (NDCH) 是一种在不同模型之间进行选择时,构建最优成本-性能权衡曲线的方法。它确保随着成本的增加,性能不会下降,从而形成一条单调递增的可实现权衡前沿。给定一组具有不同成本和成功率的 LLM,NDCH 通过在模型之间进行线性插值来识别最佳分配,从而在最大化性能的同时实现高效的预算分配。
NDCH 的工作原理是过滤掉严格支配的模型——那些比其他可用模型更昂贵且准确率更低的模型。然后,它在剩余的成本-性能点上构建一个凸包,确保曲线上每个点都代表成本和准确率之间的最佳权衡。当给定固定预算时,NDCH 确定分配给不同模型的指令的理想比例,平衡低成本、低性能模型和更昂贵、高准确率的模型。这种方法保证任何额外的支出都能带来有意义的性能提升,避免低效分配。
使用非递减凸包方法,我们
- 为所有 LLM 构建成本-性能权衡曲线。
- 确定廉价和昂贵 LLM 的最佳组合,以在满足预算约束的同时最大化性能。
伪代码:
convex_hull = compute_non_decreasing_convex_hull(cost_performance_points)
budget_allocation = determine_mixture(convex_hull, total_budget)
for instruction in batch:
model = sample_from_allocation(budget_allocation)
assign(instruction, model)
原始的 NDCH 方法没有考虑每个 LLM 对单个指令的成功概率。相反,它根据计算出的比率随机将指令分配给模型。然而,由于我们有一个分类器可以预测不同 LLM 上每个指令的成功概率,我们引入了一种名为 NDCH(P) 的改进变体。这种方法在预测成功概率高时优先将指令路由到更便宜的 LLM,从而确保更明智的分配。通过利用这些概率,NDCH(P) 通过将指令分配给最便宜但仍可能表现良好的模型来提高成本效率,从而在保持准确性的同时更好地利用整体预算。
3. 使用混合整数规划的预算分配
在必须严格遵守固定预算但同时需要最大化性能的场景中,我们将分配问题构建为混合整数规划 (MIP) 优化任务。MIP 提供了一种结构化的方法,通过同时考虑成本和性能来实现全局最优分配。
此方法的关键目标是在分配的预算内最大化所有已分配指令的平均精度。与基于启发式的方法不同,MIP 全面考虑所有可能的分配,确保每个指令都以平衡成本效益和性能的方式路由到最合适的 LLM。
为了实现这一点,我们定义了一组约束条件
- 预算约束: 所有选定模型分配的总成本不得超过给定预算。
- 独占分配: 每条指令必须恰好分配给一个 LLM。
- 性能意识: 每条指令-LLM 对的预测成功概率应指导选择过程,确保在预算限制下优先选择性能最高的分配。
通过利用 MIP,我们有效地确定了一种最优的路由策略,该策略可以动态调整分配给不同模型的指令比例。这确保了预算的利用方式能够最大化整体性能,而不是随意分配。此外,MIP 允许纳入额外的约束,例如延迟要求或不同 LLM 之间的公平性考虑,使其成为一个高度灵活的优化框架。
伪代码:
m = create_optimization_model()
for i, instruction in enumerate(batch):
for M_k in models:
x[i, M_k] = add_binary_variable(m) # Decision variable
m.setObjective(sum(x[i, M_k] * get_precision(M_k, P(i, | M_k)) for i in batch for M_k in models), MAXIMIZE)
m.addConstr(sum(x[i, M_k] * cost[M_k] for i in batch for M_k in models) <= budget)
for i in batch:
m.addConstr(sum(x[i, M_k] for M_k in models) == 1)
solve(m)
结果与性能比较
为了评估我们路由策略的有效性,我们对 LLM-Uncertainty-Bench 和 Router R1 Math Reasoning 两个数据集进行了实验,比较了不同分配方法下的准确率-成本权衡。我们使用 70% 的数据集作为训练集,20% 作为测试集;其余样本用于校准和验证。我们对不同的随机种子重复实验五次,并在图中显示了结果的方差。
LLM-Uncertainty-Bench 数据集
该数据集包含 50,000 个样本,涵盖五项关键 NLP 任务,每项任务有 10,000 个样本。每项任务旨在评估模型推理和不确定性量化的不同方面。任务包括
- 问答 (QA): 基于 MMLU 评估通用知识检索和推理。
- 阅读理解 (RC): 使用 CosmosQA 评估上下文理解。
- 常识推理 (CI): 使用 HellaSwag 测试逻辑推理。
- 对话响应选择 (DRS): 使用 HaluEval 衡量对话上下文意识。
- 文档摘要 (DS): 评估摘要能力,也源自 HaluEval。
此数据集包含 16 个 LLM 的标签,包括 Qwen、Yi、MPT、Llama、Internlm、Mistral、Falcon 和 DeepSeek 等模型家族。这些模型的成功率和成本如下表所示。我们根据模型大小设置这些模型的成本,确保它们与类似大小的公共 API 模型的成本大致相当。然而,对于 Yi-34B,我们特意将其成本提高到超过该集合中最大模型的成本,因为它是性能最佳的模型。这种调整给路由器引入了更具挑战性的分配问题,迫使其在成本和性能之间进行更细微的权衡。
模型 | 成功率 | 成本(百万 token) |
---|---|---|
Yi-34B | 0.824 | 2.20 |
Qwen-72B | 0.785 | 2.19 |
Qwen-14B | 0.740 | 0.90 |
Llama-2-70b-hf | 0.730 | 2.16 |
deepseek-llm-67b-base | 0.728 | 2.00 |
Yi-6B | 0.694 | 0.60 |
Mistral-7B-v0.1 | 0.643 | 0.69 |
Llama-2-13b-hf | 0.606 | 0.90 |
Qwen-7B | 0.603 | 0.69 |
internlm-7b | 0.485 | 0.69 |
Llama-2-7b-hf | 0.463 | 0.69 |
deepseek-llm-7b-base | 0.440 | 0.69 |
Qwen-1_8B | 0.418 | 0.18 |
falcon-40b | 0.345 | 1.24 |
mpt-7b | 0.263 | 0.69 |
falcon-7b | 0.250 | 0.69 |
此数据集的实验结果如下图所示,其中我们将预算表示为可用最昂贵模型成本的百分比,基于上表中的值。研究结果表明,基于 MIP 的路由方法以最佳模型 50% 的成本实现了其 95% 的性能。虽然 MIP 显著优于 NDCH 和 NDCH(P),但这些方法在降低成本方面也表现出有效性,在相同预算内实现了最佳模型 90% 的性能。这突出表明,即使是基于启发式的方法也能带来可观的成本节约,同时保持高性能。
路由器 R1 数学推理数据集
路由器 R1 数学推理 数据集是一个公开可用的数据集,包含 3,234 个样本。它专门用于测试数学推理和判断任务,在评估合成数据生成策略和模型鲁棒性方面发挥着关键作用。该数据集包含以下模型的标签:R1-1.5B、R1-7B、R1-14B 和 R1-32B。这些模型的成功率和成本如下所示。
模型 | 成功率 | 成本(百万 token) |
---|---|---|
R1 | 0.951 | 2.19 |
R1-32B | 0.936 | 1.18 |
R1-14B | 0.915 | 0.90 |
R1-7B | 0.900 | 0.69 |
R1-1.5B | 0.819 | 0.18 |
此数据集的实验结果如下图所示,我们观察到所有路由方法都表现相似,其中 NDCH(P) 略有优势。它在不同预算比例下相对于最佳模型始终优于其他方法。例如,在最佳模型 50% 的预算下,其他路由方法达到其准确率的约 97%,而 NDCH(P) 达到近 99%。这表明,利用预测成功概率进行分配在不牺牲性能的情况下显著提高了成本效率。值得注意的是,使用预测概率进行路由使整体成功率超过了最佳模型,而成本仅为其 70%。
结论与未来工作
在这项工作中,我们探索了在预算约束下将指令路由到 LLM 的不同策略。虽然我们引入了基于阈值的分配作为一种直接的方法,但它需要一个预定义的性能阈值,由于缺乏对可实现准确率上限的了解,这通常不切实际。在大多数实际场景中,预算限制比预定义的性能目标更相关,这促使我们专注于预算感知的分配策略。
为了解决这个问题,我们评估了 NDCH、NDCH(P) 和基于 MIP 的路由,每种方法都旨在优化 LLM 选择,同时保持在固定预算内。基于 NDCH 的分配构建了一个有效的成本-性能权衡曲线,确保资源在可用模型之间得到优化分配。NDCH(P) 通过纳入预测成功概率来完善此方法,进一步提高效率。另一方面,基于 MIP 的优化将分配问题表述为受约束的优化任务,在给定预算内产生最佳性能。
我们的结果表明,MIP 路由始终提供最高的效率,在显著降低成本的同时实现近乎最优的准确率。同时,NDCH 和 NDCH(P) 提供了具有竞争力的替代方案,尤其是在计算效率优先的设置中。
我们将在未来几周内发布代码以重现所呈现的结果。
未来的研究将探索其他数据集并纳入新约束,例如延迟感知路由和级联决策。在延迟感知路由中,模型将根据性能和推理速度进行选择,以满足延迟要求。另一方面,级联路由引入了一种自适应选择过程,其中初始响应使用更便宜的 LLM 生成,其质量通过单独的验证机制进行评估。如果响应被认为不足,则指令将升级到更强大但成本更高的模型。这种方法确保了高效的资源利用,同时保持了高质量的输出。