Title: Fewer is More: Boosting LLM Reasoning with Reinforced Context Pruning

Institute: Hong Kong University of Science and Technology, Microsoft Research

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

Arxiv Link: https://arxiv.org/abs/2312.08901

Date: 2024.2.15

1. Abstract

Even with CoT prompting, LLMs still fall short on mathematical reasoning problems. The authors propose CoT-Influx, a method that optimizes both the exemplars and the content of chain-of-thought prompts to improve the reasoning ability of LLMs. Its core idea is to use pruning to maximize the amount of useful information fed into the model.

2. Motivation

1. The potential of chain-of-thought (CoT) exemplars is under-exploited:
   Prior work has shown that supplying step-by-step chain-of-thought exemplars helps models perform complex reasoning. However, a limited context window caps the number of CoT exemplars that can be fed in, so the full potential of CoT goes unrealized.

2. Limitations of existing solutions:
   Extending the context window does make room for more CoT exemplars, but it drives up inference cost and complexity. Existing compression- and retrieval-based techniques, meanwhile, perform poorly on mathematical reasoning tasks, and in particular fail to pick out the CoT exemplars and tokens that matter most for a specific task.

These gaps motivate CoT-Influx, which improves the quality of the input text through a pruning strategy without increasing computational cost.

3. Observations

1. More CoT exemplars improve LLM reasoning performance: Adding CoT exemplars lifts how LLMs perform on mathematical reasoning tasks, but the context window caps how many can actually be supplied.

2. The choice of CoT exemplars is critical: Not every CoT exemplar helps reasoning; carelessly adding incorrect, redundant, or misleading exemplars can even degrade performance.

3. CoT exemplars contain redundant tokens: CoT exemplars often carry redundant tokens. Pruning them frees space for more useful content and thereby improves reasoning.

4. Method: CoT-Influx

The short version, skipping the details

CoT-Influx treats the LLM as a black box. Its key module is a coarse-to-fine pruner that produces high-quality, compact CoT exemplars in two steps:

1. Shot Pruner: select from a large pool of CoT exemplars the k exemplars most helpful for the target problem.

2. Token Pruner: prune the retained CoT exemplars further, stripping out unimportant tokens to yield a condensed version.

The condensed CoT exemplars are then prepended to the question, and that’s it.

The detailed version

1. Problem definition

Given a set of chain-of-thought exemplars $\mathcal{D} = \{x_{\text{cot}}^i\}_{i=1}^{n}$, where each exemplar $x_{\text{cot}}^i$ consists of a question, reasoning steps, and an answer, and whose token count exceeds the LLM’s context window length $T$, we would like to optimize $\mathcal{D}$ through a two-stage pruning process to produce an effective input whose total token count $t(x_{\text{input}})$ satisfies:

$$ t(x_{\text{input}}) \leq T $$

while still allowing the LLM to generate the correct reasoning result from that input.

2. The two-stage pruning process

CoT-Influx’s two-stage pruning process uses policy networks to select the useful chain-of-thought exemplars and, within them, the relevant tokens:

• Stage one: Shot Pruner
  This stage uses a multi-layer perceptron (MLP) with two hidden layers to select the $k'$ most useful exemplars from a batch of chain-of-thought exemplars. The MLP takes as input the text embedding of each exemplar, denoted $H_{\text{shot}} \in \mathbb{R}^{k \times N \times D_{\text{BERT}}}$, where $k$ is the batch size, $N=512$ is the number of tokens per exemplar, and $D_{\text{BERT}}$ is the dimensionality of the BERT embedding. The MLP then outputs a probability distribution:

$$ \pi_{\theta_1}(a_{\text{shot}} | s_{\text{shot}}) = \sigma(\text{MLP}(H_{\text{shot}})) $$

Here $\sigma$ is the sigmoid activation, and the action $a_{\text{shot}}$ indicates whether the CoT exemplar is kept. The selected exemplar set is:

$$ \mathcal{D'} = \{x_{\text{cot}}^j \in \mathcal{D} : a_{\text{shot}}^j = 1\} $$

• Stage two: Token Pruner
  For the retained exemplar set $\mathcal{D'}$, a second two-hidden-layer MLP prunes the tokens within each exemplar. It likewise takes as input the embeddings of the retained exemplars $H_{\text{token}} \in \mathbb{R}^{k' \times N \times D_{\text{BERT}}}$ and outputs a probability distribution that decides whether each token is kept:

$$ \pi_{\theta_2}(a_{\text{token}} | s_{\text{token}}) = \sigma(\text{MLP}(H_{\text{token}})) $$

The pruned token set is:

$$ \hat{x}_{\text{cot}}^j = \{ \text{token} \in x_{\text{cot}}^j : a_{\text{token}} = 1 \} $$

3. Optimization objective: a multi-objective reward function

CoT-Influx optimizes a multi-objective reward function that keeps reasoning accurate while cutting away as many redundant exemplars and tokens as possible. The reward combines the LLM’s reasoning loss, the accuracy of the reasoning, and the number of input tokens:

$$ R(x_{\text{input}}) = \left( \frac{1}{1 + L_{\text{LLM}}(x_{\text{input}})} + R_{\text{Acc}} \right) \times \left( \frac{t(x_{\text{input}})}{T} \right)^w $$

where:

• $L_{\text{LLM}}(x_{\text{input}})$ is the LLM’s reasoning loss;
• $R_{\text{Acc}}$ is the reasoning accuracy (1 if correct, 0 if wrong);
• $t(x_{\text{input}})$ is the number of input tokens;
• $T$ is the maximum token length of the context window;
• $w$ is a hyperparameter that tunes the influence of the token count.

This reward steers the pruner toward maximum reasoning accuracy while squeezing the total token count down far enough that the input fits inside the context window.

4. Reinforcement learning optimization

To optimize the pruning policy networks, CoT-Influx uses REINFORCE, adjusting the parameters of the two-stage pruner by maximizing the reward function.

$$ \nabla_\theta J(\theta) = R(x_{\text{input}}) \cdot \left( \nabla_\theta \log \pi_{\theta_1}(a_{\text{shot}} | s_{\text{shot}}) + \nabla_\theta \log \pi_{\theta_2}(a_{\text{token}} | s_{\text{token}}) \right) $$

5. Results

The figure below shows how CoT-Influx performs on the GSM8K dataset.

Notably, as the figure below shows, LLaMA2-70B paired with CoT-Influx surpasses larger LLMs without any fine-tuning at all. For instance, LLaMA2-70B improves on GPT-3.5 by 2.5%.

6. Findings

1. More capable LLMs prefer harder CoT exemplars, whereas smaller LLMs go for simpler ones.
2. Numbers and formatting tokens are essential for mathematical reasoning. Function words such as with, the, and then, along with background context irrelevant to the reasoning — theater, for example — can be pruned away without hurting reasoning.

7. Commentary

1. Although the idea is a natural one, the paper is still fairly inspiring, particularly in its architectural design and optimization.

2. The paper tells its story extremely well, laying out the background and motivation in a way that keeps you engaged — the Pilot Study section is a good example. It makes something simple sound profound and yet remains easy to follow.

3. There is a lot to learn from its experiments, such as the ablations and the many small studies; the logic is rigorous throughout.

4. Its practicality is open to debate. For one thing, a reinforcement learning approach carries real cost. For another, I read this as an end-to-end design, so switching datasets would mean retraining. And perhaps what the trained network ends up deleting is exactly words like with, the, and then — if it really is that rigid, a vocabulary list would do just as well. That last point is, of course, only speculation.