Cut Your Losses! Learning to Prune Paths Early for Efficient Parallel Reasoning

# Cut Your Losses! Learning to Prune Paths Early for Efficient Parallel Reasoning

Source: https://arxiv.org/html/2604.16029

Jiaxi Bi1,3, Tongxu Luo1,2*, Wenyu Du4, Zhengyang Tang1, Benyou Wang1,2

1The Chinese University of Hong Kong, Shenzhen
2Shenzhen Loop Area Institute
3USTB
4DualityRL

[email protected] [email protected] [email protected]

†Equal contribution; alphabetical by last name.
‡Work done during interning at CUHK-Shenzhen.
*Corresponding author.

## Abstract

Parallel reasoning enhances Large Reasoning Models (LRMs) but incurs prohibitive costs due to futile paths caused by early errors. To mitigate this, path pruning at the prefix level is essential, yet existing research remains fragmented without a standardized framework. In this work, we propose the first systematic taxonomy of path pruning, categorizing methods by their signal source (internal vs. external) and learnability (learnable vs. non-learnable). This classification reveals the unexplored potential of learnable internal methods, motivating our proposal of STOP (SuperTOken for Pruning). Extensive evaluations across LRMs ranging from 1.5B to 20B parameters demonstrate that STOP achieves superior effectiveness and efficiency compared to existing baselines. Furthermore, we rigorously validate the scalability of STOP under varying compute budgets—for instance, boosting GPT-OSS-20B accuracy on AIME25 from 84% to nearly 90% under fixed compute budgets. Finally, we distill our findings into formalized empirical guidelines to facilitate optimal real-world deployment. Code, data and models are available at https://bijiaxihh.github.io/STOP.

## 1 Introduction

Parallel reasoning has established itself as a standard paradigm for solving complex problems. The core principle is to sample multiple independent reasoning paths and subsequently aggregate them to derive a robust consensus. However, this accuracy gain comes at a prohibitive cost. Generating dozens or even hundreds of trajectories per query increases computational overhead by orders of magnitude and escalates inference costs to nearly $6 per query.

**Figure 1:** The necessity of pruning early. Early errors often lead to irreversible failure. Pruning these futile paths early not only saves computation but also purifies the candidate set for better consensus.

### Why Prune Early in Parallel Reasoning?

Crucially, recent studies reveal that this extensive computation is largely squandered: not every path contributes to the solution. Many trajectories are flawed from inception, yet they consume equal resources to generate and subsequently pollute the final answer aggregation. As illustrated in Figure 1, once a reasoning path begins with a flawed prefix, the LRM struggles to self-correct, inevitably spiraling into a futile trajectory. Consequently, identifying and terminating these unpromising paths at the **prefix level**—a technique known as **path pruning** (or prefix rejection)—is essential.

### A Unified Taxonomy

While existing methods attempt to filter paths using auxiliary reward models, internal confidence, or semantic redundancy, they lack a standardized evaluation protocol, leading to fragmented research. So first, we propose the first systematic taxonomy of path pruning, classifying methods based on the **source** (internal vs. external) and **learnability** (learnable vs. non-learnable) of their signals (see Figure 2). This taxonomy reveals a significant research gap: the unexplored potential of **learnable internal methods**. Conceptually, learnable internal methods offer unique advantages, as learning enables task-specific accuracy gains, while internal signals provide early, fine-grained indicators of reasoning failure without incurring extra computational overhead. To bridge this gap, we introduce **STOP** (SuperTOken for Pruning), the first efficient instantiation of this paradigm. Extensive evaluations demonstrate that STOP outperforms existing baselines in both effectiveness and efficiency.

### Further Evaluation and Empirical Analysis

Despite the promise of path pruning, its widespread adoption is currently hindered by unverified scalability across varying computational budgets and model sizes; and the absence of empirical guidelines for determining optimal pruning configurations in real-world scenarios. To overcome them, we rigorously validate the utility of path pruning in practical settings. We conduct extensive experiments across diverse model sizes (1.5B to 20B) and compute budgets, confirming that STOP exhibits robust scalability. Moreover, we distill our empirical analysis into actionable guidelines, providing a formalized method to determine the optimal retention ratio for varying resource constraints.

### Contributions

In summary, this work makes four primary contributions:
(1) We present the first systematic investigation and taxonomy of path pruning.
(2) We propose STOP, a novel pruning method based on learnable internal signals.
(3) We provide a comprehensive evaluation demonstrating STOP's superior scalability and effectiveness.
(4) We establish empirical guidelines to support the practical implementation of path pruning.

**Figure 2:** The proposed taxonomy of path pruning.

## 2 A Unified Taxonomy of Path Pruning

### 2.1 Problem Definition

Consider a LRM Θ and an input query x. Parallel reasoning improves accuracy by generating N independent trajectories T = {τi}_{i=1}^N, where τ_i ∼ P_Θ(x), and aggregating them through a consensus strategy, such as majority voting. The final prediction ŷ is typically computed as:

ŷ = vote({τ_i}_{i=1}^N)   (1)

However, generating N complete trajectories incurs a linear computational cost (C ∝ N). To mitigate this cost, path pruning aims to identify and discard unpromising trajectories early in the decoding process.

#### The Path Pruning Formulation

Formally, we define a checkpoint at length L_prefix where the generation is paused. At this stage, the model has produced a set of prefixes P = {p_i}_{i=1}^N. The core of path pruning is a **pruning signal generator** S, which maps each prefix to a scalar score representing its potential correctness:

s_i = S(p_i | x, Θ)   (2)

where s_i ∈ [0,1] denotes the pruning signal. Based on these signals, we retain only the top-k promising paths (where k ≪ N) for full completion, discarding the rest. The final aggregated answer is then derived exclusively from this pruned subset:

ŷ_pruned = vote({finish(p_i) | s_i ∈ {s_j}_{j=1}^k})   (3)

So, the objective of path pruning is to design an S that maximizes ŷ_pruned's accuracy while minimizing the computational cost (the number of generated tokens). Therefore, the design of S dictates the effectiveness of the entire framework.

### 2.2 A Unified Taxonomy of Pruning Signal Generators

**Table 1:** A Unified Taxonomy of Path Pruning Methods. We categorize methods based on the pruning signal source and learnability. Type IV satisfies both Desideratum 1 (Internal) and Desideratum 2 (Learnable).

As defined in Section 2.1, the efficacy of path pruning hinges entirely on the quality of the pruning signal generator S. While the function of S is consistent—scoring prefixes—existing methods differ fundamentally in **how** this signal is produced. To systematically evaluate these approaches, we categorize them based on two critical dimensions: the **source** of the signal (External vs. Internal) and the **learnability** of the generator (Learnable vs. Non-learnable), as summarized in Table 1.

#### Two Desiderata for Signal Generators

Before categorizing specific methods, we establish two desiderata for an ideal signal generator:

**Desideratum 1. Internal Source**

An ideal S should leverage the rich, high-dimensional internal states of the LRM. Internal signals contain fine-grained information about uncertainty and reasoning dynamics that are often lost in the final text output used by external methods.

**Desideratum 2. Learnability**

An ideal S should be trainable to adapt to specific data distributions. Learnable parameters allow the generator to capture complex, non-linear patterns of error that rigid, pre-defined heuristics cannot model.

Based on these axes, we classify existing works into four distinct types.

#### External Signal Source

Methods in this category derive pruning signals from the generated textual output or by querying separate models. They fail to satisfy Desideratum 1.

**Type I. Surface Heuristics**

These methods rely on human-designed rules (e.g. similarity) applied to the surface form of the generated text. While computationally cheap, these heuristics are rigid and blind to the model's actual confidence. To overcome these, the next type introduces learnability into the external evaluation process.

**Type II. External Judges**

These approaches employ a separate, trained model to evaluate the reasoning path. Although they satisfy Desideratum 2, they incur significant computational overhead due to the need for additional model inference and fail to access the LRM's internal certainty. To overcome this rigidity, the next category introduces learnability into the external evaluation process.

#### Internal Signal Source

Methods in this category extract signals directly from the LRM's internal states, accessing richer information (satisfying Desideratum 1).

**Type III. Raw Confidence**

This paradigm utilizes intrinsic metrics directly derived from the decoding process, such as perplexity or token probability. However, these methods rely on fixed definitions of confidence, violating Desideratum 2; raw probability does not always correlate with reasoning correctness.

**Type IV. Learned Intuition**

The final category represents the intersection of both desiderata: a trainable module inserted into the LRM to process internal states. This approach can leverage rich hidden representations (Internal) while adapting to the specific error patterns of the task (Learnable).

## 3 Methodology: Super Token for Pruning

As established in our taxonomy, Type IV represents the ideal pruning paradigm but remains unexplored. In this section, we introduce **STOP** (SuperTOken for Pruning), the first efficient instantiation of this paradigm. We delineate the motivation in Section 3.1, followed by the architectural design and workflow in Section 3.2.

### 3.1 Motivation for Type IV Pruning

As illustrated in Figure 2, prior methods compromise on either information richness or adaptability. Type II suffers from high latency, while Type III lacks the capacity to model complex error patterns. Type IV represents an ideal optimum: it combines the **efficiency** of accessing internal states with the **adaptability** of learnable parameters. However, this type remains unexplored due to the challenge of designing a module that extracts these signals without disrupting the LRM's generative capabilities.

**Figure 3:** The inference process comprises three stages: caching initial prefixes (Launch), scoring them via the STOP module (Check), and completing only the top-ranked candidates (Resume).

### 3.2 Instantiation of Type IV Pruning: STOP

To instantiate this type, we design STOP as a lightweight, non-invasive module that integrates seamlessly with the backbone LRM.

#### Components

We augment the fixed LRM Θ with three learnable components:
(1) A **Super Token** ([STOP]) added to the vocabulary, acting as a specialized query vector to aggregate information;
(2) A **Critique Adapter LoRA** (θ_LoRA), activated only when processing the [STOP] token to extract error-specific features without altering the LRM's general reasoning capabilities;
(3) A **Classification Head** (W_cls), which projects the hidden state of the [STOP] token to a scalar probability.

This design ensures **modularity**: the original parameters Θ remain frozen, preserving the LRM's generative capability while enabling efficient parameter-efficient fine-tuning (PEFT).

#### Training: Learn to Use Internal Information

The goal of training is simple: teach the model to distinguish promising prefixes from futile ones. Formally, for a prefix p_i, we derive a soft label s_i^{mc} ∈ [0,1] via Monte Carlo estimation (details in Appendix B). The training process involves two steps: First, we compute the KV cache of the prefix using the frozen LRM: C_{p_i} = LRM(p_i; Θ). Second, we append a sequence of learnable [STOP] tokens, denoted as T_s, and process them using the LoRA-augmented model. The final hidden state h_i is fed into the classifier to minimize the soft binary cross-entropy loss:

ℒ = −[s_i^{mc} log σ(W_{cls}h_i) + (1 − s_i^{mc}) log(1 − σ(W_{cls}h_i))]   (4)

where h_i = LRM(T_s | C_{p_i}; Θ, θ_LoRA)_{-1}.

#### Training Cost

Constructing the MC supervision requires sampling multiple continuations per prefix to estimate s_i^{mc} (e.g., K = 32), which introduces an upfront computational cost during data construction. However, this cost is incurred only once, and the