突破大语言模型推理成本困境！TALE框架大幅削减Token冗余

📖阅读时长：19分钟

🕙发布时间：2025-02-10

近日热文：全网最全的神经网络数学原理（代码和公式）直观解释
欢迎关注知乎和公众号的专栏内容
LLM架构专栏
 知乎LLM专栏
 知乎【柏企】
公众号【柏企科技说】【柏企阅文】

在大语言模型（LLM）的推理过程中，思维链（Chain-of-Thought, CoT）推理这类方法虽然能通过将问题分解为中间步骤来提升模型性能，但却会在token使用上带来巨大开销，进而导致成本大幅增加。为了解决这一问题，一篇论文提出了一种具有token预算感知能力的LLM推理框架。它能依据推理复杂度，动态地为不同问题预估token预算，并利用这些预估值来引导推理过程。

LLM推理中的Token冗余现象

Token预算

论文表明，LLM有能力在提示（prompt）中遵循长度约束。如下表展示了原版CoT和带有token预算的CoT之间的差异。例如，在提示中设定50个token的预算，LLM会调整输出长度（如生成86个输出token），努力契合指定预算。这充分证明了LLM在遵循明确token预算提示方面具备一定能力。

Token冗余现象

从下图呈现的Token弹性现象可以看出，x轴代表预算搜索迭代，y轴在不同子图中分别表示搜索的预算（图a和图c）以及每个搜索预算的实际token成本（图b和图d）。通过图表和表格结果能够确认，在当前最先进的LLM推理过程中，存在着大量的冗余token。

搜索最优Token Budget

用于最佳预算搜索的Vanilla方法

搜索最优token预算的原版方法，首先会运用原版CoT为每个问题生成答案，接着计算答案中的token数量，并将其设定为搜索的右边界。函数isFeasible用于判断预算的可行性，即带有该预算的CoT提示能否确保答案的正确性。在整个算法流程中，以0为左边界，当前可能的预算β被计算为0和右边界的中点，并不断循环调整β的值，直至找到最终合适的预算。

graph TD;
    A[开始] --> B[应用原版CoT生成答案];
    B --> C[计算答案中token数作为右边界right];
    C --> D[左边界设为0，计算当前预算β=(0+right)/2];
    D --> E[记录之前搜索的预算β0];
    E --> F{β是否可行};
    F -- 是 --> G[更新β=(β0+right)/2];
    G --> H[调整搜索边界];
    H --> F;
    F -- 否 --> I[返回最终预算β];
    I --> J[结束];

观察代币弹性

在执行算法1的过程中，研究人员发现当预算减少到一定范围之外时，token成本反而会上升。这意味着过度削减预算，不仅无法降低成本，还会造成token消耗的增加。如下面图表所示，x轴是预算二叉搜索的迭代次数，随着迭代，预算值逐渐降低，而图b和图d的y轴展示了每次迭代时对应的token成本变化。

基于代币弹性的最优预算搜索

根据token弹性的观察，最低预算虽然能保证答案正确，但不一定能使token成本最小化。所以在搜索最优预算时，除了要确保答案正确，更新后的预算还需比之前搜索的预算产生更低的token成本。算法2详细列出了搜索过程中使用的可行性函数，其会同时计算当前和之前评估预算的实际token成本，并依据答案正确性和token成本降低这两个标准来评估可行性，只要有一个条件不满足，搜索就会停止。

graph TD;
    A[开始] --> B[计算当前和之前预算的实际token成本];
    B --> C{答案是否正确且token成本降低};
    C -- 是 --> D[继续搜索，更新预算β];
    D --> E[重复计算和评估];
    E --> C;
    C -- 否 --> F[停止搜索];
    F --> G[结束];

方法论

概述

基于上述观察与分析，研究团队设计了用于LLM中token预算感知推理的TALE方法。TALE的工作流程是：先利用预算估算器预估token预算，再将问题与预估预算相结合生成具有token预算感知的提示，最后将该提示输入LLM以得到最终答案。

预算估算

零镜头估计器：TALE利用推理LLM本身作为零镜头机制的预算估计器。在向LLM提出问题前，先让LLM预估回答该问题所需的输出token数量，这一设计思路源自人类的思考模式，就像人们看到数学题，虽计算答案需要时间，但快速浏览题目后就能大致估计所需时间和精力。
回归估计器：基于回归的估计器旨在训练或微调另一个LLM$f(\theta)$作为预算估计器。研究人员先设定指令提示$p$为“Estimate the token budget for the following question”，用预训练的LLM（如LLaMA 3 - 8B）初始化$f(\theta)$，并根据搜索到的最优预算$\beta^{*}_{i}$为每个问题$x_{i}$制作目标输出$y_{i}$。通过最小化损失函数，让模型输出尽可能接近目标输出$y_{i}$。

损失函数定义为：$L(\theta) = -\sum_{i = 1}^{N} \log P(y_{i} | p_{i}, x_{i}; \theta)$
其中，$\theta$代表模型参数，$\theta^{*}$是优化后的参数。该损失函数的目的是最大化在所有$N$个训练样本中，给定输入$x_{i}$和指令提示$p$时，目标输出$y_{i}$的概率$P(y_{i} | p_{i}, x_{i}; \theta)$。

代币预算意识内化：微调LLM$M(\theta)$，使其能生成具有token预算意识的答案，这一过程分为目标输出生成和LLM微调两个关键阶段。在目标输出生成阶段，用结合了最优token预算的思维链（CoT）提示$M(\theta)$，生成受预算约束的目标输出$y_{i}$。在LLM微调阶段，使用第一阶段生成的目标输出来训练$M(\theta)$，让LLM在微调过程中内化token预算约束，倾向于更紧凑的推理过程。

评估

实验设置

实验使用的评估指标有：

准确性（Accuracy, Acc）：计算方式为$Acc = \frac{1}{|X|}\sum_{(x_{i}, y_{i}) \in X} I\{M(x_{i}) = y_{i}\}$。其中，$(x_{i}, y_{i}) \in X$，$x_{i}$是数据集中的数学问题，$y_{i}$是正确答案，$M(·)$用于返回给定问题的答案，$I\{·\}$是指示函数，条件成立时返回1，不成立返回0。
输出Token的数量：输出token成本的衡量公式是$Cost = \sum_{x_{i}} T(x_{i})$。$x_{i}$表示给定问题，$T$是用于测量token数量的函数，输出token越多，成本越高。

TALE的有效性

对比TALE和“直接回答”“Vanilla CoT”等其他提示工程方法在7个不同数据集上的表现，结果显示：TALE在保持较高准确性的同时，显著提升了效率。直接回答的输出token和费用最低，但准确率也最低；Vanilla CoT的准确率最高，然而token成本和费用也最高。TALE则在性能和成本之间取得了平衡，准确率达到81.03%，同时将token成本降至Vanilla CoT的32%，费用降至41%。在GSM8K数据集上，TALE的准确率更是提升至84.46%，超过了Vanilla CoT；在GSM8K-Zero数据集上，TALE在保持98.72%高准确率的同时，将输出token成本从252.96大幅降低至22.67；在MathBench-Arithmetic数据集上，TALE将费用从78.58降至18.62，同时保持了73.67%的准确率。总体而言，TALE平均降低了68.64%的输出token成本。

预算估算的有效性

理想的预估预算应处于最佳搜索预算附近，并且位于Token弹性图的底部区域。评估时使用范围内准确率和范围外距离这两个指标，范围内准确率用于判断预测预算$\hat{\beta}$是否在理想预算范围内，范围外距离用于量化预测预算$\beta^{*}$超出理想预算范围$W^{*}_{k}$时与$W^{*}_{k}$的距离。评估结果显示，范围内准确率为60.61%，范围外距离为109.64，这表明超过三分之二的预估预算处于理想范围内。

TALE的泛化

在MathBench-College数据集上，对TALE在Yi-lightning、GPT-4o-mini和GPT-4o等不同模型上的泛化能力进行测试。结果表明，与Vanilla CoT相比，TALE在不同模型上都实现了显著的token节省，其中Yi-lightning的输出token减少了62.6%，GPT-4o-mini减少了61.5%，GPT-4o减少了69.8%，同时还能保持较高的准确率。

结论

研究团队推出的TALE框架，通过融入token预算意识，有效减少了思维链推理中的token冗余。它能够根据任务复杂度动态预估token预算，在推理过程中平衡效率和准确性。实验数据显示，TALE平均减少了68.9%的token使用量，准确率损失不到5%。在成本效益方面，TALE优于Vanilla CoT，并且在各种LLM中都展现出良好的通用性。

论文链接：https://arxiv.org/abs/2412.18547
代码链接：https://github.com/GeniusHTX/TALE

本文由mdnice多平台发布

突破大语言模型推理成本困境！TALE框架大幅削减Token冗余