FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation

2025-05-29

AI Summary

Federated learning (FL) requires clients to train full models, making it infeasible for large-scale models; split learning (SL) alleviates client memory constraints but suffers from high communication latency due to sequential training. Existing parallelization approaches are accuracy-limited by the absence of server-side gradient feedback. This paper proposes a novel federated split learning framework. It introduces the first gradient estimation mechanism based on smashed activations, employing a lightweight auxiliary model to dynamically approximate server gradient behavior—enabling efficient local updates without explicit gradient backpropagation. The framework integrates local-loss-driven parallel training, knowledge distillation, and adaptive model synchronization. We theoretically establish an $O(1/sqrt{T})$ convergence rate, matching that of FedAvg. Experiments demonstrate substantial reductions in both communication overhead and client memory usage, achieving optimal trade-offs between accuracy and efficiency.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Collaborative training methods like Federated Learning (FL) and Split Learning (SL) enable distributed machine learning without sharing raw data. However, FL assumes clients can train entire models, which is infeasible for large-scale models. In contrast, while SL alleviates the client memory constraint in FL by offloading most training to the server, it increases network latency due to its sequential nature. Other methods address the conundrum by using local loss functions for parallel client-side training to improve efficiency, but they lack server feedback and potentially suffer poor accuracy. We propose FSL-SAGE (Federated Split Learning via Smashed Activation Gradient Estimation), a new federated split learning algorithm that estimates server-side gradient feedback via auxiliary models. These auxiliary models periodically adapt to emulate server behavior on local datasets. We show that FSL-SAGE achieves a convergence rate of $mathcal{O}(1/sqrt{T})$, where $T$ is the number of communication rounds. This result matches FedAvg, while significantly reducing communication costs and client memory requirements. Our empirical results also verify that it outperforms existing state-of-the-art FSL methods, offering both communication efficiency and accuracy.

Problem

• Reducing network latency in split learning
• Improving accuracy with server feedback
• Lowering client memory requirements

Innovation

• Estimates server gradients via auxiliary models
• Reduces communication costs significantly
• Maintains accuracy while lowering memory needs