A Decision-Language Model (DLM) for Dynamic Restless Multi-Armed Bandit Tasks in Public Health

1. Simulated Environment Reliance: “The paper relies on simulations for validation. Real-world trials would strengthen findings. A plan for field testing DLM in actual public health settings, including partnerships for pilot studies and addressing ethical and logistical challenges, is needed.”

This work will be followed by a real world evaluation. Domain experts from the NGO played an integral role throughout the model iteration and design process. And we would also like to re-emphasize (see Section 5 and Appendix A) that any subsequent real-world evaluations are only undertaken after securing ethics approval from the NGO's ethics board which is registered with the ICMR (Indian Council of Medical Research). Please see the global Note on Ethical Considerations above for more information. Obtaining approvals and designing and conducting real world evaluation require significant effort from the NGO and may take up to several months to complete due to complexities involved in real-world trials. Success in simulations are required for any subsequent approvals, and our simulations indicate substantial potential gains from the system (as demonstrated in this paper).

2. Field Trials and Testing: “Transitioning to real-world testing is essential. This includes expanding language support, handling larger datasets, and addressing ambiguous prompts.”

We thank the reviewer for this comment and agree that these are critical steps for future envisioned deployment. We discuss potential strategies to handle ambiguous prompts below. As referenced in Sec. 5.6, we agree also that expanded testing across different languages is critical. This may intersect with ongoing research in multilingual LLM evaluation [10], especially in LLM reasoning from Indic languages [11] which are highly represented in the ARMMAN deployment regions. Please see above global comment for additional details in consideration of future envisioned real-world testing.

3. Ethical Considerations: “Address potential negative societal impacts such as data privacy, bias, and effects on vulnerable groups. Tackling these issues and outlining a clear path for real-world use will enhance the practical and ethical value of the work.”

We thank the reviewer for the comment. We summarize the ethical considerations discussed in our paper in the global comment above, and highlight these key sections below.

We include a discussion on societal impacts in Appendix A and a discussion on data usage in Appendix C. The real world dataset is completely anonymized by ARMMAN. Data exchange from ARMMAN to the researchers was regulated through clearly defined exchange protocols including anonymization, read-only researcher access, restricted use of data for research purposes only, and approval by the ARMMAN ethics review committee. Our method helps ARMMAN in training dynamic policies to support vulnerable groups and allows ARMMAN program managers to monitor state-feature distributions over underrepresented demographic groups (see Figure 1 description and Section 5). We will further emphasize above mentioned points on ethical considerations and societal impacts in the paper.

4. Prompt Design Foundation: “Is there a theoretical foundation for designing the prompts used in the model? Referring to sociological theories or frameworks could help in creating more effective and relevant prompts.”

The prompts in our model are primarily designed to capture the practical, real-world scenarios that could be encountered through callership patterns amongst mothers in India and in collaboration with ARMMAN. While there isn't a direct theoretical foundation in prompt design, the prompts are inspired by common decision-making processes and challenges faced by ARMMAN. In addition, we design them in such a way that they become progressively more difficult to reason about, where we include a mixture of tasks that require just recovering one feature, multiple features, and inferring the correct due to ambiguity. Thus, a substantial amount of time was spent thinking about designing the prompt, although the current prompts could benefit from principled approaches.

For example, we believe that integrating sociological theories or frameworks is a wonderful idea, and could indeed enhance the contextual relevance and effectiveness of prompts. For instance, a list of criteria such as the Social Determinants of Health [12] could be used to help guide the development of prompts that would likely improve community needs and values. In addition, a large body of literature is now emerging along the direction of “prompt engineering” [13, 14], or finding methods for improving prompt efficacy.

5. Ambiguous Language Prompts: “The paper notes issues with ambiguous prompts leading to misaligned policies. Are there methods or plans to quantify and address these ambiguities to improve reliability?”

One simple idea is to reject ambiguous prompts and ask for further clarifying information, which is tackled in other work that uses the disagreement between sampled outputs of the model [15] to quantify ambiguity. This is a well-studied area [16, 17], and there should be several strategies available that effectively mitigate ambiguous questions.

1. Abstract Citations: “Avoid citing sources in the Abstract.”

We thank the reviewer for the suggestion. We will remove this citation from the abstract. We include this citation at the end of the introduction.

2. Real-World Validation: “The approach provided in this paper lacks real-world validation.”

This work will be followed by a real world evaluation. Domain experts from the NGO played an integral role throughout the model iteration and design process. And we would also like to re-emphasize (see Section 5 and Appendix A) that any subsequent real-world evaluations are only undertaken after securing ethics approval from the NGO's ethics board which is registered with the ICMR (Indian Council of Medical Research). Please see the global Note on Ethical Considerations above for more information. Obtaining approvals and designing and conducting real world evaluation require significant effort from the NGO and may take up to several months to complete due to the complexities involved in real-world trials. Success in simulations are required for any subsequent approvals, and our simulations indicate substantial potential gains from the system (as demonstrated in this paper).

3. Training Process Explanation: “In some parts of the training process, it is crucial to explain why specific numbers are chosen. For instance, the downstream policy is trained for 5 epochs, why 5? Why are 100 simulation steps selected?”

We thank the reviewer for this suggestion and will include additional details on the selection of training hyperparameters in the Appendix. Choices for these hyperparameters were made following quantitative assessments over a hyperparameter search which showed improved stability and low overfitting for the selected values for our given dataset. Our choices were additionally qualitatively guided by prior works studying network-based multi-armed bandit policy training, which similarly use 100 steps in buffer for simulated settings [9]. We will include this information on hyperparameter selection in the Appendix.

4. Ethical Implications: “This paper should also discuss the ethical implications of using AI for decision-making in health resource allocation.”

We discuss this extensively in Appendix A, Social Impact Statement, as well as Appendix B, Dataset Description and Appendix C, Dataset Consent for Collection and Analysis. We summarize these points discussed in our paper in the global comment above, including an overview of the ethical considerations taken prior to and during this work with our partner NGO.

5. Dataset Usage: “Why use the ARMMAN dataset exclusively in this study? What language is used in this dataset? Given that the test involves prompts in English, was there any preprocessing of the dataset required?”

We focus on the ARMMAN dataset because it is vetted, real-world data collected with consent ethically. Such a dataset, created from a study over 7,668 mothers, has the potential for great impact, helping us gain insight into key policy design techniques that can help improve the deployment of public health resources across millions of beneficiaries. Furthermore, the partner organization ARMMAN has an operational AI/RMAB deployment that works to allocate resources. The collected dataset is also rich and not sparse, allowing a diverse set of tasks to be tested within it. Thus, there are multiple important reasons to work with this dataset. The dataset was already collected in English. Thus, no preprocessing of the dataset was required. Due to these key strengths and benefits, we focus our analysis entirely on this highly valuable source of data.

1. “The related works section is not comprehensive…”

We thank the reviewer for the comment. There is growing research in LLMs for healthcare (see Sec. 1). In particular, a question summarization framework in healthcare domains has been proposed [1, 2]. Methods based on contrastive language image pretraining (CLIP) and LLMs to identify medical disorders, generate relevant context, and craft summarizes with visual information have been studied [3, 4]. ChatGPT applications in medical, dental, pharmacy, and public health have been investigated [5, 6]. However, the potential of LLMs to dynamically adapt resource allocation using language prompts, potentially enabling automated policy tuning using expert and community health insights, remains unstudied (see Sec. 1). We will highlight healthcare related works in the paper.

2. Clarify "policy and critic", "features (z)", and "buffer"

We describe the details of actor-critic algorithms in Appendix F. To improve clarity, we will move this explanation to earlier in Section 3. Our novel reward proposal loop enhances LLM-generated reward functions using RMAB simulation feedback, and is compatible with various policy learning algorithms, including the widely used actor-critic algorithms. The features describe demographic information such as age range, education level, and income; examples of features (z) are shown in Sec. 4.2, Sec. 4.4 and Sec. 5. We will provide examples earlier at the beginning of Section 4. We train each policy network using RMAB simulations following the concept of experience replay buffers first introduced by [7], widely used to train RL agents using previously observed transitions [8]. We record per-arm state, action, reward, next state, and lambda over n_s timesteps (see Alg. 1 line 11), and store these “trajectories” of arms in the buffer D, which is then used to update our actor-critic algorithms (see Alg. 1 line 13, App. F).

3. “The DLM with reflection works better, but the comparison is not a good fair comparison. Justifying why without reflection works good or bad is necessary.”

We appreciate the reviewer’s feedback here, but are unclear on why the comparison appears unfair. Could the reviewer clarify their concerns? We refer to DLM (No Reflection) as DLM-NR below. As highlighted (see Sec. 4.5, App. C), LLMs are trained on real-world data, and can therefore recognize and retrieve features relevant to natural language prompts. This is confirmed by our empirical findings: DLM-NR outperforms Default reward policy (see Fig. 2 and Sec. 5.4b), achieving a mean normalized reward (MNR) score of 0.87 +- 0.008 vs. Default score 0.57 +-0.027, with DLM-NR significantly outperforming Default in 11/16 tasks (see Sec. 5.4c). We justify the strong performance of DLM-NR in Table 1, demonstrating more than 70% feature recall in 9/16 tasks; thus, while DLM with Reflection potentially helps fine-tune feature selection and usage (see Fig. 2), DLM-NR still captures relevant features for reward functions without reflection.

4. “Figure 2 does not clearly summarize results… not clear why Default method sometimes matches the performance of DLM (no reflection).”

We present numerical results in Tab. 4, and observe that DLM-NR significantly outperforms Default in 11/16 tasks. In the remaining tasks, while DLM-NR still performs well in feature selection for reward proposals (see Tab. 1), reflection iterations may be required to fine-tune the usage of these features (see Fig. 3). In some cases, a prompt may have a ground truth Base reward that is close to the original Default reward, yielding higher Default performance; yet even in these cases (see Tab. 4, Idx. 2 “Hindi Speakers”) we find that DLM (Reflection) can help fine-tune features and feature weights to improve over DLM-NR and Default.

5. “The baselines are not fair … [why not use] CoT, ReAct, or Reflexion”

We note that DLM is a full framework for the use of LLMs in the reward design of an RMAB. Other cited methods such as Chain-of-Thought Reasoning (CoT) improve base LLM reasoning and, unlike our proposed model, are not a full framework. The listed methods are in fact compatible to use with our framework; we indeed attempted using CoT, but found it did not improve DLM reasoning by any significant margin (see attached PDF for results). We further emphasize that we compare against random (Fig. 2 red line) and no-action (Fig. 2 purple bar) allocation policy baselines. Thus, we compare against five other baselines in total. ReAct and Reflexion frameworks do not specify how to choose the state space, action space, and reward function for the LLM, which are all significant design decisions that DLM solves. It is unclear how they would be directly comparable to our multi-agent RMAB reward design loop.

6. “More baseline comparisons are needed”

See (5). We respectfully disagree. As no existing RMAB reward design loop exists before our method, all other comparisons would be additions on top of the proposed DLM framework. We also would emphasize to clarify that in Figure 2, our main results, we indeed compare against 5 other baselines:“No Action”, “Random”, “No Reflection”, “Base”, and “Default”. We clarify that the “Random” baseline is the red line in Fig. 2.

7. “The level of automation is not clear… the need for human prompts reduces its practicality”

The level of automation is akin to interacting with a chatbot, requiring only prompting and evaluation. Our method is significantly more automated than human design, allowing tasks to be specified via natural language. This significantly reduces the burden that nonprofits face supporting low-resource communities by enabling rapid policy adaptability without significant additional human effort. Note: in an actual deployment (see Sec. 5.5a, Sec. 5.6 and global comment above) a human operator would approve DLM-generated policy outcomes using output state-feature distributions, providing an additional layer of system scrutiny and safety.