Tryage：实时、智能地将用户提示路由至大语言模型

23年8月来自加州理工的论文“Tryage: Real-time, Intelligent Routing of User Prompts to Large Language Models”。

Transformer 架构和自注意机制的引入，使针对特定下游任务和数据域训练的语言模型数量激增。Hugging Face 生态系统中有超过 200,000 个模型，用户在解决计算、安全和近因问题的同时，还要努力选择和优化模型以适应多方面的工作流程和数据域。迫切需要机器学习框架来消除模型选择和定制的负担，并为最终用户发挥庞大的涌现模型库惊人功能。在这里提出一种上下文-觉察路由系统 Tryage，利用语言模型路由器对单个输入提示的分析，它从模型库中最佳地选择专家模型。 Tryage 受到大脑中丘脑路由器的启发，采用感知路由器预测下游模型在提示下的性能，然后使用目标函数做出路由决策，该目标函数将性能预测与通过flag (例如模型大小、模型近因) 纳入的用户目标和约束相结合。Tryage 允许用户探索帕累托前沿（Pareto front），并自动在任务准确性和次要目标 (包括最小化模型大小、近因、安全性、详细程度和可读性) 之间进行权衡。从概念上讲，Tryage 展示如何应用路由模型来编程和控制多模型 LLM 系统的行为，最大限度地提高对不断扩大和发展的语言模型生态系统有效利用。

语言模型在从 CommonCrawl 等通用数据源以及特定域数据（Pile 子集、财务数据、临床数据）中采样的数据上进行训练，并使用不同的范式（掩码语言建模、因果语言建模）进行训练和微调，因此，现有语言模型在下游任务和从域中提取数据上表现出不同的性能。

如图展示一组流行的 Transformer 在一系列不同数据域的掩码语言建模上的性能。虽然 Roberta 模型在 Pile 数据集所有数据域中表现出最高的平均准确率，但当输入数据来自 GitHub 或专利数据库等特定领域时，Codebert 和 Patent Bert 等专家模型的表现明显优于 Roberta。经验结果促使开发一种多模型架构，该架构可以根据特定输入动态选择最佳模型。该系统的目标是在从特定数据分布中提取提示或任务时保持较高的一般准确率，同时获得专家级性能。

为了能够将用户提示和任务动态路由到模型库，设计一个感知路由系统。该系统包含一个执行路由功能的 Tryage（路由）节点和一个下游/叶节点库。Tryage 节点接受用户提示并根据提示选择下游模型。路由功能结合下游模型的性能，并将其预测准确性与用户提供的约束/目标flags集成到单个目标函数中。路由功能通过权衡预测准确性和一组参数加权约束来选择特定模型。路由损失函数的优化会导致模型选择决策，提示将被路由到下游模型。如图所示：在 Tryage 系统中，系统会收到提示和flag，系统会根据flag和提示找到执行 MLM 任务的最佳模型。

考虑一个语言模型库 M = (M1 . . . Mn )，其中 Mi 是单个模型。每个模型 Mi 都接受输入 z，并将输入映射到输出 Mi(z)，例如，输出 Mi(z) 可以是下一个单词预测概率。一般来说，由于模型架构、训练数据和训练范式的不同，M 的不同元素在提示和任务上的表现会有所不同。目标是开发一个模型选择/提示路由策略，其中路由模型 R 可以为特定提示选择最佳专家模型 Mi，同时将用户约束作为次要损失函数纳入路由决策。

在 Tryage 中，给定一个输入 z 和一个损失函数为 L(z, Mi) 的任务，该任务以用户提示 z 和模型 Mi 作为输入，构建一个 Q 函数，将动作 πi 映射到奖励 ri。具体来说，动作表示将提示 z 路由到每个可能的专家模型 Mi 。在通用实现中，Q(z, πi) 将估计模型 Mi 在提示 z 上获得的奖励 Q(z, Mi) = L(z, Mi, )，Q 函数通过语言模型实现，语言模型本身经过训练可以预测下游的损失值。Q 函数将估计的损失值发送到路由函数，然后路由函数将损失估计与加权约束函数实现的目标/约束相结合，做出路由决策。

为了描述路由决策，首先考虑一个理论上的（但不切实际的）Oracle Router 模型，其中路由器模型可以访问所有提示中每个专家模型的整个 Q 性能指标表（损失值）。然后，路由器模型可以选择最大化性能的模型，同时以权重 为 λj 的约束函数 Cj(Mi) 形式结合次要用户目标。在完全了解 Q 的情况下，路由器优化在损失和约束 Ci 之间进行权衡的函数来选择模型。

路由目标函数允许路由模型在准确性目标和约束之间进行权衡。对于 λj 的任何设置，路由器模型都能够组合 Q 值和约束值，并选择最小化组合损失的模型。在 Tryage 中，用户约束被纳入输入提示本身，例如“加利福尼亚州的首府是 [空白] [Flag：最小模型]”。这两个成本结合起来给出路由损失 LR，Oracle Router 模型 RO 将其最小化。

给定 一个 Q 的 oracle，路由目标函数可以在 Q 和 Cj 上最小化。可以对不同的 λ 值执行最小化，例如将模型大小的权重从不重要的 λ ≈ 0 调节到主导的 λ > 0。路由函数可以通过合并约束函数 Cj(Mi) 来合并多个目标函数，该约束函数根据近因、冗长、幻觉概率等因素对每个 Mi 进行评分。

虽然 oracle 策略可能会根据提供的数据选择下游模型，但在所有可能的提示实例上运行所有 |M| 模型来为每个传入提示创建 Q 表，在计算上效率低下。因此，目标是训练一个估计下游损失的 Q 函数，以便可以对各个输入提示的动态分析返回估计损失。局部分析提示的一种方法，是在提示和专家模型损失的数据集上以监督学习范式训练路由模型来预测 Q 表。

使用监督训练范式，其中构建一个可微分函数 R(z, Mi; W )，权重（模型参数）为 W，其中 R 将提示 z 作为输入并返回 R 中估计模型 Mi 损失值的标量值。对于 R，可以使用任何在 W 中可微分的函数。在实践中，用现有语言模型作为基础架构来实现 R，并修改架构，以便 R 在实数上输出 n-维输出向量。例如，用 Bert-tiny 实现了良好的损失预测准确率。

包括路由器和下游模型在内的整个架构都可以进行端到端训练。具体来说，路由器通过最小化函数来选择模型。路由模型和专家模型一起训练，这样每个专家模型都可以提高路由器发送给特定提示集的性能。该范例与自组织映射的训练具有概念上的相似性。在每个训练步骤中，都会更新路由模型 R 以及专家模型 M ，以便每个专家模型都能发展自己的专业知识。对整个系统进行端到端训练，并将更新与路由器和专家模型分离，以使训练更加模块化。

除了在跨领域实现专家级模型性能外，Tryage 还在无监督训练期间构建语言的隐式潜表征，如图所示。具体来说，可以在 Tryage 生成的 UMAP（Uniform Manifold Approximation and Projection）中看到来自 Github、CommonCrawl、Arxiv、StackExchange、PubMed 和 Wikipedia 的数据聚类（图 b）。相比之下，Tryage 提供的潜空间分离比 GPT-2 更好，后者的嵌入如图 a 中的 UMAP 所示，因为 GPT-2 生成的嵌入是分散的，而不是聚类的。具有默认参数的 UMAP 空间提供了可解释性，因为人们可以推断模型已经学会了数据子集之间的分离。