No Train Still Gain. Unleash Mathematical Reasoning of Large Language Models with Monte Carlo Tree Search Guided by Energy Function (2309.03224v3)

Abstract: LLMs demonstrate impressive language understanding and contextual learning abilities, making them suitable for NLP tasks and complex mathematical reasoning. However, when applied to mathematical reasoning tasks, LLMs often struggle to generate correct reasoning steps and answers despite having high probabilities for the solutions. To overcome this limitation and enhance the mathematical reasoning capabilities of fine-tuned LLMs without additional fine-tuning steps, we propose a method that incorporates Monte Carlo Tree Search (MCTS) and a lightweight energy function to rank decision steps and enable immediate reaction and precise reasoning. Specifically, we re-formulate the fine-tuned LLMs into a Residual-based Energy Model (Residual-EBM) and employ noise contrastive estimation to estimate the energy function's parameters. We then utilize MCTS with the energy function as a path verifier to search the output space and evaluate the reasoning path. Through extensive experiments on two mathematical reasoning benchmarks, GSM8k and AQUA-RAT, we demonstrate the exceptional capabilities of our method, which significantly improves the pass@1 metric of the fine-tuned model without requiring additional fine-tuning or reinforcement learning with human feedback alignment.

Overview of Enhanced Mathematical Reasoning

LLMs have transformed natural language processing, offering advanced contextual learning and language understanding. Despite these advancements, LLMs sometimes struggle with generating accurate reasoning steps and solutions for mathematical tasks, despite seemingly high probabilities for correct answers. A paper presents a strategy to surpass this hurdle, proposing a union of Monte Carlo Tree Search (MCTS) and an energy function to refine decision-making processes, steering LLMs toward precise mathematical reasoning.

Residual Energy-Based Model and MCTS

The paper introduces a revised mechanism that transforms fine-tuned LLMs into what is known as a Residual-based Energy Model (Residual-EBM). This model, equipped with an energy function, acts as a ranking criterion pivotal for the MCTS algorithm, which, in turn, searches for the optimal reasoning path. Extensive testing on two mathematical benchmarks—the GSM8k and AQUA-RAT—showcases that this approach significantly enhances the fine-tuned model's performance without additional training phases, such as reinforcement learning or alignment with human feedback.

Methodology in Detail

The methodology consists of several key steps. It commences with fine-tuning a LLM or employing a pre-existing specifically tailored model. Following this, the paper dives into formulating a Residual EBM, where an energy function is introduced as a means to coerce the model towards a more desired output distribution. The energy function itself is optimized using Noise Contrastive Estimation (NCE), a process benefiting from noise samples generated by the model. This synergy between the Residual EBM and noise generation marks a significant deviation from methodologies requiring elaborate training datasets or expert knowledge.

Efficacious Use of MCTS

MCTS, an algorithm adept at balancing between exploratory and exploitative decision-making, is then employed to decode complex reasoning tasks. Guided by the energy function from the Residual EBM as a heuristics measure, MCTS systematically searches across sentence-based tree nodes—rather than individual words—for the most probable reasoning steps. The performance improvements observed with this approach are compelling, especially when considering the model's ability to surpass the pass@1 accuracy metrics of previously released models without intensive additional fine-tuning.

Concluding Thoughts

The results gleaned from this research are both remarkable and promising, showing a clear path towards improving LLMs' performance on math reasoning tasks. With enhanced model decision-making facilitated by the combined efforts of MCTS and an energy function, LLMs can more accurately navigate the complexities of mathematics. The versatility of the proposed methods—negating the need for task-specific adjustments or extensive model retraining—marks a significant advancement in our tools for unleashing the potential of LLMs for analytical reasoning.