Federated Learning for ICD Classification with Lightweight Models and Pretrained Embeddings

Technical Summary: Federated Learning for ICD Classification with Lightweight Models and Pretrained Embeddings

Problem Statement
Automated classification of clinical notes into International Classification of Diseases (ICD) codes is critical for epidemiological research and healthcare administration but remains a bottleneck due to the manual effort required. While Large Language Models (LLMs) have shown promise in automating this task, their practical deployment in healthcare is hindered by privacy regulations (e.g., GDPR) and ethical concerns that prevent the centralization of sensitive patient data. Traditional Federated Learning (FL) approaches often require fine-tuning massive transformer models across distributed nodes, which is computationally expensive, communication-heavy, and difficult to deploy in resource-constrained hospital environments. Furthermore, existing FL studies in healthcare often focus on simpler tasks (e.g., mortality prediction) rather than complex, multi-label ICD coding.

Methodology
The authors propose a modular, two-stage pipeline that decouples semantic representation from classification to address privacy and scalability challenges:

1. Offline Embedding Extraction: Instead of fine-tuning LLMs across institutions, the study utilizes six state-of-the-art, open-source embedding models (selected from the Massive Text Embedding Benchmark leaderboard) to encode raw clinical discharge summaries into fixed-dimensional dense vectors (768–1536 dimensions). This process occurs locally and offline at each node, ensuring no raw text leaves the facility. The selected models include Qwen3-Embed-0.6B, gte-Qwen2-1.5B-instruct, gte-multilingual-base, snowflake-arctic-embed-l/m-v2.0, and nomic-embed-text-v1.5.
2. Lightweight Classification: The resulting embeddings serve as input to simple, lightweight Multi-Layer Perceptron (MLP) classifiers. Three architectures were evaluated: a standard MLP, a DeepMLP (three stacked layers with batch normalization and dropout), and a DeepResMLP (incorporating residual connections).
3. Federated Training Setup: The study simulates a horizontal FL environment using the Flower framework across 20 virtual clients. Data is partitioned from the MIMIC-IV dataset (specifically discharge summaries labeled with ICD-9 and ICD-10 codes). The training strategy employs FedAvg for aggregation over 30 communication rounds.
4. Experimental Design: The authors conducted extensive experiments comparing centralized training (baseline) against federated training. They evaluated performance across six embedding models and three classifier architectures under two coding standards (ICD-9 and ICD-10). Additionally, an ablation study using 10 random stratified splits was performed to assess stability.

Key Contributions

• Privacy-Preserving Architecture: The paper demonstrates a viable pipeline where raw clinical text never leaves the local node. By freezing the embedding models and only training lightweight classification heads in a federated manner, the approach adheres to strict data sovereignty requirements.
• Decoupling Representation and Classification: The study challenges the necessity of end-to-end fine-tuning of massive models for ICD coding. It posits that high-quality, frozen embeddings can provide sufficient semantic richness that simple classifiers (like MLPs) can achieve competitive performance.
• Comprehensive Benchmarking: The work provides a rigorous evaluation of six distinct embedding models and three classifier depths under both centralized and federated settings, offering a detailed analysis of how embedding quality versus classifier complexity impacts multi-label performance.
• Scalability Analysis: The proposed architecture uses models with significantly fewer parameters (under 7.1M for the largest classifier) compared to state-of-the-art end-to-end transformers (e.g., PLM-ICD with ~139M parameters), making it feasible for deployment on standard hospital hardware.

Results

• Embedding Quality Dominance: The experiments reveal that the choice of embedding model is the primary determinant of predictive performance, outweighing the complexity of the classifier. For instance, Qwen3-Embed-0.6B paired with a simple MLP achieved the highest micro F1 score for ICD-10 (0.3942), while gte-Qwen2-1.5B-instruct with DeepResMLP performed best for ICD-9.
• Classifier Complexity: Deeper architectures (DeepMLP, DeepResMLP) did not consistently outperform the simple MLP. In many configurations, the added complexity yielded diminishing returns or even performance drops, suggesting that the semantic burden is effectively handled by the pre-trained embeddings.
• Federated vs. Centralized: Federated learning results closely matched centralized training baselines. For example, the Qwen3-Embed-0.6B + MLP configuration achieved a micro F1 of 0.3944 in FL compared to 0.3942 in centralized training for ICD-10, indicating that data decentralization does not significantly degrade performance when using standardized embeddings.
• Performance on Frequent Codes: On the top 10 and top 50 most frequent ICD codes, the models achieved significantly higher F1 scores (often >0.60 for Top-10), demonstrating strong discriminative capacity for common diagnoses.
• Stability: The ablation study across 10 random splits showed low variance in performance (standard deviations often <0.003 for F1 scores), confirming the robustness of the approach.

Significance and Claims
The authors claim that this work offers a "viable way toward scalable, privacy-conscious medical coding systems." They argue that high-performing medical coding systems can be built without relying on costly, rigid, end-to-end architectures that require centralizing sensitive data. The study suggests that the "paradox of abundant but isolated data" in healthcare can be addressed by leveraging frozen, standardized embeddings locally and training only lightweight heads collaboratively.

The paper maintains a modest tone regarding its limitations. It acknowledges that the reliance on static, non-fine-tuned embeddings may limit performance on rare or highly domain-specific codes. Furthermore, the FL setup assumes idealized conditions (homogeneous data distribution, synchronous updates) that may not fully reflect the heterogeneity of real-world hospital networks. The authors conclude that while their approach achieves competitive results with a fraction of the parameters of SOTA models, future work should explore adaptive federated strategies, joint fine-tuning, and the integration of domain-specific medical knowledge to further enhance robustness.