Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Achieving Linear Speedup for Composite Federated Learning

Kun Huang, Shi Pu

TL;DR

This work addresses federated learning with composite objectives of the form $\psi(x)=f(x)+\varphi(x)$, where $\varphi$ is nonsmooth and $\rho$-weakly convex. It introduces FedNMap, a normal-map-based update with a local correction to counteract drift from multiple local steps, achieving linear speedup in both the number of clients $n$ and local updates $Q$ for nonconvex losses, and improving rates under the proximal-PL condition. The analysis relies on a multistep Lyapunov function to control consensus, local-update, and multistep errors, under standard assumptions such as $L$-smoothness of $f_i$, unbiased gradient estimates with variance $\sigma^2$, and weak convexity of $\varphi$, without requiring heterogeneity bounds. Empirical results on MNIST and SVHN show FedNMap attaining lower stationarity, stable convergence, and better training loss/test accuracy compared to baselines, while reducing uplink communication. Overall, the paper delivers the first linear-speedup results for nonconvex composite FL under mild assumptions and demonstrates practical communication-efficiency advantages.

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-Łojasiewicz (PL) condition. To our knowledge, this is the first result establishing linear speedup for nonconvex composite federated learning.

Achieving Linear Speedup for Composite Federated Learning

TL;DR

This work addresses federated learning with composite objectives of the form , where is nonsmooth and -weakly convex. It introduces FedNMap, a normal-map-based update with a local correction to counteract drift from multiple local steps, achieving linear speedup in both the number of clients and local updates for nonconvex losses, and improving rates under the proximal-PL condition. The analysis relies on a multistep Lyapunov function to control consensus, local-update, and multistep errors, under standard assumptions such as -smoothness of , unbiased gradient estimates with variance , and weak convexity of , without requiring heterogeneity bounds. Empirical results on MNIST and SVHN show FedNMap attaining lower stationarity, stable convergence, and better training loss/test accuracy compared to baselines, while reducing uplink communication. Overall, the paper delivers the first linear-speedup results for nonconvex composite FL under mild assumptions and demonstrates practical communication-efficiency advantages.

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 and the number of local updates for nonconvex losses, both with and without the Polyak-Łojasiewicz (PL) condition. To our knowledge, this is the first result establishing linear speedup for nonconvex composite federated learning.
Paper Structure (22 sections, 7 theorems, 129 equations, 3 figures, 1 table, 1 algorithm)

This paper contains 22 sections, 7 theorems, 129 equations, 3 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Main Contribution
  4. Notation and Assumptions
  5. FedNMap: A Federated Learning Method with Normal Map-Based Update
  6. Connection with Existing Algorithms
  7. Convergence Results
  8. Proof Sketch
  9. Numerical Experiments
  10. A Simple Neural Network
  11. Deep Learning for Image Classification
  12. Conclusion
  13. Derivations for Related Methods
  14. The Method in zhang2024composite
  15. SCAFFOLD
  16. ...and 7 more sections

Key Result

Theorem 3.1

Let Assumptions as:abc, as:smooth, and as:phi hold. Denote $\Delta_\psi:= \psi(x_0) - \psi^*$. Set $\gamma\leq 1/[5(\rho + L)]$ and Then, the iterates generated by FedNMap satisfy In particular, if we set $\gamma = 1/[5(\rho + L)]$ and then

Figures (3)

  • Figure 1: An illustration of the subsequence $\{\boldsymbol{t}_j\}_{j=0}^R$.
  • Figure 2: Comparison among FedNMap, Zhang zhang2024composite, and FedCanon zhou2025fedcanon for training a one-hidden-layer neural network with an elastic net regularizer on the MNIST dataset. The stepsizes are set to $\eta_a = 0.5/Q$ for local updates and $\eta_s = 1$ for the outer loop across all methods. The parameter $\gamma$ in FedNMap is set to $1$.
  • Figure 3: Comparison of FedNMap, Zhang zhang2024composite, and FedCanon zhou2025fedcanon for training VGG-16 with an elastic net regularizer on the SVHN dataset. The stepsizes are set to $\eta_a = 1 / Q$ and $\eta_s=0.5$. The number of local updates is $Q\in\{10,20\}$, the number of clients is $n\in\{20,50\}$, and the parameter $\gamma$ in FedNMap is set to $2$.

Theorems & Definitions (9)

  • Theorem 3.1
  • Remark 3.1
  • Theorem 3.2
  • Remark 3.2
  • Lemma B.1
  • Lemma B.2
  • Lemma B.3
  • Lemma B.4
  • Lemma B.5