Variational Consensus Monte Carlo for Bayesian Mixture

Résumé Technique : Variational Consensus Monte Carlo pour les Modèles de Mélange Bayésiens

Énoncé du Problème
L'article traite du défi que représente l'inférence bayésienne pour les modèles de mélange dans un contexte d'apprentissage fédéré (Federated Learning - FL), spécifiquement le scénario « cross-silo » où les données sont naturellement partitionnées entre plusieurs nœuds (par exemple, des centres médicaux) et ne peuvent être regroupées en raison de contraintes de confidentialité ou de réglementation. Bien que le Consensus Monte Carlo (CMC) offre un cadre pour l'inférence bayésienne distribuée en agrégeant des échantillons de la distribution postérieure locale, les méthodes existantes présentent des limitations importantes lorsqu'elles sont appliquées au partitionnement non supervisé (clustering). Plus précisément, l'approche Variational Consensus Monte Carlo (VCMC) proposée par Rabinovich, Angelino et Jordan (2015) suppose que le nombre de clusters et les paramètres clés du mélange sont connus et que chaque cluster est présent dans chaque fragment de données. Ces hypothèses sont rarement vérifiées en pratique, particulièrement dans les applications de données de santé où la prévalence des clusters varie selon les sites et où le nombre réel de sous-groupes latents est inconnu. De plus, l'agrégation CMC standard (par exemple, les moyennes pondérées) ne garantit pas une approximation optimale de la distribution postérieure globale, et les méthodes de partitionnement fédéré existantes reposent souvent sur des approximations de type champ moyen (mean-field) nécessitant des priors conjugués ou échouent à gérer les contextes non conjugués sans accéder à l'ensemble du jeu de données.

Méthodologie
Les auteurs proposent un pipeline complet étendant le VCMC aux modèles de mélange sur-ajustés (over-fitted) de distributions catégorielles. Le cadre opérationnel se déroule en deux phases principales :

1. L'étape d'Application (Inférence Locale) :

   • Un algorithme MCMC (spécifiquement un échantillonneur de Gibbs) est exécuté indépendamment dans chaque fragment de données.
   • Pour gérer le nombre inconnu de clusters, le modèle est « sur-ajusté » en fixant le nombre de composantes du mélange ($C'$) plus grand que le nombre anticipé de vrais clusters ($C$).
   • Un prior fractionné est utilisé, où le prior global est divisé en $p(\theta)^{1/S}$ à travers les fragments. Cela garantit que le produit des distributions postérieures locales approxime la distribution postérieure globale.
   • Les auteurs implémentent un échantillonneur de Gibbs parallélisé en C++17 pour accélérer l'inférence locale.

2. L'étape d'Agrégation (Inférence Globale) :

   • Appariement de Clusters (Cluster Matching) : Puisque les modèles de mélange sont invariants aux permutations d'étiquettes, une étape critique consiste à apparier les clusters entre les fragments. L'article introduit et évalue trois algorithmes :
     • Appariement Hongrois (Hungarian Matching) : Impose une correspondance un-à-un mais échoue si tous les clusters ne sont pas présents dans chaque fragment.
     • Appariement par Divergence Minimale (Minimum Divergence Matching) : Minimise la divergence de Kullback-Leibler (KL) entre les distributions postérieures marginales locales.
     • Appariement par Balle (Ball Matching) : Un nouvel algorithme qui regroupe les clusters qui sont « proches » (à l'intérieur d'un rayon $r$) dans l'espace des paramètres, sur la base des taux de convergence des mélanges sur-ajustés. Cette méthode permet de fusionner des clusters au sein d'un fragment s'ils sont indiscernables, atténuant ainsi la surestimation locale du nombre de clusters.
   • Agrégation Variationnelle : L'agrégation est formulée comme un problème d'inférence variationnelle visant à minimiser la divergence KL entre l'approximation et la véritable distribution postérieure globale. La fonction d'agrégation est restreinte à une somme pondérée des paramètres locaux ($\sum \upsilon_s \theta_s$).
   • Optimisation : Les poids d'agrégation optimaux sont trouvés à l'aide de la descente de gradient stochastique projetée. L'article détaille quatre stratégies pour cette optimisation afin de s'adapter à différentes contraintes de FL :
     1. Partage des données brutes avec le nœud global (le moins restrictif).
     2. Descente de gradient fédérée (partage des gradients uniquement).
     3. Communication limitée (exécutions parallèles avec peu d'itérations et de multiples points de départ).
     4. Partage d'une matrice de contingence (statistiques de synthèse) avec le nœud global.

Contributions Clés
L'article présente quatre contributions méthodologiques majeures :

1. Extension aux Modèles Sur-ajustés : Il étend le VCMC pour inférer simultanément le nombre de clusters et tous les paramètres du modèle sans nécessiter de conjugaison, levant ainsi l'hypothèse que le nombre de clusters est connu.
2. Nouveaux Algorithmes d'Appariement de Clusters : Il introduit des stratégies d'appariement (spécifiquement la Divergence Minimale et la Balle) adaptées aux contextes de cross-silo où certains clusters peuvent être absents de certains fragments, un scénario où l'appariement hongrois standard échoue.
3. Stratégies d'Inférence Flexibles : Il propose plusieurs stratégies d'agrégation adaptées à différentes contraintes de l'apprentissage fédéré (par exemple, la bande passante de communication limitée ou les exigences strictes de confidentialité).
4. Directives Pratiques : Il fournit des directives empiriques pour la sélection des algorithmes d'appariement et des stratégies d'inférence en fonction des caractéristiques des données (par exemple, la séparabilité des clusters et l'hétérogénéité des fragments).

Résultats
Le cadre a été évalué via une étude de simulation exhaustive et une application aux données réelles de dossiers de santé électroniques (EHR).

• Étude de Simulation :

  • Performance de l'Appariement : Dans des contextes hétérogènes où tous les clusters n'apparaissent pas dans chaque fragment, l'algorithme hongrois a échoué à récupérer l'appariement correct (0 % de succès dans les scénarios non imbriqués). L'appariement par Balle et l'appariement par Divergence Minimale ont été nettement plus performants. L'appariement par Balle s'est avéré particulièrement efficace pour atténuer la surestimation du nombre de clusters causée par le bruit de l'échantillonnage MCMC local.
  • Qualité du Partitionnement : Le VCMC a obtenu des scores de l'indice de Rand ajusté (ARI) comparables ou légèrement supérieurs aux alternatives de pointe (FedMerDel et Federated K-modes) et au « gold standard » du MCMC sur données complètes. Notamment, le VCMC a surpassé le MCMC sur données complètes pour la récupération de petits clusters lorsque la distribution des données entre les fragments était hétérogène.
  • Estimation des Paramètres : Le VCMC a fourni un faible biais relatif pour les poids des clusters. Dans les scénarios de clusters mal séparables et de distributions de fragments hétérogènes, le VCMC a estimé les poids des petits clusters plus précisément que le MCMC sur données complètes, probablement parce que ces petits clusters étaient suffisamment importants au sein de fragments spécifiques pour être identifiés avec une incertitude moindre.
  • Comparaison avec les Alternatives : Bien que FedMerDel soit plus rapide, il repose sur des approximations de champ moyen qui sous-estiment les variances postérieures. Federated K-modes présentait un biais élevé et de faibles performances dans l'estimation des structures de clusters pour les données catégorielles.

• Application Réelle (Cohorte THIN) :

  • La méthode a été appliquée à 289 821 patients gériatriques de la base de données britannique THIN pour identifier des modèles de comorbidités.
  • L'algorithme a identifié 27 clusters. Il a caractérisé avec succès des groupes cliniquement significatifs, tels que des clusters dominés par l'AVC, la démence ou des types de cancers spécifiques.
  • Crucialement, il a identifié de petits clusters distincts (par exemple, un groupe de seulement 31 individus caractérisé par la pancréatite, la polyarthrite rhumatoïde et la dysfonction érectile) qui s'alignaient avec les corrélations médicales documentées, démontrant la capacité de la méthode à détecter des sous-groupes rares.

Signification et Revendications
Les auteurs affirment que leur cadre constitue une solution robuste pour le partitionnement non supervisé dans des environnements sensibles à la confidentialité. La portée de ce travail réside dans la capacité à :

• Inférer le nombre de clusters et les paramètres du modèle sans regroupement de données.
• Récupérer des clusters petits et rares avec une plus grande précision que le MCMC standard sur données regroupées, lorsque les fragments de données locaux contiennent un signal suffisant pour ces sous-groupes spécifiques.
• Opérer sous des contraintes de confidentialité strictes (FL cross-silo) sans nécessiter de priors conjugués ou d'accès à l'ensemble du jeu de données.

L'article note avec modestie que, bien que le cadre soit théoriquement extensible à d'autres distributions de mélange, la validation actuelle est spécifique aux distributions catégorielles. Il souligne également que le choix de l'algorithme d'appariement dépend de l'objectif spécifique : l'appariement par Balle est préférable pour l'estimation du nombre de clusters, tandis que l'appariement par Divergence Minimale peut être préférable si l'on souhaite éviter la fusion de clusters distincts. Les auteurs concluent que, bien que le VCMC puisse être plus intensif en calcul que les méthodes variationnelles comme FedMerDel, il offre une précision supérieure dans l'estimation des paramètres et la récupération des clusters dans des contextes fédérés complexes et hétérogènes.