Retrieval Augmented Thought Process for Private Data Handling in Healthcare

Authors: Thomas Pouplin, Hao Sun, Samuel Holt, Mihaela van der Schaar
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Privacy, Information-retrieval, Large Language Model, Question-Answering, retrieval based reasoning

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.

Decision: reject
Reviewer scores: 6, 3, 3, 3, 6

Reviews and discussion:
<thread>
creation date:
1729065520693

thread_type:
official_review

<thread_object>
creator: reviewer_vqsa

<content>
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.

strengths:
According to the authors the method greatly increases the performance over the tested dataset.

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?

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730153134177

thread_type:
official_review

<thread_object>
creator: reviewer_xajf

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730425011405

thread_type:
official_review

<thread_object>
creator: reviewer_7jne

<content>
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.

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

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730704388660

thread_type:
official_review

<thread_object>
creator: reviewer_sfub

<content>
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.

strengths:
+ The authors propose a novel approach that integrates RAG and MCTS into a single framework.

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

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730772526870

thread_type:
official_review

<thread_object>
creator: reviewer_cfv5

<content>
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.

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.

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

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732558477919

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We sincerely thank the reviewers for their valuable and constructive feedback. We are encouraged by the reviewers’ recognition of the novelty and effectiveness of our proposed Retrieval-Augmented Thought Process (RATP) framework in addressing key challenges related to private data handling and reasoning in healthcare applications.

$\color{red} CfV5$: “The method requires no LLM fine-tuning, preserving privacy and allows easy deployment. The papers show strong improvements against baseline experiments. ”

$\color{blue} 7Jne$: “The RATP addresses a particular type of privacy concerns which are highly relevant in the healthcare domain: the use of sensitive data during training. By avoiding this, the approach enhances the applicability of LLMs in handling sensitive healthcare information safely and efficiently.”

$\color{magenta} XAJF$: “This paper takes the privacy question in healthcare into consideration, which is very critical in the healthcare domain, and proposes corresponding methods to enhances the interaction between LLMs and private data.”

$\color{magenta} vqsA$: “This method is extremely useful in the scenarios where private data is needed to answer a question.”

---

Reviewers also had questions about the cost efficiency of the method. To address this, we provide a detailed analysis of token costs for each method below and respond to individual concerns in the respective rebuttals. We believe this point, in combination with the reviewer’s addressed concerns, significantly strengthens our paper. Thank you for your valuable feedback.

## P1 Cost efficiency analysis

We thank the reviewers for suggesting an analysis of the cost efficiency of the methods. We agree that presenting inference costs in terms of tokens provides valuable guidance for practitioners to assess the practicality of RATP.

In the table below, we provide a benchmark of the average token cost (± standard error) per question. Due to time constraints and limited computational resources, we reran all methods on the public setting dataset and one of the two private datasets. We observe that RATP, which utilizes Monte Carlo Tree Search (MCTS) incurs higher token costs than most baseline methods on average with the exception of Tree of Thought (ToT) with Information Retrieval (IR). Notably, the cost increase varies significantly by task, for example, the cost increase for the open-domain public setting is relatively modest compared to the private closed-domain setting, where costs rise by a factor of 10. This aligns with the performance gaps observed, as RATP’s advantages over baselines are more pronounced in the private setting. This indicates that RATP is particularly beneficial for private closed-domain tasks.

Finally, RATP is more cost-efficient than ToT with IR, underscoring the advantages of using MCTS over tree search.

| **Dataset** |  | **LLM** | **RAG** | **RATT** | **Self-RAG** | **ToT w/ IR** | **RATP** |
| --- | --- | --- | --- | --- | --- | --- | --- |
| emrQA | Input tokens | 70 (± 0) | 186 (± 9) | 4914 (± 130) | 2199 (± 27) | 12511 (± 71) | 9416 (± 82) |
|  | Output tokens | 4 (± 0) | 4 (± 0) | 1048 (± 16) | 530 (± 14) | 1449 (± 20) | 1180 (± 24) |
| BoolQA | Input tokens | 66 (± 0) | 226 (± 2) | 5262 (± 150) | 1929 (± 12) | 5481 (± 30) | 662 (± 60) |
|  | Output tokens | 1 (± 0) | 1 (± 0) | 1330 (± 31) | 395 (± 24) | 935 (± 9) | 76 (± 11) |

We also note that this analysis focuses solely on inference costs and does not account for training costs, which are particularly relevant for methods like Self-RAG. Fine-tuning large language models can make adapting to new tasks prohibitively expensive, especially in the healthcare domain, where the diversity of tasks (described in Appendix A of the manuscript) necessitates constant retraining.

**Update:** the cost efficiency analysis discussed above will be integrated into a new Appendix in the revised paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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