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

1. Il Problema: La Selezione degli Strumenti nei Sistemi Agentic Sanitari

I sistemi sanitari basati su agenti (agentic systems) stanno emergendo come soluzione per rispondere a query cliniche complesse, integrando modelli linguistici (LLM) con strumenti specializzati esterni. Tuttavia, il paper identifica un collo di bottiglia critico: non esiste un singolo modello "migliore" per ogni compito clinico.

• Variabilità e Spostamento di Dominio: In ambito reale, i dati variano notevolmente tra ospedali, scanner e protocolli di acquisizione. Un modello addestrato su un dataset specifico può fallire su dati provenienti da un'altra fonte (domain shift).
• Spazi di Etichetta Eterogenei: I dataset di addestramento differiscono per definizioni, granularità e prevalenza delle malattie, rendendo impossibile un approccio "taglia unica".
• La Sfida: Per ogni query clinica, un agente deve selezionare in modo affidabile il modello specializzato più adatto da un pool eterogeneo di candidati. Una selezione casuale o basata su metriche medie può degradare le prestazioni rispetto all'uso di un singolo modello robusto.

2. Metodologia: ToolSelect

Gli autori propongono ToolSelect, un framework di selezione basato sull'apprendimento che utilizza i Neural Process Attentivi (ANP) per adattare dinamicamente la scelta del modello in base alla query e al comportamento storico dei candidati.

Formulazione del Problema

Il problema è definito come la minimizzazione del rischio di popolazione su un pannello di candidati di strumenti.

• Input: Una query multimodale $p = (x, q)$ (immagine + testo) e un pannello di strumenti $S_t$ disponibili per un compito $t$.
• Obiettivo: Assegnare una probabilità di selezione $\pi_\theta(j | p, S_t)$ a ogni strumento $j$ nel pannello, massimizzando l'accuratezza prevista.
• Gestione del Supporto Parziale: Gli strumenti possono non supportare tutte le sottotask (es. un modello di rilevamento malattie potrebbe non essere addestrato su una specifica patologia). ToolSelect utilizza un indicatore di validità per mascherare gli strumenti non supportati durante la selezione e il calcolo della perdita.

Architettura di ToolSelect

1. Sintesi Comportamentale (Reference Sets): Ogni modello specializzato è rappresentato da un piccolo set di riferimento (few-shot) che ne riassume il comportamento empirico su compiti specifici (input, ground truth, e previsioni allineate).
2. Codifica e Attenzione:
   • La query e i set di riferimento vengono codificati.
   • Viene applicato un meccanismo di Cross-Attention tra la rappresentazione della query e i set di riferimento di ogni strumento. Questo permette al selettore di identificare quali strumenti hanno performato bene su esempi simili alla query corrente.
3. Testata di Selezione (Selector Head): Un MLP leggero calcola un punteggio per ogni strumento combinando la rappresentazione della query, il descriptor comportamentale (ottenuto dall'ANP) e metadati dello strumento.
4. Funzione di Perdita (Loss Function): Poiché la selezione è un processo discreto e non differenziabile, gli autori utilizzano una surrogata comp-sum (basata su lavoro di Mao et al.) per approssimare la minimizzazione del rischio, rendendo il training possibile tramite discesa del gradiente.

3. Contributi Chiave

Il paper apporta contributi significativi in tre aree:

• Contributo Concettuale e di Dataset:

  • È il primo lavoro a evidenziare la necessità di mantenere multipli modelli specializzati come candidati strumenti e di un modulo di selezione condizionato alla query.
  • Introduce ToolSelectBench, un benchmark composto da 1448 query (triplette immagine-testo-risposta) su raggi X toracici.
  • Crea un ambiente agentic con un "zoo" di 55 modelli specializzati distribuiti su quattro famiglie di compiti:
    • 17 modelli per il rilevamento di malattie.
    • 19 modelli per la generazione di referti.
    • 6 modelli per il grounding visivo (localizzazione).
    • 13 modelli per VQA (Visual Question Answering).

• Contributo Tecnico:

  • Formalizza la selezione degli strumenti come un problema di minimizzazione del rischio di popolazione.
  • Propone l'uso degli Attentive Neural Processes (ANP) per l'adattamento condizionale senza bisogno di ottimizzazione dei parametri a tempo di inferenza (i modelli sottostanti rimangono "congelati").

• Contributo Empirico:

  • Rilascia le previsioni di tutti i 55 modelli per facilitare la riproducibilità.
  • Confronta ToolSelect contro 10-12 metodi basati sullo stato dell'arte (SOTA), inclusi router basati su LLM, metodi euristici e approcci di apprendimento supervisionato.

4. Risultati Sperimentali

I risultati dimostrano che ToolSelect supera costantemente i metodi di base e si avvicina significativamente al limite superiore teorico (Oracle).

• Diagnosi di Malattie: ToolSelect raggiunge un F1-score del 43,80% su Open-I e 59,88% su VinDr, migliorando di oltre 30 punti F1 rispetto al miglior baseline e chiudendo il divario con l'Oracle (che raggiunge ~58-66%). I modelli singoli mostrano prestazioni variabili e spesso opposte (es. alta precisione vs alta recall) a seconda del dataset.
• Generazione di Referti: ToolSelect ottiene un F1-RadGraph del 24,17%, superando il miglior singolo modello (16,33%) e i router basati su LLM. Le analisi qualitative mostrano che ToolSelect seleziona modelli che producono referti clinicamente più allineati, evitando allucinazioni di gravità.
• Grounding Visivo: Il sistema migliora l'IoU medio di circa 10 punti rispetto al miglior modello singolo, dimostrando che la scelta dello strumento dipende fortemente dalla frase specifica e dall'immagine, non solo dalla performance media.
• VQA: ToolSelect raggiunge il 72,01% di accuratezza, contro il 63,22% del miglior modello singolo e il 46,49% della selezione casuale, evidenziando l'importanza di una selezione guidata dalla query.

5. Significato e Impatto

Questo lavoro cambia il paradigma di progettazione degli agenti sanitari:

1. Dall'uso di un singolo modello all'orchestrazione dinamica: Dimostra che la combinazione e la selezione intelligente di più modelli specializzati sono superiori all'uso di un singolo modello "tuttofare" o di un modello gigante.
2. Efficienza e Affidabilità: Permette di utilizzare modelli più piccoli e specializzati, riducendo i costi computazionali (selezionando solo lo strumento necessario) e aumentando l'affidabilità clinica adattandosi al contesto specifico.
3. Fondamenta per l'Agentic AI: Fornisce un framework metodologico solido per gestire l'eterogeneità degli strumenti in ambienti reali, dove i dati non sono mai perfettamente allineati con i modelli di addestramento.

In sintesi, ToolSelect risolve il problema critico di "quale strumento usare quando", trasformando un pool di modelli eterogenei in un sistema coeso e adattivo capace di fornire risposte cliniche più accurate e sicure.