Variational Consensus Monte Carlo for Bayesian Mixture

Resumen Técnico: Monte Carlo de Consenso Variacional para Modelos de Mezcla Bayesianos

Planteamiento del Problema
El artículo aborda el desafío de realizar inferencia bayesiana para modelos de mezcla en un entorno de Aprendizaje Federado (FL), específicamente en el escenario "cross-silo" donde los datos están naturalmente particionados entre múltiples nodos (por ejemplo, centros médicos) y no pueden agruparse debido a restricciones de privacidad o regulatorias. Si bien el Monte Carlo de Consenso (CMC) ofrece un marco para la inferencia bayesiana distribuida mediante la agregación de muestras de la posterior local, los métodos existentes enfrentan limitaciones significativas cuando se aplican al agrupamiento (clustering) no supervisado. Específicamente, el enfoque de Monte Carlo de Consenso Variacional (VCMC) propuesto por Rabinovich, Angelino y Jordan (2015) asume que el número de clústeres y los parámetros clave de la mezcla son conocidos y que cada clúster está presente en cada fragmento (shard) de datos. Estas suposiciones rara vez se cumplen en la práctica, particularmente en aplicaciones de datos de salud donde la prevalencia de los clústeres varía entre los sitios y el número real de subgrupos latentes es desconocido. Además, la agregación estándar de CMC (por ejemplo, medias ponderadas) no garantiza una aproximación óptima de la posterior global, y los métodos de agrupamiento federados existentes suelen depender de aproximaciones de campo medio que requieren prioris conjugadas o fallan al manejar configuraciones no conjugadas sin acceder al conjunto de datos completo.

Metodología
Los autores proponen un flujo de trabajo integral que extiende el VCMC a modelos de mezcla sobreajustados de distribuciones categóricas. El marco opera en dos fases principales:

1. El Paso de Aplicación (Inferencia Local):

   • Se ejecuta un algoritmo MCMC (específicamente un muestreador de Gibbs) de forma independiente dentro de cada fragmento de datos.
   • Para manejar el número desconocido de clústeres, el modelo se "sobreajusta" estableciendo el número de componentes de la mezcla ($C'$) mayor que el número previsto de clústeres reales ($C$).
   • Se utiliza una prior fraccionada, donde la prior global se divide como $p(\theta)^{1/S}$ a través de los fragmentos. Esto asegura que el producto de las posteriores locales aproxime la posterior global.
   • Los autores implementan un muestreador de Gibbs en C++17 paralelizado para acelerar la inferencia local.

2. El Paso de Agregación (Inferencia Global):

   • Emparejamiento de Clústeres (Cluster Matching): Dado que los modelos de mezcla son invariantes a las permutaciones de etiquetas, un paso crítico es el emparejamiento de clústeres entre fragmentos. El artículo introduce y evalúa tres algoritmos:
     • Emparejamiento Húngaro (Hungarian Matching): Fuerza una correspondencia uno a uno, pero falla si no todos los clústeres aparecen en cada fragmento.
     • Emparejamiento por Divergencia Mínima (Minimum Divergence Matching): Minimiza la divergencia de Kullback-Leibler (KL) entre las marginales locales de la posterior.
     • Emparejamiento de Bolas (Ball Matching): Un algoritmo novedoso que agrupa clústeres que están "cerca" (dentro de un radio $r$) en el espacio de parámetros, basándose en las tasas de convergencia de las mezclas sobreajustadas. Este método permite fusionar clústeres dentro de un fragmento si son indistinguibles, mitigando la sobreestimación local del recuento de clústeres.
   • Agregación Variacional: La agregación se enmarca como un problema de inferencia variacional para minimizar la divergencia KL entre la aproximada y la verdadera posterior global. La función de agregación se restringe a una suma ponderada de los parámetros locales ($\sum \upsilon_s \theta_s$).
   • Optimización: Los pesos de agregación óptimos se encuentran utilizando el descenso de gradiente estocástico proyectado. El artículo detalla cuatro estrategias para esta optimización para adaptarse a diferentes restricciones de FL:
     1. Compartir los datos brutos con el nodo global (el menos restrictivo).
     2. Descenso de gradiente federado (compartiendo solo gradientes).
     3. Comunicación limitada (ejecuciones paralelas con pocas iteraciones y múltiples puntos de partida).
     4. Compartir una matriz de contingencia (estadísticas de resumen) con el nodo global.

Contribuciones Clave
El artículo realiza cuatro contribuciones metodológicas primarias:

1. Extensión a Modelos Sobreajustados: Extiende el VCMC para inferir el número de clústeres y todos los parámetros del modelo simultáneamente sin requerir conjugación, relajando la suposición de que el número de clústeres es conocido.
2. Nuevos Algoritmos de Emparejamiento de Clústeres: Introduce estrategias de emparejamiento (específicamente Emparejamiento por Divergencia Mínima y de Bolas) adecuadas para entornos cross-silo donde los clústeres pueden estar ausentes en algunos fragmentos, un escenario donde el emparejamiento Húngaro estándar falla.
3. Estrategias de Inferencia Flexibles: Propone múltiples estrategias de agregación adaptadas a diferentes restricciones federadas (por ejemplo, ancho de banda de comunicación limitado o requisitos estrictos de privacidad).
4. Guías Prácticas: Proporciona guías empíricas para la selección de algoritmos de emparejamiento y estrategias de inferencia basadas en las características de los datos (por ejemplo, separabilidad de clústeres y heterogeneidad de los fragmentos).

Resultados
El marco fue evaluado mediante un estudio de simulación exhaustivo y una aplicación a datos reales de registros de salud electrónicos (EHR).

• Estudio de Simulación:

  • Desempeño del Emparejamiento: En entornos heterogéneos donde no todos los clústeres aparecen en cada fragmento, el algoritmo Húngaro falló en recuperar el emparejamiento correcto (0% de éxito en escenarios no anidados). El emparejamiento de Bolas y el de Divergencia Mínima funcionaron significativamente mejor. El emparejamiento de Bolas fue particularmente efectivo para mitigar la sobreestimación del número de clústeres causada por el ruido del MCMC local.
  • Calidad del Agrupamiento: El VCMC logró puntuaciones de Índice Rand Ajustado (ARI) comparables o ligeramente superiores a las de las alternativas de vanguardia (FedMerDel y Federated K-modes) y al "estándar de oro" del MCMC de datos completos. Notablemente, el VCMC superó al MCMC de datos completos en la recuperación de clústeres pequeños cuando la distribución de los datos entre los fragmentos era heterogénea.
  • Estimación de Parámetros: El VCMC proporcionó un bajo sesgo relativo para los pesos de los clústeres. En escenarios con clústeres pobremente separables y distribuciones de fragmentos heterogéneas, el VCMC estimó los pesos de los clústeres pequeños con mayor precisión que el MCMC de datos completos, probablemente porque estos clústeres pequeños eran lo suficientemente grandes dentro de fragmentos específicos para ser identificados con menor incertidumbre.
  • Comparación con Alternativas: Aunque FedMerDel era más rápido, dependía de aproximaciones de campo medio que subestimaban las varianzas de la posterior. Federated K-modes mostró un alto sesgo y un desempeño pobre en la estimación de estructuras de clústeres para datos categóricos.

• Aplicación en el Mundo Real (Cohorte THIN):

  • El método se aplicó a 289,821 pacientes geriátricos de la base de datos británica THIN para identificar patrones de multimorbilidad.
  • El algoritmo identificó 27 clústeres. Caracterizó con éxito grupos clínicamente significativos, como clústeres dominados por accidentes cerebrovasculares, demencia o tipos específicos de cáncer.
  • Crucialmente, identificó clústeres pequeños pero distintos (por ejemplo, uno con solo 31 individuos caracterizado por pancreatitis, artritis reumatoide y disfunción eréctil) que se alineaban con correlaciones médicas documentadas, demostrando la capacidad del método para detectar subgrupos raros que podrían diluirse en un análisis de datos agrupados.

Significancia y Reivindicaciones
Los autores afirman que su marco proporciona una solución robusta para el agrupamiento no supervisado en entornos sensibles a la privacidad. La significancia radica en la capacidad de:

• Inferir el número de clústeres y los parámetros del modelo sin la agrupación de datos.
• Recuperar clústeres pequeños y raros con mayor precisión que el MCMC de datos agrupados estándar cuando los fragmentos de datos locales contienen suficiente señal para esos subgrupos específicos.
• Operar bajo estrictas restricciones de privacidad (FL cross-silo) sin requerir prioris conjugadas ni acceso al conjunto de datos completo.

El artículo señala modestamente que, si bien el marco es teóricamente extensible a otras distribuciones de mezcla, la validación actual es específica para distribuciones categóricas. También destaca que la elección del algoritmo de emparejamiento de clústeres depende del objetivo específico: el emparejamiento de Bolas es preferible para estimar el número de clústeres, mientras que el emparejamiento por Divergencia Mínima puede ser mejor si la prioridad es evitar la fusión de clústeres distintos. Los autores concluyen que, aunque el VCMC puede ser computacionalmente intensivo en comparación con los métodos variacionales como FedMerDel, ofrece una precisión superior en la estimación de parámetros y la recuperación de clústeres en entornos federados complejos y heterogéneos.