## Diagram: MCTS* and LLM Self-Training Workflow ### Overview The image depicts a technical workflow for an iterative LLM self-training process utilizing Monte Carlo Tree Search (MCTS*). The diagram illustrates how a prompt ($Q_1$) is processed through a search tree to generate trajectories (traces). These traces are then evaluated and used to train two distinct components: a Process Reward Model (PRM) and a Policy Model. The process is iterative, indicated by the $j$-th iteration notation. ### Components/Axes **1. Left Column: Input and Definitions** * **Input Data:** A blue cylinder labeled $\{(Q_i, y_i)_{i=1}^M\}$ representing a set of prompts and their corresponding correct answers. * **Sampling Equations:** * $s_1 \sim \pi(\cdot | Q_1)$ * $s_2 \sim \pi(\cdot | Q_1, s_1)$ * $s_3 \sim \pi(\cdot | Q_1, s_1, s_2)$ * ... (ellipsis indicating continuation) * **Variable Definitions:** * $M$: number of prompts * $N$: number of samples * $j$: $j$-th iteration * $T$: length of a trace **2. Center: MCTS* Tree** * **Structure:** A tree diagram starting from $Q_1$, branching into states ($s$) and actions ($A$). * **Pathing:** * A primary path is highlighted with red arrows: $s_1 \rightarrow s_2 \rightarrow s_{3,1} \rightarrow s_{4,1} \rightarrow A_1$. * Each step in the red path is annotated with a weight $w = 0.25$. * Nodes are annotated with values ($v$): * $s_1: v = 0$ * $s_2: v = 0.25$ * $s_{3,1}: v = 0.5$ * $s_{4,1}: v = 0.75$ * $A_1: v = 1$ * **Outcomes:** * $A_1$: Marked with a green checkmark, labeled $A_1 = y_1$ (Correct). * $A_2$: Marked with a red 'X', labeled $A_2 \neq y_1$ (Incorrect). * $v_5^+$: Associated with the successful $A_1$ branch. * $v_5^-$: Associated with the unsuccessful $A_2$ branch. **3. Right Side: LLM Self-Training Components** * **PRM (Process Reward Model):** A box receiving data from the MCTS tree. * **Policy Model:** A box receiving data from the MCTS tree. * **Mathematical Objective:** $\mathbb{E}_{Q \in D_{S_0}} [\mathbb{E}_{s \in \pi(s|Q)} [r(Q, s)]] - \mathbb{E}_{(Q,s) \in D_{S_0}} [\sum_{t=1}^T \log \pi(s_t|s_{<t}, Q)]$ * **Data Storage (Blue Cylinders):** * **PRM Dataset:** Contains $(Q_i, s_{<t}, v_i^-)$ and $(Q_i, s_{<t}, v_i^+)$. * **Policy Model Dataset:** Contains $(Q_1, \tau_1^+), (Q_2, \tau_2^+), \dots$ * **Feedback Loop:** A line at the bottom labeled $\tau = (s_1, s_2, \dots, A)$ connects the MCTS output back to the training datasets. ### Detailed Analysis * **MCTS* Tree Logic:** The tree represents a search process where the model explores different reasoning paths. The red path represents a successful trajectory where the model reaches the correct answer ($A_1 = y_1$). The values ($v$) increase as the model gets closer to the correct answer, suggesting the PRM is learning to assign higher rewards to states that lead to the correct solution. * **Data Flow:** * The MCTS tree generates traces ($\tau$). * Successful traces ($\tau^+$) are fed into the Policy Model dataset. * State-value pairs (both positive $v_i^+$ and negative $v_i^-$) are fed into the PRM dataset. * **Training Objectives:** * The PRM is trained to predict the value $v$ of intermediate states. * The Policy Model is trained to maximize the expected reward (the first term of the equation) while minimizing the log-probability of the traces (the second term, likely a KL-divergence or entropy regularization term). ### Key Observations * **Asymmetric Outcomes:** The diagram explicitly contrasts a successful path ($A_1$) with an unsuccessful path ($A_2$), which is critical for training the PRM to distinguish between good and bad reasoning steps. * **Iterative Nature:** The inclusion of $j$ in $V_\theta^j$ and $\pi_\phi^j$ indicates that the PRM and Policy Model are updated iteratively. The MCTS uses the current versions of these models to generate new data, which then improves the models in the next iteration. * **Process Supervision:** By training on $(Q_i, s_{<t}, v_i)$, the PRM is performing "process supervision," evaluating the quality of the reasoning steps rather than just the final answer. ### Interpretation This diagram illustrates a sophisticated Reinforcement Learning from AI Feedback (RLAIF) pipeline. The core logic is that the MCTS acts as a "reasoning engine" that explores the state space of possible answers. Because the model has access to the ground truth ($y_i$), it can label the final actions as correct or incorrect. The innovation here is the dual-training approach: 1. **The PRM** learns to assign values to intermediate states ($s_{<t}$), effectively learning a "value function" for reasoning steps. This helps the model identify *why* a path is good or bad before it reaches the final answer. 2. **The Policy Model** learns to generate the actual sequences of states and actions ($\tau$) that lead to success. By feeding the results of the MCTS back into the training datasets, the system creates a self-improving loop where the models become better at reasoning (Policy Model) and better at evaluating that reasoning (PRM) over successive iterations ($j$).