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

1. Problem Formulation

The paper addresses a critical bottleneck in agentic healthcare systems: the selection of the optimal "specialist" model from a heterogeneous pool of task-specific tools for a given clinical query.

• Context: While Large Language Models (LLMs) can plan and reason, they often lack up-to-date medical knowledge or the ability to process high-dimensional data (e.g., medical images). Consequently, agentic systems rely on external task-specialized models (tools) for tasks like disease diagnosis, report generation, and visual grounding.
• The Challenge: In real-world clinical settings, no single "best" model exists for every scenario due to:
  • Domain Shifts: Variations in scanners, acquisition protocols, and patient populations across hospitals.
  • Label-Space Mismatches: Differences in label definitions, granularity, and prevalence across training datasets.
  • Partial Support: A model might excel at detecting one disease but fail on another, or lack support for specific output components.
• The Goal: Develop a mechanism to adaptively select the most appropriate specialist model for a specific input query $(x, q)$ from a pool of candidates $S_t$, minimizing the risk of selecting a model that performs poorly on that specific instance.

2. Methodology: ToolSelect

The authors propose ToolSelect, a learning-based framework that performs query-conditioned model selection without fine-tuning the underlying specialist models (which remain frozen).

A. Core Architecture: Attentive Neural Process (ANP)

ToolSelect utilizes an Attentive Neural Process to condition the selection decision on the query and the behavioral history of each candidate tool.

1. Input Representation:
   • Query: The multimodal prompt (image $x$ and text instruction $q$) is encoded into a feature vector $u(p)$.
   • Tool Behavioral Summaries: Each candidate tool $E_j$ is represented by a small reference set (few-shot examples) $D_E = \{(x_b, y_b, \hat{y}_b)\}$ summarizing its performance on the task.
2. Encoding & Attention:
   • The reference set is embedded and processed via self-attention to create a rich representation of the tool's behavior.
   • Cross-attention is applied between the query representation $u(p)$ and the tool's reference set. This allows the selector to weigh which historical examples of the tool are most relevant to the current query.
   • The output is a query-dependent tool descriptor $\psi_E(p)$.
3. Scoring:
   • A lightweight MLP scores each tool based on the concatenated query features, the tool descriptor, the tool's current prediction, and optional metadata.
   • A masked softmax converts scores into selection probabilities, ensuring tools that do not support the specific task or label space are excluded.

B. Training Objective: Comp-Sum Surrogate Loss

Directly minimizing the selection loss (choosing the tool with the lowest error) is intractable because the argmax operation is non-differentiable.

• Surrogate Loss: The authors adapt the comp-sum (comparative sum) surrogate loss (from Mao et al., 2023). This loss function encourages the selector to assign higher scores to tools with lower costs (errors) relative to others in the panel.
• Regularization: The training includes panel entropy regularization to prevent over-confident early selections and a coverage head to approximate the tool's success probability.

3. Key Contributions

1. Theoretical & Framework Contribution

• Formalization: The paper formalizes tool selection as a population-risk minimization problem over candidate models with partial task support.
• Novelty: It is the first work to highlight the necessity of maintaining multiple specialist models as tool candidates in agentic healthcare and introducing a query-conditioned routing module to select among them.

2. Dataset & Benchmark (ToolSelectBench)

To enable rigorous evaluation, the authors created a comprehensive testbed:

• Environment: An agentic Chest X-ray environment with an LLM orchestrator.
• Tool Zoo: A diverse pool of 55 task-specialized models:
  • 17 Disease Detection models (DenseNet, ResNet, ViT variants).
  • 19 Report Generation models (CheXAgent, LLaVA, Qwen, etc.).
  • 6 Visual Grounding models.
  • 13 VQA models.
• Benchmark: ToolSelectBench, comprising 1,448 queries across four task families, designed to simulate domain shifts and label mismatches.

3. Empirical Contribution

• Comprehensive evaluation of 55 individual specialists, revealing significant performance variability and complementarity (no single model dominates all cases).
• Benchmarking ToolSelect against 10 State-of-the-Art (SOTA) baselines, including heuristic routers, supervised classifiers, and LLM-based routing strategies.

4. Results & Analysis

Experiments were conducted across four task families: Disease Diagnosis, Report Generation, Visual Grounding, and VQA.

• Performance Gains: ToolSelect consistently outperforms all 10 baselines and random selection across all tasks.
  • Disease Diagnosis: On the VinDr dataset, ToolSelect achieved an F1 score of 59.88%, significantly surpassing the best single model (24.81%) and the best baseline (24.44%). It closed a large portion of the gap to the "Oracle" (theoretical upper bound of 66.22%).
  • Report Generation: ToolSelect achieved an F1-RadGraph of 24.17%, outperforming the best single generator (16.33%) and the best baseline (22.40%).
  • Visual Grounding: ToolSelect improved Mean IoU to 50.08%, a ~10 point gain over the best single model (40.66%).
  • VQA: ToolSelect reached 72.01% accuracy, compared to the best single model's 63.22%.
• Oracle Gap: In all tasks, the "Oracle" (selecting the perfect model per instance) significantly outperformed any single fixed model, proving that per-instance specialization is crucial. ToolSelect successfully captures a large fraction of this potential gain.
• Qualitative Analysis: Visualizations show that ToolSelect correctly routes queries to specialists that handle specific anomalies or report styles better, whereas baselines often default to globally strong but contextually inappropriate models, leading to hallucinations or missed findings.

5. Significance and Impact

• Paradigm Shift: The paper argues that the future of agentic healthcare lies not in finding a single "super-model," but in orchestrating a heterogeneous ensemble of specialized tools.
• Reliability: By adaptively selecting tools based on query-specific characteristics, the system mitigates the risks of domain shifts and label mismatches inherent in clinical deployment.
• Efficiency: The approach allows agents to invoke only the necessary specialist, reducing computational costs compared to running all models or relying on a single, potentially over-parameterized model.
• Clinical Relevance: The work provides a pathway to more robust, auditable, and accurate clinical decision-support systems that can handle the variability of real-world medical data.

In summary, ToolSelect demonstrates that learning to route queries to the right specialist model is a solvable and highly effective strategy for building reliable agentic healthcare systems, significantly outperforming static model selection or naive routing heuristics.