Achieving Linear Speedup for Composite Federated Learning

2026-02-03

AI Summary

This work addresses the convergence bottlenecks in composite federated learning caused by nonsmooth regularizers and data heterogeneity by proposing FedNMap. The method leverages the normal map to handle nonsmooth regularization and introduces a local correction strategy to mitigate the adverse effects of heterogeneity. Under the assumptions of weakly convex regularization and bounded stochastic gradient variance, FedNMap achieves, for the first time, linear speedup with respect to both the number of clients and the number of local updates in nonconvex composite federated learning, thereby overcoming existing theoretical limitations. Notably, this linear speedup is guaranteed regardless of whether the Polyak–Łojasiewicz condition holds.

Technology Category

• Search and Optimization: Non-convex Optimization
• Machine Learning: Distributed Machine Learning & Federated Learning
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

This paper proposes FedNMap, a normal map-based method for composite federated learning, where the objective consists of a smooth loss and a possibly nonsmooth regularizer. FedNMap leverages a normal map-based update scheme to handle the nonsmooth term and incorporates a local correction strategy to mitigate the impact of data heterogeneity across clients. Under standard assumptions, including smooth local losses, weak convexity of the regularizer, and bounded stochastic gradient variance, FedNMap achieves linear speedup with respect to both the number of clients $n$ and the number of local updates $Q$ for nonconvex losses, both with and without the Polyak-{\L}ojasiewicz (PL) condition. To our knowledge, this is the first result establishing linear speedup for nonconvex composite federated learning.

Problem

• composite federated learning
• linear speedup
• nonsmooth regularizer
• data heterogeneity
• nonconvex optimization

Innovation

• composite federated learning
• normal map
• linear speedup
• nonconvex optimization
• data heterogeneity