Federated Learning for ICD Classification with Lightweight Models and Pretrained Embeddings

TL;DR

This paper explores federated learning with lightweight models and pretrained embeddings for ICD code classification from clinical notes, demonstrating comparable performance to centralized methods while preserving privacy.

Contribution

It introduces a scalable, privacy-preserving pipeline combining frozen embeddings with simple MLP classifiers for federated clinical NLP tasks, contrasting with large language models.

Findings

Embedding quality impacts performance more than classifier complexity.

Federated learning closely matches centralized results in ideal conditions.

Lightweight models achieve competitive F1 scores with significantly reduced size.

Abstract

This study investigates the feasibility and performance of federated learning (FL) for multi-label ICD code classification using clinical notes from the MIMIC-IV dataset. Unlike previous approaches that rely on centralized training or fine-tuned large language models, we propose a lightweight and scalable pipeline combining frozen text embeddings with simple multilayer perceptron (MLP) classifiers. This design offers a privacy-preserving and deployment-efficient alternative for clinical NLP applications, particularly suited to distributed healthcare settings. Extensive experiments across both centralized and federated configurations were conducted, testing six publicly available embedding models from Massive Text Embedding Benchmark leaderboard and three MLP classifier architectures under two medical coding (ICD-9 and ICD-10). Additionally, ablation studies over ten random stratified splits assess performance stability. Results show that embedding quality substantially outweighs classifier complexity in determining predictive performance, and that federated learning can closely match centralized results in idealized conditions. While the models are orders of magnitude smaller than state-of-the-art architectures and achieved competitive micro and macro F1 scores, limitations remain including the lack of end-to-end training and the simplified FL assumptions. Nevertheless, this work demonstrates a viable way toward scalable, privacy-conscious medical coding systems and offers a step toward for future research into federated, domain-adaptive clinical AI.