SplitGP: Achieving Both Generalization and Personalization in Federated Learning

TL;DR

SplitGP is a split learning approach for federated learning that balances personalization and generalization, enabling efficient inference on resource-constrained edge devices by splitting the model into client and server components.

Contribution

It introduces a novel split learning framework that separately optimizes for personalization on the client side and generalization on the server side, with theoretical convergence analysis and practical bounds.

Findings

Outperforms baselines in inference time and accuracy

Balances personalization and generalization effectively

Provides theoretical analysis of convergence and split ratios

Abstract

A fundamental challenge to providing edge-AI services is the need for a machine learning (ML) model that achieves personalization (i.e., to individual clients) and generalization (i.e., to unseen data) properties concurrently. Existing techniques in federated learning (FL) have encountered a steep tradeoff between these objectives and impose large computational requirements on edge devices during training and inference. In this paper, we propose SplitGP, a new split learning solution that can simultaneously capture generalization and personalization capabilities for efficient inference across resource-constrained clients (e.g., mobile/IoT devices). Our key idea is to split the full ML model into client-side and server-side components, and impose different roles to them: the client-side model is trained to have strong personalization capability optimized to each client's main task, while the server-side model is trained to have strong generalization capability for handling all clients' out-of-distribution tasks. We analytically characterize the convergence behavior of SplitGP, revealing that all client models approach stationary points asymptotically. Further, we analyze the inference time in SplitGP and provide bounds for determining model split ratios. Experimental results show that SplitGP outperforms existing baselines by wide margins in inference time and test accuracy for varying amounts of out-of-distribution samples.