SMoFi: Step-wise Momentum Fusion for Split Federated Learning on Heterogeneous Data

TL;DR

SMoFi is a novel framework for split federated learning that enhances convergence and accuracy across heterogeneous data by synchronizing momentum buffers and controlling gradient divergence.

Contribution

It introduces Step-wise Momentum Fusion (SMoFi), a lightweight method to mitigate gradient divergence in split federated learning with data heterogeneity.

Findings

Improves global model accuracy by up to 7.1%.

Speeds up convergence by up to 10.25 times.

More effective with increased clients and deeper models.

Abstract

Split Federated Learning is a system-efficient federated learning paradigm that leverages the rich computing resources at a central server to train model partitions. Data heterogeneity across silos, however, presents a major challenge undermining the convergence speed and accuracy of the global model. This paper introduces Step-wise Momentum Fusion (SMoFi), an effective and lightweight framework that counteracts gradient divergence arising from data heterogeneity by synchronizing the momentum buffers across server-side optimizers. To control gradient divergence over the training process, we design a staleness-aware alignment mechanism that imposes constraints on gradient updates of the server-side submodel at each optimization step. Extensive validations on multiple real-world datasets show that SMoFi consistently improves global model accuracy (up to 7.1%) and convergence speed (up to 10.25$\times$). Furthermore, SMoFi has a greater impact with more clients involved and deeper learning models, making it particularly suitable for model training in resource-constrained contexts.