Adaptive Split Federated Learning: Efficient Training and Robust against Distribution Shifts at Test Time

Collaborative distributed training approaches, such as Federated Learning (FL) and Split Learning (SL), are hindered by substantial communication and computation overheads, as well as limited robustness to test-time distribution shifts, posing significant challenges for deployment on resource-constrained edge devices in real-world settings. Personalized models finetuned on client-specific data often exhibit poor generalization to out-of-distribution (OOD) test samples, while global models typically underperform personalized models on in-distribution (ID) data. This highlights the need for a balanced integration of global and personalized models during inference. To address these challenges, this study proposes an efficient adaptive split federated learning framework that uses gradient-norm-based evaluation at inference time to effectively balance global generalization with local personalization. Additionally, this study integrates Supervised Contrastive and Cross-Entropy loss functions to enhance feature representations, improving robustness to test-time distribution shifts. The proposed split architecture enables efficient ondevice computation, while activation caching significantly reduces communication overhead. Extensive experiments on real-world medical imaging data, ISIC2019, and heterogeneously sampled benchmarks, including CIFAR10, FMNIST, VLCS, PACS, and OfficeHome, demonstrate that the approach achieves a 96× reduction in computation cost and a 4.5× reduction in communication overhead. Moreover, the adaptive inference strategy yields an average accuracy improvement of 29% over strong baselines in OOD test-time settings, with consistent generalization gains observed across all personalized clients.