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

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

1. Problem Statement

In low- and middle-income countries (LMICs), medical certification of causes of death is often unavailable, necessitating the use of Verbal Autopsy (VA) to infer causes through caregiver interviews. While probabilistic algorithms have replaced expert review for analyzing VA data, they face two critical limitations:

1. Data Access and Privacy: State-of-the-art models typically require centralized training on large, labeled reference datasets. However, privacy regulations and logistical constraints often prevent the consolidation of data from different regions, countries, or time periods. Consequently, analysts are often forced to use rigid, fixed-parameter models (e.g., InterVA, Tariff) or rely on non-representative local data.
2. Distribution Shift: VA data collected from different populations exhibit systematic differences in symptom distributions ($p(X|Y)$) and cause-of-death prevalences ($p(Y)$), even for the same underlying causes. These shifts, driven by epidemiology, healthcare access, and cultural norms, degrade the performance of models trained on external data. Existing solutions to address shift often require either diverse pooled reference data or sufficient local labeled data, both of which are frequently unavailable.

2. Methodology: Bayesian Federated Learning (BFL)

The authors propose a Bayesian Federated Learning (BFL) framework that enables the training of robust VA models without sharing individual-level data across domains. The framework operates under the assumption that reference death data exists across multiple domains (training sites) but can only be analyzed locally, while the target domain may have limited or no labeled data.

Core Architecture:

• Base Models: Each training domain $m$ independently trains a "base model" (e.g., a latent class model like LCVA) on its local labeled data to estimate the conditional likelihood of symptoms given a cause, $\hat{p}_m(X|Y)$. Only these model summaries (parameters) are shared, not the raw data.
• Global Aggregation: In the target domain, the BFL model treats the conditional distribution of symptoms as a weighted mixture of the base models:
  $$p_0(X|Y=c) = \sum_{m=1}^M \lambda_{cm} \hat{p}_m(X|Y=c)$$
  Here, $\lambda_{cm}$ represents the weight of the $m$-th source domain for cause $c$. These weights are learned via a Bayesian latent-class aggregation model.
• Generative Approach: Unlike discriminative federated learning that aggregates prediction functions $p(Y|X)$, BFL aggregates the generative distribution $p(X|Y)$. This aligns with the natural data-generating process of VA (symptoms are consequences of disease) and is empirically more robust to distribution shifts.
• Handling Local Labels: The framework offers three strategies to incorporate optional local labeled data in the target domain:
  1. BFL-domain: Treats local labeled data as a new domain, training an additional base model to expand the convex hull of conditional distributions.
  2. BFL-partial: Uses local labels directly to update the mixing weights ($\lambda$) and cause fractions ($\pi$) within the global likelihood, without retraining base models.
  3. BFL-mix: A hybrid approach splitting local data to train a local model and update global weights simultaneously.

Inference:
Posterior inference is conducted efficiently using a Gibbs sampler (or Stan), estimating the Cause-Specific Mortality Fractions (CSMFs) $\pi$ and the domain weights $\lambda$. The method supports both individual-level cause assignment and population-level quantification.

3. Key Contributions

• Novel Framework: Introduction of the first VA-specific federated learning framework that eliminates the need for centralized labeled data while addressing both label shift ($p(Y)$) and conditional shift ($p(X|Y)$).
• Modularity and Efficiency: The framework is modular, compatible with various existing VA algorithms (e.g., LCVA, InSilicoVA) as base models. Once base models are trained, the global aggregation step is computationally negligible (seconds on a laptop).
• Robustness to Shift: By modeling the conditional likelihood rather than the posterior probability, the method explicitly decouples source domain CSMFs from the symptom-cause relationships, allowing for robust transfer even when source and target populations have vastly different mortality profiles.
• Flexible Fine-tuning: The provision of three distinct strategies (domain, partial, mix) to leverage local labeled data depending on the severity of within-target label shift and the representativeness of the local sample.

4. Experimental Results

The authors validated BFL using two real-world datasets: the PHMRC gold-standard adult VA dataset (6 sites) and the CHAMPS neonatal dataset (7 countries).

• PHMRC Experiments:

  • Scenario: Leave-one-domain-out experiments under varying degrees of within-target label shift (none, mild, severe).
  • Performance: BFL variants consistently outperformed single-domain base models (local-self, local-avg).
  • Comparison with Joint Modeling: BFL performed comparably to, and in some cases better than, the "gold standard" joint modeling approach (LCVA with full data pooling), despite lacking access to the raw pooled data.
  • Label Shift Sensitivity: Under severe label shift, BFL-partial (direct weight updating) outperformed BFL-domain (local model training), as the latter struggled when local labeled data did not represent the causes prevalent in the unlabeled target population.
  • Calibration Methods: Standard calibration methods (GBQL) performed well under random labeling but degraded significantly under severe label shift or with strong shrinkage priors, whereas BFL remained robust.

• CHAMPS Experiments:

  • Performance: In the neonatal dataset with high class imbalance, BFL achieved CSMF accuracy comparable to full-data LCVA and significantly better than single-domain models.
  • Weight Interpretation: The learned domain weights ($\lambda$) effectively identified which source domains provided the most informative symptom patterns for specific causes, even without direct access to source CSMFs.

5. Significance and Claims

The paper claims that BFL offers a practical solution to the "data silo" problem in global health surveillance. By allowing analysts to leverage diverse reference datasets without violating privacy constraints, BFL bridges the gap between simple, rigid models and complex, data-hungry joint models.

• Practical Impact: The framework facilitates the deployment of advanced, flexible VA models in real-world mortality surveillance systems where data sharing is restricted.
• Robustness: The results demonstrate that federated aggregation of generative models can effectively handle distribution shifts that typically cripple standard transfer learning approaches.
• Caveats: The authors maintain a modest stance, noting that no single method dominates universally; performance depends heavily on the specific nature of the distribution shift and the availability/representativeness of local labeled data. They emphasize that while BFL improves upon existing methods, the fundamental challenge of validating cause-of-death assignment without sufficient reference deaths remains, and claims of "solving" the problem should be viewed with caution.

The work concludes that BFL is a scalable, computationally efficient, and robust approach for modernizing cause-of-death surveillance in resource-constrained settings.