基于Token显著性压缩思维链:降低大模型CoT成本的关键技术

基于Token显著性压缩思维链:降低大模型CoT成本的关键技术 做大模型应用开发时最纠结的问题往往是我们希望模型“按步骤思考”于是让模型输出一长串 Chain-of-Thought思维链简称 CoT。这串思维链经常能显著提升推理准确率却也带来了更高的 token 消耗、更长的首字延迟、更大的输出成本。尤其当业务需要处理大量复杂问题时token 账单和响应时间都会被迅速拉高。最近在浏览推理侧压缩相关的论文时看到一个非常形象的标题Every Token Leaves a Ripple in the Stream of Thought: Eliciting Model-Internal Token Saliency for Chain-of-Thought Compression这篇标题的核心意思是在“思维流”中每一个 token 都会留下“涟漪”。如果我们能发现哪些 token 的涟漪更强、对最终答案贡献更大就可以只保留这些关键 token从而对思维链做压缩。需要说明的是本文不声称复现论文原文的具体网络结构与实验细节而是从标题给出的关键词出发做一次偏工程视角的技术拆解。对论文方法感兴趣的读者最好再回到原文对照实现。对于普通 LLM 应用开发者这篇文章也可以帮助你理解“思维链压缩”和“token 级重要性分析”到底在解决什么问题。1. 先理解痛点思维链又香又贵1.1 思维链为什么有效Chain-of-Thought 最早让人印象深刻的效果是在数学推理、逻辑推理等任务上大幅提升了大模型的表现。让模型先输出中间推导过程再给出最终结论本质上是在模仿人类解决复杂问题时的思考方式。从建模角度看思维链相当于把“一步到位”的高维映射拆成了若干个中间步骤。模型每一步只需要完成一个相对简单的子任务难度降低了因此更不容易出错。这也是如今各类推理模型普遍采用“长思维链”策略的原因之一。复杂题目的推理过程可以很长模型需要反复检查、修正、枚举可能性。1.2 思维链的成本也随之上升思维链的代价非常直观token 消耗成倍增加。延迟明显变大。输出被截断的概率更高。多轮对话中占用的上下文空间变大。如果在生产环境直接开启“无限思考”你很快会遇到两类典型问题一类是成本问题另一类是 API 层的限制例如响应文本超过了 provider 设定的输出 token 上限接口报错“已达到输出 token 上限回答被截断”。这就自然产生一个需求能否在保持高准确率的前提下把冗长的 CoT 压缩成一段“精简但命中要害”的文本这就引出了标题里提到的 “Token Saliency”。1.3 为什么“压缩思维链”比“压缩普通文本”更难普通文本压缩追求的是保留信息密度例如摘要、抽句、关键词提取。但思维链并不只是“信息”它更是模型推理过程中的“计算中间状态”。一个在最后答案中完全不出现的 token也可能对推理方向产生关键推动作用。例如中间不小心写错的一个数模型靠自检修正了又例如一个看似废话的“换个思路”可能让模型从错误的推理分支里跳了出来。所以对思维链做压缩不能只站在“文本语义”层面看而要站在“模型内部信号”层面看。这就是标题里 Model-Internal Token Saliency 想表达的方向。2. 四个关键词决定文章主线2.1 Chain-of-Thought不只是“步骤化提示”如果你接触过提示词工程大概率用过类似写法Lets think step by step.它的作用机制不只是“让模型多说话”。更准确地说这行文字会把模型的后缀生成过程切换成更细粒度的中间推演模式。由于每一步只需在短距离依赖上做预测模型犯大错的概率会下降。但 Chain-of-Thought 在生产环境并不廉价。工程上我们也见过很多变体少样本 CoT给模型几个带推导的示例。零样本 CoT直接让模型思考。CoT 集成生成多条思维链后投票。结构化 CoT要求模型输出 JSON 步骤。每一种方案都在质量和成本之间做取舍。本文讨论的是更进一步对已经生成的 CoT 再加工压缩后再用于模型推理或人类审计。2.2 Token从字符串到模型的最小原子Token 是模型处理文本的最小单位。文本进入模型前会先被 tokenizer 切分成一组 token模型每次预测一个 token生成结果后再拼接回文本。不同分词器对不同语言的切分策略不同。一个英文单词可能是一个 token也可能是两个 token。中文通常按字或按 subword 切分不同词表差异比较大。要精确知道文本的 token 数最可靠的方法是直接调用对应模型的分词器统计。在工程语境中token 还常被用来表达“用量”和“计费单位”。比如某个推理模型每百万输出 token 的价格会直接决定你的成本。你甚至会遇到“credits 怎么换算成 token”之类的问题。不同平台的换算规则不同没有统一公式只能以官方文档为准。2.3 Saliency找出哪些 Token 更关键Saliency 的中文可以翻译为“显著性”或“重要性”。它来自可解释性研究表示输入中的某个部分对模型输出结果的影响程度。在文本分类时代Saliency 通常表现为给每个输入词打分哪些词对判断为正面情感贡献最大哪些词导致模型判断为负面到了大模型时代这个思想可以被迁移到每个推理 token 上。对于一串 CoT每个房间有 4 扇窗户。已知 3 个房间所以窗户总数是 3×412。答案12。我们可以给每个 token 打分。数字“4”“3”“×”“12”通常会有较高的显著性而“每个”“已知”“所以”等词可能显著性偏低。如果能自动获得这种打分就可以把低显著性的 token 去掉。Saliency 这个方向和“注意力权重”并不完全等同。注意力权重描述的是模型在计算时“看了哪里”而 Saliency 更关心“改哪里会造成结果明显变化”。二者经常一起使用但背后的含义有差异。2.4 Model-Internal不只看输出文字标题中的 Model-Internal 强调了一个重要观点判断 token 重要性不能只看语言层面的词频或语法角色更要看模型内部状态。所谓内部状态包括但不限于每一层的隐藏状态hidden states。注意力矩阵。token 对应的梯度。残差流中累积的信息。这就像一条河流。水面上的文字是看得见的但真正推动推理方向的是水面下的暗流。一个 token 是否重要取决于它在模型内部激起了多大的涟漪。例如在数学题推理中某个数字 token 可能在语义上只占很小比例但如果把它换成别的数字最终回答会完全不同。这种 token 就是高显著性的 token。只看表面文本很容易低估它的价值。3. 为什么压缩方案要把 Token 级重要性作为核心3.1 从“粗粒度摘要”到“细粒度保留”最常见的思维链压缩方式是直接让另一个模型把长推导过程总结成短内容。这种方案实现简单但问题也很明显摘要会丢失推导细节。摘要过程本身要消耗额外 token。摘要模型并不清楚最终答案依赖哪些中间变量。有可能把关键数字改错。Token Saliency 方案则不同。它先保留完整的 CoT再通过模型内部信号给每个 token 计算贡献分。很多原始 token 可以直接保留不需要重写。这种思路在信息保真度上更友好因为“保留原 token”永远比“让另一个模型改写”更不容易发生语义偏移。3.2 什么样的 CoT 最值得压缩不是所有场景都需要压缩 CoT。比较适合的方向通常包括复杂数学题步骤多、token 长。代码生成前的自我规划模型可能先输出一堆计划。Agent 推理轨迹模型在调用工具前会输出大量观察和思考。多步检索问答中间步骤包含检索结论但最终回答只需要关键依据。在这些场景里关键步骤可能只占整个推理轨迹的一半甚至四分之一。把剩余部分丢弃理论上可以把 token 成本压到原来的 50% 以下同时保留大部分推理能力。3.3 一般化流程先测重要性再选择保留基于 Token Saliency 的压缩通常可以抽象成四步让模型生成一段完整的 CoT。让这段 CoT 的每个 token 都“连接”到最终答案通过内部信号计算显著性分数。根据分数选择保留哪些 token或者把低显著性 token 改写为占位符。把压缩后的完整序列重新输入模型让模型输出最终答案。这个流程可以在推理阶段做一次性优化也可以在推理时实时计算。整体思路与“先长篇思考再精简回答”非常相似只是它不需要让模型额外学习一套摘要提示而是用内部信号自动判断。4. 模型内部有哪些信号可以用来计算 Token 显著性4.1 注意力分布模型更“关注”谁Transformer 每一层都有多头注意力。计算下一个 token 时模型需要从前面所有 token 中聚合信息。如果某个 token 被很多后续 token 高权重关注它往往承担了更重要的信息中转职责。因此注意力矩阵是最容易获取的内部信号之一。通过多层注意力的组合计算每个 token 收到的注意力概率可以给显著性一个粗估计。但它也有缺点注意力高不等于因果贡献大有时模型关注某些 token 只是为了“忽略它”。4.2 Logits 与预测概率变化模型每一步都会输出一个 logits 向量表示下一个 token 的候选分布。要判断某个历史 token 是否重要可以做一个反事实实验如果把某个 token 替换成占位符或同类型随机 token再进行一次前向计算观察最终预测概率的变化幅度。变化越大说明该 token 越重要。这种反事实法在逻辑上最直观但计算成本较高。因为每个 token 都要单独跑一次前向序列长的时候开销会很大。因此研究中往往会使用近似手段例如只用部分层做重计算。4.3 梯度最经典的 Saliency 来源之一梯度是另一种计算显著性的方式。我们把模型最终的损失或目标 token 的概率作为目标对输入 token 的 embedding 求梯度。梯度的绝对值越大说明该 token 的 embedding 对结果的影响方向越敏感。数学上可以写成类似这样的理解Saliency(token) ≈ 目标函数对 token embedding 的梯度范数具体公式可能是 L1、L2 或者积分梯度Integrated Gradients。这类方法在图像分类的可解释性里被广泛使用迁移到文本 token 上也成立。它的优势是能给出每个 token 对特定答案的因果性影响劣势是为了计算梯度需要模型内部访问权。这也是为什么标题强调 Model-Internal闭源 API 往往不提供梯度接口你很难用这种方法做细致分析。4.4 隐藏状态与残差流观察更抽象的“涟漪”相比于只看最后一层输出我们还可以研究模型中间层的隐藏状态。在大模型的残差流里信息会沿着层不断被读写。每个 token 在每一层都会携带并更新一个状态向量。如果某一个 token 在某一层的状态向量与最终答案向量高度相似说明它携带了被最终决策使用的关键信息。常见的做法是提取某个目标位置的输出向量再计算它和每一个历史 token 向量的余弦相似度。相似度高的 token可以认为是“思维流”里被最终结论读取最深的 token。如果要给它一个直观理解那就是标题中的“Ripple”。散落在思维流中的 token 并不是等价的它们有的轻轻漂过有的激起了能被后续层一直读取的涟漪。5. 从显著性到压缩一条完整的处理管线5.1 管线总览在实际实现中Code 到 Token Saliency 再到压缩最好拆成独立模块便于调试。模块输入输出CoT 生成器原始问题完整思维链Tokenizer思维链文本token id 序列Saliency 计算器token id 模型内部状态每个 token 的显著性分数Token 选择器显著性分数保留 token 的位置集合结果生成器原始问题 压缩后的 CoT最终答案5.2 两个容易混淆的设计选择第一种压缩后直接输出给用户看。如果思维链本身要用于展示那么压缩结果仍然需要保持语言通顺。此时不能简单删除 token因为删除会造成语法断裂。比较稳妥的方式是保留高显著性 token并在关键位置让模型重新生成过渡语句。第二种压缩后继续作为内部上下文。如果压缩后的 CoT 只是为了让模型继续推理不一定非要是完整句子。可以把低显著性 token 变成特殊标记例如drop或者干脆只保留 token id 序列。模型只要能获得下一步预测所需的关键信息即可。这些选择会影响压缩率与最终效果之间的平衡也需要在具体实验中验证。5.3 压缩率不是越高越好压缩率过高会带来两个问题信息丢失模型推理能力下降。文本连贯性被破坏即使压缩文本模型仍然可能“读不懂”。压缩率过低则节省 token 的效果有限。一个好的算法应当在“节省多少 token”和“准确率下降多少”之间寻找平衡。通常论文中会用“准确率-压缩率曲线”来展示这种权衡。6. 用开源模型理解 Token Saliency可直接跑的教学示例为了更直观理解上述概念下面用一个开源 CausalLM 模型来演示基础逻辑。需要注意本文示例只是为了解释通用思路并不是论文代码。论文的正式实现结果要以作者公开的代码和模型为准。6.1 基础环境建议环境如下Python 3.10 或更高版本。PyTorch 2.x。Hugging Face Transformers 4.x。可以运行 CPU 推理的小模型例如 0.5B 级别的模型。如果你的机器无法访问模型下载地址可以把示例代码中的模型替换成本地已有的模型路径或者只从代码结构上理解原理。6.2 示例 1统计文本与 Token 的关系先看最基础的 token 化过程。from transformers import AutoTokenizer model_name Qwen/Qwen2.5-0.5B-Instruct tokenizer AutoTokenizer.from_pretrained(model_name) text Lets think step by step: 3 rooms, each has 4 windows. tokens tokenizer.tokenize(text) print(原始文本:, text) print(Token 列表:, tokens) print(Token 数量:, len(tokens))这个例子用来建立直观认知同一段文字在不同分词器下会得到完全不同的 token 切分结果。因此讨论 CoT 长度和成本时最好直接使用模型对应的 tokenizer 统计而不是肉眼数单词。6.3 示例 2用梯度计算每个 Token 的显著性接下来我们演示最典型的 Saliency 计算方式之一输入 embedding 梯度。核心思想是将问题与 CoT 文本编码为 token。把 token id 映射成 embedding。通过 embedding 前向传播计算损失。反向传播后查看每个 token 位置的梯度范数。梯度大的位置通常更敏感。import torch from transformers import AutoTokenizer, AutoModelForCausalLM model_name Qwen/Qwen2.5-0.5B-Instruct tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name) model.eval() text A farmer has 12 apples. He gives 4 to his friend. Then he buys 3 more. How many apples does he have? Lets think step by step. enc tokenizer(text, return_tensorspt) input_ids enc[input_ids] attention_mask enc[attention_mask] # 拿到输入 embedding并让梯度保留在这个中间张量上 input_embeds model.get_input_embeddings()(input_ids) input_embeds.retain_grad() # 用 input_embeds 做一次带梯度的前向 outputs model( inputs_embedsinput_embeds, attention_maskattention_mask, labelsinput_ids, ) # 因果语言模型已经完成了内部 shiftloss 表示模型预测下一个 token 的效果 loss outputs.loss loss.backward() # 梯度形状: (batch, seq_len, hidden_size) # 对 hidden 维求和得到每个 token 的重要性分数 saliency input_embeds.grad.abs().sum(dim-1)[0] print(Sequence length:, input_ids.size(1)) print(Saliency shape:, saliency.shape) # 展示前 20 个 token 对应的文本和分数 tokens tokenizer.convert_ids_to_tokens(input_ids[0][:20]) scores saliency[:20].tolist() for tok, score in zip(tokens, scores): print(f{tok:12} {score:.4f})这段代码存在一个需要理解的点模型前向时计算的是“根据前文预测下一个 token”的语言模型损失。输入 embedding 的梯度能间接反映“当前 token 对后续预测有多重要”。严格地说这种方法估算的是“预测当前语料本身时的敏感度”并不是工业级最终答案显著性。若想要针对“最终答案”的显著性可以只把最后一小段答案文本对应的 token 作为 label而不是把所有 token 都当作 label。你还可以进一步把损失替换成“目标答案的负对数似然”。整体思路是相通的。6.4 示例 3基于隐藏状态做余弦相似度显著性梯度计算需要反向传播如果只想快速观察内部状态也可以用 hidden states 做近似。思路如下获取最后一层所有 token 的隐藏状态。以最后一个 token 的隐藏状态作为“汇总信息向量”。计算每个历史 token 的隐藏状态跟这个向量的余弦相似度。相似度高说明该 token 的信息和最终聚合状态更接近。import torch from transformers import AutoTokenizer, AutoModelForCausalLM model_name Qwen/Qwen2.5-0.5B-Instruct tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name) text Each room has 4 windows. If there are 3 rooms, windows in total? Lets think step by step. enc tokenizer(text, return_tensorspt) with torch.no_grad(): outputs model( input_idsenc[input_ids], attention_maskenc[attention_mask], output_hidden_statesTrue, ) # 取最后一层隐藏状态: (batch, seq_len, hidden) hidden_states outputs.hidden_states[-1][0] # 用最后一个 token 的向量作为全局汇总向量 query hidden_states[-1].unsqueeze(0) # (1, hidden) # 与所有 token 向量计算余弦相似度 cos torch.cosine_similarity(hidden_states, query, dim-1) tokens tokenizer.convert_ids_to_tokens(enc[input_ids][0]) for tok, score in zip(tokens, cos.tolist()): print(f{tok:12} {score:.4f})这个方法的缺点是最后一个 token 并不一定等于“最终答案向量”。尤其在生成还没完成时最后位置可能只是一个普通 token。教学场景里用它来观察趋势很好但做正式实验时需要把目标向量设计得更准确。6.5 示例 4按显著性做简易删除拿到每个 token 的显著性分数后我们可以按分数做 top-k 保留再把其余 token 丢掉。为了保持文本可读下面用一个简单的“删除后拼接”示例演示。def compress_text_by_scores( tokens, saliency_scores, keep_ratio0.4, ): n len(tokens) keep_k max(1, int(n * keep_ratio)) # 按显著性降序得到索引并取前 keep_k 个 indices sorted(range(n), keylambda i: saliency_scores[i], reverseTrue) keep indices[:keep_k] # 保留位置集合 keep_set set(keep) # 按原始顺序输出未删除的 token kept_tokens [] for i in range(n): if i in keep_set: kept_tokens.append(tokens[i]) return kept_tokens # 假设 tokens 与 saliency_scores 来自前两个示例 # compressed_tokens compress_text_by_scores(tokens, cos.tolist(), keep_ratio0.3) # print(tokenizer.convert_tokens_to_string(compressed_tokens))真实压缩方法不会只做这种硬截断通常还会加入位置约束、连续片段约束、以及重生成润色。上面的代码是为了让读者直观感受“显著性排序 剪枝”的最小实现。7. 高频问题与排查思路在工程或论文复现中你很可能遇到下面几类问题。问题现象常见原因解决思路输出超过 token 上限回答被截断模型生成了过长 CoT 或过长答案先对 CoT 压缩再生成答案同时降低 max_tokens反向传播时报显存不足模型过大或序列过长使用更小模型、缩短序列、采用梯度 checkpoint压缩后准确率反而大幅下降只按表面文本选择 token忽略了内部信号改用梯度、隐藏状态等模型内部信号删除 token 后文本破碎缺少连续性约束按片段保留或加入模型重写模块不同模型之间显著性规律不一致分词器与训练目标不同不能简单跨模型迁移应重新计算压缩耗时比生成 CoT 还长逐 token 反事实计算导致开销过大改用近似梯度或只对关键片段计算“输出 token 上限”相关错误是生产中很常见的。处理思路通常是三层检查 max_tokens 参数是否设置得过小。检查 prompt 是否要求模型输出过长内容。考虑对 CoT 做压缩或让模型分多轮输出。当你看到一个错误信息类似API error: Invalid request: your request exceeded model token limit不必直接怀疑是 Saliency 方法本身的问题先确认输入 prompt 和输出 max_tokens 是否符合模型限制。在压缩场景中如果压缩后的 CoT 仍然超限很可能是保留比例设置得过高。8. 工程化建议如何把 CoT 压缩落地到业务中8.1 区分“最终答案展示”和“内部推理缓存”如果压缩后的 CoT 只是作为模型内部上下文那么它在格式上可以更自由。即使文本不连贯也不影响模型读取关键 token。但如果用户会在 UI 上看到思维链你就需要额外的改写模块。这时候要权衡收益改写模块本身也要消耗 token只有改写长度远小于原始 CoT整体才划算。8.2 用“最小可用闭环”验证方案在正式接入业务前不要马上做复杂网络设计。可以先用一个粗糙版本验证链路固定 100 道题。记录全量 CoT 下的准确率和 token 成本。用最简单的梯度显著性保留 50% token。记录压缩后的准确率和 token 成本。对比两者差异。如果最简单的梯度法都无法保住准确率说明你的任务可能更依赖那些“表面不关键但实际重要的 token”需要更复杂的内部信号组合。8.3 需要保留高价值 Token 类型从工程经验看以下内容通常要尽量保留所有数字和符号。等式、比例、单位。步骤之间出现的结论性短句。模型在自检时发现的错误。与最终答案强相关的专有名词。而以下内容可以优先丢弃重复的连接词。 -自我重复的过渡句。无意义的语气词。被模型中途否决的错误假设中的大段解释。不要完全按语言习惯来判断因为有些看似“废话”的内容可能影响模型推理状态。最终判断标准应当来自消融实验。8.4 注意安全与权限边界使用透明模型内部信号的方法时需要清楚知道梯度、隐藏状态在闭源 API 中通常拿不到。如果你只能调用 API可能需要把显著性问题转化为“多次采样投票”或“让模型自己标注关键行”。在分析思维链的过程中还要注意用户输入可能包含隐私或敏感信息。不要在日志中完整记录思维链也不要轻易将这些数据投喂给第三方摘要服务。8.5 成本治理需要全局监控token 压缩做得再好也不能忽略整体费用控制。建议建立指标监控单次请求平均输入 token。单次请求平均输出 token。CoT 占输出 token 的比例。压缩后 token 节省率。压缩后准确率变化。当准确率下降不超过某个阈值时节省率越高越好一旦超过阈值说明压缩过度。每个业务都需要自己定义这个平衡点。9. 顺着这个问题还能继续学什么如果对标题里的技术方向感兴趣下一步可以按下面顺序探索第一可解释性工具。例如学习 attention rollout、Integrated Gradients、局部可解释模型等概念。它们能帮你理解“为什么某个 token 重要”。第二LLM 推理加速。除了压缩 token还有投机采样、KV Cache 复用、结构化输出等方法。它们解决的是不同层面的性能问题可以与 CoT 压缩叠加使用。第三数据蒸馏与模型训练。当压缩后的 CoT 效果稳定后可以把它整理成新的训练集让模型在微调阶段就学会“少说废话、直击要害”。这也是一种更长期的成本治理方案。第四评估方法。不要只拿一两个样例判断效果。最好准备一个覆盖不同难度、不同长度的测试集分别统计准确率、token 数和耗时。从这篇论文标题还可以引申出一个问题我们在让模型“一步一步思考”时那些被丢弃的 token 真的没有价值吗显然不是。它们的价值可以很大但也可以被压缩。真正困难的地方在于如何在“保留足够计算信息”和“减少冗余 token”之间找到最优路径。回到文章的标题Every Token Leaves a Ripple in the Stream of Thought。每一次生成每个 token 都会在模型内部留下影响但影响有大有小。Token Saliency 要做的就是顺着这些细小的涟漪找到真正改变方向的那些浪花。理解这一点之后你在设计提示词、控制系统输出长度、优化 token 成本时都会多一层更准确的判断。如果你对这段技术拆解有不同理解欢迎按自己的实验数据来验证。毕竟这类方法还处在快速演进阶段不同任务、不同模型的结论可能差异很大。