AutoMix：自动混合语言模型

24年6月来自CMU、谷歌、IIT和USC的论文“AutoMix: Automatically Mixing Language Models”。

现在，云 API 提供商提供各种大小和配置的大语言模型 (LLM)。虽然这种多样性提供了广泛的选择，但有效利用这些选项来优化计算成本和性能仍然具有挑战性。这项工作提出 AutoMix，基于从一个较小 LM 输出的近似正确性，将查询策略性地路由到较大 LM。AutoMix 的核心是两个。首先，它有一个少样本自我验证机制，可以估计自身输出的可靠性，而无需大量训练。其次，鉴于自我验证可能会有噪声，它采用基于 POMDP 的路由器，可以根据答案置信度有效地选择适当大小的模型。在五种语言模型和五个具有挑战性的数据集上进行实验，AutoMix 始终超越强基线，在同等性能​​的情况下将计算成本降低 50% 以上。

如图所示是AutoMix介绍，它包含 3 个步骤：解决方案生成（小模型生成初始答案）、自我验证（相同的小模型评估难度）和选择性路由（根据自我验证的建议路由到更大的模型）。从高层次上讲，这个过程反映了人类解决问题的过程，它本质上遵循一个多步骤的过程：生成解决方案、验证其有效性，并根据验证结果进一步完善它。

有几项研究试图通过模型切换来优化 LLM 推理成本，并使用经过专门训练的验证器 [Chen et al., 2023, Zhu et al., 2023, vSakota et al., 2023, Ding et al., 2024]。AutoMix 通过少样本小语言模型（SLM）提示消除了昂贵的验证器训练的需要，并且不需要预先访问所有输入查询。使用少至 50 个样本进行训练的路由器性能优于专门的模型。一些努力是组合不同的模型和外部工具来改进语言模型的推理时间 [Khattab et al., 2023, Press et al., 2022, Yao et al., 2022, Zhou et al., 2022]。

如图所示，验证过程被设计为自然语言蕴涵任务，其中模型根据上下文和问题，确定模型生成的答案的有效性。对所有任务都使用通用的少样本提示。

如图所示，使用 LLAMA2-13B 进行上下文为基础的自我验证。该示例展示验证器使用与答案生成器相同的模型，通过有效利用上下文来识别和拒绝不准确的答案（“He took it in 1990”）。

为了评估 As 的可信度，AutoMix 采用了一个少样本验证器 V，它可确定 LMi 输出的有效性。与通过创建新问题进行验证的现有研究不同 [Weng et al., 2022, Jiang et al., 2023b]，验证视为一项蕴涵任务 [Dagan et al., 2005, Poliak, 2020, Dagan et al., 2022]，旨在确定 LMi 生成的答案是否与提供的上下文一致。具体而言，验证器测量 v = p(correct = 1 | As, C, q)，其中 correct = 1 表示 As 是正确的。为了估计概率，在高采样温度下从验证器 (LMi) 中采样 k > 1 次，然后计算概率。对所有任务使用相同的 4 次验证提示，并且不训练验证器。

路由在解决方案生成和自我验证之后。路由器决定是否接受 LMi 的输出或将查询路由到某个 LMj (j > i) 以提高性能。路由器也可以解释为元验证器，为少样本验证器的评估提供额外的置信度评估层。具体来说，V 确定 LMi 的答案是否由上下文决定，在不考虑问题固有难度的情况下做出决定。例如，在处理无法解决的查询时，调用更大的 LM 将浪费资源，并且不会提高性能。一个好的路由器可以解决这个问题，只要不进一步路由这样的查询，并且需要使用训练数据的验证概率和趋势来做出这个决定。

为应对大语言模型中自我校正的挑战 [Madaan et al., 2023, Huang et al., 2023]，AutoMix 采用非 LLM 设置进行路由，并避免幻觉和推理错误等问题升级 [Dziri et al., 2023]。原则上，路由器可以采用各种学习策略，包括监督学习、强化学习和符号推理。

在这种针对双-模型情况 (N = 2) 的简单路由方法中，路由到 LM2 的决定基于 LM1 验证器的概率 v 和阈值 t。如果 v ≥ t，则返回 LM1 的答案，否则将查询路由到 LM2。直观地说，高概率表示验证器对其决定有信心并且可以信任。改变 t 可以帮助探索成本性能权衡。

只有当性能差证明成本与质量权衡是合理的，路由器才应将查询定向到更大的 LM。鉴于系统性能的真实状态固有的不确定性（无法观测），将路由器表述为部分可观察的马尔可夫决策过程 (POMDP) [Åström, 1965]。POMDP 特别适合于观察结果（例如自我验证概率）可能不完全可靠的场景。

POMDP 的特征是 (S, A, T , R, Ω, O)。在应用中，状态 S 表示当前选定的 LMi 和各种 LM 在数据点上的性能指标（例如准确度或 F-分数），表示为 S = ⟨i, PerfLM1, PerfLM2, …, PerfLMN ⟩。这些动作包括要么保留当前 LM（LMi）的答案，要么路由到其中一个更大的 LM。观测值 Ω 以来自 LMi 的验证器输出 v 形式出现，使 POMDP 能够确定其信念状态 b：一个在 S 上的概率分布。观测概率 P (o|s) 表示在给定状态 s 情况下观测 o（验证输出）的似然，对于定义 POMDP 模型至关重要。例如，高的验证器置信度可能表明当前 LM 的性能 PerfLMi 足够高，从而减少了切换到更昂贵 LM 的必要性。观测概率直接在训练集上估计，计算每个状态的验证概率期望，即

但是，由于状态可以是连续的，用非参数高斯核密度估计来估计一个新状态的观测概率。不是直接估计 P (o|s)，而是首先通过在每个训练点计算 KDE 来学习联合分布 P (S, O) 和 P (S)。然后计算条件概率 P (o|s) 如下：P(o|s) = P (s,o) / P(s)。

POMDP 的目标是优化奖励函数 R = P − λ · C，其中 P 是整体性能，C 是整体成本，λ 是根据用户偏好平衡两个标准的调整参数。用 AdaOps POMDP 求解器 [Wu et al., 2021]，它使用粒子滤波器来表示信念，其中每个粒子对应一个特定状态。在推理时，POMDP 求解器从初始信念（均匀分布）开始，根据验证者的观测更新其信念状态，并根据更新后的信念状态 b′，采取行动最大化预期奖励。

如图左边是AutoMix算法总结，图右边是一个性能-成本图，考虑连接小语言模型 (SLM) 和大语言模型 (LLM) 数据点的线段。此线段的斜率表示每增加一个成本单位，性能的基本提升率。任何方法 M 的增量成本收益 (IBC) 都是从 SLM 点到代表 M 的点之间连线的斜率。位于 SLM-LLM 线段上方的方法 M 提供更陡的斜率，表示有利的 IBC（和正的 ∆IBC）。相反，如果 M 位于线段下方，则表明 IBC 不利或为负。主要目标是确定或开发能够产生持续正 IBC 的方法，从而最大限度地提高每增加一个成本单位的性能。