Fewer is More: Boosting LLM Reasoning with Reinforced Context Pruning

TL;DR

This work tackles math reasoning in LLMs by introducing CoT-Influx, a two-stage coarse-to-fine pruner that selectively retains informative Chain-of-Thought examples and prunes tokens to fit within the model’s context window. Trained with reinforcement learning on a GPT-4-evolved math dataset MRD$^3$, the pruner learns to maximize reasoning accuracy while respecting a token budget, enabling far more effective CoT prompts without fine-tuning the LLM. Across multiple LLaMA2 models and five math datasets, CoT-Influx yields significant improvements and even allows a 70B model to outperform GPT-3.5 and other larger models without additional cost. The approach is designed as a plug-and-play module that complements existing reasoning prompts like self-consistency and self-verification, highlighting the value of optimized prompt composition for advancing mathematical reasoning in LLMs.

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. 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.

Figures (7)

• Figure 1: LLaMA2-7B reasoning accuracy under an increasing number of TopK retrieved CoT examples.
• Figure 2: A compressed CoT example using the prompt compression tool of LLMLingua llmlingua. The pruned tokens contain truly redundant tokens (colored in blue) and crucial tokens (colored in red).
• Figure 3: Above: The overview procedure of CoT-Influx; Below: an example illustrating the use of CoT-Influx to first prune entire CoT examples and then prune tokens.
• Figure 4: EM(%) accuracy on GSM8K with inputting different number of CoT examples for CoT-Influx.
• Figure 5: Token length after each stage of our pruner.