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

AI summary

Overview

Research area: Reinforcement learning for large language models, applied to multi-turn medical dialogue (clinical information seeking and diagnosis support).

Technical level: Advanced. The paper assumes familiarity with reinforcement learning for LLMs (PPO, GRPO), value functions and critics, Markov decision processes, tree search, and KV caching in inference engines.

Scope: The paper proposes ATPO, an uncertainty-aware adaptive tree-search reinforcement learning algorithm that allocates rollout budget to high-uncertainty dialogue states, and evaluates it on three public medical dialogue benchmarks using Qwen3 models of 1.7B, 4B, and 8B parameters.

What This Paper Is About

In real medical conversations, patients give vague or incomplete information at first, so a good AI assistant must ask clarifying questions across several turns before committing to an answer. The paper argues that existing methods for training this behavior fall short: prompt engineering does not fundamentally improve multi-turn ability, supervised fine-tuning merely imitates training data, and RL methods such as GRPO struggle with long-horizon credit assignment while PPO suffers from unstable value estimation. The goal is a reinforcement learning algorithm that trains a language model to ask better questions over multiple turns, using an adaptive tree search to spend computation where the model is most uncertain.

Key Contributions

1. Adaptive Tree Policy Optimization (ATPO): An uncertainty-aware tree-search RL algorithm that allocates rollout budgets per turn based on a composite uncertainty score combining Bellman error (U₁) and action-value variance (U₂), aiming to improve both sampling diversity and critic accuracy.

2. Efficiency optimizations: An uncertainty-guided pruning mechanism to reduce the number of rollouts, plus an asynchronous search architecture that reuses shared prefixes to exploit KV cache, targeting higher inference throughput.

3. Hierarchical MDP formulation: Multi-turn dialogue is modeled as a Hierarchical Markov Decision Process where a macro-action is one full assistant response (a turn) and a micro-action is a single token, with turn-level advantages uniformly distributed across that turn's tokens.

4. Empirical validation: Experiments across three Qwen3 model sizes and three medical dialogue benchmarks show ATPO outperforming strong RL baselines such as TreePO and GRPO, with a Qwen3-8B model surpassing GPT-4o by 0.92% accuracy on MedQA, and with greater sample efficiency.

Main Findings

• ATPO beats tree and non-tree RL baselines across model sizes. On Qwen3-8B, ATPO (U₁+U₂) reaches 65.87 ± 3.72 on MedicalExam, 64.07 ± 0.43 on MedQA, and 53.66 ± 1.52 on MedMCQA, compared with TreePO at 65.33 ± 3.09, 61.81 ± 0.90, and 54.74 ± 1.99 respectively.

• Surpassing GPT-4o on MedQA at 8B scale. GPT-4o under the MEDIQ prompt scores 64.00 ± 3.53 on MedicalExam, 63.15 ± 0.82 on MedQA, and 53.03 ± 0.89 on MedMCQA; the paper reports ATPO (with Qwen3-8B) exceeding GPT-4o on MedQA by +0.92%.

• Both uncertainty terms matter and are complementary. ATPO (U₁+U₂) generally outperforms ATPO (U₁), which in turn exceeds TreePO. Absolute gains over TreePO on MedQA are 0.82%, 1.73%, and 2.26% for the 1.7B, 4B, and 8B models.

• Higher sample efficiency. On MedQA with Qwen3-4B, ATPO (U₁+U₂) reaches approximately 52.7% accuracy using only 55% of the training turns required by TreePO.

• Better exploration balance across dialogue depth. TreePO's fixed branching causes exponential node growth concentrated in early turns, whereas ATPO's uncertainty-based pruning enables deeper and more balanced exploration across dialogue depths.

• U₁ alone behaves differently from U₁+U₂. U₁ alone drives aggressive early exploration, concentrating expansions at shallow depths of 3–4 and causing steep local drops in node values; U₁+U₂ achieves deeper coverage and a more uniform value variance. U₁+U₂ also produces high sample-return variance (comparable to GRPO, higher than TreePO) while yielding substantially lower critic value loss than PPO (both MDP and H-MDP), with U₁ alone ranking second.

• Visit-count down-weighting is critical for stability. Removing policy down-weighting (variant EXP₁) causes uncontrolled entropy growth and excessive policy clipping. Down-weighting both policy advantage and value loss (EXP₂) causes rapid entropy collapse, after which the policy abandons multi-turn exploration and regresses to single-turn strategies.

• Hierarchical modeling helps modestly. PPO (H-MDP) slightly but consistently surpasses PPO (MDP), scoring higher in 5 out of 9 evaluation settings.

• Zero-shot prompting to ask questions can hurt. The MEDIQ prompting strategy performed worse than a Direct single-turn prompt, consistent with the finding reported in MEDIQ; SFT gave only modest accuracy gains while being important for enabling multi-turn information seeking.

• Distillation from larger models did not help. Distilling from GPT-4o and Gemini-2.5-Pro (Appendix A.3, Table 2) also failed to yield significant improvements, supporting the paper's claim that imitating expert trajectories is insufficient for generalization.

• Generalization to an unseen user simulator. Replacing the Qwen3-8B user simulator used in training with Llama-3.3-70B-Instruct at evaluation left assistant performance similar across test sets (Appendix A.6, Table A.6).

Methodology in Plain English

The researchers frame a multi-turn medical conversation as a hierarchy: at the top level, each "move" is an entire assistant reply (a turn), and at the bottom level, each move is a single generated token.

To train the model, they grow a search tree of possible conversations. The root is the user's initial query. At each node, the assistant can produce several candidate responses (the expansion size N = 4). Each candidate response leads to a next state, either because the user answers a clarifying question (creating a child node) or because the assistant gives a final answer (ending that branch).

For each node the method computes how "uncertain" the state is, using two signals. The first, U₁, is the Bellman error: the gap between the critic's current value estimate for the state and the average one-step lookahead value across the candidate actions, treated as a proxy for aleatoric uncertainty. The second, U₂, is the variance of the candidate action-values, capturing a blend of epistemic and aleatoric uncertainty. These are combined as U = αU₁ + (1−α)U₂, with U₂ normalized by Z-score scaling against historical samples, and α ∈ [0,1] balancing the two (α = 0.3 in the main configuration).

Expansion is threshold-based. If U exceeds a threshold τ, all N branches are kept. If not, usually only one randomly chosen branch is kept, although with a small probability (e.g., 10%) pruning is bypassed and all N branches are expanded to preserve baseline diversity. Expansion stops when dialogues terminate or the tree reaches a leaf-node (width) budget.

Afterward, target values are computed by a backward pass from the leaves: a terminal state's target is its immediate reward, and other states average the one-step TD targets over their child branches. Advantages use a one-step temporal-difference formulation with the critic's values rather than the target values, because keeping a single branch would otherwise yield zero advantage. The tree is then decomposed into root-to-leaf trajectories and optimized with a PPO-style clipped objective, with the KL penalty against a reference policy (β = 0.01) and normalization by state visit count and response length. The critic is the LLM backbone plus a linear value head, predicting from the average over the final h special tokens and trained by mean squared error against the target values.

Efficiency comes from reusing shared prefixes via the KV cache and running answer generation, user-model interaction, and critic value estimation asynchronously within the sampling phase. The implementations build on the VeRL Agentic RL framework, and the paper reports decoding speeds up to 2,500 tokens/sec/GPU on a 1.7B model with TreePO.

Environment details: the User Simulator is Qwen3-8B, instructed to answer strictly from a set of atomic facts and to refuse out-of-scope queries; the authors report verifying it with GPT-4o at 100% accuracy in following instructions and rejecting irrelevant queries, with a 1.2% hallucination rate. Training uses 14,256 samples, 66% (9,400) from the MEDIQ training dataset and 34% (4,856) constructed from the MedMCQA training data. Rewards are +3 for a correct final answer, 0 for incorrect, and −1 for invalid format. Policy learning rate is 1×10⁻⁶, critic learning rate 1×10⁻⁵, γ = 1, the critic is initialized from the actor weights and warmed up for 5 steps, GRPO uses group size 32, ATPO uses N = 4 with total expansion budget 128, ATPO (U₁) uses τ = 0.5, and ATPO (U₁+U₂) uses α = 0.3 and τ = 1.5. The SFT baseline used Gemini-2.5-Pro in self-play to generate 1,269 multi-turn dialogues from the MEDIQ validation dataset. Evaluation reports the mean and standard deviation of five independent runs, with accuracy defined as the percentage of test cases where the chosen option matches the ground truth.

Why This Matters

Impact on research. The paper argues that token-level tree-search RL methods (VinePPO, SPO, TreePO and entropy-based exploration) do not naturally transfer to macro-level decisions in multi-turn dialogue, and offers a turn-level uncertainty measure as an alternative. It also provides evidence that tree-structured credit assignment outperforms a single trajectory-level advantage (TreePO substantially outperforming GRPO in these experiments), and that pure imitation—SFT or distillation from GPT-4o and Gemini-2.5-Pro—does not generalize as well as goal-driven RL.

Real-world applications:

• Clinical intake and triage assistants that must ask clarifying questions before suggesting a course of action.
• Diagnostic decision-support tools that work from vague or incomplete patient descriptions.
• Telemedicine and symptom-checker chatbots that need to gather history efficiently rather than answer prematurely.
• Any multi-turn information-gathering agent, including tool-use assistants, where the paper says ATPO could also be applied.

Industry relevance. The efficiency mechanisms (uncertainty-guided pruning, prefix sharing with KV cache reuse, asynchronous search) matter for the cost of running tree-based RL at scale, and the reported decoding throughput of up to 2,500 tokens/sec/GPU on a 1.7B model is a practical data point. A relatively small 8B model outperforming GPT-4o on MedQA by 0.92% is relevant for deployment cost in healthcare settings. Code is released at https://github.com/Quark-Medical/ATPO, and the paper states the associated datasets are also available there.

Future Directions

• Replacing the fixed-threshold expansion rule with a learnable, soft control policy, which could reduce hyperparameter tuning and let the expansion strategy adapt as the policy evolves.
• Adaptively determining how many nodes to expand based on the uncertainty metrics, rather than choosing the number randomly.
• Refining credit assignment inside the Hierarchical MDP, moving beyond uniformly cloning a turn-level advantage onto every token of that turn, to allow more precise policy optimization.
• Extending ATPO beyond medical dialogue to multi-turn open-ended dialogue and tool use, as the authors suggest.

Target Audience

This paper is most useful to reinforcement learning researchers and practitioners working on LLM post-training, particularly those interested in tree-search RL, credit assignment over long horizons, and uncertainty-driven exploration. It will also interest applied researchers and engineers building multi-turn conversational agents for healthcare who need to evaluate whether a small open model trained with RL can rival large proprietary models. Readers without a background in RL, value functions, and policy optimization will find the methodology section demanding, though the problem framing and results are accessible.

Authors’ abstract

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