阅读下面的内容，用python来完全复现论文内容:
ABSTRACT This paper presents an advanced mathematical problem-solving framework, LLaMA-Berry, for enhancing the mathematical reasoning ability of Large Language Models (LLMs). The framework combines Monte Carlo Tree Search (MCTS) with iterative Self-Refine to optimize the reasoning path and utilizes a pairwise reward model to evaluate different paths globally. By leveraging the self-critic and rewriting capabilities of LLMs, Self-Refine applied to MCTS (SR-MCTS) overcomes the inefficiencies and limitations of conventional stepwise and greedy search algorithms by fostering a more efficient exploration of solution spaces. Pairwise Preference Reward Model (PPRM), inspired by Reinforcement Learning from Human Feedback (RLHF), is then used to model pairwise preferences between solutions, utilizing an Enhanced Borda Count (EBC) method to synthesize these preferences into a global ranking score to find better answers. This approach addresses the challenges of scoring variability and nonindependent distributions in mathematical reasoning tasks. The framework has been tested on general and advanced benchmarks, showing superior performance in terms of search efficiency and problem-solving capability compared to existing methods like ToT and rStar, particularly in complex Olympiad-level benchmarks, including AIME24 and AMC23. 1 INTRODUCTION Mathematical reasoning represents a great challenge in artificial intelligence, with broad applications across automated theorem proving, mathematical problem solving, and scientific discovery (Ahn et al., 2024). Recently, significant strides have been made by Large Language Models (LLMs) like GPT-4 (Achiam et al., 2023) in general mathematical tasks involving arithmetic or geometric problem-solving (Cobbe et al., 2021; Sun et al., 2024; Ying et al., 2024). However, complex mathematical reasoning remains challenging, especially at the Olympiad-level benchmarks such as AIME (MMA). One intuitive way is to break the complex problem-solving solutions into step-bystep thoughts, known as reasoning paths, as shown in Chain-of-Thought (CoT) (Wei et al., 2022). Although prompt-based methods can directly construct such reasoning paths, they remain limited in efficiency and accuracy. As a result, researchers have begun exploring search-based optimization algorithms to enhance the generation of effective reasoning paths. One line of research focuses on generating solutions based on complete reasoning paths. Tree search techniques with backtracking capabilities are proposed to address the inefficiencies of naive methods, such as Tree-of-Thought (ToT) (Yao et al., 2024) and Breadth-First Search (BFS) based on greedy strategies (Zhang et al., 2024). Heuristic tree search methods like RAP (Hao et al., 2023) and Q* (Wang et al., 2024) treat each step in a reasoning path as a state and the generation of steps as actions, framing the search for reasoning paths as a Markov Decision Process (MDP) problem.
Monte Carlo Tree Search (MCTS) methods have further improved the efficiency of path searches in reasoning paths. Yet, the generation of steps in reasoning paths faces challenges due to the infinite action space of thought steps, particularly under limited search conditions, which constrains the performance of step-wise search methods. Unlike methods that generate a single step simultaneously, another approach considers the entire CoT solution as an independent state with iterative optimization actions. This method, known as rewriting, includes techniques such as reflection and self-refine, where optimization of a solution depends on its previous contents and typically involves only local corrections, ensuring that the new solution remains within the “trust region” of the original. This strategy offers the iterative optimization methods a proximal optimization characteristic, making the action space more compact. However, this continuous and greedy iterative approach also suffers from low search efficiency and the risk of local optima, limiting the performance ceiling of rewriting methods. Another line of research focuses on evaluating the quality of reasoning paths, which includes outcome and process reward models (Uesato et al., 2022). The former focuses on whether the final answer of a reasoning path is correct, while the latter emphasizes the correctness of each step and the logical connections between steps. Although both approaches enable models to score using a scalar, labelled data for training these reward models are scarce. Furthermore, the scoring standards in mathematical reasoning tasks are highly variable, and the reward distribution for different problems is typically not i.i.d. (independently and identically distributed), with each problem presenting unique scoring characteristics. This non-independence complicates the scaling up of single-scalar score-based reward models and their abilities to detail local preference relations between answers. To address the above challenges, we design the LLaMA-Berry framework (in Figure 1) based on LLaMA-3.1-8B Instruct (Meta, 2024b), which includes two novel methods called Self-Refine applied to Monte Carlo Tree Search (SR-MCTS) and Pairwise Preference Reward Model (PPRM). Specifically, SR-MCTS is a novel Markov Decision Process (MDP) framework focusing on a pathwise problem-solving search where a complete solution is defined as an independent state, and selfrefine is considered an optimization action. Then, Monte Carlo Tree Search (MCTS) is leveraged for the exploration of optimal solutions with guiding of PPRM evaluations. PPRM is inspired by Reinforcement Learning from Human Feedback (RLHF) (Christiano et al., 2017) techniques; we have transformed the task of reward models into an RL problem that models preferences between solutions, then aggregating local preferences into a global ranking score with an Enhanced Borda Count (EBC) method. Finally, evaluations on general benchmarks such as GSM8K (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021) have demonstrated that our method outperforms baseline approaches like ToT (Yao et al., 2024) and rStar (Qi et al., 2024) in both search efficiency and performance. Moreover, in challenging Olympiad-level benchmarks such as AIME24 and AMC23 and in the Graduate-level GPQA Diamond (Rein et al., 2023) benchmark, our approach has elevated the performance of LLaMA-3.1-8B without any further training to a level comparable to proprietary solutions including GPT-4 Turbo. These encouraging results suggest that our method can effectively improve the LLM’s reasoning ability with only a small amount of data, and this capability can be generalized to other fields, such as physics and chemistry. Furthermore, it can enable multimodal reasoning, allowing the model to achieve even greater reasoning power.
To summarize, our contributions in LLaMA-Berry framework integrate three essential components: • SR-MCTS establishes a novel MDP framework that treats complete path-wise solutions as independent states and utilizes Self-Refine as an optimization action, balancing efficiency and performance in search. • Inspired by Reinforcement Learning from Human Feedback (RLHF), PPRM leverages the preference relationship between solutions to evaluate their quality, which avoids the volatility of absolute scores while providing a more guided exploration of optimal paths. • Finally, EBC method aggregates the local preferences into a global ranking score, allowing the model to identify the overall patterns and enhance the decision-making process.
Selection Using UCT Expansion with Self-Refine Evaluation Backpropagation Figure 1: The main pipeline of LLaMA-Berry, where Si stand for problem-solving solutions and Ci stands for critiques. This Pipeline consists of four phases detailed in 2.2, including Selection, Expansion, Evaluation and Backpropagation.
2 METHODOLOGY 2.1 PRELIMINARY In the realm of mathematical problem-solving, we face significant challenges that require systematic approaches to derive accurate and efficient solutions. A core issue in this field is generating and optimizing reasoning paths, which can be formalized as a path-wise Markov Decision Process (MDP) framework. This framework allows us to define the state space S, where each state s represents a complete solution to a given problem, and the action space A consists of all feasible rewriting actions a that make transitions between states. The central issues we aim to address include the development of the value function Q(s, a), which quantifies the expected rewards from executing action a at state s, defined as: Q(s, a) = E[R(s ′ )|s ′ = T(s, a)] (1) Where T(s, a) indicates the transition to another solution s ′ from solution s with rewriting operation denoted as a. Our primary objective is to identify the optimal state s ∗ , which represents the best solution to the problem. This optimal state can be achieved through the selection of actions that maximize the value function, guiding us toward the most desirable outcomes: s ∗ = arg max s ′∈S E[Q(s)] (2) 2.2 SELF-REFINE APPLIED TO MONTE CARLO TREE SEARCH Monte Carlo Tree Search (MCTS) is an effective reinforcement learning method within the Markov Decision Processes (MDP) framework, employing states, actions, and cost functions through sampling. The algorithm follows four key steps: selection, expansion, simulation/evaluation, and backpropagation. In the selection phase, the root node is expanded using the Upper Confidence Bound applied to Trees (UCT) algorithm, which selects a node s by balancing exploration and exploitation, that is, a = arg max a∈A(s) Q(s, a) + c · s ln N(s) N(s, a) ! , (3) where N(s) is the visitation count of node s, N(s, a) is the action frequency, and c is a parameter controlling exploration. In the expansion phase, node s generates subsequent states s ′ ,
Figure 2: (Left) Self-Refine Process and (Right) Preference Prediction process of PPRM new nodes in tree T . The evaluation phase typically uses simulations or heuristics to estimate the Q-values for these nodes. Finally, during backpropagation, the estimated Q-values are updated retroactively from the leaf nodes to the root. This iterative process allows MCTS to explore and exploit the state space effectively, refining decision-making by balancing the exploration of new options with the exploitation of known high-value paths. As shown in Figure 1, the SR-MCTS method combines Monte Carlo Tree Search (MCTS) with Self-Refine, which optimizes the solution search via a structured tree. This method continuously evaluates and improves solutions by combining the iterative nature of MCTS with the self-critique and rewriting capabilities of LLMs as following sections: Selection Phase: The selection phase identifies a node si from the search tree T for expansion, where each node represents a complete solution state. The Upper Confidence Bound applied to Trees (UCT) algorithm is employed to select the optimal node, with dynamic pruning used to avoid local optima. A node si is considered fully explored when its child nodes reach a predefined limit, and at least one child node’s Q value exceeds the Q value of si . Expansion Phase: In the expansion phase, as shown in Figure 2, the selected answer si is expanded by generating successor answers through a Self-Refine process, which includes a Critizing and Rewriting process. The Critizing generates a critique ci = C(si) that identifies drawbacks (e.g., mathematical wrongs or logical fault) in the current chosen answer Si , and then Rewriting process generates a new answer si+1 = R(si , ci). In practice, to simplify the problem, we assume this process is deterministic, ensuring that the same original state of solutions si consistently produces the same successor state of solution si+1. The new state of solution s ′ is then added to the search tree T as a new node. Evaluation Phase: The evaluation phase calculates the value Q(s ′ ) of the newly generated node s ′ using the Pairwise Preference Reward Model (PPRM). The evaluation involves two steps: global and local value assessments. The global value Qg(s ′ ) is determined by the quantile of s ′ in a winloss matrix M, which reflects the win-loss relationships between nodes. The local value Ql(s ′ ) is derived from comparisons with adjacent nodes in the search tree T . The total value Q(s ′ ) is then computed as a weighted combination of global and local values: Q(s ′ ) = αQg(s ′ ) + (1−α)Ql(s ′ ), where α controls the relative influence of each component. Backpropagation Phase: In the backpropagation phase, the value Q(s ′ ) of the new node is propagated back to its parent node si , updating Q value of si as a discounted sum of its child nodes’ Q values: Q(si) = (1 − γ)Q(si) + γQ(s ′ ). The discount factor γ represents the importance of future rewards. This iterative update mechanism ensures that the values of parent nodes are progressively refined, enhancing the guidance for future selections. Additionally, to control the growth of the search tree, the SR-MCTS method restricts the maximum number of rollouts Nmax . The search process terminates when the restriction is reached, imposing limits on the unbounded expansion of the tree. The overarching objective of SR-MCTS is to maximize the expected highest Q value of all existing nodes S, guiding us towards the most desirable outcomes s ∗ , ensuring that the search method efficiently converges to high-quality solutions.
2.3 PAIRWISE PREFERENCE REWARD MODEL In mathematical problem-solving tasks, reliable evaluation of solutions is crucial, as it leads to better estimation of Q-values, thereby offering improved guidance. Existing reward models typically evaluate solutions by giving absolute scores, such as (PRM) (Lightman et al., 2023) and outcome reward mode (ORM) (Yu et al., 2023a). However, the score-based reward models may fall short in leveraging the instruction following capability of LLMs or effectively handling the variations in data distributions and exhibit volatility, especially when the differences between different solutions are subtle. To address this issue, we propose the Pairwise Preference Reward Model (PPRM), which leverages a comprehensive preference dataset incorporating substantial samples from both approaches to learn preference relationships among mathematical solutions. Please refer to Appendix for details about training and inference of PPRM. For two answers to a given mathematical problem, denoted as a1 ≻ a2|ϕ where a1 and a2 denote the preferred and dispreferred completion, PPRM predicts their relative quality using a pairwise partial ordering, represented by the following probability formula: P(a1 ≻ a2 | ϕ) = e ϕ(a1) e ϕ(a1) + e ϕ(a2) (4) Where ϕ(a1 ≻ a2) denotes the feature representation of a partial ordering relation between solution a1 and a2, with ϕ representing the parameters of the model. In our method, a1 ≻ a2 are represented by tokens of the LLM, and P(a1 ≻ a2 | ϕ) is estimated using the logits value of tokens calculated by the LLM. Then, inspired by advancements in Language Interface Fine-Tuning (LIFT) (Dinh et al., 2022), we frame the training process of PPRM as a question-answering task. The model is tasked with answering the question, ”For Question Q, is solution a1 better than solution a2?” as shown in Figure 2. To form a robust training objective, the predicted token labels yˆ (’Yes’ or ’No’) are evaluated using the indicator function I: I(ˆy, y) =  1, if yˆ = y 0, if yˆ ̸= y (5) Finally, a pairwise preference dataset D that contains millions of mathematical problem-solving solution pairs is converted into a dataset D′ suitable for a question-answering task. We employ RLHF techniques to train the model to improve its performance in the partial-order prediction questionanswering task. Subsequently, the Direct Preference Optimization (DPO) (Rafailov et al., 2024) method is utilized to find the optimal Pϕ by maximizing the objective arg maxϕ EP [I(ˆy, y)].
2.4 ENHANCED BORDA COUNT METHOD Although PPRM allows us to directly compare the quality of two answers, we still need to convert these local partial order relations into a cohesive global ranking to gain a comprehensive evaluation for the answers. This conversion process can be formalized as the global optimal ranking aggregation (GORA) problem. To address this problem, we propose the Enhanced Borda Count (EBC) method based on a transitivity assumption of mathematical problem-solving solutions (detailed in Figure 3), which integrates the naive Borda Count Algorithm with a transitive closure of partial order relations calculated by the Floyd-Warshall (Cormen et al., 2009) Algorithm. For formalized proof, please refer to Appendix E. Local Preference Calculation: First, as shown in Figure 3 (1), the PPRM generates a win-loss matrix M ∈ R n×n for all n problem-solving solutions, where M[i, j] = 1 indicates that solution ai is superior to solution aj , and M[i, j] = 0 otherwise. This process can be represented as: M[i, j] =  1, if P(ai ≻ aj ) ≥ 0.5 0, if P(ai ≻ aj ) < 0.5
This matrix can be viewed as an adjacency matrix of a directed graph G = (V, E), where each solution ai corresponds to a vertex vi , and an edge e = (vi , vj ) exists if M[i, j] = 1, indicating that solution ai is preferred over aj . Transitive Closure: As shown in Figure 3 (2), to simplify the problem, we adopt the assumption of transitivity for mathematical solutions, that is, if vi ≻ vk and vk ≻ vj , then vi ≻ vj . Under this assumption, the transitive closure C of a preference matrix can be computed. Specifically, if M[i, k] = 1 and M[k, j] = 1, then M[i, j] = 1. This is achieved by updating the transitive closure matrix C using Floyd-Warshall algorithm with dynamic programming rule: C[i, j] = _n k=1 (C[i, k] ∧ C[k, j]) (7) where W denotes the logical OR operator and ∧ denotes the logical AND operator. Borda Count-Based Global Ranking: Next, based on the transitive closure matrix C, we apply the Borda Count method for global ranking. The Borda Count determines the ranking of each node by calculating its out-degree, which corresponds to the number of nodes it defeats. For each node vi , the Borda Count Borda(vi) is defined as Pn j=1 C[i, j], like the ranked node list in Figure 3 (3). Nodes with higher Borda counts are ranked higher. However, in practice, cyclic preferences can cause efficiency issues with the naive Borda Count method. To further refine the ranking, as shown in Figure 3 (4), a re-ranking stage is introduced, where the logits generated by the PPRM are used for soft comparison among nodes with equal Borda counts. Specifically, for two nodes vi and vj with equal Borda counts, the soft comparison rule can be denoted as vi ≻ vj ⇐⇒ P(ai ≻ aj ) > P(aj ≻ ai). This process ensures that even in the presence of cycles or local ambiguities, the ranking remains consistent and reasonable. Global Quantile Score of Solutions: Finally, as shown in Figure 3 (5), the ranking is converted into a global quantile quality score Qg for each solution v is Qg[v] = 1 − rank(v)−1 |V |−1 , where rank(v) is the position of v in the ranking based on Borda counts, and |V | is the total number of nodes. To reflect local advantages in the search tree structure, the local win rate Ql(x) for a node v is calculated with its children nodes Children[v] as follows: Ql [v] = P u∈Children[v] C[u, v] |Children[v]| (8)
The final score Q(x) for a solution is a weighted sum of local win rate Ql(x) and global ranking quantification Qg, computed as: Q(x) = α · Ql(x) + (1 − α) · Qg(x) (9) where α controls the balance between local win probabilities and global ranking. The softened re-ranking process ensures that even in cyclic preferences, the local win probabilities and global ranking results are consistent, ultimately producing a reasonable and reliable ranking.