Technische Samenvatting: MedLLM – Een Open Medisch Taalmodel op de Sub-Billion Schaal

Probleemstelling

Het veld van medische large language models (LLM's) is geconvergeerd naar een enkele schaal: veelgebruikte systemen opereren bij 7 miljard parameters of meer. Deze trend laat het sub-billion regime (modellen met <1B parameters) onkarakteriseerd. Huidige medische LLM's zijn ofwel closed-source met niet-openbaar gemaakte trainingsdetails (bijv. GPT-4, Med-PaLM-2), ofwel open maar groot (7B–540B), wat toegang, reproduceerbaarheid en het auditeren van trainingsprocessen beperkt.

Er bestaat een kritieke kloof in het begrip van welke medische capaciteiten overleven wanneer de modelcapaciteit met een orde van grootte wordt verminderd. Het aanpassen van algemene modellen aan de geneeskunde is uitdagend vanwege de distributionele mismatch tussen webtekst en medische tekst, wat vereist dat men omgaat met long-tail terminologie, gestructureerde redenering en hoge precisie-eisen. Eerder werk adresseert dit via voortgezette pretraining op biomedische corpora en instructie-tuning, maar al deze inspanningen opereren bij 7B parameters of hoger. Het blijft onbekend of deze adaptatie-recepten overdraagbaar zijn naar kleinere modellen of dat een verminderde capaciteit ervoor zorgt dat specifieke medische capaciteiten eerder instorten dan andere.

Methodologie

Modelarchitectuur: MedLLM

MedLLM is een decoder-only transformer met 0,1 miljard (100M) parameters. Belangrijke architecturale keuzes zijn:

• Tokenisatie & Embedding: Gebruikt een SentencePiece vocabulaire ($V=32.000$) en een embedding dimensie $d=512$. Gezien de hoge vocabulaire-tot-breedte ratio ($V/d = 62,5$), hebben de auteurs de input embeddings gekoppeld aan de output projectie. Dit verwijdert een aparte $V \times d$ matrix (bespaart 14% van de totale parameters) en fungeert als een regularisator.
• Transformer Lagen: Stapelt 20 lagen met RMSNorm geplaatst binnen de residual branch (post-attention en post-FFN) in plaats van pre-norm. Dit houdt de residual stream op een gecontroleerde schaal, wat de stabiliteit verbetert voor diepe, smalle modellen.
• Attention: Gebruikt 8 koppen met Rotary Position Embeddings (RoPE) en QK-normalisatie om de magnitudes van de attention logits te begrenzen, wat instabiliteiten tijdens de training voorkomt die gebruikelijk zijn bij kleine modellen.
• Feed-Forward Network: Maakt gebruik van SwiGLU met een expansiefactor van 4 ($d_{ff} = 2048$). Dit wijst 62,7% van de parameters van het model toe aan de feed-forward sublaag, waarbij de prioriteit ligt bij diepte en compositionele transformatie boven breedte.

Datapijplijn: MedFineWeb

Om het model aan te passen zonder afhankelijk te zijn van handmatig gecureerde medische bronnen (die vaak schaars zijn of een bias vertonen naar abstract-stijl schrijven), introduceren de auteurs MedFineWeb, een medisch plain-text corpus dat via referentie-gestuurde selectie is geselecteerd uit algemene webdata:

1. Kandidaat-pool: Algemene webdata (RedPajama-V2, C4, OpenWebText, Wikipedia) is gesplitst in chunks van ~1.000 woorden.
2. Referentiesets: Drie medische QA-datasets (MedMCQA, MedQA, PubMedQA) dienen als semantische referenties. Hun prompts en antwoorden zijn geëmbed.
3. Selectie: Elke kandidaat-chunk wordt gescoord op basis van de semantische gelijkenis (cosine similarity via FAISS) met de referentiesets. De top 50% van de kandidaten voor elke referentieset wordt geselecteerd, samengevoegd en gededupliceerd.
4. Resultaat: Een 3B-token corpus dat de medische stijl en onderwerpen vangt zonder direct benchmark-vragen en -antwoorden te kopiëren.

Trainingspijplijn

De training volgt een volledig open, driefasig proces:

1. Algemene Pretraining: Het model wordt gepretraind op ~29,5B tokens van algemene webdata ($D_{gen}$) met causal language modeling met een curriculum sequence-length schema (128 $\to$ 256 $\to$ 512 tokens) en een z-loss regularisator om de output logits te stabiliseren.
2. Domain Fine-Tuning (DFT): Voortgezette pretraining op MedFineWeb voor 25.000 stappen om de distributie te verschuiven van algemene webtekst naar medische tekst.
3. Preference Alignment: Een tweestaps proces per benchmark:
   • Supervised Fine-Tuning (SFT): Fine-tuning op benchmark training-sets met label-predictie doelstellingen.
   • Direct Preference Optimization (DPO): Het optimaliseren van het model om de voorkeur te geven aan correcte antwoorden boven incorrecte antwoorden zonder een apart reward model. De KL-penalty sterkte wordt geschaald met $\beta=0,1$ om de preference margin te vergroten.

Kernbijdragen

1. MedLLM: Het eerste open 0,1B-parameter medisch taalmodel, dat een reproduceerbare baseline biedt die een orde van grootte kleiner is dan bestaande open medische LLM's.
2. MedFineWeb: Een nieuw medisch pretraining corpus geconstrueerd uit algemene webdata middels embedding-gelijkenis met medische QA. Het vereist geen gecureerde medische bronnen en kan worden aangepast aan elk domein mits er een kleine set target-domein QA-voorbeelden beschikbaar is.
3. Taak-gestructureerde Dissociatie: De identificatie dat onder compressie medische competentie niet uniform degradeert. In plaats daarvan splitst het zich per taaktype: het model bereikt een 7B-concurrerende nauwkeurigheid bij context-gegronde QA, maar blijft capaciteitsgebonden bij kennis-recall QA.

Resultaten

Pretrained Evaluatie

Zelfs vóór taakspecifieke fine-tuning bereikte het MedLLM-model (geëvalueerd via zero-shot top-token selectie) de hoogste gemiddelde nauwkeurigheid (35,1%) onder alle baselines, inclusclusief 7B instruction-tuned modellen (Zephyr-7B-$\beta$: 32,4%) en medisch aangepaste 7B modellen (PMC-LLaMA-7B: 30,6%). Het presteerde beter dan grotere modellen op recall-georiënteerde benchmarks (MMLU en MedMCQA), ondanks het feit dat het 70x kleiner is.

Fine-Tuned Evaluatie

Na SFT en DPO onthulden de resultaten een duidelijke splitsing op basis van taaktype:

• Knowledge-Recall Benchmarks (MedMCQA, MMLU): MedLLM (DPO) overtrof aanzienlijk alle 7B en sub-7B baselines. Op MedMCQA bereikte het 34,9% (vs. 33,0% voor Zephyr-7B-$\beta$ en 26,6% voor PMC-LLaMA-7B). Op MMLU bereikte het 32,4%.
• Context-Grounded QA (PubMedQA): MedLLM bereikte 58,2%, waarmee het de instruction-tuned 7B baselines (Zephyr-7B-$\beta$: 43,0%) versloeg, maar onder de continued-pretraining baselines bleef (PMC-LLaMA-7B: 72,9%).
• Complexe Klinische Casussen (MedQA): Het model had moeite en bereikte slechts 28,1% (nabij de 25% vier-opties vloer), aanzienlijk achter de instruction-tuned 7B modellen.

Analyse van Trainingsdynamiek

• Perplexity vs. Nauwkeurigheid: Tijdens de domain fine-tuning daalde de validatie perplexity op medische tekst met de helft (27,3 $\to$ 13,6), maar de benchmark nauwkeurigheid bleef vlak. Dit geeft aan dat het model op 0,1B schaal snel de medische vorm (taaldistributie) leert, maar faalt in het verwerven van de opvraagbare feiten die nodig zijn voor recall benchmarks.
• DPO Effect: DPO verhoogde de nauwkeurigheid slechts marginaal (bijv. +1,0 pp op MedMCQA) maar vergrootte de log-waarschijnlijkheidsmarge tussen correcte en incorrecte antwoorden met 5,2x. Dit suggereert dat DPO op deze schaal bestaande discriminatie verfijnt in plaats van nieuwe kennis injecteert.

Betekenis en Claims

Het artikel claimt dat het sub-billion regime niet louter een gedegradeerde versie is van grotere modellen, maar een uniek gedrag vertoont waarbij medische bekwaamheid het best wordt verkregen uit retrieval en geleverde context in plaats van parametrische recall.

• Capaciteit vs. Adaptatie: De resultaten demonstreren dat waar recall slaagt (MedMCQA), de beperkende factor de modelcapaciteit is, niet de adaptatiemethode. Waar recall faalt (MedQA), is de capaciteitslimiet absoluut.
• Taakspecifieke Dissociatie: De "split" in prestaties wordt gemaskeerd bij 7B (waar beide capaciteiten aanwezig zijn) en komt pas aan het licht wanneer de capaciteit schaars is.
• Reproduceerbaarheid: Door MedLLM en MedFineWeb vrij te geven, bieden de auteurs een pad om medische modellen te trainen, te auditeren en te implementeren op bescheiden hardware, wat de noodzaak van enorme parameteraantallen voor specifieke medische taken uitdaagt.

De auteurs concluderen dat op sub-billion schaal de breedte van de kennis inherent beperkt is, en dat prestaties sterk afhangen van de kwaliteit van de geleverde context in plaats van geïnternaliseerde parametrische kennis.