Reasoning Planning for Language Models

TL;DR

This work tackles the challenge of selecting the most effective reasoning method for a given query to balance accuracy and computational cost. It introduces EPIC, a contrastive learning framework that jointly learns method embeddings and a question-to-method mapping, guided by theoretical probabilistic bounds on aggregation strategies to regularize learning. EPIC demonstrates consistent improvements in accuracy-cost trade-offs on mathematical reasoning benchmarks, transfers across datasets and model scales, and competitive performance on code-generation tasks, all while reducing inference cost. These results highlight the value of principled method-selection in LLM reasoning and provide a practical pathway for adaptive, resource-aware deployments.

Abstract

Selecting an appropriate reasoning method for a given query remains a key challenge in language model generation. Existing approaches typically generate multiple candidate responses and use an aggregation strategy to select the output answer, often assuming that more candidate answers yield higher accuracy. We revisit this assumption through a rigorous theoretical analysis, deriving accuracy bounds for standard aggregation methods under fixed generation distributions and candidate sizes. Building on these insights, we introduce EPIC, an Ensemble Planning with Contrastive learning framework to learn a shared representation space that captures both model reasoning abilities and query-method compatibility. EPIC incorporates our probability bounds as a regularizer in a utility-driven optimization that balances accuracy and computational cost. Experiments on diverse mathematical reasoning tasks show that EPIC consistently selects optimal reasoning methods, improving accuracy while reducing computational overhead. Our code can be found at https://github.com/nguyenngocbaocmt02/EPIC.

Figures (3)

• Figure 1: Our method employs the regularized representation learning loss \ref{['eq:training']} to learn both the reasoning method representation vectors, denoted as $v_1, \ldots, v_M$, and the question embedding network parameters $\theta$. During inference, we route suitable math questions to the appropriate reasoning method by computing the similarity between the input questions and the learned method representations. Color codes on problem difficulty levels are provided for illustration purposes only.
• Figure 2: Average success rate and token counts on the test set with embedding dimension $d = 64$. Our ensemble planner performances with varying $\lambda \in \{0, 0.25, 0.5, 0.75, 1\}$ are highlighted in red, and individual reasoning models in $\mathcal{M}$ are plotted in blue. The boundary of the ensemble planners covers the individual models in the universe $\mathcal{M}$. The Upper Bound (UB) under $\mathcal{M}$ is the proportion of questions that at least one method in $\mathcal{M}$ could successfully solve.
• Figure 3: PCA visualization of question (grey) and reasoning method (colored crosses) embeddings for three different settings of the utility trade-off $\lambda$. Each method is labeled by index:num, where num is the number of generated answers.

Theorems & Definitions (12)

• Theorem 3.1: Majority voting
• Theorem 3.2: Voting with score sum
• Theorem 3.3: Voting with score maximum
• Proposition A.1: Count upper-bound
• proof : Proof of Proposition \ref{['prop:upper']}
• proof : Proof of Theorem \ref{['thm:majority-voting']}
• Theorem A.2: Sum upper-bound
• proof
• proof : Proof of Theorem \ref{['thm:prm-vote']}
• Proposition A.3: Max PRM upper bound