MedFeat: Model-Aware and Explainability-Driven Feature Engineering with LLMs for Clinical Tabular Prediction

MedFeat is een feedbackgedreven framework dat Large Language Models combineert met modelbewustzijn en SHAP-waarden om in de gezondheidszorg interpreteerbare en robuuste kenmerken te ontwerpen die de prestaties van klinische voorspellingsmodellen verbeteren.

De kern

MedFeat: De Slimme Architect voor Medische Voorspellingen

Stel je voor dat je een zeer ervaren arts bent die probeert te voorspellen of een patiënt binnen 24 uur ernstig ziek zal worden. Je hebt een enorme stapel papieren dossiers (de data) voor je liggen met duizenden details: leeftijd, bloeddruk, medicijnen, en hoe vaak ze hebben gekoesterd.

Vroeger deden computers dit door gewoon naar de cijfers te kijken. Maar zoals de auteurs van dit paper uitleggen, zijn de beste 'ouderwetse' computers (zoals beslissingsbomen) vaak beter in het vinden van patronen dan de supermoderne 'diepe neurale netwerken' die we zien in AI-films. Het probleem? Die oude computers hebben hulp nodig om de juiste details te vinden. Dat heet feature engineering (het maken van nieuwe, slimme variabelen).

Hier komt MedFeat om de hoek kijken. Het is een nieuw systeem dat een Super-Artiest (een Large Language Model of LLM) gebruikt om die nieuwe variabelen te bedenken, maar dan op een heel slimme manier.

Hier is hoe het werkt, vertaald naar alledaagse analogieën:

1. Het Probleem: De Zoektocht in de Donkere Kamer

Stel je voor dat je in een donkere kamer bent met duizenden lampen. Je moet de ene lamp vinden die het licht aandoen.

• Oude methoden: Ze proberen willekeurig lampen aan te doen (wiskundige formules). Soms werkt het, maar vaak zijn ze domme combinaties die geen zin hebben.
• Eerdere AI-methoden: Ze hebben een boek met medische kennis, maar ze schieten blindelings in het donker. Ze vragen de AI: "Bedenk iets!" zonder te weten welke lamp de computer al kent en welke niet.

2. De Oplossing: MedFeat als een Slimme Architect

MedFeat is als een architect die samenwerkt met een superintelligente consultant. Ze hebben drie slimme trucs:

A. De "Spiegel" (Model Awareness)

Stel je voor dat je een schilderij maakt voor een specifieke kijker.

• Als de kijker een kunstliefhebber is (een simpele lineaire model), moet je de details heel duidelijk en rechtstreeks uitleggen.
• Als de kijker een abstracte kunstenaar is (een complexe boom-structuur zoals XGBoost), hoeft je niet elke lijn te tekenen; die ziet de patronen al.

MedFeat kijkt eerst naar de "kijker" (het computermodel). Als het model al slim is in het zien van patronen, vraagt MedFeat de AI: "Bedenk iets dat dit model niet goed kan zien, zoals een complexe tijdslijn of een combinatie van factoren die samen een gevaar vormen." Zo wordt er geen tijd verspild aan dingen die het model al kent.

B. De "Zetel" (Explainability & SHAP)

De AI werkt niet blind. MedFeat gebruikt een techniek genaamd SHAP.

• Denk aan SHAP als een spotlight. Het licht op de belangrijkste details in de patiëntendossiers.
• In plaats van de AI te vragen om alles te bedenken, zegt MedFeat: "Kijk eens naar deze drie belangrijke lampen (bijvoorbeeld: leeftijd, bloeddruk en een medicijn). Bedenk een nieuwe, slimme combinatie van deze drie."
• Dit zorgt ervoor dat de AI zich richt op wat echt belangrijk is, in plaats van ruis te creëren.

C. De "Island Hopper" (Island Sampling)

Stel je voor dat je een gigantische bibliotheek hebt. Als je de AI alle boeken tegelijk geeft, raakt hij in de war en wordt hij traag.

• MedFeat maakt kleine eilanden. Het pakt een klein groepje van de belangrijkste variabelen (bijvoorbeeld alleen de hartslag en de ademhaling) en vraagt de AI: "Wat kun je met deze twee bedenken?"
• Dit houdt de AI scherp, bespaart tijd en zorgt voor betere ideeën.

3. De Leercyclus: Probeer, Faal, Herinner

MedFeat is niet alleen slim, het leert ook van zijn fouten.

• Het systeem probeert een nieuw idee.
• Als het werkt, wordt het onthouden als een "Gouden Idee".
• Als het faalt, wordt het onthouden als een "Vuilnisbak-Idee".
• De volgende keer vraagt de AI: "Geef me geen ideeën die lijken op die in de vuilnisbak, maar bouw voort op de gouden ideeën."

Waarom is dit belangrijk? (De Resultaten)

De auteurs hebben MedFeat getest op echte medische data (zoals patiënten in het ziekenhuis en ouderen).

• Betrouwbaarheid: Het werkt beter dan de oude methoden, zelfs als de data rommelig is of als er weinig zieke mensen zijn (een zeldzame gebeurtenis).
• Toekomstbestendig: Het beste deel? De nieuwe variabelen die MedFeat bedenkt, werken ook als de patiëntenpopulatie verandert (bijvoorbeeld van intensieve zorg naar algemene ziekenhuizen). Het is alsof je een sleutel maakt die niet alleen voor één deur werkt, maar voor een hele buurt.
• Privacy: De AI ziet nooit de echte namen of medische dossiers van patiënten. Hij krijgt alleen de "samenvatting" van wat belangrijk is. Dit is cruciaal voor privacy.

Samenvattend

MedFeat is als een medische detective die een super-intelligente assistent heeft. In plaats van willekeurig te gissen, kijkt de detective eerst wat de assistent al weet, focust hij op de belangrijkste bewijsstukken, en bedenkt hij slimme nieuwe hypotheses die de computer zelf niet zou zien. Het resultaat is een voorspellingssysteem dat nauwkeuriger is, makkelijker te begrijpen voor artsen, en beter bestand tegen veranderingen in de wereld.

Technische samenvatting

Probleemstelling

In de gezondheidszorg presteren klassieke machine learning-modellen (zoals boomgebaseerde modellen en logistische regressie) vaak beter dan diepe leerbenaderingen op tabulaire data. Dit komt door uitdagingen zoals ernstige class-imbalance, heterogene feature-types, complexe tijdspatronen en onregelmatige missingness. Hoewel automatische feature engineering (AFE) methoden zoals AutoFeat en OpenFE bestaan, hebben deze beperkingen:

1. Ze vertrouwen op vooraf gedefinieerde operatoren (bijv. wiskundige combinaties) die vaak klinisch niet onderbouwd zijn.
2. Ze zijn computatie-intensief en missen inzicht in de specifieke inductieve bias van het downstream model.
3. Bestaande LLM-gebaseerde methoden (zoals CAAFE, FeatLLM) behandelen features vaak als uitwisselbaar, negeren feature-importance signalen en sturen lange prompts die de kwaliteit van de generatie verminderen.
4. Privacyproblemen: Veel methoden sturen patiëntdata naar LLM's, wat in strijd is met klinische privacyregels.

Het doel is een schaalbaar, interpreteerbaar en privacy-bewust framework te creëren dat LLM's gebruikt voor feature engineering, maar wel rekening houdt met de beperkingen van het downstream model en gebruikmaakt van explainability-signalen.

Methodologie: MedFeat

MedFeat is een iteratief, feedback-gedreven framework dat Large Language Models (LLMs) combineert met model-awareness en explainability (SHAP-waarden). Het proces verloopt als volgt:

1. Initiële Training & Explainability:

   • Een baseline model (bijv. XGBoost of Logistische Regressie) wordt getraind op de oorspronkelijke data.
   • SHAP-waarden worden berekend om de relevantie en het belang van elke feature te bepalen.

2. Importance-Weighted Island Sampling:

   • In plaats van alle features in één prompt te sturen, worden er "eilanden" (subsets van features) gegenereerd.
   • Deze eilanden worden gesampleerd op basis van de SHAP-rangschikking (hoge prioriteit voor belangrijke features) en feature-typen (statisch vs. temporair).
   • Dit beperkt de prompt-lengte, verlaagt de token-kosten en focust de LLM op lokale, informatieve regio's van de feature-ruimte.

3. Model-Aware Feature Generatie:

   • De LLM (bijv. GPT-4o) ontvangt een prompt die specifiek is afgestemd op het downstream model (de "hypothesis class").
   • Voorbeeld: Voor Logistische Regressie (lineair) vraagt MedFeat de LLM om niet-lineaire transformaties en interacties te genereren die het model zelf niet kan leren. Voor XGBoost (boomgebaseerd, dat al veel niet-lineariteiten kan leren) vraagt het om complexe tijdspatronen, globale statistieken of context-gedreven interacties die moeilijk te leren zijn voor bomen.
   • De prompt bevat ook een "Memory Bank" met eerdere succesvolle en gefaalde transformaties om herhaling te voorkomen en de zoekrichting te sturen.

4. Validatie & Selectie:

   • De gegenereerde features worden lokaal uitgevoerd (geen data naar de LLM gestuurd).
   • Een nieuw model wordt getraind op de verrijkte dataset.
   • Als de validatieprestatie (bijv. AUC) een bepaalde drempel ($\beta$) overschrijdt, worden de features geaccepteerd en toegevoegd aan de dataset voor de volgende iteratie. Anders worden ze verworpen.

5. Privacy:

   • Er wordt nooit ruwe patiëntdata naar de LLM gestuurd. Alleen metadata, feature-beschrijvingen en geaggregeerde SHAP-scores worden gebruikt.

Belangrijkste Bijdragen

1. Eerste Model-Aware LLM Framework: MedFeat is het eerste framework dat feature-proposals aanpast aan de representatieve beperkingen van het downstream model, waardoor verspilde evaluatiebudgetten worden voorkomen.
2. Explainability-Driven Feedback: Het gebruik van SHAP-waarden als sturingssignaal (in plaats van alleen validatiescores) zorgt voor een stabielere zoekrichting, vooral bij ruis en class-imbalance.
3. Island Search Strategie: Door features te groeperen in "eilanden" gebaseerd op belangrijkheid, wordt de prompt-efficiëntie verbeterd en wordt gefocust op de meest informatieve predictors.
4. Privacy & Robuustheid: Het framework is ontworpen voor klinische inzet zonder patiëntdata te lekken en toont robuustheid onder distributiewisselingen (bijv. overdracht van ICU-data naar algemene ziekenhuisdata).

Resultaten

Het framework is geëvalueerd op drie grote klinische datasets (IORD, MIMIC-IV, HRS) voor diverse taken (bijv. 24-uurs mortaliteit, hartfalen, 10-jaar mortaliteit).

• Prestatieverbetering: MedFeat presteerde consistent beter dan baselines (AutoFeat, OpenFE) en andere LLM-methoden (CAAFE, OCTree, FeatLLM).
  • Bijvoorbeeld: Voor 24-uurs inpatient mortaliteit steeg de AUC van 0.686 naar 0.740 (een verbetering van 7.87%) met XGBoost.
  • Voor Logistische Regressie waren de verbeteringen nog sterker, omdat de gegenereerde features de lineaire beperkingen van het model direct opvulden.
• Robuustheid: MedFeat bleef prestaties leveren zelfs na uitgebreide hyperparameter-optimalisatie (HPO), terwijl andere methodes vaak instabiel waren of trade-offs maakten tussen AUC en F1-score.
• Generalisatie: Features gegenereerd op ICU-data (MIMIC) verbeterden de prestaties op algemene ziekenhuisdata (IORD), wat aantoont dat de gegenereerde features klinisch betekenisvol en generaliseerbaar zijn.
• Ablatie Studies: Het verwijderen van "model-awareness" of "feature importance guidance" leidde tot significante prestatiedalingen (tot -35% F1-verlies in extreme onbalans), wat de noodzaak van deze componenten bevestigt.

Significantie

MedFeat biedt een praktische route naar schaalbare en interpreteerbare feature engineering in de gezondheidszorg. Het lost het probleem op dat klassieke modellen vaak beter presteren dan deep learning op tabulaire data, maar dat handmatige feature engineering te duur en tijdrovend is. Door LLM's te koppelen aan model-inzicht en explainability, creëert MedFeat features die:

1. Klinisch relevant zijn: Gebaseerd op domeinkennis en niet alleen op statistische correlaties.
2. Model-specifiek zijn: Ontworpen om de zwaktes van het specifieke leermodel aan te vullen.
3. Privacy-bewust zijn: Geschikt voor real-world deployment zonder inbreuk op patiëntgegevens.

De studie suggereert dat toekomstige systemen in de gezondheidszorg niet alleen moeten vertrouwen op ruwe data, maar op slimme, door LLM's gegenereerde features die de inductieve bias van het model versterken, wat leidt tot stabielere en betrouwbaardere voorspellingen onder tijds- en locatieverschillen.