Reasoning Language Models: A Blueprint

🔼 This figure shows the relationship between different types of language models and the three key components that enable reasoning language models (RLMs): Large Language Models (LLMs), Reinforcement Learning (RL), and High-Performance Computing (HPC). The left side illustrates how LLMs, which are adept at generating human-like text but lack structured reasoning, and RL agents, which can find optimal solutions but lack world knowledge, are combined with the power of HPC to create RLMs. The right side presents a hierarchy showing how LLMs are a subset of language models, with RLMs being a more advanced type of LLM that incorporates reasoning capabilities.

read the caption

Figure 3: Hierarchy of language models (right) and the three pillars of RLMs (left).

🔼 Figure 4 illustrates the architecture of a Reasoning Language Model (RLM). It starts with a high-level overview showing the three main pipelines: inference, training, and data generation. The medium-level overview zooms in on the inference and training processes, distinguishing between implicit and explicit RLMs. The bottom part provides a very detailed view of the processes, including the reasoning scheme (how the reasoning process evolves), operators, models, and various training data types (output-based, process-based, and trace-based supervision). The figure details the flow of data and how the different components interact to achieve reasoning. Appendices C and D provide further details on the algorithms and data generation.

read the caption

Figure 4: Overview of a general RLM design and core concepts. We provide a high-level overview (the top-left part), a more detailed medium-level overview (the top-right part), and a very detailed diagram showing the inference and training pipelines (the bottom part). A detailed specification of the inference pipeline can be found in Appendix C.1 and in Algorithm 1. Pipelines for different training phases and paradigms can be found in Appendix C.2–Appendix C.4, and in Algorithms 2–7. The data generation pipeline is detailed in Appendix D.

🔼 Figure 5 presents a comprehensive blueprint for building reasoning Language Models (RLMs). It’s structured around four key component categories: Reasoning Schemes (defining the overall reasoning strategy and structure), Operators (actions that modify the reasoning process), Models (neural networks used to implement operators), and Pipelines (workflows for inference, training, and data generation). The figure visually organizes these components, highlighting their interrelationships and providing references to detailed descriptions and algorithms found in the paper’s appendices.

read the caption

Figure 5: A blueprint for reasoning LMs. It consists of four main toolboxes: the reasoning scheme (the top part), operators (the bottom-left part), and models (the bottom-right part); pipelines are mentioned in the center and detailed in Appendix C.1 and in Algorithm 1 (the inference pipeline), Appendix C.2–Appendix C.4, and in Algorithms 2–7 (the training pipelines), and in Appendix D (the data generation pipeline).

🔼 The figure illustrates the x1 framework’s two-phase training process. Phase 1 initializes the models (policy and value models) using supervised fine-tuning. Phase 2 iteratively refines these models using reinforcement learning. This involves alternating between generating many Monte Carlo Tree Search (MCTS) trees and training the models using data derived from the MCTS process. This iterative refinement helps improve both model accuracy and efficiency.

read the caption

Figure 6: An overview of the x1 framework is presented, highlighting its two-phase training process. In phase 1, the models are initialized, while in phase 2, the models are iteratively refined by alternating between constructing a sufficient number of MCTS trees and training the models on data derived from these trees.

🔼 This figure displays an example of model output, highlighting the probabilities associated with individual tokens. The color-coding helps visualize the model’s confidence in its predictions: purple indicates low confidence (highest probability below 0.8), blue signifies high confidence with a potential for alternative predictions (second highest probability above 0.1), and red represents uncertainty (both conditions are met). This exemplifies how the model’s certainty is reflected in its token generation process.

read the caption

Figure 7: Example model output with highlighted tokens. The output has been colored in the following way: Purple if the highest probability is below 0.8 (not so certain but no contention), blue if the second highest probability is above 0.1 (very certain, but maybe another one) and red if both are true (uncertain).

🔼 This figure visualizes the token probability distributions generated by a language model during the reasoning process. For each token in the model’s output, it displays the probability of the two most likely tokens and the combined probability of all other less-likely tokens. This allows for a detailed analysis of the model’s certainty and uncertainty in each step, revealing potential areas of confidence and ambiguity in its reasoning.

read the caption

Figure 8: Token probabilities of example model outputs We show the two highest probabilities as well as the sum of the other probabilities.

🔼 This figure visualizes the uncertainty of token probabilities in a model’s output sequence, using four metrics: variance, entropy, VarEntropy, and the Gini coefficient. These metrics help analyze the model’s confidence in its predictions and can reveal where the model is most uncertain. The x-axis represents the position in the token sequence, and the y-axis shows the values of each metric. High values for variance, VarEntropy, and the Gini coefficient indicate high uncertainty, suggesting areas where the model may be making less informed predictions.

read the caption

Figure 9: Uncertainty metrics (variance, entropy, VarEntropy, and the Gini coefficient) plotted against the output token sequence.

🔼 This figure displays the relationship between the size of a question set and the length of the 95% confidence interval for the accuracy of large language models. The data was generated by sampling responses from a subset of 1000 questions, generating 8 answers for each question at a temperature of 1. The confidence interval was calculated across 32 randomly sampled sets with replacement for each question set size. The chart visually demonstrates how increasing the number of questions reduces the uncertainty in the model’s accuracy.

read the caption

Figure 10: Estimated 95%-confidence interval length for different question set sizes using sampled generated answers from a subset of 1000 questions with eight generated answers per question at temperature 1. The confidence interval is calculated over the eight different pass@1 subsets of each question with 32 sets randomly sampled with replacement for each set size.

🔼 This figure illustrates the different types of reward models used in reinforcement learning for reasoning tasks, specifically in the context of Monte Carlo Tree Search (MCTS). It compares Outcome-Based Reward Models (ORMs), which only consider the final outcome, with Process-Based Reward Models (PRMs), which evaluate intermediate steps. The figure also introduces Outcome-Driven Process-Based Reward Models (O-PRMs), a hybrid approach that combines elements of both ORMs and PRMs. Gray nodes in the diagrams represent terminal nodes (i.e., end of the reasoning process). The visual representation helps clarify the differences in how these reward models evaluate the quality of reasoning paths. This visualization helps to understand why PRMs are usually preferred, as they provide richer feedback signals during the training process.

read the caption

Figure 11: Comparison of Outcome vs. Process-Based label generation, and the introduction of Outcome-Driven Process Based Reward Models (O-PRMs). Gray nodes mark terminal nodes.

🔼 This figure illustrates the differences between reward models, V-value models, and Q-value models, specifically in scenarios with sparse rewards (only terminal states receive non-zero rewards). The reward model only provides information at the terminal state, indicating whether the solution is correct or incorrect. V-value models estimate a global value for each state, while Q-value models provide a more granular evaluation, estimating a value for each state-action pair. This more detailed information is especially useful for search algorithms like Monte Carlo Tree Search (MCTS) which make decisions based on the values of individual actions. Gray nodes in the figure represent terminal states.

read the caption

Figure 12: Comparison of reward, v-value and q-value models in a sparse reward setting (only terminal states receive non-zero rewards). Gray nodes mark terminal nodes. The reward model should predict the rewards for transitioning from one state to another which is 0 for non-terminal states and not providing information. V-VMs and Q-VMs however, predict a global value and are therefore informative for non-terminal states.

🔼 This figure illustrates an example of a tree structure generated by the Monte Carlo Tree Search (MCTS) algorithm used in the paper’s reasoning language model (RLM). Each node represents a reasoning step in the process of solving a problem, and the edges represent the transitions between steps. The tree expands by exploring multiple possible paths, guided by a policy model that predicts the likelihood of each next step and a value model that estimates the overall quality of each path.

read the caption

Figure 13: Example MCTS generated tree of reasoning sequences.

🔼 This figure illustrates the two-phase training pipeline for the Reasoning Language Model (RLM). Phase 1 focuses on supervised fine-tuning using Chain-of-Thought (CoT) and Monte Carlo Tree Search (MCTS) to initialize the policy and value models. Phase 2 uses reinforcement learning with Process-Based Reward Models to refine the models through iterative improvement based on data generated via the MCTS process. The process involves alternating between model training and MCTS simulations, each iteration enhancing model performance and data quality. Implicit and Explicit RLMs are shown to illustrate the generalizability of the pipeline.

read the caption

Figure 14: The two phases of the training pipeline.