Retrieval Augmented Thought Process for Private Data Handling in Healthcare

## Abstract

Abstract Large Language Models (LLMs) have demonstrated the strong potential to assist both clinicians and the general public with their extensive medical knowledge. However, their application in healthcare is constrained due to concerns about the privacy of data used in training, which prevents the integration of private and personal information because of security and ethical issues. Moreover, if their capabilities can be enhanced with information retrieval to access up-to-date knowledge, the current integration of LLMs with Information retrieval lacks robustness to imperfect retrieval, which can hinder their effectiveness and even reduce overall performance. In this work, we address this challenge by introducing the Retrieval-Augmented Thought Process (RATP). Given access to external knowledge, RATP formulates the thought generation of LLMs as a multiple-step decision process. To optimise such a thought process, RATP leverages Monte-Carlo Tree Search and learns a proxy reward function that permits cost-efficient inference. On a private dataset of electronic medical records, deliberately excluded from any LLM training set, RATP achieves 35% additional accuracy compared to in-context retrieval-augmented generation for the question-answering task.

## 1 Introduction

Recent advancements in Large Language Models (LLMs) trained on extensive datasets have showcased their enhanced capabilities in diverse tasks such as question-answering [], conversational abilities with humans [] and notably providing medical knowledge [, ]. Yet, their usage in healthcare is hindered by their limited proficiency in accessing and accurately handling private data.

Private data is any sensitive knowledge deliberately kept unavailable to LLMs during training due to ethical or business considerations. This includes, for example, medical records [, , ] or banking details. These data must be excluded from large language model training sets because once learned, their privacy cannot be ensured [, , ]. This creates a significant barrier to LLM usage in healthcare as organisations cannot risk potential Personal Identifiable Information (PII) leaks, especially under regulations such as GDPR which protect closely the usage of such data. Furthermore, private databases are frequently updated, sometimes daily, making the knowledge embedded within a model’s parameters quickly outdated. Continuously retraining the model with new data involves significant computational and financial resources [] which can be afforded by very few organisations. Hence, relying solely on LLMs, regardless of their increasing size, is impractical.

Thus, in scenarios such as administrative tasks (e.g., letters, discharge summaries) [] or decision aids (semi-autonomous diagnostic)[], where accessing information from a private database is required, the fusion of LLMs’ capabilities with external knowledge sources becomes crucial. While Retrieval-Augmented Generation (RAG) [] has introduced the LLM and information retrieval pairing with the embedded cross-product between a query and the documents, recent works have shown its lack of robustness to unsuccessful retrieval [], to the extend that it can negatively impact performance [].

Additionally, beyond the privacy issue, external knowledge can potentially address some of the inherent limitations of LLMs []. Importantly, it enables the LLMs to recall information from outside of their limited context window which is particularly beneficial in multi-turn dialogues []. Furthermore, grounding LLMs with factual external knowledge reduces the production of factual inaccuracies and hallucinations [, ], which are major issues impacting their performance, especially in less common domains or when dealing with private data. Moreover, using external knowledge allows the clinician to trace back the source of the information, unlike the implicit knowledge stored in model parameters. Hence, it has the potential to improve the accuracy and interpretability of responses, addressing two noted limitations of LLMs that hinder their impact in healthcare []. Consequently, this approach could support various medical applications detailed in Appendix A.

However, current literature still faces several desiderata to be fulfilled simultaneously for LLM and external knowledge pairing:

- nosep
Guaranteeing data privacy.
- nosep
Processing batches of documents beyond the LLM’s context window.
- nosep
Exploit reasoning capabilities of LLMs to filter out irrelevant or noisy information.
- nosep
Ensuring interpretability of the retrieval-augmented thought process.

To fulfil these desiderata we introduce the Retrieval-Augmented Thought Process (RATP) in Figure 1. RATP enhances the thought generation capabilities of LLMs by treating it as a multi-step decision-making process utilising external knowledge sources.

Our contributions include:

- wide
Formally, we formulate the open-book question-answering task as a sequential decision-making problem. In light of this new formalism, we compare existing methods and, in collaboration with clinicians from different specialities and countries, highlight the significance of retrieval-augmented thought processes for healthcare applications.
- wiide
Methodologically, we introduce the RATP, which leverages Monte-Carlo Tree Search to combine the reasoning capabilities of Large Language Models and the access to external knowledge.
- wiiide
Empirically, using publicly available LLMs, we evaluate RATP’s effectiveness on the emrQA dataset, a private QA dataset on real electronic medical records (EMRs).

Figure: Figure 1: Retrieval-Augmented Thought Process overview. \raisebox{-0.9pt}{1}⃝ The frozen LLM $l_{thought}$ given an answer $\hat{y}$ to the question $x$ by using the extra context $s_{T}$. \raisebox{-0.9pt}{2}⃝ The thought process starts from the question $x$ and outputs the best thought found $s_{t}$ to help answering $x$. The actions $\{a_{i}\}$ are decided by the MCTS with feedback from the scoring model. This component is detailed in Figure 2. \raisebox{-0.9pt}{3}⃝ The information retrieval system interacts with the thought process by answering its queries with retrieved documents $\{I_{i}\}$.
Refer to caption: x1.png

## 2 Preliminaries

The Question Answering Task
We consider the general setting of question-answering where
an autonomous machine learning model provides answer $\hat{y}\sim\mathcal{Y}$ given an input query $x\sim\mathcal{X}$. Given the state-of-the-art performance achieved by LLMs, in our work we consider the machine learning model to be a general-purpose LLM $\ell$.
The loss of the task can be evaluated through a metric $R$ with $\mathbb{E}_{(x,y)\sim(\mathcal{X},\mathcal{Y})}R(\hat{y},y)$, where $\hat{y}=\ell(x)$

Open-Book QA
In healthcare applications, there are many scenarios where an external knowledge database is a must to answer questions as privacy concerns and/or computational costs prevent this knowledge from being encoded in the LLM’s weight. For instance, redacting discharge summaries or pre-screenings of patients are tasks that require the information contained in private electronic medical records. Without loss of generality, we use $\mathcal{K}$ to denote the space of external knowledge, the LLM $\ell$ then answers the query with a subset of the knowledge, such that
$\hat{y}=\ell(K,x)$ — as in the most general cases, the knowledge database can be too large to be fed to the $\ell$.
The problem of finding the most appropriate piece of information is known as the information retrieval (IR) problem:

$$ $K^{*}=\arg\max_{K\in\mathcal{K}}R(\hat{y},y)~{}~{}~{}\mathrm{where}~{}~{}~{} \hat{y}=\ell(x,K).$ (1) $$

## 3 Retrieval-Augmented Thought Process

In this section, we first use formal language to define the retrieval-augmented thought process as an MDP in Sec. 3.1; we then introduce Monte-Carlo Tree Search as an efficient and effective planning algorithm to solve the MDP in Sec. 3.2; finally, we discuss practical implementation choices of the scoring model in MCTS in Sec. 3.3.

### 3.1 Multi-Step Thought Generation as a Markov Decision Process

Figure: Figure 2: Modeling the thought process. Each thought is generated from previous thoughts and/or documents, effectively creating a graph. The planning policy controlling the construction of this graph is detailed in Figure 3.
Refer to caption: extracted/5777609/images/fig2_ver_h2.png

Formally, we define the multi-step thought generation as a Markov Decision Process (MDP), denoted as $\mathcal{M}=(\mathcal{S},\mathcal{A},P,r,\mathcal{X},T)$. $\mathcal{S}$ is the state space. In our context, the state space is the space of thought process — the current reasoning graph leading to the answers. A state $s$ is composed of previous thoughts $\phi$ and previous actions. $\mathcal{A}$ is the action space; in our context, the action space contains all possible combinations of thoughts in the current thought process $s_{t}$. For ease of notation, any external document $K$ is also considered to be a thought. $P$ is the dynamics, in our context, the transition dynamics is known and instantiated as feeding the combined thoughts to a properly prompted LLM $\ell_{\mathrm{thought}}$ to generate a new thought, i.e., $s_{t+1}=(s_{t},a_{t},\phi_{t+1})\text{ with }\phi_{t+1}=\ell_{\mathrm{thought}
}(a_{t})$.
$R$ is the reward function that evaluates the final answer. We denote the query $x$ as the initial thought $\phi_{0}=x$ (hence the thought process also initializes with $s_{0}=x$). Therefore, the query distribution $\mathcal{X}$ is the initial state distribution. $T$ is the problem horizon — the maximal number of reasoning steps.

Starting the query $s_{0}=x$, an external policy $\pi:\mathcal{S}\mapsto\mathcal{A}$ determines whether the LLM continues to develop the existing thought with either external knowledge or another existing thought, i.e., $a_{0}=\pi(s_{0})$. Then, the LLM $\ell_{\mathrm{thought}}$ transits the thought process (state) to $s_{1}=(s_{0},a_{0},\phi_{1})\text{ with }\phi_{1}=\ell_{\mathrm{thought}}(a_{0})$, and an action $a_{1}$ will be generated by $a_{1}=\pi(s_{1})$, consequently, we have $s_{2}=(s_{1},a_{1},\phi_{2})\text{ with }\phi_{2}=\ell_{\mathrm{thought}}(a_{1})$, $a_{2}=\pi(s_{2})$, …
Such a thought process terminates until the maximal reasoning timestep $T$ is reached, and the final thought $\phi_{T}=\ell_{\mathrm{thought}}(a_{T-1})$ will be used to generate the final answer of the initial query: $\hat{y}=\ell_{\mathrm{out}}(x,\phi_{T})$,
where $\ell_{\mathrm{out}}$ denotes a frozen LLM.
The reward function $R$ provides an episodic scalar feedback only at the end of a thought process: $r_{T}=R(\hat{y},y)$, and $r_{t<T}=0$.
In our work, we use $s_{\pi}=(\phi_{0},a_{0},\phi_{1},a_{1},...,\phi_{T})$ to denote the above thought process generated with policy $\pi$, and consider the policy optimization problem:

$$ $\pi^{*}=\arg\max_{\pi}\mathbb{E}_{(x,y)\sim(\mathcal{X},\mathcal{Y}),\hat{y} \sim s_{\pi}}\left[R(\hat{y},y)\right]$ $$

### 3.2 Planning with Monte Carlo Tree Search

Given the vast action space in the above problem and the fact that the system dynamics model (i.e., the thought generation LLM $\ell_{\mathrm{thought}}$) is accessible during inference, we choose to use Monte Carlo Tree Search (MCTS) [, ] for effective and efficient decision optimization. MCTS explores the tree of decisions and returns the best action found. In our case, we use MCTS to build the graph of thoughts and documents. Guided by our scoring system, MCTS realizes the trade-off between exploiting the best thought and exploring other thoughts. In the tree-search, the root $\phi_{0}$ is the initial query $x$ and each node is a thought $\phi$. The children of a given thought are the thoughts generated from it.
The 4 key steps of MCTS are
Selection, Expansion, Simulation, and Backpropagation. Figure 3 illustrates one step of the MCTS decision process by summarising the action of these four functions. This step is repeated until the answer is found or we reach the thought process decision step limit $T$.

Figure: Figure 3: One complete step from our MCTS decision process. It is divided into four functions, which are repeated until we find the answer or the thought process size limit is reached. The Selection, Expansion, Simulation, and Backpropagation functions are described in section 3.2. Their associated algorithm can be found in Appendix G.
Refer to caption: x2.png

Selection. In the selection, we decide which thought to be expanded. From the root $\phi_{0}$ of our graph, we explore the graph until finding a thought $\phi_{i}$ that has never been used to generate another thought: a node without children.
The exploration is guided by the Upper Confidence bounds applied to Trees (UCT) formula [, ]. When we meet a node with children, we choose to explore the child with the highest UCT value. For the $i^{th}$ thought that has been visited $n_{i}$ times, with a score $q_{i}$, and whose parents have been visited $N_{i}$ times, the UCT value expression with an exploration parameter $c$ is
$\frac{q_{i}}{n_{i}}+c\sqrt{\frac{ln(N_{i})}{n_{i}}}$

Expansion. The thought $\phi_{i}$ selected in the previous step is paired with other thoughts (including external documents) to generate a new thought. The choice to choose documents or thoughts for the pairing is decided by the probability $p_{\mathrm{doc}}$ and the scoring model (see section 3.3). Then, the chosen thoughts and documents are combined with a prompt template (see examples in Appendix B.1) and given to $l_{\mathrm{thought}}$ to generate the new thought.

Simulation. In our case, the simulation function computes the score of the new thought by using one of the scoring models described in section 3.3.

Backpropagation. This function recursively updates the UCT value of every node that has led to the new thought. Starting from this new node, it uses the new thought’s score and the updated visit counts in the UCT formula to modify the UCT value of each parent until the root or documents are reached.

Documents retrieval. As the number of documents in the collection can be very large, we perform an initial document retrieval step using traditional retriever models Additional retrieval queries are made when needed. The initial implementation of the framework used Contriever [] but other retrieval models, in particular for different modalities (images, tables), can be used to broaden the applicability of the RATP framework.

### 3.3 Scoring models

Oracle Score To score the value of each thought, a naive way is to use the QA LLM $\ell_{\mathrm{out}}$, which takes the query $x$ and any given thought $\phi$ as inputs and outputs the answer to the question.
An oracle score can be defined by comparing this generated answer to the true answer with the metric score $R$ of the QA task. However, such a perfect signal is by definition not available at inference time. Therefore, two scoring models are introduced as proxies of such signals. Details on the implementation and comparison of the scoring models can be found in Appendix C.

Model-Based Estimation. In the first approach, we propose a model-based method to learn a proxy reward function $R_{\theta}$ that estimates the oracle scores using a training dataset.
Specifically, when running MCTS on the training data, we would be able to generate an offline dataset that contains various thought combinations $(\phi^{(i)},\phi^{(j)})$, and whether those thoughts are informative enough to solve the original query $x$:

$$ $r^{(i,j)}=R(\ell_{\mathrm{out}}(\ell_{\mathrm{thought}}(\phi^{(i)},\phi^{(j)}) ,x),y).$ $$

Given such an offline dataset $\mathcal{D}=\{\phi^{(i)},\phi^{(j)},r^{(i,j)}\}_{(i,j)}$,
$R_{\theta}$ can be optimized through

$$ $\theta\leftarrow\arg\min_{\theta}||R_{\theta}(\phi^{(i)},\phi^{(j)})-r^{(i,j)} ||^{2}$ (2) $$

and used as a proxy of the Oracle score during inference time when we do not have access to the true answers. $R_{\theta}$ is typically instantiated as an MLP or XGboost model.

Self-critics score. However, acknowledging that such an offline dataset may not always exist. We also introduce the second approach based on LLM self-critic. In the literature, using an LLM as a self-criticism agent is widely applied to reflect on a previous LLM answer and improve it [, , ]. It has been reported that LLM can detect their own mistake when prompted with their previous outputs. Thus, we leverage this ability by asking the LLM to predict if a thought is accurate and contains all the information to answer the query.

## 4 Related Work

Table: Table 1: Comparison with Related work. The desiderata laid down in Section 1 are considered. RATP is the only method that fulfills them all.

Private data and LLM
Recent studies have extensively demonstrated the vulnerability of private data processed during LLM training [, , , , , ]. The primary issue is the LLM’s memorisation of data, which cannot be prevented during pre-training [] or fine-tuning []. This memorisation makes LLMs susceptible to malicious extraction attacks. Since no method fully prevents this memorisation [], the safest approach is to exclude private data from training datasets and use RAG instead []. Despite recognition of the problem, few solutions have been proposed. To the best of our knowledge, this is the first paper to propose and benchmark a solution using a private dataset.

Thought Processes as New Prompting Strategies
Recent advancements in prompting strategies have greatly improved the reasoning capabilities of Large Language Models. This progress began with the Chain of Thought (CoT) approach by [], which breaks down complex problems into simpler steps, known as thoughts. Building on CoT, subsequent research has focused on generating, combining, and selecting these thoughts. Notable methods include the Tree of Thought (ToT) [] and Graph of Thought (GoT) [], which offer greater flexibility in thought combination. Similar to our approach, other strategies employ Reinforcement Learning techniques, such as Monte Carlo Tree Search, to explore the thought space in constrained reasoning problems like the Game of 24 and the 8-puzzle []. However, these methods rely heavily on LLM-generated content and are susceptible to hallucination, limiting their use in real healthcare applications. Therefore, integrating these strategies with information retrieval is essential for grounding them in reality. Additional presentation and comparison of popular thought processes can be found in Appendix F.

Table: Table 2: Theoretical comparison between RATP and other thought processes with the lens of the Information Retrieval as Multi-step decision-making problem formalism.

Pairing Information Retrieval and LLM
Integrating Information Retrieval with Large Language Models significantly enhances the efficiency of accessing and processing vast amounts of information [, , ].
One method involves pre-training LLMs with IR models like Contriever [, , ]. Another method fine-tunes LLMs with task-specific tokens [] or adds layers to the model structure [], though these approaches require significant computational resources and are vulnerable to extraction attacks. Alternatively, Retrieval-Augmented Generation [, ] leverages LLMs’ in-context learning abilities to integrate IR data without extensive additional training, thus avoiding sensitive data leaks. Iterative prompting techniques like RAT or RATT [, , , ] enhance their reasoning capabilities by integrating external knowledge from documents in a chain-of-thought manner and/or with LLM fine-tuning [], addressing the vulnerability of LLMs to misleading documents []. For the same reason, HipporRAG []uses an LLM preprocessing of the document collection. Table 2 illustrates that the MDP formalism developed in section 3.1 can analyse and compare all the popular iterative prompting strategies.

Our approach uniquely enhances RAG by safely and accurately handling healthcare-sensitive data, utilising frozen LLMs to prevent training data leaks. We are the first to propose a method that grounds a complete thought process, planned via reinforcement learning solutions, in factual documents for private information retrieval.

### 4.1 Experimental setup

Main experimental setup: private medical datasets We analysed and benchmarked our method on two real, private, and sensitive healthcare data. First, we consider unstructured EMRs. These EMRS have been gathered from 2004 to 2014 at the Partners HealthCare System, Boston, during the i2b2 project (examples in Appendix B.1). We evaluated our method on the Question Answering task using the emrQA dataset [], which consists of open-ended medical questions based on patient records. These questions were created in coordination with physicians and medical experts. To evaluate our performance on open-ended questions, we used the Exact Match metric from SQuAD [, , ]. (Details in Appendix B.1.) Second, we focus on discharge summaries. The EHRQA dataset [] contains multiple-choice questions derived from patients’ discharge summaries from the MIMIC-IV [] database, which includes anonymized patient records from Beth Israel Deaconess Medical Center between 2008 and 2019. The questions are categorized into eight types: Treatment, Assessment, Problem, Etiology, Sign/Symptom, Vitals, Test Results, History, Instruction, and Plan. Each question has been reviewed and refined by three clinicians. A crucial property of both these datasets is that their access is controlled preventing its use for LLM training.

Additionnal experimental setup: Open-Domain Question Answering. We add a second experimental setup to highlight the difference between the private and public knowledge settings. For the latter, we consider the open-domain question-answering task using the Boolq dataset []. This widely-used dataset consists of closed-ended questions based on Wikipedia common knowledge. (The analysis of this setting is in Appendix D).

### 4.2 Analysis

Baselines. Our methods are compared against two prevalent zero-shot open-domain question-answering approaches. The first baseline, LLM ($\hat{y}=\ell_{\mathrm{out}}(x)$), involves directly prompting an LLM with the question. The second baseline, RAG ($\hat{y}=\ell_{\mathrm{out}}(x,k)$), represents the in-context IR method [, ]. Here, the LLM generation process is augmented with one document $k$ retrieved from the knowledge base by Contriever [].

Oracles methods use the oracle scoring system that requires the answers to the questions. As such, these methods do not apply to answering new questions. They are included to demonstrate the importance of information retrieval and the potential of our methods with an ideal scoring model. MCTS oracle is the full method with the oracle score, and MCTS oracle w/o IR is identical but with the information retrieval function deactivated. The former can generate thoughts from documents and previous thoughts, while the latter can only access previous thoughts ($\mathcal{K}=\varnothing$).

RATP. Two versions are introduced : MCTS self-critic and MCTS estimation are our MCTS algorithms using the self-critic and offline model-based estimation scores, respectively. Tree-of-Thought with IR (ToT \w IR) is an enhanced version of Tree-of-Thought [] which integrates Information Retrieval as delineated in Figure 2.

Table: Table 3: Ablation study for the Private Knowledge Question Answering task. We report the Exact Match accuracy (%). The margins of error are computed by bootstrapped student-t tests (95%). The thought process size is 25.

Even without access to patient records, the LLM correctly answers 34% of the time by leveraging its prior knowledge of standard prescriptions. For instance, when asked, ”What is the drug strength of aspirin prescribed to the patient?", the LLM correctly infers the default dosage of 1mg. Table 3 indicate that adding a document to the prompt (RAG) leads to worse performance. This aligns with previous studies showing that RAG can be unreliable, depending on the dataset, retriever model, and knowledge base[, ]. We identify two main reasons for this performance decline: current retriever models often fail to find relevant documents in large knowledge bases [], and LLMs can be confused by retrieved information, particularly when it is irrelevant or noisy []. Takeaway: Our experiment confirms that adding information in context can reduce LLM performance.

The enhanced performance of our methods, underscores that our thought process can guide an LLM from an initially incorrect answer to a correct one : $\mathbb{E}_{(\mathcal{X},\mathcal{Y})}[R(\ell_{\mathrm{out}}(s_{0}),y)]\geq
\mathbb{E}_{(\mathcal{X},\mathcal{Y})}[R(\ell_{\mathrm{out}}(s_{T}),y)]$. This phenomenon is evidenced by the increase in accuracy correlating with the size of the thought process, as depicted in Figure 4: the LLM’s performance improves after processing multiple documents and generating several thoughts. Notably, the thought process not only allows the viewing of more documents compared to other methods but also strengthens resilience against both unsuccessful information retrieval and misleading documents, as exemplified in Figure 16 in the appendix. In this example, our MCTS method is not misled by irrelevant document 1 and 3 and eventually succeeds in extracting pertinent information from note 2. Takeaway: the knowledge provided by RATP increases substantially the number of correct answers.

Figure: Figure 4: Evolution of the accuracy and the number of LLM queries. When we increase the thought process size (i.e. the number of thoughts generated), the accuracy increases but the number of LLM queries too.
Refer to caption: extracted/5777609/images/accuracy_speed_plot.png

Even with the oracle score, RATP is not flawless. Incorrect answers occur when the LLM cannot revise its stance within the thought process size limit, $T$. This happens if all retrieved documents are irrelevant, the LLM cannot use them effectively, or it fixates on an incorrect answer. The limit $T$ is set to prevent excessively long inferences. As shown in Figure 4, the number of LLM queries increases linearly with the number of thoughts, leading to significant financial or time costs for extensive thought processes. Typically, our methods find the correct answer within 20 thoughts, as indicated by the accuracy plateau in Figure 4. In contrast, the oracle score accuracy rises to 88% with 25 thoughts. Takeaway: LLMs may need an excessive number of thought generation steps to pivot, exceeding an acceptable cost.

Ablation study:
Comparing the MCTS oracle with and without IR shows that access to retrieved documents accounts for a 37% increase in final performance. Surprisingly, a similar experiment in an Open-Domain public data setting (see Appendix D) also demonstrates that IR boosts performance. This result is unexpected, as the LLM has been extensively trained on Wikipedia articles, suggesting little benefit from retrieving documents from this source. However, this can be explained by the imperfect memory of LLMs [] and their propensity for hallucination []. It is widely recognised that prompting LLMs with relevant factual documents helps reduce hallucinations [, ], thereby enhancing performance. In Appendix E, we present an additional experiment demonstrating the impact of IR on improving the general thought quality.
Takeaway: IR improves performance, even when the LLM has been trained on the knowledge base. Moreover, we use the ToT with IR method to highlight the performance gains brought by the MCTS algorithm. Unlike standard thought processes, which solely follow the scoring model’s guidance, MCTS optimizes an exploration-exploitation trade-off modelled by the UCT equation. This allows it to develop poorly scored thoughts as well, leading to two interesting properties: a) Robustness against scoring system flaws, as a thought mistakenly flagged as irrelevant has a non-zero probability of being developed. b) It fosters more "creativity" by allowing concurrent and/or connected lines of thought. Takeaway: the MCTS algorithm demonstrates superior performance compared to other thought process algorithms.

Table 3 demonstrates the effectiveness of the offline model-based estimator. It matches the performance of the self-critic score while being significantly faster, requiring fewer LLM queries compared to its self-critic equivalent (see Figure 4).

High Interpretability. Finally, we wish to highlight the high interpretability of RATP. As demonstrated in the inference example in Figure 16, for each answer provided by our system, the entire thought process $s_{\pi}=(\phi_{0},a_{0},\phi_{1},...,\phi_{T})$ is accessible. The final thought $\phi_{T}$ provides the rationale behind the generated answer, and additionally, we can trace the source of the information through the retrieved documents.

### 4.3 Benchmark

Self-RAG[]. This method relies on teaching the model to use a new <retrieval> token and a self-reflection thought process to assess the relevancy of the retrieved documents. Self-RAG resorts to LLM fine-tuning but can be evaluated on documents outside of its training corpus.

RATT[]. This method generates multiple thoughts from the initial query, retrieves a document for each thought, and combines all documents and thoughts into a prompt to generate the final thought. This process iterates multiple times.

Table: Table 4: Benchmark of LLM and Information Retrieval pairing methods on both Public and Private dataset. We report the binary accuracy (%) for the BoolQA dataset and the Exact Match accuracy (%) for the emrQA dataset and EhrQA datasaet. The margins of error are computed by bootstrapped student-t tests (95%). Additional details on the implementation of each method are given in Appendix B. This experiment has been conducted with Mixtral8x7B deployed locally.

Difference between the private and public knowledge setting : This benchmark highlights the unique aspects of the private knowledge setting, which is rarely studied. As expected, the LLM baseline performance is much lower for questions requiring knowledge it hasn’t been trained on, emphasising the importance of IR in the private setting. Additionally, the LLM shows greater confusion with RAG in the private setting, indicated by a larger drop in accuracy. We interpret this as a greater reliance on retrieved documents in the private setting, where the LLM cannot rely on its prior knowledge. This increased confusion may also be due to out-of-distribution formatting, such as the technical jargon, abbreviations, and layouts of EMRs and discharge summaries, which differ significantly from the text the LLM has been trained on. Similarly, Self-RAG performs poorly because the documents to be retrieved and processed differ greatly from its fine-tuning dataset. Additionally, while Self-RAG realises multiple thought processes of depth 1, RATP performs one deep thought process which provides additional context and enables queries on information that might not be explicitly stated in the initial question. Finally, we observe that the final accuracy is similar in both settings, suggesting that performance depends more on the reasoning abilities of the LLMs than on the knowledge setting, as shown by a comparison of different LLM sizes in Appendix I.
This benchmark validates the effectiveness of our approach on both public and private data. For the Question Answering task on Electronic Medical Records, RATP demonstrates over 70% accuracy improvement.

Limitations. One significant limitation of our work is the difficulty in obtaining publicly available large-scale datasets that are guaranteed not to have been part of the LLMs’ training set. This increasingly challenging condition poses a problem for machine learning researchers aiming to conduct research on LLM in Healthcare. Furthermore, while the performance of our two scoring systems is similar, each has distinct limitations. The self-critic score is computationally intensive, effectively doubling the number of LLM calls required (see Figure 4). In contrast, the estimation score is more computationally efficient but requires a training dataset composed of questions and samples from the knowledge base. We also want to clarify that our focus is on one specific privacy concern: sensitive data leakage from the training set. Other threats, such as interception of retrieved data [], profiling individuals [], or contextual privacy [], are beyond the scope of this paper. Finally, while our method may appear complex with several components, it is actually easy to apply and generalise to new datasets. We provide a guide for practitioners in Appendix H, demonstrating that the method only requires two steps and minimal tuning to be implemented

## 5 Conclusion

In this work, we introduced RATP, a new framework that pairs LLMs with private information using a Markovian multi-step decision process. Our method meets the four criteria outlined in the introduction for such approaches. 1. It does not rely on LLM training, ensuring that patient-sensitive data will not be leaked. 2. By dividing the retrieval into a multi-step process, it can handle knowledge well beyond its context window size. 3. It leverages MCTS to enhance the self-reflection and self-critique abilities of LLMs across many documents, making it robust to unsuccessful or noisy retrievals. This allows it to outperform other methods by 50% in navigating real EMRs. 4. The entire process occurs in natural language, enabling each step to be verified by the practitioner.

Futur work: Applying the RATP framework directly to other data modalities, such as images or tables, using multimodal LLMs and alternative retriever models would broaden the method’s applicability. Table 3 highlights the importance of guiding thought generation with a scoring model trained on medical data. Thus, developing robust scoring models using larger fine-tuned LLMs, while preserving data privacy, is a significant research direction.Moreover, clinicians might still be the best scoring model, making a "human in the loop" approach particularly valuable. This method, being transparent and providing reasoning steps in natural language, is well-suited for such an approach.

Impact statement: Our manuscript contributes positively to both the research community and healthcare applications by presenting a method to safely integrate new sensitive private knowledge into LLMs. This inclusion of private data, previously excluded due to ethical or safety reasons, facilitates personalised applications of LLMs tailored to each patient’s medical context. Moreover, its transparent reasoning across multiple documents of diverse natures and origins supports multifactorial decision-making, which is crucial for complex medical scenarios: additional healthcare applications enabled by RATP can be found in Appendix A. Furthermore, by eliminating the need for extensive training, our approach reduces barriers to LLM utilisation, making them more accessible for various scenarios, especially for organisations, institutions, or countries with limited resources. However, given the medical context, it is crucial to emphasise that the use of LLMs must be carefully controlled and validated to prevent errors and ensure patient safety.

## Appendix A Large Language Applications in Healthcare enabled by RATP

Table: Table 5: Examples of current healthcare issues that could be addressed using large language model applications with access to private and sensitive data. RATP enables such applications.

## Appendix B Experimental details

For all our experiments, the large language model used is Mixtral8x7B [] locally deployed on 4 A100 GPU accelerators with 80GB VRAM. The batch of document retrieval from the knowledge base is performed by a Contriever model [] in its standard configuration. This model has been pre-trained on CC-net and English Wikipedia.

### B.1 Private knowledge: Unstructured Electronic Medical Records

To realise our experiment on private knowledge, we use the question with fine-grained answers from the emrQA [] dataset. These are questions of the type “What is the dosage of |medication|?". We only keep the questions whose answers are of the form "X mg" with X a number. The 4436 questions filtered are split into a training set (3110 questions) and a test set (1000 questions). The knowledge base for this dataset is composed of unstructured electronic medical reports which are gathered in patient records. For each patient, we re-split this patient record (Figure 5 provides examples) into 100-word chunks. During the retrieval step, the patient associated with the query is identified and the retrieved documents are these 100-word passages from its medical record. For this experiment, we directly embed the initial question as the input for Contriever.

Figure: Figure 5: Examples of Unstructured Electronic Medical Records. For privacy reasons, we present simulated EMRs resembling the actual dataset.

The prompts used to build the thought process in this experiment are the following.

Figure: Figure 6: This prompt template is employed with the final thought to obtain the answer to the question. This is the same prompt template used in the RAG baseline.

Figure: Figure 8: This prompt template is employed to generate new thoughts for the amrQA experiment.

The MCTS is configured with the following hyperparameters: exploration rate: $\sqrt{2}$, pick document probability: 1, thought sample size: 5, maximum thought process size: 25, score threshold to stop the thought process : (oracle: 0.5, self-critic: 0.49, model-based estimation: 0.9), retrieved document batch size: 5.

### B.2 Public Knowledge : BoolQ dataset

The experiments on public knowledge have been realised by running our methods on the validation split from the Hugginface Boolq dataset. The knowledge base includes all the articles from the English Wikipedia. To create this database, we filter the most recent Wikipedia dump (December 2023) to only keep the body of the articles. Then, we merge and split these articles to create 500-word chunks of text. Thus, a retrieved document in this experiment is one of this 500-word passage. Moreover, the embedded query used as input to Contriever to retrieve a new batch of documents is the LLM’s output from the "query" prompt template with the best-scored thought (see Figure 11).

Figure: Figure 10: "This prompt template is employed to generate queries to retrieve documents with Contriever."

Figure: Figure 12: This prompt template is employed with the final thought to obtain the answer to the question. This is the same prompt template used in the RAG baseline.

The MCTS is configured with the following hyperparameters: exploration rate: $\sqrt{2}$, pick document probability: 0.5, thought sample size: 5, maximum thought process size: 10, score threshold to stop the thought process : (oracle: 0.5, self-critic: 0.49, model-based estimation: 0.21), retrieved document batch size: 2.

### B.3 Additional methods

Self-RAG: the Self-RAG implementation relies on the official Self-RAG codebase. The largest model offered by the authors (Llama-13B) is being used on both our public and private dataset. The two settings proposed by the self-RAG authors (short-form and long-form text generation) have been tested with the default parameters. Only the best results are reported in Section 4.3. Full results are presented in table 6.

Table: Table 6: Complete results for the Self-RAG method in both settings.

ToT w\ IR: We use standard hyperparameters that match the size of our thought processes (25) for ToT: step limit: 6, breadth limit: 2, size limit: 2. $P_{doc}$ that controls the integration of information retrieval is set to 1 for EmrQA and 0.7 for BoolQ.

## Appendix C Scoring Models

### C.1 Scoring model implementation

Training of the offline model-based score. For both datasets the method to build a dataset and train the score estimation model is similar. We collect complete thought processes by answering questions from the training split of the dataset with our MCTS-oracle method. Thus, we are able to link every thought $\phi_{t}$ collected with its parents i.e. the thoughts or documents that have been used to generate $\phi_{t}$. The parents of $\phi_{t}$ are indicated by the action $a_{t-1}=(\phi^{t,1},\phi^{t,2})$ because $\phi_{t}=\ell_{\mathrm{thought}}(a_{t-1})$. As the thoughts processes have been generated by the MCTS-oracle policy, we also have the true score $R(\ell_{\mathrm{out}}(\phi_{t}),x_{i}),y_{i})=R(\ell_{\mathrm{out}}(\ell_{
\mathrm{thought}}(\phi^{t,1},\phi^{t,2}),x_{i}),y_{i})=r^{(t,1,t,2)}$. Hence, we effectively have a dataset $\mathcal{D}=\{\phi^{(i)},\phi^{(j)},r^{(i,j)}\}_{(i,j)}$.

In practice, in the dataset $\mathcal{D}$, the feature vectors are the concatenations of the embedded versions of thoughts $\phi^{(i)}$ and $\phi^{(j)}$. These embeddings are generated by Contriever []. The size of the dataset $\mathcal{D}$ is typically around 4000 samples. We trained gradient-boosting or MLP models on these datasets to perform the score estimations. The thresholds to compute the accuracy are chosen to maximize the precision on the training set. Finally, they are tested on the test split of their respective dataset.

Self-critic scoring model. The self-critic method entails prompting the LLM to reflect on its previous outputs. In our case, we combine the initial question and a given thought with a prompt template (see Figure 9) to query a frozen LLM about the accuracy and adequacy of the thought.

In our framework, the LLM is constrained to respond with either 0(indicating the response is insufficient or inaccurate) or 1 (indicating adequacy and accuracy).To derive a more nuanced self-critic score, we examine the relative softmax probabilities assigned to the tokens “0" and “1" by the LLM. The self-critic score is calculated using the formula:

$$ $\text{Self-Critic Score}=\frac{\text{Softmax Score of Token 1}}{\text{Softmax Score of Token 1}+\text{Softmax Score of Token 0}}.$ (3) $$

### C.2 Scoring Model Comparison

As detailed in section 3.3, we proposed two different scoring models to estimate the potential of the thoughts. These models should predict the oracle score which is not available at inference time. Thus, to assess the performance of these scores, we answer a subset of queries from the Boolq test set using our MCTS with the oracle score as the policy $\pi$. By doing so, we generate a set of pairs (question, thought-process) : $\{x_{i},s_{\pi,i}\}$. For every generated thought of these thought processes, we compute its score using the two scoring models. Finally, we compare how well these models predict the oracle score in Table 7.

In Table 7, it is evident that the model-based estimators outperform the self-critic model. The main reason for this disparity is that the estimation score model is specifically trained for this task while the self-critic score prompts a frozen LLM.

However, we also want to highlight that the weak performance of the self-critic model can be explained by its over-confidence. As depicted in Figure 14, the distribution of self-critic scores is more skewed towards 1 (indicating confidence that the thought contains the required information to answer the query) compared to the oracle score, which is skewed towards 0 (suggesting the thought is inaccurate or lacks information).

Figure: Figure 14: Histograms of Score Distributions. Each scoring system’s score value is on the horizontal axis, and the vertical axis is the number of samples for the value. The dashed lines correspond to the thresholds used to compute the accuracy .
Refer to caption: extracted/5777609/images/hist_scores.png

Table: Table 7: Comparison of Scoring Models. These scoring models predict the oracle score for each thought from multiple runs of our MCTS method. The queries are from the Boolq test set, and we retrieved documents from Wikipedia. Accuracy is computed using thresholds computed on the training set (0.9 for self-critic, 0.21 for the model-based estimation).

The unbalanced distribution of oracle scores is a consequence of our MCTS implementation: we stop the generation of the thought process $s_{\pi}=(\phi_{0},a_{0},\phi_{1},a_{1},...\phi_{T})$ when the oracle score indicates that a thought $\phi_{T}$ has been generated which successfully answers the query. Therefore, in each run, only the last thought $\phi_{T}$ receives a high score, while preceding thoughts $\phi_{0},\phi_{1},...,\phi_{T-1}$ that led to this final thought – either partially answering the query or being irrelevant – receive low scores.

However, if using LLMs as self-critics causes hallucinations in evaluating thoughts, the self-critic performance should improve with higher quality LLMs (more parameters). To verify how the quality of the LLM impacts the self-critic score’s tendency to hallucinate and affects RATP’s performance, we experimented with GPT-3.5-turbo and GPT-4, which are stronger general-purpose LLMs known to suffer less from hallucinations. This experiment was limited to the BoolQA dataset, as feeding sensitive data from emrQA to a non-locally deployed LLM raises ethical issues.

Table: Table 8: Comparison of RATP (MCTS self-critic) performance using different LLMs as critic models. The thought generation LLM is consistently Llama 70B.

Table 8 supports our conclusion that the performance of RATP improves with stronger LLMs and fewer hallucinations in acting as self-critics. In addition, if LLMs that are better aligned with humans yield better results, it shows that leveraging human feedback would be an interesting way of improving the method.

## Appendix D Additional ablation study : Public Knowledge Setting

Table: Table 9: Ablation study for the Open Domain Question Answering task. This comparison was conducted on the Boolq test set. The knowledge base was a Wikipedia dump from 2023. We report the binary accuracy. The margins of error are computed by bootstrapped student-t tests (95%). This experiment has been conducted with Llama-2 70B.

As seen in Table 9, IR still enhances performance even though the LLM has been trained on the knowledge base. There are two main differences compared to the private knowledge setting. First, the limited performance gain from RAG can be attributed to the nature of the retrieved documents, which are Wikipedia articles. Both the retriever model and the LLM have been trained on Wikipedia, leading to more successful retrievals and better handling of the documents by the LLM. Second, the performance gain from IR is much lower; the difference for the Oracle method with and without IR is only 3%, which is expected given the nature of the knowledge base.

## Appendix E Hallucination experiment

To demonstrate the impact of IR (even for public knowledge base), an additional experiments has been conducted to explicitly quantify the extent of hallucination in the generated thoughts. In these experiments, we evaluate the quality of thoughts produced by the MCTS+oracle scoring method, both with and without the integration of IR. The assessment uses the BoolQA test split as the dataset.

To effectively measure hallucination, we establish three proxy scores: A. The oracle score, as defined in the paper, which pertains to the downstream QA task performance of a frozen LLM prompted with the generated thought. B. F1 and ROUGE-L similarity metrics, which compare the generated thoughts with the gold standard paragraphs associated with the queries in the BoolQA dataset.

Table: Table 10: Assessement of the thought quality with and without IR. The frozen LLM used is Mixtral8x7B.

Table 10 shows that across all these measures, the inclusion of IR in our method demonstrates a marked improvement in the quality of thought generation. There is a clear indication that IR contributes to generating higher-quality thoughts with reduced instances of hallucination and enhanced accuracy.

## Appendix F Comparison with existing thoughts processes and baselines

Figure: Figure 15: Comparison with other popular types of thought process
Refer to caption: extracted/5777609/images/Figure11.png

In this section, we contrast different thought processes with their graphical models.

#### Vanilla LLM for QA

The naive approach is to directly query the LLMs for the answers. In this approach, there is no decision, and the answer is generated with a single-round dialogue with the LLMs.

#### Retrieval Augmented Generation (RAG)

In RAG, an external document can be used to alleviate the hallucination of the LLMs. The limitation of such an approach is the length of the document can be restricted by the context window size. In RAG, the LLMs answer the query within a single interaction. By formulating RAG as an MDP, the decision is 1-step, and the action can be defined as selecting the most relevant document form an external database.

#### Tree of Thought (ToT)

In ToT, the thought process is generated in a tree structure with specifically designed prompts that ask the LLMs to generate multiple thoughts at each timestep. To search over the generated tree-of-thought structure, either DFS or BFS can be used. Such a search is over the entire thought generation process, hence can be extremely high-dimensional and computationally challenging.

#### Graph of Thought (GoT)

In GoT, the thought process is generated as a directional graph. If we consider the thought generation process in GoT as an MDP, the action space is defined as the graph operators, such as aggregating or reflecting, and the exploration over such actions is conducted by LLMs.

#### RATP

in RATP, we solve the decision-making problem with an external planning policy learned with MCTS. Moreover, RATP has a more general State space where external documents are also considered to be thoughts. In addition, RATP can work with model-based return estimation such that in inference time the value of thoughts can be effectively estimated with light machine learning model — without using LLMs for self reflection.

## Appendix G Monte-Carlo Tree Search implementation

Figure: Algorithm 1 Overview of the MCTS algorithm that builds our thought process. The algorithms of the functions Selection, Expansion, Simulation and Backpropagation can be found in section G.

Figure: Algorithm 2 Selection

Figure: Algorithm 3 Expansion

Figure: Algorithm 4 Simulation

Figure: Algorithm 5 Backpropagation

## Appendix H Guide for practionner

In this section, we present a concise guide for implementing our method with new datasets.

1. Choosing a Retriever Mode and LLM: In our experiments, we use a pre-trained Contriever [4] with default parameters from its official GitHub repo. This model has been pre-trained on Wikipedia but can be used for any textual database as shown by our results on the electronic medical record. Moreover, as our method is built to be robust to unreliable retrievals, the choice of the retriever model should have a low impact.
Similarly, while we provide the implementation for Llama-2, Gemma, and Mixtral, our method is directly compatible with any Large Language Models.

2. Customizing Prompts: Our methods use different prompt templates for the different steps (final answer, thought generation, self-critic). While we provide general-purpose prompts (see Appendix B.2), you can customize these prompts by incorporating expert knowledge (see Appendix B.1). This is particularly useful to enforce a desired shape for the output.

3. (Optional) Training an Oracle Score Estimator: The practitioner might want to use an estimator of the oracle score as a scoring model because it is cheaper in LLM queries with similar performance (see section 4.2). They would only need to provide a left-out dataset. The methods can then be run with the oracle scoring method and save every score thought. It will then train a model on such collected data points. We provide the training script for XGboost models.

4. (Optional) Fine-Tuning MCTS Parameters: To improve the performance of the method on a particular dataset, it is possible to adapt the values of some parameters.

- •
$P_{doc}$ : the probability of pairing a thought with a retrieved document or another previous thought.
- •
$c$ : the exploration parameter in the UCT formula, a higher value will favor the exploration of unseen thoughts, lower value favors the exploitation of highly scored thoughts.
- •
Early stopping : There are two mechanisms to stop the thought process. First a threshold, we stop the process when a thought has a score higher than a given value. Second, a finite size stops the thought process when a given amount of thoughts is reached.

## Appendix I Generalisation to other Large Language Models

This section provides empirical evidence to support that RAPT can use any LLM models by performing a complete re-run of our experimental setup for Gemma 2B, Llama-2 70B and Mixtral8x7B.

Table 11 findings align with our expectations that larger models exhibit enhanced performance. Interestingly, the Gemma 2B model, being relatively smaller, does not exhibit significant gains from the RAPT method on the BoolQA dataset. This observation is attributed to the limitations of smaller language models in executing complex reasoning processes akin to self-reflection or chain-of-thought, which are critical for RATP’s efficacy.

Table: Table 11: Comparison of RATP performance for different Large Language Models. Unlike table 8, in this experiment the different LLMs are also used to generate the thought. We report the binary accuracy (%) for the BoolQA dataset and the Exact Match accuracy (%) for the emrQA dataset.

## Appendix J Analysis of the data bias

In this section, we analyse bias in the EMRQA dataset. Our method relies on LLMs, which have been demonstrated to have biases such as geographic bias [] or gender bias [].
From the EmrQA case study, we identify three potential biases:

- •
Geographic bias: EMRs were collected at Partners HealthCare System in Boston.
- •
Temporal bias: EMRs were collected between 2004 and 2014.
- •
Age bias: The dataset excludes EMRs concerning minors.

Additionally, we provide further analysis of the data used regarding two criteria: age and sex (as the EMRs have been pre-processed to prevent identification, these are the only criteria we can reliably infer from the data). The results of these analyses are presented in Tables 12 and 13 13 of the supplementary PDF.

Table: Table 12: Accuracy of the method depending on the sex of the patient. The sex of a subset of patients was manually retrieved from their EMRs. We report the accuracy (ExactMatch metric) of the answers to questions related to these patients.

Table: Table 13: Accuracy of the method depending on the age of the patient. The age of a subset of patients was manually retrieved from their EMRs. We report the accuracy (ExactMatch metric) of the answers to questions related to these patients.

From this analysis, we conclude that the data used for the case study are slightly biased toward aged males. Some bias can be observed in the method’s accuracy, as the method performs slightly better for younger patients and male patients.

## Appendix K Example

Figure: Figure 16: Example of question answering by RATP. It has been realized by the self-critic MCTS method on a question from the Boolq test set.
Refer to caption: extracted/5777609/images/exemple_graph.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. How does the method compare to ToT in terms of computational costs such as  LLM calls, total input + output tokens, and total time to answer a question?
2. How does the method compare to methods such as ReACT or DSP?
3. The RAG baseline implementation lacks clarity - was chain-of-thought prompting used? How much of RATP's improvement comes from better reasoning versus better retrieval?
4. How does RATP's document selection compare to retrieval metrics (precision@k, recall@k)? Is MCTS actually finding better documents or just processing more of them?
soundness: 2
strengths: The paper is well written, it improves upon ToT by unifying document retrieval and thought generation in a single MDP framework, allowing for dynamic graph expansion rather than fixed-tree exploration. The method requires no LLM fine-tuning, preserving privacy and allows easy deployment. The method scales across different LLMs. The papers shows strong improvements against baselines experiments.
summary: Large Language Models (LLMs) two major challenges in healthcare such as:
1. Privacy concerns with sensitive medical data
2. Difficulty in reliably accessing and using external knowledge

RATP addresses these by formulating thought generation as a markov decision process where states are reasoning graphs and actions are combinations of thoughts with retrieved documents. The privacy challenge is addressed by keeping sensitive data separate from LLM training, using only frozen inference with retrieved chunks. For retrieval integration, RATP treats documents as potential thoughts in the state space, allowing MCTS to optimize the document selection and reasoning paths, unlike standard RAG which uses fixed retrieval.

The paper shows strong results on benchmarks against their baselines.
weaknesses: My main concerns with the paper is that it doesn’t compare against strong methods that have been established in the literature such as ReACT [1] or DSP [2] with the only one being ToT. Compared to ToT, how does it compare on LLM calls, total input + output tokens, and time?  The evaluation of retrieval quality is incomplete - there's no analysis of whether MCTS truly finds optimal documents or just compensates for poor retrieval through more reasoning. Does the RAG baseline incorporate Chain of Thought? 

[1] Yao et al. ReAct: Synergizing Reasoning and Acting in Language Models
[2] Khattab et al. DEMONSTRATE–SEARCH–PREDICT: Composing retrieval and language models for knowledge-intensive NLP

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: + Is $p_{doc}$ a predefined constant? If so, this should be clarified in line 205. Additionally, the pseudocode in the Appendix should be moved to the main article to improve understanding of the expansion phase.

+ Is there a reason why some of the methods listed in Table 2 are not included as baselines?
soundness: 2
strengths: + The authors propose a novel approach that integrates RAG and MCTS into a single framework.
summary: The authors addressed privacy data leakage in language model training and inference. They proposed a novel retrieval-augmented thought process that formulates multi-step thought generation as a Markov Decision Process, effectively solving it using Monte Carlo Tree Search. They evaluated the proposed method and baselines on two biomedical and two open-domain QA tasks, demonstrating the effectiveness of their approach.
weaknesses: + The experimental results are not sufficient to demonstrate the effectiveness of the proposed method. 
For instance, S. Soni and K. Roberts [1] showed that they achieved an EM accuracy of 70.56 on emrQA using the training dataset and ClinicalBERT. 
Since one of the proposed reward models requires the training data, the authors should include additional baselines that utilize the training data of each dataset.
Furthermore, X. Zhang et al. [2] proposed a pipeline that improves SLM performance by incorporating medical contexts from LLMs in privacy-restricted scenarios. 
Including such baselines would strengthen the authors' contribution.

+ The authors do not conduct any privacy analysis. 
Although they used a frozen LLM, the reward model was trained on a privacy-sensitive dataset, which methods like the model inversion attack [3] could extract training data. 
The authors should include a privacy analysis of the generated output, as done in [2].

+ The strength of the proposed method is unclear. 
Based on the results in Table 9, this approach does not appear to generalize well to open-domain QA datasets. 
Additionally, it is not specifically designed for the biomedical domain. 
For example, the authors could more carefully design the self-critic prompt to incorporate essential domain knowledge for each biomedical task. 
Another concern is the computational efficiency of the proposed method. 
While MCTS is effective for complex reasoning tasks, it is known to be computationally costly. 
I suggest that the authors include an efficiency analysis to compare the performance gains against the computational cost.


[1] S. Soni and K. Roberts, Evaluation of Dataset Selection for Pre-Training and Fine-Tuning Transformer Language Models for Clinical Question Answering, LREC 2020 

[2] X. Zhang et al., Enhancing Small Medical Learners with Privacy-preserving Contextual Prompting, ICLR 2024

[3] S.V. Dibbo et al., SoK: Model Inversion Attack Landscape: Taxonomy, Challenges, and Future Roadmap, CSF 2023

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. There appears to be an inconsistency in Figure 3 regarding document selection. In the "Expansion" step, there is no arrow indicating that \( I_4 \) is selected, yet it is included in the prompt for generating \( \phi_{10} \). This raises questions about how \( I_4 \) was chosen for the prompt without a connecting arrow in the illustration. Clarification on this selection process and whether \( I_4 \) should indeed be part of the prompt would improve the figure’s accuracy and help clarify the decision-making flow in the MCTS process.
2. As mentioned in the weaknesses, it would be beneficial to include a comparison of inference costs. Given the high computational demands of the multi-step thought process, a detailed analysis of inference costs compared to other methods would provide valuable insights into the efficiency of the RATP approach.
soundness: 1
strengths: 1. **Novelty**: The authors introduce RATP, which combines retrieval with a thought process through MCTS, enhancing LLM reasoning and robustness to imperfect retrievals.
2. **Focusing on Healthcare problem**: The RATP addresses privacy concerns which is highly relevant in the healthcare domain. By avoiding the use of sensitive data in training, the approach enhances the applicability of LLMs in handling sensitive healthcare information safely and efficiently.
3. Empirical Results: RATP demonstrates significant improvements in accuracy (35%) over standard retrieval-augmented generation methods on private datasets, underscoring its utility in healthcare QA task
summary: This paper presents the Retrieval-Augmented Thought Process (RATP), an approach designed to enhance large language models (LLMs) in healthcare by allowing them to interact with external knowledge sources while maintaining data privacy. RATP leverages Monte-Carlo Tree Search (MCTS) to support LLM decision-making in a multi-step reasoning framework, specifically addressing privacy concerns with private datasets excluded from LLM training. The approach yields substantial accuracy improvements over baseline retrieval-augmented generation (RAG) methods, notably achieving a 35% gain in accuracy on question-answering tasks over EMR datasets.
weaknesses: 1. **Privacy Concerns Still Existing**: While the training-free method mitigates some privacy risks by excluding sensitive data from training, it limits applicability to proprietary or on-premise LLM setups. Additionally, it raises questions about whether training-free alone is sufficient to address all privacy concerns, as data handling during inference still poses potential risks.
2. **Appropriateness of Experimental Setup**: The experimental setup for emrQA raises questions about the feasibility of accurately retrieving the "**oracle**" EMR. Since emrQA queries are designed to be answered directly from specific patient records, it is unclear how the system reliably identifies and retrieves the relevant EMR document (oracle) to support accurate answers. This critical detail is not fully explained in the paper, complicating understanding of the retrieval process and the model's effectiveness in real-world scenarios.
3. **High Inference Cost**: The RATP framework involves a multi-step thought generation process using Monte-Carlo Tree Search, which increases inference time and computational cost. Given the linear increase in LLM queries with the number of thought steps, extensive thought processes can become financially and temporally demanding, particularly in real-time applications.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Could you clarity what LLM indicates in the paper.

2. Could you compare with simple baselines such as COT with or without self-consistency.

3. Could you should the effectiveness of your model across basemodel family and size?

4. Could you show the computational cost comparison comparing to baselines?
soundness: 2
strengths: This paper take the privacy question in healthcare into consideration, which is very critical in the healthcare domain,  and propose corresponding  methods to enhances the interaction between LLMs and private data.

The experimental results comparing to proposed baselines looks promising in the EHR QA datasets.
summary: This paper introduces a framework, Retrieval-Augmented Thought Process (RATP), designed to integrate external information retrieval with LLMs to safely and effectively use private medical data. It addresses challenges in privacy, reasoning, and robustness when LLMs access sensitive information for healthcare applications.
weaknesses: 1. Incremental Contribution:
While the paper addresses several critical issues in the healthcare domain—such as privacy, hallucination, and reasoning—it relies heavily on existing approaches, including RAG, multi-step thinking paths, and scoring systems, without introducing fundamentally new ideas.

2. Ambiguity in Baseline Models:
The authors mention using "LLM" as a baseline but fail to specify which LLM or group of LLMs they experimented with. This lack of clarity makes it difficult for readers to accurately assess the reproducibility and comparability of their results.

3. Lack of Comparison with Simpler Methods:
The paper does not compare the proposed method with simpler and widely adopted reasoning frameworks, such as Chain of Thought (CoT) with or without self-consistency. Including such comparisons would have provided better insights into the benefits of using the proposed Monte Carlo Tree Search (MCTS) framework.

3. High Computational Cost:
Although the use of MCTS improves the reasoning capacity of the LLM, the computational cost is likely much higher than that of simpler baselines. A detailed comparison of computational overhead, including the number of LLM queries and time or resource consumption, would strengthen the evaluation and help justify the trade-off between performance and efficiency.

4. Insufficient Coverage of Related Work:
The related work section lacks coverage of recent advancements in Medical LLMs and privacy in healthcare. Here are some works are notably missing:

Medical LLMs: 
Li et al. (2023): ChatDoctor: A Medical Chat Model Fine-Tuned on a Large Language Model (LLaMA) Using Medical Domain Knowledge
Zhang et al. (2023): AlpaCare: Instruction-Tuned Large Language Models for Medical Applications
Kim et al. (2024): Health-LLM: Large Language Models for Health Prediction via Wearable Sensors

Health-domian privacy:  
Zhang et al. (2022): Enhancing Small Medical Learners with Privacy-Preserving Contextual Prompting

Including these references would provide a more comprehensive understanding of related work in the healthcare domain, particularly concerning privacy and medical LLMs.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: None
soundness: 3
strengths: According to the authors the method greatly increases the performance over the tested dataset.
summary: The paper describe a novel method to to expand LLM knowledge using recursive thought process that includes retrieving relevant documents. This method is extremely useful in the scenarios where private data is needed to answer a question.
weaknesses: - I think the paper is written in too abstract and over formulated manner, a concrete step-by-step walkthrough example of the process would be very useful to better understand  the process. 

- What is the reason that RAG underperforms LLM in some cases, it is especially peculiar when we know the answers requires private data. Is the retrieval wrong? Does the model has issue utilizing the retrieved data? An analsys of the reason for failure would be useful and can contribute to the understanding of why your method is successful

- Figure 4 seems to be missing the MCTS oracle #LLM queries line, generally this figure is hard to read to the fact it has two y values, splitting the figure and adding the missing line would be useful

- For easier understanding of the method cost I think total tokens used would be easier to understand. Can total tokens per method be added to the analsys?

Meta-review: The paper introduces Retrieval-Augmented Thought Process (RATP), leveraging Monte Carlo Tree Search to enhance privacy-preserving reasoning in healthcare, achieving superior accuracy and robustness against imperfect retrieval in question-answering tasks.

Strengths
- RATP combines retrieval-augmented reasoning and Monte Carlo Tree Search to optimize decision-making in LLMs while ensuring privacy, addressing critical healthcare problems.
- RATP achieves up to 35% accuracy improvements on private datasets, demonstrating its effectiveness compared to baseline methods like RAG and Tree-of-Thought.

Weaknesses
- The paper relies  heavily on existing methods (e.g., retrieval-augmented generation, multi-step reasoning) without introducing fundamentally new ideas, significantly reducing the originality of its contribution.
- The lack of comparisons with well-established methods like Chain-of-Thought and other relevant baselines (e.g., ReAct, ClinicalBERT) weakens the evaluation of the proposed framework's benefits
- RATP’s multi-step reasoning with MCTS significantly increases inference costs, raising doubts about its practicality for real-time or resource-constrained applications

Decision: Reject