arXiv:2409.08561cs.CLcs.AI2024-09被引 18

用压缩思想让大模型推理更快,不丢精度。

Expediting and Elevating Large Language Model Reasoning via Hidden Chain-of-Thought Decoding

  • 用辅助模型将完整推理过程压缩成一个语义对齐的特殊标记。
  • 在数学、问答等任务上实现1.5倍以上推理加速,性能不降反而提升。
  • 适合追求高效推理的工业应用或资源受限场景使用。

大语言模型在需要多步推理的任务中表现出色,但链式思维(CoT)提示会导致输出序列过长,增加计算开销和推理延迟。为此,我们提出一种通过语义对齐压缩CoT过程的新方法,实现更高效的解码同时保留CoT优势。该方法引入一个辅助的CoT模型,学习将完整的思考过程压缩为与原始输出语义对齐的紧凑标记表示,并将其作为隐藏链式思维(HCoT)模型的输入。训练分两阶段:首先用对比损失优化辅助模型生成对齐的压缩标记;随后冻结辅助模型参数,微调HCoT模型以基于指令前缀和压缩标记生成准确预测。在数学推理、智能体调用和问答三个挑战性领域上的实验表明,该方法在保持或优于全量CoT基线性能的同时,推理速度提升至少1.5倍。引入对比学习进一步提升了压缩表示质量,改善了推理效果和任务准确率。本工作为大模型在广泛应用场景中更高效地利用多步推理能力开辟了新路径。

原文摘要 · Abstract (English)

Large language models (LLMs) have demonstrated remarkable capabilities in tasks requiring reasoning and multi-step problem-solving through the use of chain-of-thought (CoT) prompting. However, generating the full CoT process results in significantly longer output sequences, leading to increased computational costs and latency during inference. To address this challenge, we propose a novel approach to compress the CoT process through semantic alignment, enabling more efficient decoding while preserving the benefits of CoT reasoning. Our method introduces an auxiliary CoT model that learns to generate and compress the full thought process into a compact special token representation semantically aligned with the original CoT output. This compressed representation is then integrated into the input of the Hidden Chain-of-Thought (HCoT) model. The training process follows a two-stage procedure: First, the CoT model is optimized to generate the compressed token representations aligned with the ground-truth CoT outputs using a contrastive loss. Subsequently, with the CoT model parameters frozen, the HCoT model is fine-tuned to generate accurate subsequent predictions conditioned on the prefix instruction and the compressed CoT representations from the CoT model. Extensive experiments across three challenging domains - mathematical reasoning, agent invocation, and question answering - demonstrate that our semantic compression approach achieves competitive or improved performance compared to the full CoT baseline, while providing significant speedups of at least 1.5x in decoding time. Moreover, incorporating contrastive learning objectives further enhances the quality of the compressed representations, leading to better CoT prompting and improved task accuracy. Our work paves the way for more efficient exploitation of multi-step reasoning capabilities in LLMs across a wide range of applications.

推理加速链式思维压缩表示

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。