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

TL;DR

This paper introduces ToolSelect, an attentive neural process-based method for selecting the most suitable task-specialized models in healthcare, demonstrating superior performance on a new Chest X-ray benchmark.

Contribution

The paper proposes a novel model selection approach using attentive neural processes and introduces a new agentic Chest X-ray environment with a comprehensive benchmark.

Findings

ToolSelect outperforms 10 state-of-the-art methods across multiple tasks.

The new benchmark includes 1448 queries across diverse medical tasks.

The approach effectively adapts to heterogeneous model pools in healthcare applications.

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.