CleanSurvival: Automated data preprocessing for time-to-event models using reinforcement learning

Résumé Technique : CleanSurvival

Énoncé du Problème
Le prétraitement des données est un composant critique, bien que souvent négligé, des pipelines d'apprentissage automatique, ayant un impact significatif sur les performances des modèles. Bien que les frameworks d'Automatisation de l'Apprentissage Automatique (AutoML) aient commencé à intégrer un prétraitement automatisé pour les tâches de classification et de régression, il existe un manque distinct de solutions automatisées et adaptées à l'analyse de survie (modélisation du temps jusqu'à l'événement). L'analyse de survie présente des défis uniques en raison des données censurées, où les temps d'événement ne sont que partiellement observés. Les outils AutoML existants échouent soit à gérer nativement les données censurées, soit à s'appuyer sur un prétraitement manuel, soit à convertir les tâches de survie en problèmes de classification (par exemple, via une discrétisation du temps), entraînant une perte d'information. De plus, les outils actuels de prétraitement automatisé manquent souvent de support pour les structures de récompense et les types de modèles spécifiques requis pour l'analyse de survie, tels que le modèle à risques proportionnels de Cox ou les Forêts Aléatoires de Survie.

Méthodologie
L'article présente CleanSurvival, un framework de prétraitement de données automatisé spécifiquement conçu pour l'analyse de survie, construit sur une approche d'Apprentissage par Renforcement (RL) utilisant l'apprentissage Q.

• Architecture : Le framework traite la sélection des étapes de prétraitement comme un problème de prise de décision séquentielle. Un agent navigue dans un espace d'états défini par les caractéristiques du jeu de données pour sélectionner une séquence d'actions (techniques de prétraitement) maximisant une fonction de récompense.
• Actions de Prétraitement : Le framework prend en charge diverses techniques pour :
  • Imputation des Valeurs Manquantes : Analyse des Cas Complets (CCA), imputation par Moyenne/Médiane, Imputation Multiple par Équations Chaînées (MICE) et k-Plus Proches Voisins (KNN).
  • Détection des Valeurs Aberrantes : Enveloppe Elliptique (basée sur la distance de Mahalanobis) et résidus de Martingale.
  • Sélection/Extraction de Caractéristiques : Sélection univariée par Risques Proportionnels de Cox, LASSO, Élimination Récursive des Caractéristiques (RFE) et Gain d'Information.
• Modèles de Survie : Le système est compatible avec divers modèles de temps jusqu'à l'événement, notamment le modèle à Risques Proportionnels de Cox, les Forêts Aléatoires de Survie (RSF), DeepHit (une approche d'apprentissage profond) et les modèles fournis par l'utilisateur.
• Structure de Récompense : L'agent d'apprentissage Q est guidé par un signal de récompense basé sur l'Indice de Concordance (C-index) évalué sur un jeu de données de retenue. Cela garantit que le pipeline de prétraitement sélectionné est optimisé spécifiquement pour la performance prédictive du modèle de survie en aval.
• Modes de Fonctionnement : L'outil propose quatre modes :
  1. Algorithme Principal : Optimisation par apprentissage Q.
  2. Nettoyage Aléatoire : Recherche aléatoire pour comparaison de référence.
  3. Pipeline Personnalisé : Pipelines fixes définis par l'utilisateur via des fichiers de configuration.
  4. Aucune Préparation : Entrée de données brutes pour une performance de référence.

Contributions Clés

1. Premier Prétraitement basé sur le RL pour la Survie : CleanSurvival étend le framework Learn2Clean (à l'origine pour la régression/classification) pour gérer les contraintes spécifiques des données de temps jusqu'à l'événement censurées.
2. Optimisation Spécifique au Modèle : Contrairement au prétraitement statique, le framework sélectionne dynamiquement des pipelines optimaux pour des classes de modèles spécifiques (par exemple, un pipeline optimal pour Cox peut différer de celui pour RSF).
3. Implémentation Open Source : Les auteurs fournissent un package Python disponible sur GitHub, s'intégrant à des bibliothèques standards comme scikit-survival.
4. Gestion de la Censure : Le framework intègre explicitement la pondération par l'inverse de la probabilité de censure (IPCW) dans des métriques comme le Score Brier Intégré (IBS) pour tenir compte des observations censurées lors de l'évaluation.

Résultats Expérimentaux
Le framework a été évalué sur trois jeux de données réels : Rotterdam (cancer du sein), Flchain (chaîne légère libre sérique) et GBSG (Groupe d'Étude Allemand sur le Cancer du Sein). Une absence de données synthétique (MCAR, MAR, MNAR) a été introduite pour tester la robustesse.

• Performance : CleanSurvival a constamment surpassé des basiques simples (Analyse des Cas Complets et Imputation par Moyenne) et a souvent égalé ou dépassé une recherche aléatoire non dirigée et Optuna (optimisation bayésienne) en termes de C-index.
  • Sur le jeu de données Flchain, CleanSurvival a atteint un C-index de 0,67, contre 0,39 pour la CCA et 0,62 pour l'imputation par moyenne.
  • Sur le jeu de données Rotterdam avec 50 % de données manquantes, CleanSurvival a atteint des C-index entre 0,82 et 0,83, surpassant significativement la CCA (0,77–0,79) et l'imputation par moyenne (0,78–0,79).
• Sélection de Caractéristiques : La politique d'apprentissage Q a fréquemment sélectionné des opérateurs de sélection de caractéristiques, avec LASSO apparaissant dans environ 68 % des pipelines optimisés et RFE dans environ 42 %.
• Calibration : Le Score Brier Intégré (IBS) a indiqué que la calibration était généralement stable à travers les mécanismes de données manquantes, à l'exception des conditions MCAR à 40 % qui ont montré une variance plus élevée.
• Temps d'Exécution : Le comportement du temps d'exécution dépendait des conditions. Bien que le framework ait identifié avec succès des pipelines performants, le temps nécessaire pour atteindre 95 % du meilleur C-index observé variait selon le jeu de données et le mécanisme de données manquantes.

Signification et Revendications
L'article revendique que CleanSurvival comble une lacune critique dans l'écosystème AutoML en généralisant le prétraitement automatisé à l'analyse de survie. Sa signification réside dans :

• Robustesse : Démontrer que le prétraitement automatisé et guidé par le RL peut produire des modèles de survie plus robustes face à des niveaux variables de données manquantes et de bruit par rapport aux approches manuelles ou statiques.
• Accessibilité : Rendre les études d'analyse de survie "plus faciles et plus rapides à réaliser" pour les praticiens qui pourraient manquer d'expertise métier pour ajuster manuellement des pipelines de prétraitement complexes pour des données censurées.
• Intégration : Fournir un pont transparent entre le nettoyage des données et la modélisation de survie en aval, garantissant que les décisions de nettoyage sont directement guidées par leur impact sur les prédictions de survie plutôt que par des métriques de qualité de données génériques.

Les auteurs maintiennent un ton modeste, reconnaissant que bien que les résultats soient prometteurs, la généralisation à des jeux de données cliniques multidimensionnels et multimodaux nécessite une vérification supplémentaire, et que la surcharge computationnelle de la boucle d'apprentissage Q (en particulier avec les modèles KNN et d'apprentissage profond) reste une limitation. Un travail futur est proposé pour étendre le framework à des tâches de prétraitement supplémentaires et améliorer l'évolutivité.