cs.LGMay 15, 2026

Federated Imputation under Heterogeneous Feature Spaces

Authors: Imane HocineChaimaa MedjadjiSylvain KublerGregoire DanoyYves Le Traon

Organizations: SnT, University of Luxembourg, Luxembourg · FSTM/DCS, University of Luxembourg, Luxembourg

Abstract

Federated Learning (FL) enables collaborative training across decentralized clients, but most methods assume aligned feature schemas, an assumption that rarely holds in tabular settings where clients observe only partially overlapping feature subsets. In these heterogeneous feature spaces, parameter-averaging methods (e.g., FedAvg) transfer little information across weakly overlapping or disjoint feature groups, limiting their effectiveness for federated imputation. To overcome this, we propose \textbf{FedHF-Impute}, a federated imputation framework that separates structural feature unavailability from conventional missingness and uses a shared global feature graph to propagate information across statistically related features through message passing. This enables indirect cross-client knowledge transfer, even when features are never jointly observed locally, while preserving standard federated communication. Under simulated partial schema overlap on the SECOM and AirQuality datasets, FedHF-Impute improves imputation accuracy (RMSE) over FL baselines by 26.9%, and 8.4% respectively, while achieving comparable performance on PhysioNET, with only a 0.3% difference relative to the best baseline.

Explore similar work

Apr 25, 2026cs.LG

Conditional Imputation for Within-Modality Missingness in Multi-Modal Federated Learning

Multimodal Federated Learning (MMFL) enables privacy-preserving collaborative training, but real-world clinical applications often suffer from within-modality missingness caused by sensor intermittency or irregular sampling. Existing methods implicitly represent unobserved data via architectural alignment or missing embeddings, often failing to recover the true distribution and yielding sub-optimal performance. We propose CondI, a federated framework explicitly addressing this missingness using conditional diffusion models. CondI employs a two-phase training pipeline: first, imputing unobserved temporal components using available multimodal context and conditional embeddings; second, optimizing modality-specific extractors and joint embedding spaces. During inference, imputed raw data pass through trained extractors to generate robust features, providing a holistic representation for downstream tasks. Explicit data imputation ensures models operate on complete semantic structures, significantly enhancing resilience against severe data incompleteness. Experiments on three clinical datasets (PTB-XL, SLEEP-EDF, MIMIC-IV) demonstrate CondI achieves comparable results to state-of-the-art baselines. Code: https://github.com/ZhengWugeng/CondI
Wugeng Zheng, Ziwen Kan, Katie Wang +2
Jul 5, 2026eess.IV

FedProIn: Mitigating Client Drift for Learnable Prototypes in Federated Medical Imaging

Federated learning (FL) is severely hindered by statistical heterogeneity due to variations in scanners, acquisition protocols, and patient populations. Such non-IID data induces client drift during local optimization, leading to unstable convergence and suboptimal global models when parameter-based aggregation is applied. We propose a prototype-based, influence-aware federated learning framework (FedProIn) that uses multiple learnable class prototypes to capture shared semantic structures across heterogeneous clients. We introduce feature divergence loss and prototype contrastive loss to mitigate client drift by decomposing it into feature drift and prototype drift. In addition, we propose a normalized influence aggregation strategy that adaptively weights client prototypes according to their contribution to the global representation, reducing the impact of biased or low-quality updates. Experimental results on two publicly available medical datasets, HAM10000 and Matek-19, demonstrate that FedProIn achieves accuracies of (83.5% IID, 81.1% non-IID) on HAM10000 and (96.2% IID, 95.8% non-IID) on Matek-19, respectively, outperforming existing baselines in both conditions. Our code is available at https://github.com/harsh-kmr/FedProIn.
Harsh Kumar, Tarun Kumar Garg, Vaanathi Sundaresan
Jul 23, 2024cs.LG

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation

Federated Learning (FL) faces challenges due to data heterogeneity, which limits the global model's performance across diverse client distributions. Personalized Federated Learning (PFL) addresses this by enabling each client to possess an individual model adapted to its local distribution. Many existing methods assume that certain global model parameters are difficult to train effectively in a collaborative manner under heterogeneous data. Consequently, they localize or fine-tune these parameters to obtain personalized models. In this paper, we reveal that both the feature extractor and classifier of the global model are inherently strong, and the primary cause of its suboptimal performance is the mismatch between local features and the global classifier. Although existing methods alleviate this mismatch to some extent and improve performance, we find that they either (1) fail to fully resolve the mismatch while degrading the feature extractor, or (2) address the mismatch only post-training, allowing it to persist during training. This increases inter-client gradient divergence, hinders model aggregation, and ultimately leaves the feature extractor suboptimal for client data. To address this issue, we propose FedPFT, a novel framework that resolves the mismatch during training using personalized prompts. These prompts, along with local features, are processed by a shared self-attention-based transformation module, ensuring alignment with the global classifier. Additionally, this prompt-driven approach offers strong flexibility, enabling task-specific prompts to incorporate additional training objectives (e.g., contrastive learning) to further enhance the feature extractor. Extensive experiments show that FedPFT outperforms state-of-the-art methods by up to 5.07%, with further gains of up to 7.08% when collaborative contrastive learning is incorporated.
Xinghao Wu, Xuefeng Liu, Jianwei Niu +4