Fewer is More: Boosting LLM Reasoning with Reinforced Context Pruning

📄 arXiv: 2312.08901v3 📥 PDF

作者: Xijie Huang, Li Lyna Zhang, Kwang-Ting Cheng, Fan Yang, Mao Yang

分类: cs.CL, cs.AI

发布日期: 2023-12-14 (更新: 2024-02-15)

💡 一句话要点

CoT-Influx：强化上下文剪枝提升大语言模型数学推理能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 数学推理 链式思维 上下文剪枝 强化学习

📋 核心要点

  1. 现有大语言模型在数学推理方面表现不足，尤其是在复杂问题上，需要更有效的提示方法。
  2. CoT-Influx通过由粗到精的剪枝器，选择并优化CoT示例，在有限的上下文窗口内提供更多有效信息。
  3. 实验表明，CoT-Influx显著提升了LLaMA2系列模型在多个数学数据集上的性能，甚至超越了更大的模型。

📝 摘要（中文）

大型语言模型(LLMs)展现了令人印象深刻的能力，但它们在数学推理方面仍然存在困难。本文提出了一种名为CoT-Influx的新方法，旨在突破少样本链式思维(CoT)学习的界限，以提高LLM的数学推理能力。受到在提示中添加更简洁的CoT示例可以提高LLM推理性能的观察结果的启发，CoT-Influx采用了一种由粗到精的剪枝器，以最大化有效且简洁的CoT示例的输入。该剪枝器首先选择尽可能多的关键CoT示例，然后剪除不重要的token以适应上下文窗口。使用具有不同难度级别和推理步骤的数学推理数据集来训练剪枝器，并结合数学专用强化学习方法。结果表明，通过在token数量上增加一倍的上下文窗口大小，从而启用更多的CoT示例，CoT-Influx在各种LLM（LLaMA2-7B、13B、70B）和5个数学数据集上显著优于各种提示基线，实现了高达4.55%的绝对改进。值得注意的是，在没有任何微调的情况下，具有CoT-Influx的LLaMA2-70B在GSM8K上超过了GPT-3.5和各种更大的LLM（PaLM、Minerva 540B等）。CoT-Influx作为LLM的即插即用模块，与大多数现有的推理提示技术（如自洽性和自验证）兼容。

🔬 方法详解

问题定义：论文旨在解决大语言模型在数学推理任务中，由于上下文窗口限制，无法充分利用CoT示例的问题。现有方法要么无法选择关键CoT示例，要么无法在有限的窗口内塞入足够多的示例，导致推理性能受限。

核心思路：论文的核心思路是设计一个上下文剪枝器，该剪枝器能够从大量的CoT示例中选择最关键的示例，并对这些示例进行token级别的剪枝，从而在有限的上下文窗口内最大化有效信息的密度。这样，模型可以利用更多的CoT示例进行推理，从而提高性能。

技术框架：CoT-Influx包含一个由粗到精的剪枝器。首先，粗剪枝阶段选择尽可能多的关键CoT示例。然后，细剪枝阶段对选定的示例进行token级别的剪枝，移除不重要的token。剪枝器的训练使用数学推理数据集，并结合了数学专用强化学习方法。整个框架可以作为LLM的即插即用模块，与现有的推理提示技术兼容。

关键创新：论文的关键创新在于提出了由粗到精的上下文剪枝策略，并使用强化学习方法训练剪枝器。这种方法能够有效地选择和优化CoT示例，从而在有限的上下文窗口内提供更多有用的信息。与现有方法相比，CoT-Influx能够更有效地利用上下文信息，从而提高数学推理性能。

关键设计：剪枝器的训练使用强化学习方法，奖励函数的设计考虑了数学推理的正确率和CoT示例的简洁性。具体来说，奖励函数会奖励模型选择能够提高推理正确率的CoT示例，并惩罚模型选择冗余或不相关的token。此外，论文还使用了数学专用强化学习方法，以更好地适应数学推理任务的特点。具体的网络结构和参数设置在论文中有详细描述（未知）。

📊 实验亮点

CoT-Influx在LLaMA2-7B、13B和70B等多个模型上进行了评估，并在GSM8K等5个数学数据集上取得了显著的性能提升。例如，在GSM8K数据集上，CoT-Influx使LLaMA2-70B的性能超过了GPT-3.5和PaLM等更大的模型，实现了高达4.55%的绝对改进。这些结果表明，CoT-Influx是一种有效的数学推理增强方法。

🎯 应用场景

CoT-Influx具有广泛的应用前景，可以应用于各种需要复杂推理的任务，例如科学计算、金融分析、法律推理等。该方法可以作为LLM的通用增强模块，提高其在各种领域的应用性能。此外，该方法还可以用于优化LLM的提示工程，提高LLM的效率和可靠性。

📄 摘要（原文）

Large Language Models (LLMs) have shown impressive capabilities, yet they still struggle with math reasoning. In this work, we propose CoT-Influx, a novel approach that pushes the boundary of few-shot Chain-of-Thoughts (CoT) learning to improve LLM mathematical reasoning. Motivated by the observation that adding more concise CoT examples in the prompt can improve LLM reasoning performance, CoT-Influx employs a coarse-to-fine pruner to maximize the input of effective and concise CoT examples. The pruner first selects as many crucial CoT examples as possible and then prunes unimportant tokens to fit the context window. A math reasoning dataset with diverse difficulty levels and reasoning steps is used to train the pruner, along with a math-specialized reinforcement learning approach. As a result, by enabling more CoT examples with double the context window size in tokens, CoT-Influx significantly outperforms various prompting baselines across various LLMs (LLaMA2-7B, 13B, 70B) and 5 math datasets, achieving up to 4.55% absolute improvements. Remarkably, without any fine-tuning, LLaMA2-70B with CoT-Influx surpasses GPT-3.5 and a wide range of larger LLMs (PaLM, Minerva 540B, etc.) on the GSM8K. CoT-Influx serves as a plug-and-play module for LLMs and is compatible with most existing reasoning prompting techniques, such as self-consistency and self-verification.