Do LLMs Encode Functional Importance of Reasoning Tokens?
作者: Janvijay Singh, Dilek Hakkani-Tür
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-01-06
备注: 20 pages, 8 figures, 2 tables
💡 一句话要点
提出贪婪剪枝方法,探究LLM推理token的功能重要性并提升蒸馏效果
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 推理链压缩 贪婪剪枝 蒸馏训练 功能重要性 注意力机制
📋 核心要点
- 现有方法在压缩LLM推理链时,缺乏对token级别功能重要性的深入理解,限制了压缩效率。
- 提出贪婪剪枝方法,通过迭代删除对模型似然性影响最小的token,实现推理链的压缩。
- 实验表明,使用剪枝后的推理链进行蒸馏训练,学生模型在相同推理长度下优于基线方法。
📝 摘要(中文)
大型语言模型通过生成长推理链来解决复杂任务,虽然提高了准确性,但也增加了计算成本,并降低了隔离功能相关推理的能力。现有工作通过概率抽样、启发式方法或来自前沿模型的监督来缩短这些链,但对于模型是否在内部编码了token级别的功能重要性以用于答案生成,提供的见解有限。本文通过诊断方法填补了这一空白,并提出了一种贪婪剪枝方法,该方法是一种保留似然性的删除过程,它迭代地删除那些在指定目标下删除后对模型似然性影响最小的推理token,从而产生长度可控的推理链。在蒸馏框架中评估了剪枝后的推理,结果表明,在匹配的推理长度下,用剪枝链训练的学生模型优于前沿模型监督的压缩基线。最后,分析揭示了系统的剪枝模式,并表明注意力分数可以预测贪婪剪枝的排名,进一步表明模型在推理token上编码了重要的功能重要性结构。
🔬 方法详解
问题定义:论文旨在解决大型语言模型(LLM)在推理过程中生成的冗长推理链所带来的问题。这些长链虽然提高了准确性,但也增加了计算成本,并且难以确定哪些token对于最终答案的生成至关重要。现有压缩推理链的方法,如概率抽样、启发式方法或前沿模型监督,未能充分理解模型内部对token功能重要性的编码。
核心思路:论文的核心思路是通过一种名为“贪婪剪枝”的方法,系统性地识别并删除推理链中功能重要性较低的token。该方法旨在保留模型在删除token后的似然性,确保在压缩推理链的同时,尽可能地保留模型的推理能力。通过分析剪枝模式和token的注意力分数,进一步探究LLM是否以及如何编码token的功能重要性。
技术框架:整体框架包括以下几个主要步骤: 1. 推理链生成:使用LLM生成原始的、未剪枝的推理链。 2. 贪婪剪枝:迭代地删除推理链中的token,每次删除对模型似然性影响最小的token。似然性损失基于指定的objective计算。 3. 蒸馏训练:使用剪枝后的推理链作为训练数据,对学生模型进行蒸馏训练。 4. 评估:评估学生模型在下游任务上的性能,并与基线方法进行比较。 5. 分析:分析剪枝模式和token的注意力分数,以了解LLM如何编码token的功能重要性。
关键创新:论文的关键创新在于提出了贪婪剪枝方法,这是一种保留似然性的token删除过程,能够有效地压缩LLM的推理链。与现有方法相比,贪婪剪枝方法更加注重token的功能重要性,并且能够生成长度可控的推理链。此外,论文还通过实验验证了剪枝后的推理链在蒸馏训练中的有效性,并深入分析了LLM对token功能重要性的编码方式。
关键设计: * 似然性损失函数:用于衡量删除token后对模型似然性的影响。具体形式未知,但目标是最小化删除token对模型预测结果的影响。 * 迭代删除过程:每次迭代删除一个token,并重新计算剩余token的似然性损失。迭代过程持续到达到预定的推理链长度或满足其他停止条件。 * 注意力分数分析:使用注意力分数作为token功能重要性的指标,并分析注意力分数与贪婪剪枝排名之间的关系。
🖼️ 关键图片
📊 实验亮点
实验结果表明,使用贪婪剪枝方法压缩后的推理链进行蒸馏训练,学生模型在匹配的推理长度下,性能优于前沿模型监督的压缩基线。这表明贪婪剪枝方法能够有效地保留推理链中的关键信息,并提高蒸馏训练的效果。具体性能提升数据未知。
🎯 应用场景
该研究成果可应用于降低大型语言模型推理的计算成本,提高模型在资源受限环境下的部署能力。通过压缩推理链,可以加速模型的推理速度,并减少内存占用。此外,该研究还有助于理解LLM内部的推理机制,为模型优化和改进提供新的思路。
📄 摘要(原文)
Large language models solve complex tasks by generating long reasoning chains, achieving higher accuracy at the cost of increased computational cost and reduced ability to isolate functionally relevant reasoning. Prior work on compact reasoning shortens such chains through probabilistic sampling, heuristics, or supervision from frontier models, but offers limited insight into whether models internally encode token-level functional importance for answer generation. We address this gap diagnostically and propose greedy pruning, a likelihood-preserving deletion procedure that iteratively removes reasoning tokens whose removal minimally degrades model likelihood under a specified objective, yielding length-controlled reasoning chains. We evaluate pruned reasoning in a distillation framework and show that students trained on pruned chains outperform a frontier-model-supervised compression baseline at matched reasoning lengths. Finally, our analysis reveals systematic pruning patterns and shows that attention scores can predict greedy pruning ranks, further suggesting that models encode a nontrivial functional importance structure over reasoning tokens.