Large language models (LLMs) have shown impressive capabilities in various tasks, yet they still struggle with math reasoning. Despite efforts to optimize Chain-of-Thoughts (CoT) prompts and fine-tune LLMs, the potential of few-shot learning remains unexplored. In this work, we propose CoT-Max, a novel approach pushing the boundaries of few-shot CoT learning to improve LLM math reasoning capabilities. CoT-Max addresses the challenges of the selection of useful examples and limited number of examples due to restricted context window length. Inspired by our observation that natural language inputs contain many redundancy, we propose a coarse-to-fine pruner as a plug-and-play module for LLMs, which first identifies crucial CoT examples from a large batch and then further prunes unimportant tokens. To train the pruner, we collect a math reasoning dataset with diverse difficulty and steps, introduce a reward to measure both the input's effectiveness for math reasoning and token length constraints, and propose a novel training approach with reinforcement learning. As a result, CoT-Max significantly outperforms CoT and few-shot prompting baselines across various LLMs (LLaMA2-7B, 13B, 70B) and 5 mathematical datasets, achieving up to 4.55% absolute improvements. Remarkably, without any fine-tuning, LLaMA2-70B with CoT-Max surpasses GPT-3.5 and a wide range of larger LLMs (PaLM, Minerva, etc.) on the GSM8K.

翻译：大语言模型（LLM）在各类任务中展现出卓越能力，但在数学推理方面仍存在不足。尽管已有研究致力于优化思维链（CoT）提示和微调LLM，少样本学习的潜力仍未被充分挖掘。本文提出CoT-Max这一创新方法，突破少样本思维链学习的边界，以提升LLM数学推理能力。CoT-Max有效解决了示例选择难题及上下文窗口长度限制导致的示例数量不足问题。基于自然语言输入存在大量冗余的观察，我们设计了一种粗到细的剪枝器作为LLM的即插即用模块：该模块首先从海量样本中识别关键思维链示例，继而剪除无关标记。为训练该剪枝器，我们构建了包含多难度多步骤的数学推理数据集，引入衡量输入数学推理有效性与标记长度约束的奖励函数，并创新性地提出基于强化学习的训练方法。实验表明，CoT-Max在LLaMA2-7B/13B/70B等多种LLM及5个数学数据集上显著优于CoT与少样本提示基线方法，绝对性能提升最高达4.55%。尤为值得注意的是，无需任何微调，搭载CoT-Max的LLaMA2-70B在GSM8K任务中已超越GPT-3.5及众多更大规模LLM（如PaLM、Minerva等）。