--- title: 'TreePS-RAG: Tree Supervision in RAG' url: https://www.emergentmind.com/topics/treeps-rag type: topic --- # TreePS-RAG: Tree Supervision in RAG TreePS-RAG is a reinforcement learning (RL) framework for agentic retrieval-augmented generation (RAG) that introduces tree-based process supervision. It models the entire multi-step reasoning and retrieval process as a rollout tree, enabling step-wise credit assignment and process-level supervision without intermediate annotations. TreePS-RAG achieves consistent improvements on multi-hop and general question answering (QA) tasks compared to outcome-supervised and leading process-supervised RL methods [2601.06922]. ## 1. Formalization of Agentic RAG as a Rollout Tree In TreePS-RAG, the agentic RAG paradigm is formulated using the ReAct framework, in which an LLM alternates between reasoning and issuing information retrieval (IR) actions, or emitting an answer to terminate the episode. Each RL episode consists of a sequence: - The state at step $i$: $s_i = [q,~(r_1, a_1, o_1), \ldots, (r_{i-1}, a_{i-1}, o_{i-1})]$, where $q$ is the user question, $r_j$ is the reasoning string, $a_j \in \{\textrm{search},~\textrm{answer}\}$, and $o_j$ is either search results or null. - The action space: $A = \{\textrm{search},~\textrm{answer}\}$. - The policy samples: $(r_i, a_i) \sim \pi_\theta(\cdot | s_i)$. TreePS-RAG represents possible agentic trajectories as a tree $T = (N, E)$: - Each node $n_i$ corresponds to $(r_i, a_i, o_i)$. - The root $n_{\mathrm{root}}$ is $[q]$. - Edges represent the transition to children by available actions. - Answering nodes are leaves; nodes with $a_i = \textrm{search}$ can be expanded. A root-to-leaf path denotes a complete trajectory. ## 2. Monte Carlo Process-Level Credit Assignment Standard RL for agentic RAG applies only sparse reward at the root—matching the model's final answer to ground truth. TreePS-RAG propagates this outcome reward back to all process steps using Monte Carlo estimation over the tree: - For leaf trajectory $y$, the reward is $R(y) = \textrm{EM}(a_\textrm{pred}, a_\textrm{gold}) \in \{0,1\}$ with $\textrm{EM}$ denoting exact match. - For non-leaf node $n$, let $L(n)$ be its descendant leaves. Estimated value: \[ V(n) = \frac{1}{|L(n)|} \sum_{\ell \in L(n)} R(\ell). \] - Step-wise process advantage: \[ A_{\mathrm{global}}(n) = V(n) - V(n_\mathrm{root}), ~~ A_{\mathrm{local}}(n) = V(n) - V(\mathrm{parent}(n)) \] \[ A(n) = \frac{1}{\sqrt{|L(n)|}} \left[2V(n) - V(n_\mathrm{root}) - V(\mathrm{parent}(n)) \right]. \] $A(n)$ is used as a dense RL reward for $(r_i, a_i)$ at each step, enabling more informative credit assignment than outcome-only RL. ## 3. Efficient Online Tree Construction and Pruning TreePS-RAG presents an online, generate-then-prune algorithm for tractable rollout tree construction: - Given constraints: budget $N$ (number of root-to-leaf rollouts), depth $D$, and per-node retention $N_\text{retain}$. - At each depth $d$, for each node, sample $B_d = \lceil N / n_{d-1} \rceil$ children from $\pi_\theta$. - For each "search" action, retrieve candidate passages; compute Jaccard similarity over passage sets for child nodes; cluster by hierarchical clustering with $1-$Jaccard as distance. - Retain one representative per cluster—this preserves branch diversity and keeps total node count $O(ND)$. **Algorithm 1: Online Tree Construction (abridged)** ```python for d = 1, ..., D: for parent in M(d-1): candidates = [sample child from πθ] C_search, C_ans = partition(candidates) # Add answer children as leaves # For search-children: retrieve passages, cluster, prune clusters = hierarchical_cluster(C_search, distance=1 - Jaccard) kept = [pick one per cluster] M(d).extend(kept) ``` This approach ensures computational cost remains comparable to $N$ flat rollouts, but with greater exploratory coverage. ## 4. RL Objective and Training Integration With the constructed tree, $N$ root-to-leaf trajectories are sampled. The step-level process advantage $A(n)$ is applied to each token of the corresponding $(r_i, a_i)$ generation, masking observations. Training follows a PPO-style policy-gradient: \[ \mathcal{J}(\theta) = \mathbb{E}_{y \sim \pi_{\theta_{\mathrm{old}}}} \sum_{t=1}^{|y|} A_t \log \frac{\pi_\theta(y_t | y_{