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

Le papier présente MedFeat, un cadre d'ingénierie de fonctionnalités alimenté par les LLMs qui intègre la connaissance médicale, la conscience du modèle et l'explicabilité SHAP pour améliorer de manière stable les prédictions cliniques tabulaires et découvrir des caractéristiques robustes et cliniquement pertinentes.

L'essentiel

Imaginez que vous essayez de prédire la santé d'un patient en utilisant un tableau rempli de données : l'âge, la tension artérielle, le nombre de médicaments, etc. C'est comme essayer de deviner le résultat d'une partie d'échecs en regardant seulement la position des pièces, sans comprendre les stratégies cachées.

Dans le monde médical, les ordinateurs (les modèles d'intelligence artificielle) ont du mal à faire ces prédictions avec des tableaux de données classiques. Souvent, les méthodes "anciennes" (comme les arbres de décision) fonctionnent mieux que les réseaux de neurones complexes, mais elles ont besoin d'aide pour bien voir les détails. C'est là qu'intervient MedFeat.

Voici une explication simple de ce que fait MedFeat, en utilisant des analogies du quotidien :

1. Le Problème : Le Chef Cuisinier et les Ingrédients Bruts

Imaginez un chef cuisinier (le modèle d'IA) très talentueux, mais qui ne sait cuisiner que des plats très simples. Si vous lui donnez des ingrédients bruts (les données brutes du patient), il peut faire un bon plat, mais il rate souvent les saveurs subtiles.

Traditionnellement, pour aider le chef, des experts humains (les médecins) devaient passer des heures à mélanger les ingrédients eux-mêmes pour créer de nouvelles recettes (créer de nouvelles "fonctions" ou caractéristiques). C'est long, coûteux et épuisant.

D'autres méthodes automatiques ont essayé de faire cela en mélangeant tous les ingrédients possibles au hasard, mais cela prenait trop de temps et créait souvent des plats immangeables (des données inutiles).

2. La Solution : MedFeat, le "Chef Assistant" Intelligents

MedFeat est comme un assistant de cuisine ultra-intelligent (une Intelligence Artificielle de type LLM) qui connaît parfaitement la cuisine médicale. Mais il ne travaille pas seul. Il a trois super-pouvoirs :

• Il écoute le Chef (Conscience du Modèle) :
  Au lieu de donner n'importe quel ingrédient, MedFeat demande au chef : "Qu'est-ce que tu as du mal à cuisiner ?".

  • Si le chef est un modèle simple (comme une régression logistique), MedFeat lui prépare des mélanges complexes (des interactions entre l'âge et la pauvreté, par exemple) que le chef ne pourrait pas inventer seul.
  • Si le chef est déjà très fort (comme XGBoost), MedFeat lui propose des choses très spécifiques, comme des tendances temporelles complexes, pour ne pas lui donner des choses qu'il sait déjà faire.
  • Analogie : C'est comme un assistant qui ne vous donne pas un marteau si vous avez déjà un tournevis, mais qui vous donne la vis exacte dont vous avez besoin pour finir le meuble.

• Il utilise une "Lampe Torche" (Explicabilité / SHAP) :
  MedFeat ne devine pas au hasard. Il utilise une "lampe torche" (appelée SHAP) pour voir quels ingrédients sont les plus importants pour la santé du patient. Il se concentre uniquement sur ces ingrédients clés pour créer de nouvelles recettes.

  • Analogie : Au lieu de chercher une aiguille dans une botte de foin, MedFeat utilise un détecteur de métaux pour savoir exactement où l'aiguille se trouve, puis il fabrique une nouvelle clé pour l'ouvrir.

• Il apprend de ses erreurs (Mémoire) :
  MedFeat garde un carnet de notes. Si une recette a échoué hier, il ne la propose pas aujourd'hui. Si une recette a fonctionné, il s'en inspire pour en créer de meilleures.

  • Analogie : C'est comme un joueur d'échecs qui se souvient de ses parties perdues pour ne pas refaire les mêmes erreurs, et qui répète ses coups gagnants.

3. Comment ça marche en pratique ? (Le Processus)

1. Le Début : On entraîne le modèle de base avec les données habituelles.
2. L'Inspection : MedFeat regarde ce que le modèle a bien compris et ce qu'il a raté (grâce à la "lampe torche").
3. La Création : Il demande à son cerveau (le LLM) de créer de nouvelles données basées sur ce qui manque, mais en se concentrant sur un petit groupe d'ingrédients importants (les "îles" de données) pour ne pas se perdre.
4. Le Test : Il teste la nouvelle recette. Si elle améliore le plat, elle est ajoutée au menu. Sinon, elle est jetée.
5. La Répétition : Il répète ce processus plusieurs fois jusqu'à obtenir le plat parfait.

4. Pourquoi c'est génial ?

• Confidentialité : MedFeat est très prudent. Il ne montre jamais les noms ou les visages des patients à l'IA. Il ne lui donne que des statistiques résumées (comme "l'âge moyen" ou "l'importance d'un symptôme"). C'est comme donner à un architecte un plan de maison sans jamais montrer la famille qui y habite.
• Robustesse : Même si les données changent un peu (par exemple, si les patients arrivent d'un hôpital différent ou d'une année différente), les recettes créées par MedFeat fonctionnent toujours. C'est comme une recette de gâteau qui reste délicieuse même si vous changez légèrement la marque de farine.
• Efficacité : Dans les tests, MedFeat a permis de prédire la mortalité ou les maladies avec beaucoup plus de précision que les méthodes précédentes, et ce, sans avoir besoin de régler des milliers de paramètres compliqués.

En résumé

MedFeat, c'est comme avoir un médecin-robot qui aide un ordinateur à mieux comprendre les patients. Au lieu de jeter des millions de données au hasard, il utilise son intelligence pour créer les "mots-clés" exacts dont l'ordinateur a besoin pour prendre de meilleures décisions, tout en respectant la vie privée des patients. C'est une façon intelligente, rapide et sûre d'améliorer les soins de santé grâce à l'IA.

Résumé technique

1. Problématique

Dans le domaine de la prédiction clinique sur des données tabulaires, les modèles d'apprentissage automatique classiques (notamment les modèles basés sur les arbres de décision comme XGBoost) surpassent souvent les approches d'apprentissage profond, en raison de la nature bruitée, déséquilibrée et hétérogène des données de santé. Cependant, la performance de ces modèles dépend fortement de l'ingénierie des caractéristiques (feature engineering).

Les défis actuels sont les suivants :

• Ingénierie manuelle coûteuse : L'exploration exhaustive de l'espace des transformations possibles est impraticable et nécessite une expertise clinique approfondie.
• Limites des méthodes automatiques existantes : Les outils comme AutoFeat ou OpenFE reposent sur des opérateurs prédéfinis (combinaisons arithmétiques, transformations non linéaires) qui manquent de justification clinique et peuvent être coûteux en calcul.
• Limites des approches récentes basées sur les LLM : Les méthodes utilisant les Grands Modèles de Langage (LLM) pour générer des caractéristiques (ex: CAAFE, FeatLLM) souffrent de trois défauts majeurs :
  1. Agnosticisme vis-à-vis du modèle : Elles ignorent les capacités d'apprentissage du modèle de prédiction en aval (ex: une régression logistique a besoin de non-linéarités explicites, tandis qu'un arbre de décision les apprend nativement).
  2. Négligence de l'importance des caractéristiques : Elles traitent toutes les caractéristiques comme interchangeables lors de la génération.
  3. Problèmes de contexte et de confidentialité : L'envoi de toutes les descriptions de caractéristiques aux LLM allonge les prompts (réduisant la qualité de génération) et l'utilisation d'échantillons de patients pose des problèmes de confidentialité.

2. Méthodologie : Le Framework MedFeat

MedFeat est un cadre d'ingénierie de caractéristiques itératif, piloté par l'explicabilité et conscient du modèle (model-aware), qui utilise les LLM pour découvrir des transformations cliniquement pertinentes.

Architecture et Flux de Travail

Le processus fonctionne par itérations (boucle de rétroaction) :

1. Entraînement de la ligne de base : Un modèle de prédiction (ex: XGBoost ou Régression Logistique) est entraîné sur les données initiales.
2. Calcul de l'importance (Explicabilité) : Les valeurs SHAP (SHapley Additive exPlanations) sont calculées sur l'ensemble de validation pour obtenir un score d'importance pour chaque caractéristique.
3. Échantillonnage par "Îles" (Island Sampling) :
   • Au lieu de soumettre toutes les caractéristiques au LLM, MedFeat crée des sous-ensembles appelés "îles".
   • Ces îles sont échantillonnées de manière pondérée par l'importance SHAP, se concentrant sur les prédicteurs les plus pertinents tout en explorant certaines caractéristiques moins dominantes.
   • Cela limite la longueur du prompt, réduit l'usage de tokens et évite la dilution du signal.
4. Génération Pilotée par le Modèle (Model-Aware Generation) :
   • Le LLM reçoit un prompt structuré incluant : le contexte de l'île, les métadonnées des caractéristiques, les scores SHAP, et les contraintes du modèle en aval.
   • Exemple : Si le modèle est une régression logistique, le LLM est invité à générer des interactions non linéaires et des termes quadratiques. Si le modèle est XGBoost, le LLM est invité à proposer des structures temporelles complexes ou des statistiques globales que l'arbre peine à apprendre directement.
   • Le LLM agit comme un clinicien-épidémiologiste, utilisant son savoir-faire médical pour proposer des transformations justifiées.
5. Validation et Sélection :
   • Les transformations proposées sont exécutées localement (sans envoyer de données brutes de patients au LLM).
   • Un nouveau modèle est entraîné avec les nouvelles caractéristiques.
   • Si la performance sur l'ensemble de validation dépasse un seuil de tolérance ($\beta$) par rapport à la ligne de base, les caractéristiques sont acceptées et ajoutées à l'ensemble de données. Sinon, elles sont rejetées.
6. Mémoire de Rétroaction : Le système conserve une mémoire des propositions acceptées et rejetées pour guider les itérations suivantes, évitant de répéter les échecs et renforçant les motifs réussis.

3. Contributions Clés

1. Premier cadre d'ingénierie de caractéristiques basé sur LLM "conscient du modèle" : MedFeat adapte les propositions de caractéristiques aux limites représentatives du modèle d'apprentissage en aval, évitant ainsi le gaspillage de ressources sur des motifs que le modèle peut déjà apprendre implicitement.
2. Guidage par l'explicabilité (SHAP) : L'utilisation des scores SHAP comme signal de génération permet de prioriser les caractéristiques porteuses d'information et de conditionner le prompt du LLM avec une information d'importance structurée, améliorant la stabilité dans des régimes bruyants et déséquilibrés.
3. Stratégie d'Îles et Efficacité : L'échantillonnage par îles permet de maintenir la longueur des prompts sous contrôle et de se concentrer sur des régions locales de l'espace des caractéristiques, améliorant la qualité des générations et l'efficacité computationnelle.
4. Respect de la vie privée : Aucune donnée patient au niveau individuel n'est envoyée au LLM ; seul le métadonnées et les scores d'importance agrégés sont utilisés.

4. Résultats Expérimentaux

Les auteurs ont évalué MedFeat sur trois jeux de données cliniques réels (IORD, MIMIC-IV, HRS) couvrant cinq tâches de prédiction (mortalité, réadmission, insuffisance cardiaque, etc.).

• Performance Globale : MedFeat surpasse systématiquement les modèles de base (sans ingénierie) et les méthodes automatiques existantes (AutoFeat, OpenFE, CAAFE, FeatLLM, OCTree).
  • Sur la tâche de mortalité hospitalière sous 24h (IORD), MedFeat améliore l'AUC de 0,686 à 0,740 (+7,87%) avec un XGBoost non optimisé.
  • Il démontre une stabilité supérieure, améliorant à la fois l'AUC et le F1, contrairement à d'autres méthodes qui améliorent un métrique au détriment de l'autre.
• Robustesse à l'Optimisation des Hyperparamètres (HPO) : Même après une optimisation intensive des hyperparamètres, MedFeat maintient des gains significatifs, particulièrement sur les métriques F1 pour les tâches très déséquilibrées, prouvant qu'il ajoute une structure utile que le simple réglage des paramètres ne peut pas récupérer.
• Généralisation et Robustesse Temporelle :
  • Transfert inter-cohorte : Des caractéristiques générées sur des données d'unités de soins intensifs (MIMIC) améliorent la performance sur des données de patients hospitalisés généraux (IORD), démontrant la transférabilité des connaissances cliniques.
  • Robustesse temporelle : Les modèles enrichis par MedFeat résistent mieux aux dérifts de distribution temporels (données de 2020-2024) que les modèles réentraînés annuellement, suggérant que les caractéristiques générées capturent des régularités cliniques stables.
• Analyse par Ablation : La suppression de la composante "consciente du modèle" ou de l'échantillonnage par îles entraîne une baisse significative des performances, confirmant l'importance critique de ces mécanismes.

5. Signification et Impact

MedFeat représente une avancée majeure pour le déploiement de l'IA en santé :

• Praticité : Il offre une voie scalable pour l'ingénierie de caractéristiques dans des environnements où les données sont limitées et où la distribution change (dérift temporel).
• Interprétabilité : En générant des caractéristiques basées sur des raisonnements cliniques et en les validant via des scores SHAP, le cadre rend le processus de décision plus transparent pour les cliniciens.
• Confidentialité : Il résout le dilemme de l'utilisation des LLM en santé en évitant l'envoi de données sensibles, ce qui est crucial pour la conformité réglementaire (RGPD, HIPAA).
• Synergie Homme-Machine : Le framework ne remplace pas l'expert humain mais l'assiste en automatisant la découverte de motifs complexes que les modèles standards ne peuvent pas extraire seuls, tout en respectant les contraintes de confidentialité.

En conclusion, MedFeat démontre que l'intégration judicieuse de l'IA générative (LLM) avec des signaux d'explicabilité et une conscience du modèle permet de créer des systèmes de prédiction clinique plus robustes, précis et déployables dans le monde réel.