Fewer is More: Boosting Math Reasoning with Reinforced Context Pruning (2024)

Abstract: 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. 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 is a plug-and-play module for LLMs, adaptable in various scenarios. It’s compatible with advanced reasoning prompting techniques, such as self-consistency, and supports different long-context LLMs, including Mistral-7B-v0.3-32K and Yi-6B-200K.

## introduction large language models (llms) have demonstrated remarkable capabilities across a range of † work was done during the internship at microsoft research ‡ corresponding author: lzhani@microsoft.com tasks @xcite . however, it remains a significant challenge to improve llm performance on math reasoning tasks. while existing efforts focus on optimizing chainof-thought (cot) prompts @xcite @xcite and fine-tuning llms @xcite under the zero-shot setting, the potential of few-shot learning in improving llm reasoning has not been fully explored. inspired by the human reasoning process, we propose the hypothesis: if llms are exposed to more step-by-step problem-solving examples (i.e., cots) before answering questions, it could potentially improve llms reasoning capability to generate a correct solution. this leads us to ask: what's the boundary of llm reasoning capability achievable through inputting more cot examples? we face two major obstacles. first, the limited token length of llms' context window restricts the number of few-shot examples. extending the context window is one solution, but it requires expensive fine-tuning and increases inference overhead @xcite . while prompt compression @xcite is another approach, it underperforms in math reasoning. tokens like numerical and format ones, though identified as redundant, are crucial for few-shot prompting to solve math problems. second, it's challenging to select effective cot examples. section 3 reveals that random choices can even harm reasoning performance. existing retrieval-based methods @xcite are not tailored for math reasoning, making them suboptimal. these retrieved examples are model-agnostic, while we found that different llms favor cot examples of varying characteristics (e.g., diverse difficulty levels). in this work, we propose cot-influx, which addresses all the above challenges and pushes the boundaries of utilizing few-shot learning to improve llm math reasoning capability. cot-influx is motivated by the observation that current llm context window has not been fully utilized due to redundancy at both the example and token levels in natural language input. as such, these redundant inputs can be pruned to free up space for more informative context. the central idea of cot-influx is to input long lengthy cot examples, select the crucial examples for the target llm, and then prune redundant tokens. as a result, by inputting much more helpful cot examples, each composed solely of informative tokens and with a shorter length, we greatly improve llm math reasoning ability. moreover, as all these inputs remain within the context window, we do not increase any inference overhead. this stands in stark contrast to other methods @xcite . cot-influx treats the target llm as a black box, and serves as a plug-and-play module for llms as shown in fig. 3 . the key module is a coarse-to-fine pruner consisting of two steps: (i) a shot pruner first selects the most helpful cot examples from candidates, and (ii) a token pruner then removes unimportant tokens from selected cot examples. however, training the pruner presents two challenges: (1) since we identify discrete tokens before the llm tokenizer, the loss gradient cannot be backpropagated through the tokenizer to update the pruner; (2) the high difficulty of many math problems, which consistently yield incorrect answers regardless of the quality of compressed fewshot examples, poses a challenge to the effective training of the pruner. to this end, we introduce a novel training approach with reinforcement learning to mitigate the gradient issue. we design a reward function to measure the llm loss, fewshot math reasoning effectiveness, and token length constraints. then, we propose a difficulty-aware dataloader filtering appropriate problems and introduce two techniques to stabilize the rl training. extensive experiments on various llms and diverse reasoning datasets demonstrate the effectiveness of cot-influx. cot-influx significantly boosts llm reasoning capability, achieving up to 14.09% absolute improvements over sota baselines, and establishes a new prompting-based benchmark in math reasoning accuracy without any fine-tuning or additional inference costs. remarkably, llama2-70b with cot-influx outperforms a broad range of larger llms and surpasses gpt-3.5 by 2.5% on gsm8k. moreover, we prove that cot-influx is also highly beneficial on long-context llms. for example, we achieve ∼2.5% higher gms8k accu-racy with 15× fewer input tokens on yi-6b-200k. however, most works focus on the zero-shot setting with only task instruction or cot prompts, leaving the potential of few-shot cot largely untapped. we explore leveraging few-shot cot learning to improve llms' math reasoning capabilities. prompt compression. one related works involve compressing input prompts, including (1) token pruning @xcite ; (2) soft prompt compression @xcite ; and (3) information-entropy-based approaches @xcite . however, they do not effectively solve our problem for two reasons. first, they prune tokens based on designed metrics, often failing to remove redundancy of the entire cot examples. second, some tokens such as numerical and format tokens, although redundant, are crucial for math reasoning. prompt retrieval aims to optimize llm performance by selecting effective few-shot examples. heuristic methods, such as the topk retrieval @xcite , bm25 @xcite , @xcite , and entropy @xcite , select examples based on semantic similarity. recently, supervisedbased methods like epr @xcite , llm-r @xcite , and ids @xcite have been proposed, which train a retrieval model to learn better example selection. however, these methods do not consider token redundancy, which can limit the number of retrieved examples. ## pilot study this section presents our key observations of fewshot learning in improving llms math reasoning, upon which the cot-influx design is based. note that experiments are done with our cot dataset, mrd 3 , as introduced in sec. 4.1. observation 1: llms can improve reasoning with decreased accuracy due to squeezed token limits of response 0 5 10 15 0 2 4 6 8 12 16 20 number of cot examples exceed token limits! llama2-7b on gsm8k accuracy standard 8-shot cot more helpful cot examples, but the context window length limits further accuracy gains. a standard practice for evaluating llms' math reasoning capability is the use of 8-shot manuallydesigned cots @xcite . we increase the number of cot shots to see if reasoning accuracy improves. to avoid poor-quality examples, we use the topk method @xcite to select the k most relevant cot examples for each question. given llama2's context window limit of 4096 tokens, we could only input up to 20 cot examples ‡ . as fig. 1 shows, increasing cot examples improves llama2-7b's reasoning accuracy on the gsm8k dataset, significantly outperforming the standard 8-shot setting. however, the limited llm context window hinders the full potential of fewshot cot learning for improving math reasoning. for instance, even with 20 cots not hitting the token limit, accuracy drops as the large input context limits the llm's response space. observation 2: cot example selection is crucial for math reasoning. simply adding cot examples randomly doesn't boost performance. the prior study suggests that more cot examples can improve llm reasoning performance. however, the quality of cot examples is crucial to the final performance. as shown in ## cot-influx methodology motivated by our observations, this section introduces cot-influx, which maximizes cot examples within the original llm context window by identifying the most important cot examples and tokens from a long input context. ## evaluation and analysis models, datasets and metric. we evaluate cot-influx on llama2-7b, llama2-13b, llama2-70b and two long-context llms. the math datasets for evaluation include gsm8k @xcite , addsub @xcite , multiarith @xcite , svamp @xcite , and singleeq @xcite . we also include mmlu-stem @xcite for general reasoning evaluation, as well as math @xcite for evaluation on challenging reasoning tasks. the exact match (em) accuracy of the predicted answers are reported as the metric. baselines we set three baselines for comparison: ## conclusion we present cot-influx, a plug-and-play module that improves llm math reasoning by pruning unnecessary few-shot examples at shot and token levels for a more effective input context. to train the module, we use reinforcement learning to optimize a math reasoning-specific reward. extensive experiments on various datasets and llms compared with state-of-the-art baselines demonstrate the effectiveness of our method. this paper highlights the vast potential of few-shot cot prompting in augmenting llms' math reasoning abilities.