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

Probleemstelling

In agente gezondheidszorgsystemen worden taakgespecialiseerde modellen (specialisten) gebruikt als externe hulpmiddelen om klinische vragen te beantwoorden, zoals ziektediagnose, rapportgeneratie en visuele grounding. Echter, er bestaat zelden één "beste" model voor een specifieke taak. In de praktijk presteren verschillende modellen beter op verschillende data-stalen vanwege:

• Domeinverschuivingen: Variatie in ziekenhuizen, scanners en patiëntenpopulaties.
• Labelruimte-mismatch: Verschillen in definities, granulariteit en prevalentie van labels in trainingsdatasets.
• Heterogeniteit: Modellen die goed presteren in één setting, falen vaak in een andere.

Dit creëert een kritieke bottleneck: hoe kan een agent betrouwably het juiste specialistische model selecteren uit een heterogene pool van kandidaten voor een specifieke query? Bestaande systemen gebruiken vaak één vast model of een naïeve routering, wat leidt tot suboptimale prestaties en onbetrouwbare klinische antwoorden.

Methodologie: ToolSelect

De auteurs introduceren ToolSelect, een leerbaar kader voor query-gebaseerde adaptieve modelselectie. Het systeem houdt alle onderliggende specialistische modellen "bevroren" (frozen) en leert alleen de selector.

1. Formele Probleemdefinitie:

• De selectie wordt geformuleerd als een probleem van het minimaliseren van populatierisico over een panel van kandidaat-modellen.
• Voor elke query $p$ en een panel van tools $S_t$, moet een selector $r_\theta$ een score toekennen aan elke tool.
• Het doel is om de "population selection loss" te minimaliseren, waarbij rekening wordt gehouden met de kosten (fouten) van de geselecteerde tool.

2. Architectuur (Attentive Neural Process - ANP):
De selector is gebaseerd op Attentive Neural Processes (ANP) en werkt als volgt:

• Referentie-sets: Elk specialistisch model heeft een kleine, few-shot referentie-set ($D_E$) die zijn gedrag samenvat op de taak (input, ground-truth, en voorspelling).
• Koppeling Query en Referentie: De input-query (beeld + tekst) wordt gecodeerd en gebruikt om via cross-attention te "kijken" naar de referentie-set van elk kandidaat-model.
• Adaptieve Representatie: Dit creëert een query-afhankelijke descriptor voor elk model, die benadrukt welke referentievoorbeelden het meest relevant zijn voor de huidige query.
• Selectie: Een lichtgewicht MLP (Multi-Layer Perceptron) scoret de modellen op basis van de query, de ANP-descriptor en optionele metadata. De scores worden omgezet in selectiekansen via een gemaskerde softmax (modellen die de taak niet ondersteunen, worden gemaskeerd).

3. Trainingsdoel:
Omdat de directe selectie niet differentieerbaar is, gebruiken de auteurs een comp-sum surrogate loss (gebaseerd op Mao et al., 2023). Dit stelt het systeem in staat om de selector te trainen door de verwachte kosten te minimaliseren zonder de onderliggende modellen te finetunen.

Belangrijkste Bijdragen

1. Nieuw Kader en Dataset (ToolSelectBench):

   • Dit is het eerste werk dat de noodzaak benadrukt om meerdere specialisten als tool-kandidaten te onderhouden en een query-geconditioneerde selectiemodule te gebruiken.
   • Introductie van een agente Chest X-ray omgeving met een diverse suite van 55 taakgespecialiseerde modellen:
     • 17 modellen voor ziektedetectie.
     • 19 modellen voor rapportgeneratie.
     • 6 modellen voor visuele grounding.
     • 13 modellen voor Visual Question Answering (VQA).
   • Een benchmark van 1448 queries (ToolSelectBench) die realistische uitdagingen zoals domeinverschuivingen en label-mismatch simuleert.

2. Technische Innovatie:

   • Formalisering van multi-task tool-selectie als populatierisico-minimalisatie.
   • Toepassing van ANP voor query-geconditioneerde selectie zonder test-tijd optimalisatie (geen gradiëntupdates tijdens inferentie).

3. Empirische Validatie:

   • Uitgebreide benchmarks tegen 10 state-of-the-art (SOTA) methoden, waaronder KNN, SVM, Matrix Factorization, en LLM-routers.
   • Publicatie van de outputs van alle 55 modellen om reproduceerbaarheid te garanderen.

Resultaten

De resultaten tonen aan dat ToolSelect consistent presteert beter dan individuele modellen en andere routeringsstrategieën over alle vier taakfamilies:

• Ziektediagnose: ToolSelect bereikte een F1-score van 43,80% op de Open-I dataset en 59,88% op VinDr. Dit is een enorme verbetering ten opzichte van de beste baselines (bijv. +31,82 punten op Open-I) en benadert de "Oracle" (het theoretische maximum door per geval de beste te kiezen) veel dichter dan eerdere methoden.
• Rapportgeneratie: ToolSelect behaalde een F1-RadGraph van 24,17%, wat een verbetering is van +7,84 punten ten opzichte van de beste enkele generator (CheXAgent-8B) en +1,77 punten ten opzichte van de beste baseline (MF Router). Kwalitatieve analyses tonen aan dat ToolSelect klinisch nauwkeurigere rapporten genereert die beter overeenkomen met de ernst van de afwijkingen.
• Visuele Grounding: ToolSelect verbeterde de Mean IoU met ongeveer 10 punten ten opzichte van de beste individuele grounding-modellen, door de juiste specialist te kiezen voor specifieke beeld- en frasecombinaties.
• Visual Question Answering (VQA): ToolSelect bereikte een nauwkeurigheid van 72,01%, vergeleken met 63,22% voor de beste enkele tool. Dit benadrukt dat de "beste" tool sterk varieert per vraag en dat naïeve of statische routering de voordelen van sterke specialisten kan tenietdoen.

Betekenis en Impact

Dit paper legt een fundamentele basis voor de volgende generatie agente gezondheidszorgsystemen. De belangrijkste inzichten zijn:

• Complementariteit: Geen enkel model is universeel superieur; de kracht ligt in de combinatie van meerdere specialisten.
• Selectie is cruciaal: Betrouwbare prestaties in klinische pipelines hangen niet alleen af van de kwaliteit van de individuele tools, maar vooral van het vermogen om de juiste tool te selecteren op basis van de specifieke query.
• Efficiëntie en Betrouwbaarheid: Door de juiste specialist te kiezen, kunnen systemen rekenkosten besparen (door niet alle modellen te draaien) en de nauwkeurigheid en transparantie van klinische beslissingsondersteuning verhogen.

De auteurs benadrukken dat dit systeem bedoeld is om menselijke expertise te ondersteunen, niet te vervangen, en dat verantwoord gebruik vereist dat er klinisch toezicht blijft bestaan.