A deep cut into Split Federated Self-supervised Learning

TL;DR

The paper investigates how the depth at which the neural network is split between clients and a central server affects privacy, communication overhead, and learning performance in federated self-supervised learning. It analyzes the shortcomings of the MoCo-based Split Federated Learning (MocoSFL) approach, particularly its sensitivity to deep splits due to misalignment between online and momentum models during synchronization. The authors propose MonAcoSFL, a momentum-aligned variant that synchronizes both online and momentum client models, proving that this alignment preserves the contrastive objective and yields substantial accuracy gains under deeper, more communication-efficient splits. Empirically, MonAcoSFL achieves state-of-the-art performance across ResNet-18 and MobileNetV2 backbones on CIFAR-10/100 with non-IID data, while also enhancing privacy protections and reducing communication overhead, making federated SSL more practical for real-world deployments.

Abstract

Collaborative self-supervised learning has recently become feasible in highly distributed environments by dividing the network layers between client devices and a central server. However, state-of-the-art methods, such as MocoSFL, are optimized for network division at the initial layers, which decreases the protection of the client data and increases communication overhead. In this paper, we demonstrate that splitting depth is crucial for maintaining privacy and communication efficiency in distributed training. We also show that MocoSFL suffers from a catastrophic quality deterioration for the minimal communication overhead. As a remedy, we introduce Momentum-Aligned contrastive Split Federated Learning (MonAcoSFL), which aligns online and momentum client models during training procedure. Consequently, we achieve state-of-the-art accuracy while significantly reducing the communication overhead, making MonAcoSFL more practical in real-world scenarios.

Figures (9)

• Figure 1: Communication overhead and accuracy for MocoSFL and MonAcoSFL depending on the splitting depth. In this example, the minimal communication overhead is obtained for dividing the model into 11 and 7 layers on the client and server sides, respectively. For such an optimal case, the accuracy of MocoSFL drops significantly, in contrast to the accuracy of MonAcoSFL. Notice that computational overhead consists of forward and backward propagation (marked as blue) and parameter synchronization (marked as orange).
• Figure 2: Activations obtained from successive layers of untrained ResNet-18 for a sample ImageNet image. One can observe that representations of low layers highly resemble the respective input data, increasing the privacy risks associated with broadcasting network activations.
• Figure 3: Communication overhead of a single client device for one training epoch of ResNet-18 he2015deep and MobileNet-V2 sandler2018mobilenetv2 for different splitting depths. The 11-th and 7-th layers are the most communication-efficient for ResNet-18 and MobileNetV2, respectively. Note that the training epoch corresponds to $250$ images of resolution $224 \times 224$ are processed, and $10$ synchronizations of parameters. Moreover, the blue bars correspond to communication in the optimization phase, and the orange bars correspond to parameter synchronization.
• Figure 4: Accuracy of MocoSFL drops significantly with increased splitting depth regardless of the number of clients. Here, presented for CIFAR-10 krizhevsky2009learning.
• Figure 5: Visualization of parameters changing in MocoSFL (left) and MonAcoSFL (right) for two clients. The solid and dashed lines represent the progression of online and momentum parameters, respectively. The dotted lines symbolize the synchronization of parameters. The difference between MocoSFL and MonAcoSFL lies in the synchronization procedure that, in the case of MonAcoSFL, ensures that both online and the momentum models remain aligned, preserving their ability to optimize the contrastive objective.