Riepilogo Tecnico: MedLLM – Un Modello Linguistico Medico Open a Scala Sub-Billion

Definizione del Problema

Il campo dei modelli linguistici di grandi dimensioni (LLM) in ambito medico ha converguto su una singola scala: i sistemi ampiamente utilizzati operano a 7 miliardi di parametri o più. Questa tendenza lascia il regime sub-miliardo (modelli con <1 miliardo di parametri) non caratterizzato. Gli attuali LLM medici sono o a codice chiuso con dettagli di addestramento non dichiarati (es. GPT-4, Med-PaLM-2) o open ma grandi (7B–540B), limitando l'accesso, la riproducibilità e la capacità di audit dei processi di addestramento.

Esiste un divario critico nella comprensione di quali capacità mediche sopravvivano quando la capacità del modello viene ridotta di un ordine di grandezza. L'adattamento di modelli general-purpose alla medicina è complesso a causa del mismatch distributivo tra il testo web e il testo medico, che richiede la gestione di terminologie a coda lunga (long-tail), ragionamento strutturato e requisiti di alta precisione. Il lavoro precedente affronta questo problema tramite il pre-addestramento continuo su corpora biomedici e l'instruction tuning, ma tutti tali sforzi operano a 7 miliardi di parametri o superiori. Resta da scoprire se queste ricette di adattamento si trasferiscano ai modelli più piccoli o se la riduzione della capacità costringa specifiche capacità mediche a collassare prima di altre.

Metodologia

Architettura del Modello: MedLLM

MedLLM è un decoder-only transformer con 0,1 miliardi (100M) di parametri. Le scelte architettoniche chiave includono:

• Tokenizzazione e Embedding: Utilizza un vocabolario SentencePiece ($V=32.000$) e una dimensione di embedding $d=512$. Dato l'alto rapporto vocabolario-larghezza ($V/d = 62,5$), gli autori legano l'embedding di input alla proiezione di output. Questo rimuove una matrice separata $V \times d$ (risparmiando il 14% dei parametri totali) e funge da regolarizzatore.
• Layer Transformer: Impila 20 layer utilizzando RMSNorm posizionata all'interno del ramo residuo (post-attention e post-FFN) invece di pre-norm. Ciò mantiene il flusso residuo a una scala controllata, migliorando la stabilità per modelli profondi e stretti.
• Attention: Utilizza 8 teste con Rotary Position Embeddings (RoPE) e QK-normalization per limitare l'ampiezza dei logit dell'attenzione, prevenendo instabilità nell'addestramento comuni nei modelli piccoli.
• Feed-Forward Network: Impiega SwiGLU con un fattore di espansione di 4 ($d_{ff} = 2048$). Questo alloca il 62,7% dei parametri del modello al sottolayer feed-forward, dando priorità alla profondità e alla trasformazione compositiva rispetto alla larghezza.

Pipeline dei Dati: MedFineWeb

Per adattare il modello senza fare affidamento su fonti mediche curate manualmente (che sono spesso scarse o sbilanciate verso uno stile di scrittura accademico), gli autori introducono MedFineWeb, un corpus di testo piano medico selezionato da dati web generali tramite selezione guidata da riferimenti:

1. Pool di Candidati: I dati web generali (RedPajama-V2, C4, OpenWebText, Wikipedia) vengono suddivisi in chunk di circa 1.000 parole.
2. Set di Riferimento: Tre dataset di QA medica (MedMCQA, MedQA, PubMedQA) fungono da riferimenti semantici. I loro prompt e le risposte vengono vettorizzati (embedded).
3. Selezione: Ogni chunk candidato viene valutato in base alla sua somiglianza semantica (similarità del coseno tramite FAISS) rispetto ai set di riferimento. Il top 50% dei candidati per ogni set di riferimento viene selezionato, unito e de-duplicato.
4. Risultato: Un corpus di 3 miliardi di token che cattura lo stile e gli argomenti medici senza copiare direttamente domande e risposte dei benchmark.

Pipeline di Addestramento

L'addestramento segue una pipeline completamente open in tre fasi:

1. Pre-addestramento Generale: Il modello è pre-addestrato su ~29,5 miliardi di token di dati web generali ($D_{gen}$) utilizzando il causal language modeling con uno schema di lunghezza della sequenza curricolare (128 $\to$ 256 $\to$ 512 token) e un regolarizzatore z-loss per stabilizzare i logit di output.
2. Domain Fine-Tuning (DFT): Pre-addestramento continuo su MedFineWeb per 25.000 step per spostare la distribuzione dal testo web generale al testo medico.
3. Allineamento delle Preferenze: Un processo in due fasi per ogni benchmark:
   • Supervised Fine-Tuning (SFT): Fine-tuning sui set di addestramento dei benchmark utilizzando obiettivi di predizione delle etichette.
   • Direct Preference Optimization (DPO): Ottimizzazione del modello affinché preferisca risposte corrette rispetto a quelle errate senza un modello di ricompensa separato. La forza della penalità KL è scalata con $\beta=0,1$ per amplificare il margine di preferenza.

Contributi Chiave

1. MedLLM: Il primo modello linguistico medico aperto da 0,1 miliardi di parametri, che fornisce una baseline riproducibile un ordine di grandezza più piccola degli esistenti LLM medici open.
2. MedFineWeb: Un nuovo corpus di pre-addestramento medico costruito da dati web generali tramite similarità di embedding con QA medica. Non richiede fonti mediche curate e può essere adattato a qualsiasi dominio fornendo un piccolo set di esempi di QA del dominio target.
3. Dissociazione Strutturata per Task: L'identificazione del fatto che, sotto compressione, la competenza medica non degrada uniformemente. Invece, si divide per tipo di task: il modello raggiunge un'accuratezza competitiva con i modelli da 7B nei task di QA basati sul contesto, ma rimane limitato dalla capacità nei task di richiamo della conoscenza (knowledge-recall).

Risultati

Valutazione del Pre-addestramento

Anche prima del fine-tuning specifico per il task, il modello MedLLM (valutato zero-shot tramite selezione del top-token) ha ottenuto l'accuratezza media più alta (35,1%) tra tutte le baseline, inclusi i modelli da 7B istruiti (Zephyr-7B-β: 32,4%) e i modelli da 7B adattati al dominio medico (PMC-LLaMA-7B: 30,6%). Ha superato modelli più grandi nei benchmark orientati al richiamo (MMLU e MedMCQA) nonostante fosse 70 volte più piccolo.

Valutazione del Fine-Tuning

Dopo SFT e DPO, i risultati hanno rivelato una netta divisione basata sul tipo di task:

• Benchmark di Richiamo della Conoscenza (MedMCQA, MMLU): MedLLM (DPO) ha superato significativamente tutte le baseline da 7B e sub-7B. Su MedMCQA, ha raggiunto il 34,9% (contro il 33,0% di Zephyr-7B-β e il 26,6% di PMC-LLaMA-7B). Su MMLU, ha raggiunto il 32,4%.
• QA Basata sul Contesto (PubMedQA): MedLLM ha raggiunto il 58,2%, superando le baseline da 7B istruite (Zephyr-7B-β: 43,0%), ma rimanendo al di sotto delle baseline di pre-addestramento continuo (PMC-LLaMA-7B: 72,9%).
• Vignette Cliniche Complesse (MedQA): Il modello ha faticato, raggiungendo solo il 28,1% (vicino al floor del 25% per le quattro opzioni), significativamente indietro rispetto ai modelli da 7B istruiti.

Analisi della Dinamica di Addestramento

• Perplexity vs. Accuratezza: Durante il domain fine-tuning, la perplexity di validazione sul testo medico è diminuita della metà (27,3 $\to$ 13,6), ma l'accuratezza dei benchmark è rimasta piatta. Ciò indica che, alla scala di 0,1B, il modello apprende rapidamente la forma medica (distribuzione del linguaggio) ma non riesce ad acquisire i fatti recuperabili richiesti dai benchmark di richiamo.
• Effetto DPO: Il DPO ha aumentato l'accuratezza solo moderatamente (es. +1,0 pp su MedMCQA) ma ha amplificato il margine di log-probabilità tra risposte corrette e errate di 5,2 volte. Ciò sugger lisce che il DPO a questa scala affina la discriminazione esistente piuttosto che iniettare nuova conoscenza.

Significato e Rivendicazioni

Il paper sostiene che il regime sub-miliardo non è semplicemente una versione degradata dei modelli più grandi, ma esibisce un comportamento unico dove la capacità medica è meglio derivata dal recupero e dal contesto fornito piuttosto che dal richiamo parametrico.

• Capacità vs. Adattamento: I risultati dimostrano che dove il richiamo ha successo (MedMCQA), il vincolo è la capacità del modello, non il metodo di adattamento. Dove il richiamo fallisce (MedQA), il limite della capacità è assoluto.
• Dissociazione per Task: Questa "scissione" nelle prestazioni è mascherata a 7B (dove entrambe le capacità esistono) e emerge solo quando la capacità è scarsa.
• Riproducibilità: Rilasciando MedLLM e MedFineWeb, gli autori forniscono un percorso per addestrare, auditare e distribuire modelli medici su hardware modesto, sfidando la necessità di conteggi di parametri massicci per compiti medici specifici.

Gli autori concludono che, alle scale sub-miliardo, l'ampiezza della conoscenza del modello è intrinsecamente limitata e le prestazioni dipendono fortemente dalla qualità del contesto fornito piuttosto che dalla conoscenza parametrica internalizzata.