HARMONY: Bridging the Personalization-Generalization Gap by Mitigating Representation Skew in Heterogeneous Split Federated Learning

TL;DR

HARMONY is a novel hybrid split federated learning framework that aligns heterogeneous client representations through contrastive learning, significantly improving accuracy for personalized and out-of-distribution classes without sacrificing privacy.

Contribution

HARMONY introduces a meta-learning based approach to support heterogeneous client architectures and employs contrastive learning to mitigate representation skew in federated settings.

Findings

HARMONY improves test accuracy by up to 43.0% for in-distribution classes.

HARMONY enhances out-of-distribution prediction accuracy by up to 28.3%.

The framework maintains acceptable latency across multiple datasets and model types.

Abstract

Mobile devices face diverse resource constraints and non-IID data class distributions, requiring fast on-device inference for local in-distribution (ID) classes and on-demand remote support for client-specific out-of-distribution (OOD) classes. Hybrid split federated learning (Hybrid SFL) couples personalized client-side front ends (supporting early exit) with a generalized server-side backend for fallback inference, balancing accuracy and cost. However, under client architectural heterogeneity, the existing hybrid SFL suffers from representation skew, where features from customized extractors fail to align in the shared space, leading to a sharp degradation in the server model responsible for OOD prediction. We propose HARMONY, the first hybrid SFL framework to support heterogeneous client architectures. HARMONY modifies meta-learning to simulate diverse extractors across parameters and architectures, and to learn to personalize. To mitigate representation skew, HARMONY conducts server-side contrastive learning to align extracted features, neither sacrificing clients' personalization nor sharing raw labels. Compared to the state of the art across multiple datasets and model families, HARMONY improves test accuracy by up to 43.0%/28.3% without/with OOD, respectively, while maintaining acceptable latency.