Embedding-Based Federated Learning with Runtime Governance for Iron Deficiency Prediction

Technical Summary: Embedding-Based Federated Learning with Runtime Governance for Iron Deficiency Prediction

Problem Statement
Despite the theoretical promise of Federated Learning (FL) in healthcare, recent reviews indicate that the vast majority of published studies fail to reach real-world deployment. This gap is attributed to methodological weaknesses regarding bias, privacy, generalization, communication efficiency, and governance compliance. Furthermore, clinical data is often structurally non-independent and identically distributed (non-IID) due to genuine differences in population demographics, clinical workflows, and testing indications, rather than mere sampling noise. Standard FL approaches, which often rely on sample-size-weighted aggregation, may degrade performance when sites differ significantly in task relevance and data distribution. Additionally, deploying FL in healthcare requires robust runtime governance to ensure policy compliance and auditability across institutional boundaries, a feature often absent in research prototypes.

Methodology
The authors present a cross-institutional deployment across two distinct clinical environments: Amsterdam University Medical Centre (AUMC) and NHS Blood and Transplant (NHSBT). The system addresses iron deficiency prediction using routine Full Blood Count (FBC) data.

1. Embedding-Based Two-Stage Design:

   • Representation Extraction: The pipeline utilizes a frozen, domain-specific haematology foundation model, DeepCBC, pre-trained on large-scale flow cytometry and impedance data. This model performs site-local representation extraction, converting raw FBC inputs into 256-dimensional deterministic embeddings (posterior means).
   • Federated Training: Only a compact downstream classifier (a two-hidden-layer MLP with ~41k parameters) is trained federatively. This approach restricts recurrent communication to classifier parameters, avoiding the transmission of the massive encoder weights required by full-encoder federation.

2. Runtime Governance (FLA3):

   • The deployment runs on FLA3, a healthcare-oriented FL platform.
   • Governance is enforced at runtime rather than via static agreements. The system utilizes XACML (eXtensible Access Control Markup Language) for policy-based authorization (deny-by-default) and maintains an append-only, signed audit trail.
   • This ensures study-scoped execution, preventing unauthorized parameter transmission or participation outside assigned windows.

3. Aggregation Strategies:
   The study compares three aggregation methods:

   • FedAvg: Standard sample-size-weighted averaging.
   • FedProx: Introduces a proximal penalty to stabilize local optimization under heterogeneity.
   • FedMAP: A personalized aggregation method that casts local optimization as Maximum A Posteriori (MAP) estimation. It learns a global regularizer and uses posterior-informed weights based on local data fit and proximity to the global prior, rather than relying solely on dataset size.

Key Contributions

1. Practical FL Architecture: A two-stage design separating representation learning (via a frozen foundation model) from federated task adaptation, significantly reducing communication overhead and simplifying deployment in restrictive institutional environments.
2. Characterization of Structural Heterogeneity: A detailed analysis of a clinical federation where sites differ substantially in iron deficiency prevalence (2.8% at AUMC vs. 19.5% at NHSBT), ferritin distributions, and effective positive-class volumes due to distinct clinical workflows (reactive hospital testing vs. universal donor screening).
3. Empirical Validation of Personalized Aggregation: Evidence demonstrating that sample-size-weighted aggregation (FedAvg) degrades performance in heterogeneous clinical settings, whereas personalized aggregation (FedMAP) recovers and exceeds local-only performance.
4. Operational Governance: A demonstration of how runtime controls (study scoping, policy authorization, and signed auditing) are integrated into a live healthcare FL system to satisfy regulatory requirements across international borders (Netherlands and UK).

Results
The study evaluated performance using ROC-AUC and Balanced Accuracy on held-out test sets:

• FedAvg Performance: Standard aggregation reduced ROC-AUC at both sites compared to local-only training. The global update was biased toward the larger AUMC distribution, degrading the decision boundary for both cohorts.
• FedProx Performance: Partially recovered the performance loss but did not outperform local-only training or FedMAP, as it remained insensitive to client task relevance.
• FedMAP Performance: Achieved the highest performance across all metrics:
  • AUMC: ROC-AUC increased from 0.9470 (local-only) to 0.9594.
  • NHSBT: ROC-AUC increased from 0.8558 (local-only) to 0.8671.
  • Macro Performance: FedMAP achieved the highest macro ROC-AUC (0.9133) and macro balanced accuracy (0.8389).
• Balanced Accuracy Nuance: While FedMAP improved Balanced Accuracy at NHSBT, the gain at AUMC was marginal (0.8899 vs. 0.8964 for local-only). The authors attribute this to threshold sensitivity under low positive prevalence; while FedMAP improved discrimination (ROC-AUC), the fixed WHO threshold (15 µg/L) did not shift to capture the optimal operating point for the AUMC cohort, which has higher inflammation levels.

Significance and Claims
The paper claims that successful clinical FL deployment requires moving beyond statistical optimization to address structural clinical heterogeneity and operational governance.

• Personalized Aggregation is Critical: In federations where client sample size and task relevance diverge (e.g., different clinical indications and population health statuses), standard size-weighted aggregation is suboptimal. Personalized methods like FedMAP, which account for client relevance, are necessary to improve performance across all sites.
• Embedding-Based Design is Practical: Using a frozen foundation model for local representation extraction is a viable pattern for hospitals with limited data or infrastructure, as it decouples domain knowledge transfer from the federated training loop.
• Governance Must Be Operational: Static data-sharing agreements are insufficient for multi-institutional trust. Runtime enforcement of policies and auditability are prerequisites for deploying FL across different regulatory frameworks (e.g., Dutch and UK clinical governance).

The authors conclude that these findings, derived from a two-site operational deployment, support the case for personalized aggregation and runtime governance in clinical federations characterized by substantial structural heterogeneity. They acknowledge limitations regarding the lack of formal privacy attack analysis (e.g., membership inference) and the need for future work on model calibration and temporal drift monitoring.