PaperlayerView exact revision on arXiv

Source-linked AI summary

ATPO: Adaptive Tree Policy Optimization for Multi-Turn Medical Dialogue

Ruike Cao, Shaojie Bai, Fugen Yao, Liang Dong, Jian Xu, Li Xiao

arXiv:2603.02216v1cs.LGcs.AI

TL;DR

Medical dialogue models need to gather clarifying information when users’ initial queries are incomplete, but training and evaluation often focus on single turns. ATPO uses uncertainty-guided tree search to allocate turn-level exploration, and Qwen3-8B exceeds GPT-4o accuracy on MedQA by 0.92%.

  • Problem

    Medical LLM training and evaluation often center on single-turn responses, although incomplete user queries call for proactive information gathering in multi-turn dialogue.

    Source: §1 · p. 1

  • Method

    ATPO uses uncertainty-guided tree search to allocate rollout budget across dialogue turns, combining Bellman error and action-value variance while reusing shared prefixes and asynchronous execution.

    Source: §1 · p. 2

  • Results

    0.92% higher accuracy on MedQA for Qwen3-8B than GPT-4o.

    Source: §4.2 · p. 8, §5 · p. 9

  • Takeaways & Limitations

    ATPO is presented as applicable beyond multi-turn medical dialogue, including multi-turn open-ended dialogue and tool use.

    Source: §5 · p. 9

    ATPO uses fixed-threshold expansion with random child selection, and its H-MDP credit assignment could move beyond uniform cloning.

    Source: §5 · pp. 9–10

Abstract

from arXiv · show

Effective information seeking in multi-turn medical dialogues is critical for accurate diagnosis, especially when dealing with incomplete information. Aligning Large Language Models (LLMs) for these interactive scenarios is challenging due to the uncertainty inherent in user-agent interactions, which we formulate as a Hierarchical Markov Decision Process (H-MDP). While conventional Reinforcement Learning (RL) methods like Group Relative Policy Optimization (GRPO) struggle with long-horizon credit assignment and Proximal Policy Optimization (PPO) suffers from unstable value estimation in this context, we propose a novel uncertainty-aware Adaptive Tree Policy Optimization (ATPO) algorithm. Our method adaptively allocates the rollout budget to states with high uncertainty, quantified by a composite metric of Bellman error and action-value variance. This strategy enables more accurate value estimation, while fostering more efficient and diverse exploration. To mitigate the high computational cost of tree-based RL, we introduce two key optimizations: an uncertainty-guided pruning mechanism to minimize the number of rollouts, and an asynchronous search architecture that leverages KV cache reuse to maximize inference throughput. Extensive experiments on three public medical dialogue benchmarks demonstrate that our algorithm significantly outperforms several strong baselines, culminating in Qwen3-8B model surpassing the much larger GPT-4o ($+0.92\%$ accuracy).

1. Introduction

The introduction frames multi-turn medical information gathering as an under-addressed challenge and presents ATPO as an uncertainty-aware approach to training dialogue models. It reports that ATPO outperforms strong RL baselines across datasets and model sizes, with improved sample efficiency.

1. Introduction

  • Medical LLM training and evaluation often focus on single-turn interactions, although incomplete initial queries make proactive clarification important in real-world dialogues.Prior approaches include prompting, supervised fine-tuning, and trajectory-level preference optimization; the introduction notes limitations in these approaches and in GRPO and PPO.Source: §1 · pp. 1–2
  • ATPO significantly outperforms strong RL baselines across three medical dialogue datasets and three Qwen3 model sizes.The paper also reports greater sample efficiency and robust generalization.Source: §1 · p. 2
  • ATPO combines Bellman error and action-value variance to guide turn-level rollout allocation, balancing critic training with diverse sampling.High-uncertainty nodes are fully expanded; low-uncertainty nodes are pruned to one randomly selected child before rollout.Source: §1 · p. 2, Figure 1 · p. 2
  • Shared-prefix KV-cache reuse and asynchronous execution are designed to improve training efficiency and inference throughput.The introduction identifies both as optimizations for ATPO’s training process.Source: §1 · p. 2

2. Related Work

Prior work explores reinforcement learning for proactive medical dialogue and tree-based search for LLM reasoning. The paper distinguishes these efforts from its focus on uncertainty-guided, turn-level exploration for multi-turn interactions.

2.1. Reinforcement Learning in Multi-turn Medical Dialogues

  • Medical dialogue RL studies include simulated doctor–patient agents, Shapley Information Gain rewards for question utility, and branched conversation training, but a robust algorithm for real-world dialogue remains unexplored.Examples include DoctorAgent-RL [10], ProMed [7], and Savage Conversation Forests [30].Source: §2.1 · p. 3

2.2. Tree-based Reinforcement Learning

  • Single-turn tree-search methods target value estimation, uncertainty-guided exploration, or search structure, but their token-level units and indirect or absent uncertainty measures do not naturally fit multi-turn decisions.ATPO instead uses turn-level Q-value variance to represent uncertainty over future rewards for conversational macro-actions.Source: §2.2 · p. 3

3. Methods

ATPO models multi-turn dialogue hierarchically and explores dialogue trees by allocating rollouts according to state uncertainty. It then uses the resulting tree to estimate values, calculate advantages, and optimize the policy.

3.1. Multi-Turn Dialogues as a Hierarchical MDP

  • A turn-level assistant response is a macro-action, while its individual tokens are micro-actions generated within a nested low-level MDP.The high-level state contains the prior interaction history and current user query; the low-level state adds the partial response generated so far.Source: §3.1 · p. 3

3.2. Dialogue Exploration via Tree Expansion

  • Incremental tree expansion reuses shared dialogue prefixes and directs sampling toward promising or uncertain states instead of repeatedly exploring from the root.Clarifying questions create next-turn nodes, while definitive answers terminate their branches; expansion stops at the leaf-node budget or when dialogues terminate.Source: §3.2 · pp. 3–4, §3.3 · p. 5

3.3. Uncertainty-Aware Tree Expansion Mechanism

  • Frontier uncertainty combines critic Bellman error and candidate action-value variance; high-uncertainty nodes retain all branches, while low-uncertainty nodes usually retain one.The uncertainty components are normalized using historical samples and balanced by α; low-uncertainty pruning is bypassed with a small probability, such as 10%, to preserve diversity.Source: §3.3 · pp. 4–5

3.4. Value Traceback, Tree Decomposition and Model Optimization

  • A backward pass sets terminal targets to immediate reward and other targets to mean child-branch TD targets, then derives advantages and turns root-to-leaf paths into PPO-style trajectories.When only one branch is retained, critic estimates preserve non-zero learning signals where the target value would yield zero advantage.Source: §3.4 · p. 5

4. Experiments

Across three medical-dialogue benchmarks, ATPO achieves the highest accuracy in most settings and surpasses GPT-4o on MedQA with Qwen3-8B. The experiments also examine its sample efficiency, uncertainty components, training stability, and generalization to an unseen simulator.

4.1. Experiment Setup

  • The environment pairs an Assistant Agent with a User Simulator that answers only from predefined atomic facts, while the assistant queries until answering or reaching the turn limit.Evaluation uses MedQA, MedMCQA, and MedicalExam, with final-answer accuracy reported as the mean and standard deviation over five runs.Source: §4.1 · p. 6, §4.2 · p. 8
  • The baselines include direct and interactive zero-shot prompting, SFT and dynamic fine-tuning, PPO, GRPO, and binary-tree TreePO.SFT dialogues are generated in self-play using Gemini-2.5-Pro; RL rewards depend only on final-answer correctness.Source: §4.1 · p. 6

4.2. Results

  • Zero-shot MEDIQ prompting performs worse than direct single-turn prompting, while SFT yields modest accuracy gains and enables multi-turn information seeking.Distilling dialogues from GPT-4o or Gemini-2.5-Pro also failed to yield significant improvements.Source: §4.2 · p. 8
  • ATPO achieves the highest accuracy in most settings and exceeds GPT-4o on MedQA by 0.92% with Qwen3-8B.Combining U1 and U2 yields the best results, with absolute gains over TreePO on MedQA of 0.82%, 1.73%, and 2.26% for the 1.7B, 4B, and 8B models, respectively.Source: §4.2 · p. 8
  • On MedQA with Qwen3-4B, ATPO (U1 + U2) reaches approximately 52.7% accuracy using only 55% of TreePO’s training turns.Uncertainty-based pruning supports deeper, more balanced exploration than TreePO’s fixed branching; PPO (H-MDP) also scores higher than PPO (MDP) in 5 of 9 settings, and TreePO substantially outperforms GRPO.Source: §4.2 · p. 8

4.3. Ablation and Analysis

  • Combining U1 and U2 produces higher return variance than TreePO and lower critic value loss than PPO, while U1 alone reduces diversity and ranks second on critic loss.U1 alone concentrates expansion at shallow depths (3–4); U1 + U2 explores more deeply with more uniform value variance.Source: §4.3 · p. 8
  • Without visit-count down-weighting, policy updates cause uncontrolled entropy growth and excessive clipping; down-weighting value loss as well leads to entropy collapse and weaker multi-turn exploration.The authors attribute the latter instability to underweighting value loss at high-visit nodes, which they describe as providing the most reliable value estimates.Source: §4.3 · p. 9
  • After replacing the training-time Qwen3-8B simulator with an unseen Llama-3.3-70B-Instruct simulator, the assistant maintains similar performance across all test sets.The authors interpret this as evidence of generalization beyond one simulator’s conversational patterns.Source: §4.3 · p. 9

5. Conclusions and Future Work

The paper presents ATPO as an adaptive tree-search method for multi-turn dialogue, reporting strong results on medical benchmarks. Future work could make expansion control learnable and refine how turn-level advantages are assigned to token actions.

5. Conclusions and Future Work

  • ATPO guides multi-turn dialogue exploration through state-uncertainty evaluation, outperforming strong RL baselines across clinical benchmarks and GPT-4o on MedQA with Qwen3-8B.The conclusion also identifies multi-turn open-ended dialogue and tool use as possible application areas.Source: §5 · p. 9
  • Future work could replace fixed-threshold expansion with a learnable soft control policy that adapts as the policy evolves and reduces hyperparameter tuning.The authors also suggest selecting the number of expanded nodes from uncertainty metrics rather than random selection.Source: §5 · pp. 9–10
  • Refining H-MDP credit assignment beyond uniform cloning of high-level advantages to token actions may improve policy optimization.The paper presents this as a future direction, not an established improvement.Source: §5 · p. 10

A.1. Interaction Process between Assistant Agent and User Simulator

The interaction begins with an incomplete clinical query and proceeds through targeted questions and fact-bounded simulator replies until the assistant answers or reaches the turn limit.

A.1. Interaction Process between Assistant Agent and User Simulator

  • The assistant asks targeted questions to obtain missing facts, while the simulator replies only from predefined atomic facts, until a final answer or the turn limit.The schematic starts from an incomplete user query and an answer-option set.Source: §A.1 · p. 14, Figure 3 · p. 14

A.2. Test Datasets

The three test sets are assembled from different sources and use distinct filtering or fact-generation procedures. MedicalExam contains 150 samples, MEDQA 1,268, and MedMCQA 536.

A.2. Test Datasets

  • MedicalExam combines five sources and contains 150 samples, with Gemini-2.5-pro used to structure MedMCQA and MMLU items lacking atomic facts.The sources are MedQA, MedMCQA, MMLU, SelfExam, and QMAX.Source: §A.2 · p. 15
  • MEDQA is derived from the MEDIQ medical-dialogue test set, filtered to exclude samples with empty atomic facts, leaving 1,268 samples.Source: §A.2 · p. 15
  • MedMCQA is constructed from its official validation set by selecting questions over 150 characters and synthesizing atomic facts and questions, yielding 536 samples.Source: §A.2 · p. 15

A.3. Distillation Experiment and Analysis

The distillation experiment finds little benefit from adding GPT-4o- or Gemini-2.5-Pro-generated dialogues to Qwen3-8B SFT, while identifying RL as necessary for stronger generalization.

A.3. Distillation Experiment and Analysis

  • Additional GPT-4o or Gemini-2.5-Pro dialogue data and longer SFT training produced almost no improvement in Qwen3-8B performance.Table 2 reports the comparison of Qwen3-8B distilled from each model’s dialogue outputs.Source: §A.3 · p. 15, Table 2 · p. 15
  • SFT primarily teaches the question–answer format and does not by itself provide strong generalization, even with large-scale data.The authors describe SFT as a foundation for subsequent RL training.Source: §A.3 · p. 15
  • The SFT starting checkpoint can affect RL sampling data and thereby influence subsequent RL training.Source: §A.3 · p. 15

A.4. Analysis of Model Effective Question Rate with ATPO Training Progress

As ATPO training progresses, the proportion of effective questions rises, and the authors connect this trend to more useful information gathering in fewer dialogue turns.

A.4. Analysis of Model Effective Question Rate with ATPO Training Progress

  • The proportion of effective questions increases steadily throughout ATPO training.Effective questions are defined as those that elicit valid responses from the user.Source: §A.4 · p. 15
  • The authors say more answerable questions let the assistant gather more useful information and complete tasks in fewer dialogue turns.Source: §A.4 · p. 15
  • The authors associate this change in questioning with an enhanced overall user experience.Source: §A.4 · p. 15

A.5. Comparative Analysis of the Computational Requirements of Different Algorithms

The computational comparison reports that ATPO reaches a target accuracy fastest and uses more rollout-stage value estimation, while the authors attribute its lower total resource use to higher-quality sampling and faster convergence.

A.5. Comparative Analysis of the Computational Requirements of Different Algorithms

  • ATPO reaches the target accuracy in the shortest time, while GRPO requires the longest, under identical hardware.The authors state that shorter training time translates to lower computational resource consumption and that ATPO’s gains are not attributable to greater compute.Source: §A.5 · pp. 15–16
  • About 45% of ATPO’s total training time is spent in rollout, compared with about 25% for other methods; over half of ATPO rollout computation is value estimation.Source: §A.5 · p. 16
  • The authors report that value-estimation-focused rollout produces higher-quality sampling data and faster convergence, reducing total training time and overall resource consumption.They also report PPO convergence without further improvement and entropy explosion and crashes for GRPO and TreePO during continued training.Source: §A.5 · p. 16

A.6. Tested with Different User Simulators

The appendix reports that ATPO-trained assistants retain performance with an unseen user simulator and analyzes the sampling-cost savings of tree expansion; it also provides interaction prompts and an atomic-facts example.

A.6. Tested with Different User Simulators

  • With Llama-3.3-70B-Instruct substituted for the Qwen3-8B training simulator, the three ATPO-trained Qwen3 assistants showed no significant performance differences.The authors interpret this as evidence of no overfitting to the training simulator’s tone or response style and effective performance with an unseen simulator.Source: §A.6 · p. 16, Table 5 · p. 18

A.7. Analysis of the Computational Savings of Tree Search

  • Tree expansion performs the prefill step once before decoding N samples, and the analysis says ATPO inherits this computational cost reduction.The subsection presents a theoretical analysis of potential savings from tree expansion.Source: §A.7 · pp. 18–19

A.8. Prompts

  • The assistant prompt allows a clarifying question when confidence is insufficient, limits interaction to eight turns, and requires a final answer by the last turn.Source: §A.8 · pp. 19–20

Atomic Facts Data Example

  • The atomic-facts example lists five facts, including exam-triggered symptoms, no loss-of-consciousness history, and EEG findings suggestive of epileptic spikes.Source: Atomic Facts Data Example · p. 21

Related papers

Loading 2603.02216v1…