Federated Learning over Wireless Networks: Convergence Analysis and Resource Allocation

In plain words

Federated learning trains one shared model while the training data stays on each user's device, which sends back model updates instead of data. This paper asks how to make that process converge when the devices hold very different data. Each device solves a slightly adjusted local problem rather than taking a few plain steps. The adjustment pulls the device toward the direction the whole group is moving, and the server averages the results. A device does not need to solve its problem exactly: one setting controls how precise each local solve is, and it trades extra on-device work against fewer rounds with the server.

The devices are phones and similar hardware on radio links with limited battery. How long a training round takes, and how much energy the devices use, depends on how fast each device computes and how much power it spends sending. The authors combine these costs into one problem, split it into three smaller pieces and solve each piece with a formula or a small numerical search. The method is tested on a few standard datasets and on simulated wireless setups.

TL;DR

The method, called FEDL here, changes the usual averaging scheme. Each user equipment (UE) minimises a local surrogate that adds a global gradient estimate weighted by a hyper-learning rate η, and it solves this only to a local accuracy θ §IV alg. 1. Under smoothness and strong convexity, the authors prove a linear convergence rate Θ that depends on η, θ and the condition number ρ §IV eq. 10 eq. 11. They also model the wall-clock time and UE energy of a time-shared wireless uplink, and split the resulting non-convex problem into three sub-problems §V. Author-reported PyTorch experiments show lower training loss and higher test accuracy than FedAvg on MNIST, FEMNIST and synthetic data §VI. Released code covers the learning experiments only; the wireless part has no code.

Problem

Federated learning fits one model to data spread over many UEs whose data are not identically distributed §I. FedAvg works well in such settings but, as the authors note, had no convergence analysis for heterogeneous data §I. Other analyses need an extra bound on gradient divergence or primal-dual arguments §II.

In wireless deployments the time of one round depends on UE heterogeneity in data size, CPU, transmit power and channel gain §I §V-A. The authors aim for two things: convergence under strong convexity and smoothness alone §III [assumption 1], and a choice of local accuracy, step size, CPU frequency and transmit power that trades training time against UE energy §V-B.

Assumption 1 requires each local loss F_n to be L-smooth and β-strongly convex, with condition number ρ = L/β §III [assumption 1]. Experiments use a linear regression loss with squared error §VI and a multinomial logistic loss with cross-entropy plus an L2 term (β/2) Σ_c ||w_c||² §VI.

Method

The global objective is the data-weighted sum of the local losses, with weight p_n = D_n / D §III eq. 1. The server runs K_g global rounds §IV alg. 1.

1. Each UE n receives the global model w^(t-1) and the global gradient estimate §IV alg. 1.
2. UE n approximately minimises its surrogate J_n^t eq. 2. The solver's accuracy must satisfy the test in eq. 3. The gradient of J_n^t is the local gradient at the current point, plus η times the global estimate, minus the local gradient at w^(t-1) §IV.
3. Each UE uploads its new model and its full local gradient at that model §IV alg. 1.
4. The server averages the models and the gradients with weights p_n eq. 4 eq. 5, then broadcasts them §IV.

Lemma 1 gives the number of local rounds K_l that meets the accuracy test, when the local solver is gradient descent with a linear rate §IV [lemma 1]. The local update uses the plain gradient step eq. 7.

# FEDL global rounds  [alg. 1] [§IV]
# Input: w0, theta in [0, 1], eta > 0; output: global model w

def fedl(w0, theta, eta, h, K_g, K_l, UEs):
    # Model after round 0  [alg. 1]
    w = w0
    # Global gradient estimate; its initial value is not stated  [§IV]
    grad_bar = None
    for t in range(1, K_g + 1):
        # Each UE solves its surrogate approximately  [alg. 1] [eq. 2]
        local_models = [
            local_solve(ue, w, grad_bar, eta, h, K_l) for ue in UEs
        ]
        # Each UE sends its full local gradient at the new model  [alg. 1]
        local_grads = [grad_F(ue, v) for ue, v in zip(UEs, local_models)]
        # Average with weights p_n = D_n / D over all UEs  [eq. 4] [eq. 5]
        weights = [ue.num_samples / total_samples(UEs) for ue in UEs]
        w = weighted_sum(local_models, weights)
        grad_bar = weighted_sum(local_grads, weights)
    return w

# Approximate minimiser of J_n^t  [eq. 2] [eq. 7] [§IV]
# J(v) = F_n(v) + <eta * grad_bar - grad F_n(w), v>

def local_solve(ue, w, grad_bar, eta, h, K_l):
    # Local gradient at the received model, computed once  [§IV]
    anchor_grad = grad_F(ue, w)
    # Start from the global model  [alg. 1]
    v = w
    # Run K_l local rounds; Lemma 1 says this meets the test in [eq. 3]
    for k in range(K_l):
        # Gradient of the surrogate at the current point  [§IV]
        surrogate_grad = grad_F(ue, v) + eta * grad_bar - anchor_grad
        # Plain gradient step with learning rate h  [eq. 7]
        v = v - h * surrogate_grad
    return v

Sub-problem 1 (CPU frequency) gives each UE one of three regimes: maximum frequency, minimum frequency, or an interior value set by a common deadline [lemma 2] alg. 2 eq. 26. Sub-problem 2 (time shares) is a clipped power-control rule [lemma 3]:

# Time share of UE n on the uplink  [lemma 3] [§V-C]
# g(kappa) = (s / B) / (1 + W((kappa * h / N0 - 1) / e)), W = Lambert W
# h is the average channel gain of the UE  [§V-A]

def time_share(s, B, h, N0, kappa, p_min, p_max):
    # Shortest and longest time shares, from the power limits  [§V-C]
    tau_max = s / (B * ln(h * p_min / N0 + 1))
    tau_min = s / (B * ln(h * p_max / N0 + 1))
    # g is strictly decreasing, so the three cases of lemma 3 are a clip
    g = (s / B) / (1 + lambert_w((kappa * h / N0 - 1) / e))
    return min(max(g, tau_min), tau_max)

Sub-problem 3 chooses θ and η numerically, with the stated objective divided by Θ §V-C. Sub-problems 1 and 2 do not depend on θ or η, so the authors solve them first §V-C. The combined point is a stationary point, not a proven global optimum §V-C [theorem 2].

Key equations:

J_n^t(w) = F_n(w) + \langle \eta \bar\nabla F^t - \nabla F_n(w^{t-1}),\, w \rangle eq. 2

\lVert \nabla J_n^t(w_n^t) \rVert \le \theta \, \lVert \nabla J_n^t(w^{t-1}) \rVert eq. 3

K_l = \frac{2}{\gamma} \log \frac{C}{\theta}, \qquad C = c\rho eq. 9 [lemma 1]

\Theta = \frac{\eta\bigl(2(\theta-1)^2 - (\theta+1)\theta(3\eta+2)\rho^2 - (\theta+1)\eta\rho^2\bigr)}{2\rho\bigl((1+\theta)^2\eta^2\rho^2 + 1\bigr)} eq. 11

F(w^t) - F(w^*) \le (1-\Theta)^t \bigl(F(w^0) - F(w^*)\bigr) eq. 10

K_g = \frac{1}{\Theta} \log \frac{F(w^0) - F(w^*)}{\varepsilon} eq. 12 [corollary 1]

Wireless model: CPU energy per local round is (α_n/2) c_n D_n f_n² eq. 13; the rate is B ln(1 + h̄_n p_n / N_0) eq. 14; the time share satisfies r_n = s_n / τ_n eq. 15; the uplink energy is τ_n p_n(s_n/τ_n), where p_n(r) = (N_0 / h̄_n)(e^(r/B) − 1) eq. 16 eq. 17. One global round takes T_comm + K_l T_comp seconds §V-A.

Hyperparameters

Table of symbols (no values): tab. 3.

Results

Author-reported. Setup: N = 100 UEs unless stated; MNIST and FEMNIST with multinomial logistic regression; synthetic data with linear regression; a one-hidden-layer DNN for the non-convex check; FedAvg as the baseline; hardware not stated §VI.

Table I: sub-problem 3 solutions (numerical) for five UEs §V-C. Each cell lists ρ = 1.4 / 2 / 5 in that order.

Table II: complexity of the three sub-problems, N = number of UEs: sub-problem 1 is O(N²), sub-problem 2 is O(1), sub-problem 3 is O(N) tab. 2.

Fig. 3 reports that a small enough η gives convergence for all ρ tested, a larger η converges faster, and the gap to the exact solution is negligible §VI fig. 3. Fig. 5 reports that a larger K_l speeds convergence in time but costs more local computation §VI fig. 5. Fig. 7 and Fig. 8 report total cost and the Pareto curve against L_comp and L_comm §VII.

Paper vs code

Code: CharlieDinh/FEDL_Pytorch at code_commit. The bibliography entry that the paper cites for it lists the first author. Its LICENSE is GPL-3.0, a copyleft license, so this card cites code only by path and line, in our own words. The README names a TensorFlow version at CharlieDinh/FEDL, which we did not use. The repository covers the learning experiments (Section VI). It has no code for the wireless sub-problems (Section V), so Figs. 1, 2, 7 and 8 and Table I cannot be reproduced from it. The paper says sub-problem solutions were checked with IPOPT §V-C, and that solver is not in the repository.

Discrepancies found:

• First round has no global gradient. Algorithm 1 takes only w⁰, θ and η alg. 1. In code, the server model has no gradient before the first aggregation, so the first local step drops both correction terms and is plain gradient descent on F_n userbase.py fedoptimizer.py. The correction applies from round 2 on serverbase.py. Inferred from code; not run.
• Default η is 0. The paper needs η > 0 alg. 1. The default in main.py is 0 main.py. The step then subtracts the anchor gradient but drops the global term fedoptimizer.py. The README's commands set η explicitly README.md.
• θ is never used. No code checks the accuracy test eq. 3 or computes K_l from Lemma 1. K_l is the local_epochs flag main.py userfedl.py.
• Local round means different things. The paper's local round is one gradient step eq. 7. The code runs local_epochs passes over the local data in mini-batches of batch_size userfedl.py. Batch size 0 means full batch, and then the two agree userbase.py.
• Mini-batch gradient. Each step uses a mini-batch gradient for the local term and the full anchor gradient userfedl.py. The convergence analysis covers full-batch gradient descent only §IV remark 5.
• Partial participation weights. The paper's weights are D_n/D over all N UEs eq. 4 eq. 5. The code renormalises over the selected UEs serverbase.py. Selection is uniform without replacement and seeded per round serverbase.py. The two agree only for S = N.
• Regulariser missing. The paper's logistic loss includes the L2 term §VI. The code uses plain NLL userfedl.py. The L argument is stored and never used in the loss or the step main.py fedoptimizer.py.
• Bias terms. The paper's models have no intercept §III. Both the logistic and the linear layers include a bias models.py models.py.
• ρ does not drive anything. The paper's η and θ choices depend on ρ §IV remarks 3 and 4. In code, ρ only appears in result file names serverbase.py serverbase.py.
• Learning rates for FedAvg. The paper says a grid search tuned h for both methods §VI. The README uses the same MNIST learning rate, 0.003, for FedAvg and FEDL README.md. The grid results are not in the paper.
• Synthetic data generator. The paper says σ_i ~ U(1, 10) per UE and covariance Σ_ii = i^(−p) with p = log ρ / log d §VI. The generator draws one σ ~ U(0.1, 10) for all UEs generate_linear_regession_updated.py. It uses exponent p/2 generate_linear_regession_updated.py. It also adds a random mean shift and normalises the inputs by a spectral-norm factor generate_linear_regession_updated.py generate_linear_regession_updated.py. The paper does not state the labelling model; the code uses a random weight vector and a noise term generate_linear_regession_updated.py.
• MNIST partition matches. Three labels per UE §VI matches the script generate_niid_mnist_100users.py.
• Matches. The gradient step fedoptimizer.py and the averaged full-gradient upload userfedl.py follow the paper's rules. The Fig. 4 settings (K_g = 800, S = 10, K_l = 20) match the README README.md.
• Four ρ values? Fig. 3's text says four ρ values; Table I and the README use 1.4, 2 and 5 §VI README.md. The fourth is not stated.
• Framework. The abstract says PyTorch; the conclusion says TensorFlow §VIII. The code is PyTorch.

Pitfalls

• Lemma 1 constant. The lemma states K_l = (2/γ) log(C/θ) with C = cρ eq. 9. The proof ends with cρ e^(−kγ) ≤ θ² app. B, which gives k ≥ (1/γ) log(cρ/θ²) = (2/γ) log(√(cρ)/θ). The stated K_l is larger, so it is still sufficient when cρ ≥ 1, a condition not stated. The constants c and γ are not given values eq. 8.
• Θ formula. Remark 1 writes the numerator of eq. (11) as 2η(A − B), where A and B are defined there. Eq. (11) as printed has η(A − B) over 2ρ(·). By hand, ρ = 2 and κ = 0.1 with θ = 0.015 and η = 0.177 give Θ ≈ 0.042 from eq. (11), which matches Table I. The remark's version gives about 0.084. Use eq. (11). The sign conditions in the remark are unaffected.
• Θ ∈ (0, 1) is sufficient only. The authors say it is not necessary §IV remark 1. A large η can push Θ out of (0, 1) and diverge §IV remark 3. Small θ means many local rounds and fewer global ones §IV remark 2.
• K_g is a conservative count. Eq. (12) uses 1/Θ. Since −log(1−Θ) ≥ Θ, the exact count from eq. (10) is at most this. The proof is only "similar to Lemma 1" §IV.
• Theory needs full gradients. Remark 5 limits the analysis to full-batch gradient descent, full participation and strongly convex, smooth losses §IV. Mini-batch results are experimental only, and the README warns that FEDL is unstable with mini-batch 20 on the DNN README.md §VI.
• Local round cost differs. Section V counts one local round as c_n D_n cycles, which is one full pass over the data §V-A. Section IV's local round is one gradient step eq. 7. The two agree only for full-batch gradient descent.
• Energy sum. The text writes E_iter = Σ_n E_comm + K_l E_comp §V-A. Sub-problem 1 sums E_comp over all UEs §V-C. Sum E_comp over UEs as well.
• Hidden communication cost. Each UE uploads a gradient as well as its model §IV alg. 1. Section V folds both into s_n §V-A. Budget the uplink for both.
• Units. Fig. 1 gives D_n in MB, but the CPU model counts cycles per bit, and s_n is in nats. The bit conversion is not stated. The caption's 25,000 nats ≈ 4.5 KB is consistent with 1 nat = 1/ln 2 bits.
• Noise notation. The Fig. 1 caption uses σ for the noise power, the text uses N_0 fig. 1 §V-A. Treat them as the same value; not stated.
• Channel gain. The Fig. 1 caption draws gains from an exponential distribution with mean g₀(d₀/d)⁴. The model uses one average gain h̄_n per UE §V-A. Whether the numerical results sample gains is not stated.
• κ values. The κ values for Figs. 7 and 8 are not stated. The Fig. 1 thresholds (0.004, 0.1, 1) apply to that setup only §V-C.
• Lambert W branch. The branch is not stated. The argument is at least −1/e, so two real branches exist for some arguments. Check the sign of the resulting time shares.
• Algorithm 2 needs checking. It updates T_N3 inside one pass over the sorted UEs alg. 2. Validate the three sets against the four κ regions of Corollary 2 §V-C [corollary 2].
• Only stationary points. The combined solution is a stationary point of a non-convex problem §V-C [theorem 2], not a global optimum. The sub-problem 3 solver is not named.
• Learning-rate fairness. The paper claims a tuned FedAvg §VI, but the grid results are missing, and the README uses one learning rate for both methods on MNIST README.md.
• Small gains at batch 40. The training-loss change is −0.2% at batch 40 §VI fig. 4. The gain over FedAvg is not uniform.
• Choosing η and θ in practice. Their rules depend on L and β, and so on ρ §IV. The paper does not say how to estimate them. Not stated.

Lineage

• Builds on: FedAvg §II §VI; the surrogate is motivated by DANE, whose exact global gradient, proximal term and exact local solve FEDL drops §IV; the approximation-factor idea follows CoCoA-style analysis §IV.
• Compared against: FedAvg in all experiments §VI.
• Related analyses: primal-dual convergence §II; gradient-divergence bounds §II.
• Superseded by: not stated.