PaperlayerView exact revision on arXiv

Source-linked AI summary

Fewer is More: Boosting LLM Reasoning with Reinforced Context Pruning

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

arXiv:2312.08901v3cs.CLcs.AI

TL;DR

Math reasoning remains challenging for LLMs, and few-shot CoT use is constrained by context length and example selection. CoT-Influx selects useful examples and prunes tokens, then reports improvements over prompting baselines across math datasets and LLMs.

  • Problem

    Math reasoning remains challenging for LLMs, while limited context and model-agnostic retrieval constrain the use of effective few-shot CoT examples.

    Source: §1 · p. 1

  • Method

    CoT-Influx uses a coarse-to-fine pruner to select CoT examples and remove unimportant tokens, trained with reinforcement learning and the GPT-4-evolved MRD3 dataset.

    Source: §1 · p. 2, §6 · p. 8

  • Results

    Across various LLMs and five math datasets, CoT-Influx improves over state-of-the-art prompting baselines by 1.36%-14.09% absolute.

    Source: §1 · p. 2

  • Takeaways & Limitations

    By keeping inputs within the original context window, CoT-Influx serves as a plug-and-play module without additional inference overhead.

    Source: §1 · p. 2

    Performance can depend on CoT generation quality, and the BERT example embeddings cannot handle sequences longer than 512 tokens.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

The introduction presents CoT-Influx as a way to improve math reasoning by fitting more useful chain-of-thought examples into the existing context window. It combines example selection, token pruning, and math-specialized training, and reports gains over prompting baselines.

1 Introduction

  • The method addresses context limits that restrict few-shot examples and retrieval methods that may select unsuitable, model-agnostic examples.Different LLMs favor CoT examples with varying characteristics, such as difficulty levels.Source: §1 · p. 1
  • CoT-Influx prunes redundancy at both example and token levels to fit more informative CoT examples within the original context window.The authors state that keeping inputs within that window avoids additional inference overhead.Source: §1 · pp. 1–2
  • A coarse-to-fine pruner selects helpful CoT examples for the target LLM, then removes unimportant tokens from those examples.CoT-Influx treats the target LLM as a black box and operates as a plug-and-play module.Source: §1 · p. 2
  • MRD3, created with GPT-4 and Evol-Instruct, supplies math problems with varied difficulty levels and reasoning steps for training and inference.The dataset is intended to support generalization across a wide range of math problems.Source: §1 · p. 2
  • Reinforcement learning addresses the lack of gradients through discrete token selection, using a reward for LLM loss, reasoning effectiveness, and token length.Training also uses a difficulty-aware dataloader and two techniques to stabilize reinforcement learning.Source: §1 · p. 2
  • CoT-Influx improves over state-of-the-art prompting baselines by 1.36%-14.09% absolute and surpasses GPT-3.5 by 2.5% on GSM8K.The authors also report that LLaMA2-70B with CoT-Influx outperforms a broad range of larger LLMs.Source: §1 · p. 2

2 Related Works

Related work covers CoT prompting and training, prompt compression, and few-shot retrieval. The paper identifies math-specific shortcomings in compression and retrieval that motivate its approach.

2 Related Works

  • Prior math-reasoning work enhances CoT prompts or creates CoT-based training data, while much work leaves few-shot CoT relatively unexplored.Examples include Program-of-Thoughts, Tree-of-Thoughts, and Everything-of-Thoughts.Source: §2 · p. 2
  • Prompt compression methods prune tokens using designed metrics, but may miss redundancy across whole CoT examples or remove numerical and format tokens crucial for math reasoning.The passage groups approaches as token pruning, soft prompt compression, and information-entropy-based methods.Source: §2 · p. 2
  • Retrieval methods are suboptimal for math reasoning because model-agnostic examples may not match LLM-specific preferences and retrieval does not address token redundancy.Different LLM capabilities favor CoT examples of varying complexity.Source: §2 · p. 3

3 Pilot Study

The pilot study examines how CoT example quantity and selection affect math reasoning, then considers whether redundant tokens can be removed. It finds gains from more examples but notes context and example-quality constraints.

3 Pilot Study

  • The pilot uses MRD3 and a 4096-token LLaMA2 context window; answer generation shares that limit.The study could input up to 20 CoT examples under this budget.Source: §3 · p. 3
  • Increasing TopK-selected CoT examples improves LLaMA2-7B GSM8K accuracy over the standard 8-shot setting, but context limits constrain the potential gains.With 20 examples, accuracy can drop even below the token limit because the large input context limits response space.Source: §3 · p. 3
  • Randomly selected CoT examples, even up to 16 shots, underperform the manually curated 8-shot setting.Table 1 is identified as showing that example selection heavily impacts reasoning performance.Source: §3 · p. 3, Table 1 · p. 3
  • CoT examples contain redundant tokens that can be pruned to make space for more informative content, but identifying them for math reasoning is difficult.The LLMLingua example shows that tokens it prunes can include crucial numerical and format tokens.Source: §3 · p. 3, Figure 2 · p. 3

4 CoT-Influx Methodology

CoT-Influx prunes CoT examples and then their tokens to fit more useful reasoning context within the target LLM’s window. Its methodology combines a diverse math-reasoning dataset, a coarse-to-fine pruner, and reinforcement-learning optimization.

4 CoT-Influx Methodology

  • The two-stage pruner first selects helpful CoT examples from a larger batch, then removes unimportant tokens from retained examples.The lightweight module is designed to operate alongside a target LLM.Source: §4 · p. 3

4.1 CoT Dataset Collection

  • MRD3 starts from 9.7K question-answer pairs across four math datasets, with GPT-4-generated, consistently formatted CoT reasoning and difficulty scores.GPT-4 mutations address the initial dataset’s concentration of questions at difficulty levels 2–4.Source: §4.1 · p. 4

4.2 Problem Formulation

  • CoT-Influx optimizes the pruned input to support a correct answer while keeping its token count below the LLM context-window limit.The input places retained CoT tokens before the target question.Source: §4.2 · p. 4

4.3 Coarse-to-fine Pruner Design

  • A two-stage policy network predicts whether to retain each CoT example and then each token in the retained examples.BERT-Large supplies sentence-level CoT-example embeddings, enabling features for inputs longer than the LLM context window.Source: §4.3 · p. 5

4.4 End-to-end RL Optimization

  • The multi-objective reward evaluates CoT effectiveness and context-window fit, using LLM loss and reasoning accuracy as effectiveness signals.Reasoning accuracy is included because average generated-token loss does not fully reflect whether the answer is correct.Source: §4.4 · p. 5
  • Training uses REINFORCE, filters out questions above a difficulty threshold, repeats sampled questions with different pruned prompts, and smooths accuracy rewards with an exponential moving average.The manually designed 8-shot CoTs are appended during early training to stabilize randomly pruned prompts.Source: §4.4 · p. 6
  • Figure 4 reports EM(%) accuracy on GSM8K for CoT-Influx, comparing results obtained with different numbers of chain-of-thought (CoT) examples.Source: Figure 4 · p. 6

5 Evaluation

CoT-Influx improves math reasoning across LLaMA2 models and five datasets by using more selected CoT examples with redundant tokens pruned. Its evaluation also examines model-specific shot counts, compatibility with other prompting methods, dataset and pruning ablations, and inference costs.

5 Evaluation

  • Evaluation measures Exact Match accuracy for LLaMA2-7B, 13B, and 70B on GSM8K, AddSub, Multiarith, Svamp, and Singleeq.Baselines include zero-shot and few-shot-CoT prompts, retrieval methods, and prompt-compression methods.Source: §5 · p. 6

5.1 Main Results

  • 2.13% to 4.55% absolute improvements over all baselines; with fewer input tokens, CoT-Influx also beats retrieval baselines by 1.36% to 14.09%.The comparisons cover GSM8K and four other math datasets, using CoT-Influx inputs up to twice the LLM context window.Source: Table 2 · p. 7, Table 3 · p. 7, §5.1 · p. 7
  • On GSM8K, LLaMA2-13B with 48 CoTs beats standard few-shot CoT by 4.55% and TopK by 8.72%, while peak shot counts are 40 for 7B and 48 for 13B.The tested range was 16 to 64 shots, corresponding to 0.7× to 2.8× the LLaMA2 context-window token count.Source: §5.1 · p. 7
  • Without fine-tuning, LLaMA2-70B with CoT-Influx surpasses GPT-3.5 by 2.5% on GSM8K and outperforms much larger LLMs.The comparison is under the few-shot-CoT setting.Source: §5.1 · p. 7, Table 4 · p. 7
  • Applying self-consistency or self-verification to CoT-Influx prompts further improves LLaMA2 performance on GSM8K.The evaluation sampled 20 times; Table 5 identifies the methods as maj@20 and verify@20.Source: §5.1 · p. 8, Table 5 · p. 8

5.2 Ablation Study and Analysis

  • MRD3’s GPT-4-generated and evolved CoTs are vital to reasoning performance, and removing either pruning stage reduces accuracy, with a larger drop when token pruning is removed.The dataset comparison uses the same pruner setting; the pruning-strategy comparison tests shot-only and token-only pruning.Source: §5.2 · p. 8, Table 6 · p. 8, Table 7 · p. 8
  • 4.28× is the total pruning ratio, with shot pruning contributing 3.87×; CoT-Influx adds 1.4GB memory and 1.7% latency on LLaMA2-7B.The cost measurements use GSM8K on a single NVIDIA A100 GPU.Source: §5.2 · p. 8, Table 8 · p. 8
  • More capable LLMs favor harder CoTs, smaller LLMs favor simpler ones, and numerical and format tokens are essential for math reasoning.Function words and irrelevant background context can be pruned without affecting reasoning capability.Source: §5.2 · p. 8

6 Conclusion

CoT-Influx improves math reasoning by pruning unnecessary few-shot examples at both shot and token levels, using reinforcement learning with the GPT-4-evolved MRD3 dataset. The paper reports strong performance while noting dependencies on CoT quality and BERT’s 512-token limit, and outlines additional analyses in the appendix.

6 Conclusion

  • CoT-Influx is a plug-and-play module that prunes unnecessary few-shot examples at shot and token levels, trained with a math-specific reward and the GPT-4-evolved MRD3 dataset.The authors report that experiments across datasets and LLMs demonstrate the method’s effectiveness.Source: §6 · p. 8

Limitations

  • Performance can be influenced by CoT-generation quality, and the BERT feature extractor cannot handle examples exceeding 512 tokens.The authors report strong performance on MRD3 despite the dependence on CoT quality and identify both issues as limitations for future work.Source: Limitations · p. 9

Appendix

  • The appendix covers context-window extension comparisons, finetuned models, reward ablations, hyperparameter sensitivity, MRD3 evolution, training details, related work, and prompt settings.It also includes case studies and evaluation-dataset details.Source: Appendix · p. 12

A.1 Comparison with context window extension methods

On GSM8K, CoT-Influx outperforms prompting baselines on finetuned models and is compared with context-window extension for LLaMA2-7B. The reported comparison finds weaker reasoning with extension under the same prompt and describes gains for CoT-Influx on finetuned models.

A.1 Comparison with context window extension methods

  • The context-extension comparison evaluates LLaMA2-7B against a 32K-token version using Positional Interpolation, with CoT-Influx applied when prompts exceed the original window.Table 9 presents GSM8K EM (%) accuracy for these methods.Source: §A.1 · p. 12, Table 9 · p. 12
  • With the same input prompt, the 32K context-window extension weakened LLaMA2-7B’s math reasoning, while CoT-Influx pruned over-limit prompts for the original model.The authors also report that the observation about gains from more helpful CoT examples did not hold for the extended-context model.Source: Table 9 · p. 12, §A.1 · p. 12

A.2 CoT-Influx on finetuned LLMs

  • 9.78% absolute improvement over TopK retrieval on LLaMA2-13B-Chat was achieved with CoT-Influx using 57.6% average tokens.The comparison is on GSM8K; Table 10 reports EM (%) accuracy, and the stated context limit for LLaMA2-13B-Chat and GPT-3.5-Turbo is 4096 tokens.Source: §A.2 · p. 13, Table 10 · p. 13
  • CoT-Influx also surpassed a wide range of prompting baselines on GSM8K for LLaMA2-13B-Chat and GPT-3.5-Turbo using more input shots and fewer tokens.The passage specifically quantifies the LLaMA2-13B-Chat comparison with TopK retrieval.Source: §A.2 · p. 13, Table 10 · p. 13

A.3 Ablation study on reward design

The reward ablation finds that removing any component produces suboptimal prompt selection and compression, with the context-window token-limit reward especially important because its removal prevents pruning.

A.3 Ablation study on reward design

  • The full reward combines math-accuracy, LLM-loss, and context-window token-limit components.The passage gives RLoss as 1/(1 + LLM loss); the displayed expression for RLimit is incomplete in the supplied text.Source: §A.3 · p. 13
  • Removing any reward component produced suboptimal prompt selection and compression compared with the full-reward baseline.Table 11 reports EM (%) accuracy on GSM8K for LLaMA2-7B and LLaMA2-13B with different reward functions.Source: §A.3 · p. 13, Table 11 · p. 13
  • Without the token-limit reward, the pruner did not prune shots or tokens and instead output the truncated redundant prompt.The authors identify this component as the most important of the three reward terms.Source: §A.3 · p. 13

A.4 Sensitivity analysis on hyperparameters

The hyperparameter analysis examines token target T and token penalty coefficient w, and reports robust training when the token target is not overly aggressive.

A.4 Sensitivity analysis on hyperparameters

  • The authors identify token target T and token penalty coefficient w as the most important pruner-training settings examined.The analysis investigates robustness to these hyperparameters.Source: §A.4 · p. 13
  • The sensitivity analysis tests the pruner with different token targets T and token penalty coefficients w.The results are presented in Table 12.Source: §A.4 · p. 13, Table 12 · p. 13
  • Training was reported as highly robust when the token target T was not overly aggressive.The authors caution that T should not be too small.Source: §A.4 · p. 13

A.5 Case Study on different prompt compression methods

The case study contrasts how prompt-compression methods preserve or discard CoT structure and reasoning, while describing prompt evolution strategies and the resulting difficulty mix. CoT-Influx selects representative examples and removes redundant function words; the study reports that evolution changes candidate diversity and difficulty distributions.

A.5 Case Study on different prompt compression methods

  • Prior methods can damage prompt structure at similar token lengths, whereas CoT-Influx preserves more representative examples and their useful structure.The case study notes that the examples share backgrounds but vary in difficulty and reasoning-step count, and that solutions contain substantial redundancy.Source: §A.5 · pp. 14–15
  • Selective Context and LLMLingua frequently discard question-answer markers and reasoning structure, while GPT-4 compression retains most markers but systematically removes reasoning steps.CoT-Influx selects representative examples and removes redundant function words.Source: §A.5 · p. 15

B.1 Prompt template for evolution

  • The evolution templates cover easier questions, added constraints, deeper questions, explicit multi-step reasoning, and reduced difficulty.The templates also ask for a different topic when producing an easier question.Source: §B.1 · p. 16
  • Quantitative analyses of difficulty and reasoning-step distributions indicate that GPT-4 can effectively follow evolution instructions to modify constraints or the difficulty level of input questions.Source: §B.1 · p. 16

B.2 Difficulty and Reasoning Steps Distribution of MRD3

  • Figure 6 shows the distributions of question difficulty and reasoning steps; the original questions cluster around levels 2–4, while evolution schemes shift difficulty in different directions.Add-constraints, deepening, and increase-reasoning generate more high-difficulty questions; revise-difficulty and produce-easier generate easier questions.Source: §B.2 · p. 16

B.3 Additional observation on difficulty distribution

This section reports that LLMs with different capabilities favor CoT examples of different difficulty levels. It also describes how evolved examples diversify the candidate set and how the training data and pruner settings are prepared.

B.3 Additional observation on difficulty distribution

  • LLaMA2-7B favors difficulty levels 3–4, while LLaMA2-13B favors level 4 or above; selected examples for LLaMA2-70B also have higher average difficulty and reasoning-step counts than those for 13B.The reported averages are 3.57 difficulty and 3.04 steps for 70B, versus 3.51 and 2.98 for 13B.Source: §B.3 · p. 17

B.4 Evolution example

  • Evolved prompts can change constraints, answers, or question backgrounds, adding variation to the candidate sets used for training and reasoning.The add-constraints example changes the answer from 16 to 12, while produce-easier changes an arcade-ticket question into a game-score question.Source: §B.4 · pp. 17–18

C.1 Detailed algorithm for training data preparation

  • Training-data preparation evolves CoT examples, evaluates their difficulty with GPT-4, and splits examples into question and prompt-candidate sets using a difficulty threshold.The algorithm omits reasoning-step information from the split, although difficulty and reasoning steps are both evaluated.Source: §C.1 · p. 18

C.2 Detailed settings and hyperparameters

  • Most pruner hyperparameters are shared across LLMs, and the evaluation prompt-candidate subset is selected by searching math-reasoning performance on 100 random examples.Table 14 is identified as the detailed hyperparameter listing.Source: §C.2 · p. 19, Table 14 · p. 19

C.3 Training dynamics

The appendix describes pruner training dynamics, evaluation datasets, prompt reconstruction, related work, and prompt settings. It reports that reward rises during RL training while token counts approach the target, and notes that pruning can still remove important structural tokens.

C.3 Training dynamics

  • During LLaMA2-13B pruner training, reward rises steadily and remaining tokens decrease rapidly at first before oscillating around the target.The discrete prediction reward takes values −0.1, 0, or 1; question repetition and an exponential moving average are reported to suppress its oscillations.Source: Figure 7 · p. 19, §C.3 · p. 19

C.4 Detailed introduction of dataset for evaluation

  • The evaluation datasets include GSM8K, SVAMP, MultiArith, AddSub, and SingleEq, spanning grade-school word problems, arithmetic variations, multi-step arithmetic, addition/subtraction, and single-equation algebra.The passages report 1,319 GSM8K validation examples and use all examples for evaluation in SVAMP, MultiArith, and AddSub.Source: §C.4 · pp. 19–20

C.5 Rule-based prompt reconstruction

  • Rule-based reconstruction restores markers such as Q:, A:, and “Let’s think step by step” when pruning removes them; occasional removal of important tokens can still lead to incorrect reasoning.The reconstruction is described as alleviating the effects of sub-optimal pruning.Source: §C.5 · p. 20

D Additional Related Works

  • The related-work discussion positions CoT-Influx as applicable to black-box LLMs and orthogonal to context-window extension approaches that modify models or require complex training.It also characterizes CoT-Influx as advantageous over prior prompting strategies across gradient-free, token-limit, difficulty-awareness, and dynamic-shot dimensions.Source: D Additional Related Works · p. 20, Table 15 · p. 21

E Prompt Settings

  • The GPT-4 compression prompt requests fewer than a given token count, but compression outputs still fail to meet the restriction in most cases.The passage describes uncontrollable output length as a disadvantage of this approach.Source: E Prompt Settings · p. 21
Loading 2312.08901v3…