Bayesian Federated Cause-of-Death Classification and Quantification Under Distribution Shift

Resumen Técnico: Clasificación y Cuantificación Bayesiana Federada de la Causa de Muerte bajo Desplazamiento de Distribución

1. Planteamiento del Problema

En los países de ingresos bajos y medios (PIBM), la certificación médica de las causas de muerte suele ser inexistente, lo que requiere el uso de la Autopsia Verbal (AV) para inferir las causas mediante entrevistas con cuidadores. Aunque los algoritmos probabilísticos han reemplazado la revisión de expertos para analizar los datos de AV, enfrentan dos limitaciones críticas:

1. Acceso a Datos y Privacidad: Los modelos de vanguardia suelen requerir un entrenamiento centralizado en grandes conjuntos de datos de referencia etiquetados. Sin embargo, las regulaciones de privacidad y las restricciones logísticas a menudo impiden la consolidación de datos de diferentes regiones, países o períodos de tiempo. En consecuencia, los analistas se ven obligados a utilizar modelos de parámetros fijos y rígidos (por ejemplo, InterVA, Tariff) o a depender de datos locales no representativos.
2. Desplazamiento de Distribución (Distribution Shift): Los datos de AV recolectados de diferentes poblaciones exhiben diferencias sistemáticas en las distribuciones de síntomas ($p(X|Y)$) y en las prevalencias de las causas de muerte ($p(Y)$), incluso para las mismas causas subyacentes. Estos desplazamientos, impulsados por la epidemiología, el acceso a la salud y las normas culturales, degradan el rendimiento de los modelos entrenados con datos externos. Las soluciones existentes para abordar este desplazamiento requieren, ya sea datos de referencia diversos y agrupados, o suficientes datos etiquetados locales, ambos de los cuales suelen estar ausentes.

2. Metodología: Aprendizaje Federado Bayesiano (BFL)

Los autores proponen un marco de Aprendizaje Federado Bayesiano (BFL) que permite el entrenamiento de modelos de AV robustos sin compartir datos a nivel individual entre dominios. El marco opera bajo el supuesto de que existen datos de referencia de defunciones en múltiples dominios (sitios de entrenamiento) que solo pueden analizarse localmente, mientras que el dominio objetivo puede tener datos etiquetados limitados o inexistentes.

Arquitectura Central:

• Modelos Base: Cada dominio de entrenamiento $m$ entrena independientemente un "modelo base" (por ejemplo, un modelo de clase latente como LCVA) sobre sus datos etiquetados locales para estimar la verosimilitud condicional de los síntomas dado una causa, $\hat{p}_m(X|Y)$. Solo se comparten estos resúmenes de modelos (parámetros), no los datos brutos.
• Agregación Global: En el dominio objetivo, el modelo BFL trata la distribución condicional de los síntomas como una mezcla ponderada de los modelos base:
  $$p_0(X|Y=c) = \sum_{m=1}^M \lambda_{cm} \hat{p}_m(X|Y=c)$$
  Aquí, $\lambda_{cm}$ representa el peso del $m$-ésimo dominio de origen para la causa $c$. Estos pesos se aprenden mediante un modelo de agregación de clase latente bayesiano.
• Enfoque Generativo: A diferencia del aprendizaje federado discriminativo que agrega funciones de predicción $p(Y|X)$, el BFL agrega la distribución generativa $p(X|Y)$. Esto se alinea con el proceso natural de generación de datos en la AV (los síntomas son consecuencias de la enfermedad) y es empíricamente más robusto al desplazamiento de la distribución.
• Gestión de Etiquetas Locales: El marco ofrece tres estrategias para incorporar datos etiquetados locales opcionales en el dominio objetivo:
  1. BFL-domain: Trata los datos etiquetados locales como un nuevo dominio, entrenando un modelo base adicional para expandir el envolvente convexo de las distribuciones condicionales.
  2. BFL-partial: Utiliza las etiquetas locales directamente para actualizar los pesos de la mezcla ($\lambda$) y las fracciones de causa ($\pi$) dentro de la verosimilitud global, sin reentrenar los modelos base.
  3. BFL-mix: Un enfoque híbrido que divide los datos locales para entrenar un modelo local y actualizar los pesos globales simultáneamente.

Inferencia:
La inferencia posterior se realiza eficientemente utilizando un muestreador de Gibbs (o Stan), estimando las Fracciones de Mortalidad Específicas por Causa (CSMFs) $\pi$ y los pesos de dominio $\lambda$. El método soporta tanto la asignación de causas a nivel individual como la cuantificación a nivel poblacional.

3. Contribuciones Clave

• Marco Novedoso: Introducción del primer marco de aprendizaje federado específico para AV que elimina la necesidad de datos centralizados etiquetados mientras aborda tanto el desplazamiento de etiquetas ($p(Y)$) como el desplazamiento condicional ($p(X|Y)$).
• Modularidad y Eficiencia: El marco es modular, compatible con varios algoritmos de AV existentes (por ejemplo, LCVA, InSilicoVA) como modelos base. Una vez entrenados los modelos base, el paso de agregación global es computacionalmente insignificante (segundos en una computadora portátil).
• Robustez al Desplazamiento: Al modelar la verosimilitud condicional en lugar de la probabilidad posterior, el método desacopla explícitamente las CSMFs de los dominios de origen de las relaciones síntoma-causa, lo que permite una transferencia robusta incluso cuando las poblaciones de origen y objetivo tienen perfiles de mortalidad muy distintos.
• Ajuste Fino Flexible: La provisión de tres estrategias distintas (domain, partial, mix) para aprovechar los datos etiquetados locales según la severidad del desplazamiento de etiquetas dentro del objetivo y la representatividad de la muestra local.

4. Resultados Experimentales

Los autores validaron el BFL utilizando dos conjuntos de datos del mundo real: el conjunto de datos de adultos de estándar de oro PHMRC (6 sitios) y el conjunto de datos neonatales CHAMP (7 países).

• Experimentos PHMRC:

  • Escenario: Experimentos de dejar un dominio fuera (leave-one-domain-out) bajo grados variables de desplazamiento de etiquetas dentro del objetivo (ninguno, leve, severo).
  • Rendimiento: Las variantes de BFL superaron consistentemente a los modelos base de un solo dominio (local-self, local-avg).
  • Comparación con el Modelado Conjunto: El BFL funcionó de manera comparable a, y en algunos casos mejor que, el enfoque de modelado conjunto "estándar de oro" (LCVA con agrupación total de datos), a pesar de carecer del acceso a los datos brutos agrupados.
  • Sensibilidad al Desplazamiento de Etiquetas: Bajo un desplazamiento de etiquetas severo, BFL-partial (actualización directa de pesos) superó a BFL-domain (entrenamiento de modelo local), ya que este último tuvo dificultades cuando los datos etiquetados locales no representaban las causas prevalentes en la población objetivo no etiquetada.
  • Métodos de Calibración: Los métodos de calibración estándar (GBQL) funcionaron bien bajo etiquetado aleatorio pero se degradaron significativamente bajo un desplazamiento de etiquetas severo o con priors de contracción (shrinkage) fuertes, mientras que el BFL se mantuvo robusto.

• Experimentos CHAMPS:

  • Rendimiento: En el conjunto de datos neonatal con alto desequilibrio de clases, el BFL logró una precisión de CSMF comparable al LCVA de datos completos y significativamente mejor que los modelos de un solo dominio.
  • Interpretación de Pesos: Los pesos de dominio ($\lambda$) aprendidos identificaron eficazmente qué dominios de origen proporcionaban los patrones de síntomas más informativos para causas específicas, incluso sin acceso directo a las CSMFs de origen.

5. Significancia y Reivindicaciones

El artículo afirma que el BFL ofrece una solución práctica al problema de los "silos de datos" en la vigilancia de la salud global. Al permitir que los analistas aprovechen diversos conjuntos de datos de referencia sin violar las restricciones de privacidad, el BFL cierra la brecha entre los modelos simples y rígidos y los modelos complejos que demandan grandes cantidades de datos.

• Impacto Práctico: El marco facilita el despliegue de modelos de AV avanzados y flexibles en sistemas de vigilancia de mortalidad del mundo real donde el intercambio de datos está restringido.
• Robustez: Los resultados demuestran que la agregación federada de modelos generativos puede manejar eficazmente los desplazamientos de distribución que típicamente paralizan los enfoques de transferencia de aprendizaje estándar.
• Advertencias: Los autores mantienen una postura modesta, señalando que ningún método domina universalmente; el rendimiento depende fuertemente de la naturaleza específica del desplazamiento de distribución y de la disponibilidad o representatividad de los datos etiquetados locales. Enfatizan que, si bien el BFL mejora los métodos existentes, el desafío fundamental de validar la asignación de la causa de muerte sin suficientes defunciones de referencia persiste, y las afirmaciones de "resolver" el problema deben verse con cautela.

El trabajo concluye que el BFL es un enfoque escalable, computacionalmente eficiente y robusto para modernizar la vigilancia de la causa de muerte en entornos con recursos limitados.