Picking the Right Specialist: Attentive Neural Process-based Selection of Task-Specialized Models as Tools for Agentic Healthcare Systems

AI summary

Overview

Research area: Machine learning for healthcare agents — specifically, query-conditioned selection among task-specialized clinical models used as tools by an LLM-based agent.

Technical level: Advanced. The paper combines risk-minimization theory (comp-sum surrogate losses), Attentive Neural Processes, and large-scale empirical benchmarking across 55 chest X-ray models.

Scope: The paper introduces ToolSelect, an Attentive Neural Process-based router that picks the best specialist model for each clinical query, together with a new chest X-ray agentic testbed and the 1448-query ToolSelectBench benchmark.

What This Paper Is About

Agentic healthcare systems depend on task-specialized models (disease detection, report generation, grounding, VQA) as external tools, but there is rarely one "best" model for a task: different specialists excel on different samples because of domain shift, mismatched label spaces, and differing annotation quality. The paper asks how an agent can, given a clinical query and a pool of frozen black-box specialists, reliably pick the most appropriate one for that specific input. To answer this, the authors propose ToolSelect and build the first testbed — an agentic chest X-ray environment plus ToolSelectBench — for studying this selection problem reproducibly.

Key Contributions

1. Task, framework, and dataset contribution: The authors state this is the first work to highlight (a) the importance of maintaining multiple specialist models as tool candidates in agentic healthcare systems and (b) the need for a query-conditioned tool-selection module. They build a testbed of 1448 chest X-ray and question–answer triplets and a chest X-ray agentic environment with 17 disease-detection, 19 report-generation, 6 visual-grounding, and 13 VQA tool candidates (55 specialists total).

2. Technical contribution: They formalize multi-task query-based tool candidate selection as population-risk minimization over candidate models with partial task support, and propose ToolSelect, which conditions an Attentive Neural Process selector on the query and per-model behavioral summaries. The paper states it introduces no new theorems but provides theoretical justification for the ANP-based selection and comp-sum loss in Supplements B and C.

3. Empirical contribution: They report standalone performance of all 55 specialists and release their outputs and responses so others can reuse them without re-running training or inference.

4. Benchmarking contribution: They benchmark ToolSelect across four chest X-ray task families. The abstract and Section 3.4 describe adapting 10 SOTA LLM routers; the contribution list and Section 4 describe 12 diverse baselines spanning heuristic, ML-based selectors, and LLM routing strategies.

Main Findings

• No single specialist is uniformly reliable. In disease diagnosis (Table 1), R-All reaches the highest accuracy on both Open-I (84.08%) and VinDr (75.66%) but has the lowest F1 (Open-I 6.61%, VinDr 4.70%) and very low recall (Open-I 14.49%, VinDr 6.87%), reflecting a conservative regime that misses positives. Conversely, D-PC and D-NIH achieve higher Open-I recall (42.39% / 42.93%) at low precision (approximately 10%).

• Large oracle headroom. Oracle reaches 58.00% F1 on Open-I and 66.22% on VinDr, while on VinDr the best single-model F1 is only 24.81% (D-All/D-M-NB) — close to Random (20.81% F1, 53.32% accuracy). This shows the optimal specialist varies substantially across cases.

• Random and baseline routers fall short. Random routing yields 9.54% F1 / 60.15% accuracy on Open-I and 20.81% F1 / 53.32% accuracy on VinDr. The best baseline F1 is MF Router at 11.98% on Open-I and Causal LLM at 24.44% on VinDr; some baselines achieve high accuracy but low F1 (RouterDC: 81.52% accuracy vs 8.51% F1 on Open-I).

• ToolSelect closes much of the gap. On Open-I it reaches 43.80% F1 / 87.64% accuracy (AUC 45.09), improving over the best baseline by +31.82 F1 points and over the best single-model F1 by +30.45. On VinDr it reaches 59.88% F1 / 83.39% accuracy (AUC 49.91), improving over the best baseline by +35.44 F1 points and over the best single-model F1 by +35.07.

• Report generation: modest absolute performance, consistent gains. Toolkit-wide variation is large (CheXAgent-8B: 16.33% F1-RadGraph; CheXpert Plus (C+M): 16.10%; LLaVA-1.5-13B: 1.57%; LLaVA-1.5-7B: 3.05%). ToolSelect achieves 24.17% F1-RadGraph and 57.86 RateScore, improving over Random by +12.46 points, over the best single generator by +7.84, and over the strongest baseline (MF Router, 22.40%) by +1.77. Oracle reaches 27.70% F1-RadGraph and 59.60 RateScore.

• Visual grounding is highly phrase-dependent. ChEX is the most consistently strong single tool (Atelectasis 49.44, Calcification 18.86, Cardiomegaly 71.50), but AG-KD is stronger on Aortic enlargement and Pulmonary fibrosis (22.55 vs 20.22), and RadVLM is best on ILD. ToolSelect reaches 63.98 mAP@0.25, 48.39 mAP@0.5, and 50.08 Mean IoU, versus Oracle at 78.34, 56.31, and 56.20, and Random at 50.90, 10.12, and 32.41.

• VQA: largest oracle gap. ToolSelect achieves 72.01% VQA accuracy versus Oracle 96.49% and Random 46.49%. The strongest individual candidate is NVIDIA CXR NV-Reason-3B at 63.22%, and the best baseline is Elo Router at 63.74%.

• Query-guided selection can beat globally stronger models. In the qualitative report-generation comparison, baselines all default to the globally strong CheXAgent-8B and fail on a major-anomaly case, whereas ToolSelect selects CheXpert Plus (despite its lower average performance) and succeeds.

Methodology in Plain English

The problem is framed as follows. Each user query is a multimodal prompt — an image plus a natural-language instruction — and an LLM orchestrator maps it to a task (disease detection, visual grounding, report generation, or VQA). For that task, the agent sees a panel of specialist models drawn from a population. Each specialist produces a prediction, and each prediction is scored with a bounded task-specific cost: 0/1 error (or calibrated cross-entropy clipped to [0, 1]) for detection, 1 minus IoU for grounding, 1 minus F1-RadGraph for report generation, and 0/1 error for VQA. Tools that cannot support a task or output component emit a null prediction and are masked out of both selection and loss.

Directly minimizing the resulting selection loss is intractable because the selection rule is combinatorial and non-differentiable. The authors adapt a comp-sum surrogate loss, using the logistic form where the decreasing function applied to the selection probability is −log u, weighting each tool by the sum of other tools' costs. For a panel of size 2 this reduces to a cost-sensitive binary margin form involving the score difference.

The selector itself is an Attentive Neural Process. Each tool comes with a small task-specific reference set of behavioral descriptors — examples of (image, ground truth, tool prediction) triples — with reference set size held between 16 and 64 examples in the experiments. Queries are encoded with a task-shared ViT-B/16 image encoder and a CheXbert text encoder. Reference elements are embedded, passed through self-attention, and then used in cross-attention against the query, producing a query-dependent tool descriptor that emphasizes the reference examples most relevant to the current input. This descriptor is concatenated with the query representation, the tool's aligned prediction, and optional metadata, and fed to a shared lightweight MLP that scores each tool. Scores are turned into a masked softmax distribution, and the highest-scoring supported tool is selected. Tools themselves are never fine-tuned — they stay frozen and only the router and ANP encoders receive gradients.

Training adds two regularizers: a panel entropy term to discourage overconfident early selections, and a per-tool coverage head trained to approximate the tool's task-conditional success probability. Settings used: two-layer MLP head with hidden width 512 and GELU activations, AdamW with learning rate 3 × 10⁻⁵ and weight decay 10⁻⁴, up to 50 epochs with early stopping (patience 10, min_delta 10⁻⁴), batch size 16, panel entropy weight 0.05, L2 penalty on the selector, and dropout 0.1 in the ANP. Evaluation data comes from Open-I and VinDr-CXR for diagnosis, Open-I for report generation, VinDr-CXR for phrase-level grounding, and 1% of ReX-VQA for multiple-choice VQA, with standard dataset splits used for training and testing.

Why This Matters

Impact on research. The paper reframes tool use in clinical agents from "pick the best model" to "pick the right model for this input," and supplies the first benchmark for studying that problem. It also provides 55 specialists' outputs and responses as supplementary material, lowering the cost of reproducing routing comparisons without re-training or re-running inference, and formalizes selection as population-risk minimization with partial task support — a setting existing routing baselines were not designed for.

Real-world applications:

• Clinical triage and screening: high-recall specialists such as D-RSNA (64.99% recall on VinDr) and D-M-Ch / D-M-NB (62.31% / 62.27% recall on VinDr) may suit screening, while high-precision specialists such as R-BioViL (61.65% precision on VinDr) and E-S-C14 (58.65% precision on VinDr) may suit prioritization or confirmation — a router could pick per case.
• Multi-site deployment: clinical variability across hospitals, scanners, acquisition protocols, and patient populations creates domain shift; a selection layer lets one system draw on models trained under different regimes rather than committing to one.
• Radiology report generation: routing toward clinically aligned generators (e.g., selecting CheXpert Plus when it suits the case) instead of always calling the globally strongest model.
• Phrase-level grounding for specific findings: matching the specialist for the queried phrase, given that ChEX, AG-KD, and RadVLM each win on different findings.

Industry relevance. Any agentic healthcare stack that exposes multiple models as tools faces this routing problem operationally. ToolSelect keeps all tools frozen, so adding or swapping specialists does not require retraining the underlying models, and a lightweight shared selector head plus small per-tool reference sets is a deployment-friendly design compared with fine-tuning every candidate.

Future Directions

• Closing the remaining oracle gap. ToolSelect still lags Oracle substantially on every task (e.g., VQA 72.01% vs 96.49%; disease diagnosis Open-I F1 43.80% vs 58.00%), so better query-conditioned selection remains open.

• Extending and hardening the benchmark. The environment covers four chest X-ray task families and 1448 queries; the paper does not report results for other modalities, so generalization to other imaging domains, tasks, and panel sizes is untested here.

• Scaling and maintaining reference sets. The authors note B_t should be as large as possible while staying feasible for tool maintenance, but they do not report how performance varies across the 16–64 range they used.

• Reconciling baseline counts and expanding comparison. The paper reports 10 SOTA LLM routers in Section 3.4 and the abstract, but 12 baselines in the contribution list and Section 4; clarifying this and testing additional routing paradigms would strengthen the evaluation. The paper points to Supplementary E.1 for an ablation on reference point and aggregation, and Supplementary B and C for theoretical justification.

Target Audience

Researchers and practitioners working on LLM-based agents, model routing and selection, and clinical machine learning — particularly those building healthcare agent frameworks that expose specialist models as tools. It is also relevant to medical imaging researchers interested in chest X-ray benchmarks, domain shift, and label-space mismatch, and to engineers designing deployment infrastructure where multiple competing models must be maintained and chosen among at runtime. The paper is written at an advanced level: it assumes familiarity with population risk, surrogate losses, attention mechanisms, and neural processes.

Authors’ abstract

Task-specialized models form the backbone of agentic healthcare systems, enabling the agents to answer clinical queries across tasks such as disease diagnosis, localization, and report generation. Yet, for a given task, a single "best" model rarely exists. In practice, each task is better served by multiple competing specialist models where different models excel on different data samples. As a result, for any given query, agents must reliably select the right specialist model from a heterogeneous pool of tool candidates. To this end, we introduce ToolSelect, which adaptively learns model selection for tools by minimizing a population risk over sampled specialist tool candidates using a consistent surrogate of the task-conditional selection loss. Concretely, we propose an Attentive Neural Process-based selector conditioned on the query and per-model behavioral summaries to choose among the specialist models. Motivated by the absence of any established testbed, we, for the first time, introduce an agentic Chest X-ray environment equipped with a diverse suite of task-specialized models (17 disease detection, 19 report generation, 6 visual grounding, and 13 VQA) and develop ToolSelectBench, a benchmark of 1448 queries. Our results demonstrate that ToolSelect consistently outperforms 10 SOTA methods across four different task families.