To Talk or to Work: Flexible Communication Compression for Energy Efficient Federated Learning over Heterogeneous Mobile Edge Devices

In plain words

Phones and other edge devices can train a shared model without sending their private data anywhere: each device trains on its own data and only sends model updates to a server. The catch is battery. Training costs computing energy, and every upload costs radio energy. This paper asks how to balance the two, which it calls "working" (computing) and "talking" (transmitting).

Two knobs control the balance. A device can compute several training steps before it talks to the server, which means fewer uploads. It can also send only the largest parts of its update, which makes each upload smaller. Sending less saves radio energy but can slow learning, and too few uploads can also hurt. Devices differ in their radio conditions, so one fixed setting is wasteful.

The proposed scheme lets each device choose its own compression level, and a control procedure picks these levels to minimise the total energy of all devices. Devices with poor radio links compress more. The authors report, from simulations, that this saves a lot of energy compared with simpler schemes.

TL;DR

• Federated learning with top-k gradient sparsification, local steps, error compensation and growing mini-batches, where each device has its own sparsity level §III-A.
• The paper proves a non-convex convergence bound for this algorithm and turns it into a communication-round count §III-B, Theorem 1, Corollary 1.
• A control algorithm chooses per-device sparsity and the sync interval H to minimise total computing and communication energy §IV-B, alg. 2.
• Author-reported and simulated only: 1.5 to 100 times less energy than the compared schemes to reach a target accuracy on ResNet20/CIFAR-10 and LeNet5/MNIST §V fig. 2.

Problem

Federated learning (FL) on mobile devices is limited by energy, both for local training and for sending gradients over wireless links §I. Compressed distributed SGD methods usually require the same compressor on all participants, which ignores differences in radio capacity §I §II. Reducing the number of communication rounds (local SGD, larger batches) and reducing the bits per round (sparsification, quantisation) are known, but the paper argues that they interact: compressing hard and computing many local steps can slow convergence enough to cost more energy in total §I. Heavier sparsity raises distortion and can need more communication rounds; aggressively enlarging H can also impair learning efficiency §I §III-B.

The analysis assumes the standard setting of synchronous FL with a finite sum of local losses, each device holding its own data §III-A. Two assumptions are needed: each local loss is L-smooth, and stochastic gradients have bounded variance and bounded second moment §III-B, Assumptions 1 and 2.

Method

Algorithm (FT-LSGD-DB). Each of M devices holds a local dataset. It keeps a local model copy and an error memory. In every iteration it grows its mini-batch geometrically and takes one local SGD step. Every H iterations it adds its error memory to the accumulated local change, keeps only the k_m largest-magnitude entries (its own k_m), uploads them, and stores the rest as the new error memory. The server averages the uploads, applies the average to the global model, and broadcasts it §III-A alg. 1.

Top-k sparsification keeps every entry whose absolute value is at least the k-th largest absolute value, and sets the others to zero §III-A. The paper sets the sparsity of device m to delta_m = d / k_m, where d is the model dimension, and requires delta_m >= 1 §III-A §III-B.

# FT-LSGD-DB, one run  [alg. 1]  [§III-A]
# Inputs: local data, b0, rho > 1, per-device k, sync interval H, T steps
# Output: global model w

def train(datasets, b0, rho, T, H, k):
    # Start: every device holds the global model and an error memory
    # (the paper writes e = w here; see Pitfalls)  [alg. 1, Initialization]
    w = initial_model
    w_hat = [w for each device]
    e = [w for each device]
    for t in range(T):
        # Mini-batch size grows geometrically  [alg. 1, line 4]
        b = floor(rho ** t * b0)
        for each device m in parallel:
            # Sample b examples uniformly from device m  [alg. 1, line 5]
            batch = sample(datasets[m], b)
            # One local SGD step  [alg. 1, line 6]
            w_half[m] = w_hat[m] - eta * grad_f(m, w, batch)
            if (t + 1) % H == 0:
                # Add error memory to the accumulated change  [alg. 1, line 8]
                u[m] = e[m] + w - w_half[m]
                # Keep the k_m largest entries  [§III-A, Top_k]
                g[m] = top_k(u[m], k[m])
                # Store the unsent rest as new error memory  [alg. 1, line 10]
                e[m] = u[m] - g[m]
                # Take the broadcast global model  [alg. 1, line 11]
                w_hat[m] = w_next[m]
            else:
                # No communication: keep local iterate  [alg. 1, lines 14-16]
                w_hat[m] = w_half[m]
        if (t + 1) % H == 0:
            # Server averages sparse uploads and applies them  [alg. 1]
            w_next = w - mean(g)
            w = w_next
    return w

Energy and bit model. A device needs FPP bits per kept value plus one sign bit, and log2 of the binomial coefficient binom(d, k_m) bits to encode positions, with s1 and s0 as overhead coefficients §IV-A, eqs. for S_val, S_pos, S_m. For large delta the paper approximates the per-round bits as

# Approximate bits sent per round by device m  [§IV-B, Corollary 2]
# Derived for delta_m >> 1 (Stirling approximation)  [§IV-B]

def bits_per_round(d, delta, FPP, s0, s1):
    # Sign bit plus value bits for each kept entry  [§IV-A]
    kappa = FPP + 1
    # Approximate bits, valid for delta >= e^(3/2)  [§IV-B]
    return (s1 * d / delta) * (log2(delta) + kappa) + s0

The paper states that devices with poor channels should sparsify more severely, and that devices with good channels can compress less §IV-A.

Radio energy is P_m * S_m / R_m, where R_m is the average rate over fading §IV-A. Computing energy per iteration is runtime power times runtime, modelled as a function of core voltage and frequency and of memory frequency, with constants measured on a platform §IV-A. Per global round, a device uses its per-iteration energy times H, and the paper assumes this stays constant while batches grow §IV-A.

Objective. The total energy until convergence is the per-round energy summed over devices, multiplied by the number of global rounds from the convergence bound. The number of rounds scales as O(M H delta^2) plus O(1/(sqrt(M) H)), with delta the root mean square of the device sparsities §III-B, Corollary 1 §IV-A. The problem chooses delta_m in a box [delta_lb, delta_ub] and H from an integer set, to minimise this energy §IV-A.

Control algorithm. The problem mixes continuous sparsities with an integer H, so the paper uses generalised Benders decomposition §IV-B alg. 2. The outer loop fixes H and solves a primal problem in the sparsities. The primal is solved by successive convex approximation: the product of two convex functions is replaced by a surrogate that is convex in delta, and each inner step has a closed-form best response §IV-B, Corollary 3. The master problem chooses H over the integer set with branch and bound, using cuts built from the Lagrange multipliers of the box constraints §IV-B, MasProb alg. 2.

# Compression control  [alg. 2]  [§IV-B]
# Outer loop: Benders over H. Inner loop: convex approximation over delta.

def control(H_init, I_out, eps=1e-5, iota=1e-5, xi=1e-5):
    # Starting integer sync interval, tolerances  [alg. 2, Initialization]
    H = H_init
    cuts = []
    LBD = -inf
    UBD = inf
    for i in range(I_out):
        # Inner loop starts at the lower bound of the box  [alg. 2]
        delta = [delta_lb for each device]
        gamma = gamma0  # step in (0, 1]; value not stated
        while True:
            # Closed-form best response to the convex surrogate  [§IV-B, Cor. 3]
            delta_star = best_response_all(delta, H)
            delta_next = delta + gamma * (delta_star - delta)
            gamma = gamma * (1 - xi * gamma)
            if squared_norm(delta_next - delta) <= iota:
                delta = delta_next
                break
            delta = delta_next
        # Primal value at this H is an upper bound  [alg. 2]
        UBD = primal_objective(delta, H)
        # Add a cut from the multipliers of the primal  [§IV-B, MasProb]
        cuts.append(make_cut(delta, H))
        # Master problem over H gives a lower bound  [alg. 2]
        H, LBD = solve_master(cuts)
        if UBD - LBD <= eps:
            return delta, H
    return delta, H

# Closed-form best response for all devices  [§IV-B, Corollary 3]
# Lambert W branch is not stated; see Pitfalls

def best_response_all(delta, H):
    # Shared sums over all devices at the current iterate  [§IV-B, Cor. 3]
    D = sum over m of (alpha * H * delta[m]**2 + beta / (M**1.5 * H))
    E = ln(2) * sum over m of (
        P[m] * s1 * d * (log2(delta[m]) + kappa) / (R[m] * delta[m])
        + P[m] * s0 / R[m] + E0[m] * H)
    A = alpha * H
    B = 1 - kappa * ln(2)
    out = []
    for m in devices:
        C = P[m] * s1 * d / R[m]
        # Lambert W of the argument  [§IV-B, Cor. 3]
        z = -6 * A * E * exp(3 * B) / (D * C)
        delta_bar = exp(-lambert_w(z) / 3 + B)
        # Clip to the box constraint  [§IV-B, eq. ClosedForm]
        out.append(clip(delta_bar, delta_lb, delta_ub))
    return out

Convergence result. With constant learning rate eta = theta * sqrt(M) / sqrt(T), where theta is a constant with theta * sqrt(M) / sqrt(T) <= 1 / (2L), the average expected squared gradient norm over a random previous iterate (sampled with probability 1/(MT)) is bounded by three terms: an optimisation term, a term from the batch-size growth, and a term that grows with the square of the RMS sparsity and of H §III-B, Theorem 1 § Appendix (Proof of Theorem 1).

Hyperparameters

Most training hyperparameters of the experiments are not stated in the paper.

Results

Author-reported, simulated. The text states the setups below; numerical values are only in the figures, which this card does not read off.

Paper vs code

No official code.

Searched for code in this order. (1) The paper, abstract and arXiv metadata contain no repository link. A commented-out line in the source (in the appendix, near the end of the TeX file) names a GitHub repository for a proof file. It is commented out, the paper does not use it, and its owner is not shown to be an author, so it is not recorded here. (2) No project page is linked. (3) GitHub repository searches returned no results for "FlexibleSpars", "FT-LSGD-DB", "federated topk local sgd dynamic batch" and "energy efficient federated learning compression control". A web search was not available in this run, so the search is not exhaustive.

Pitfalls

• Initial error memory. Algorithm 1 sets the error memory equal to the initial model, e = w, which is probably meant to be zero alg. 1, Initialization. The paper does not say which.
• Gradient point. Algorithm 1 evaluates the local gradient at the global model w^(t) alg. 1, line 6. The appendix analysis evaluates it at the local iterate § Appendix (Proof of Theorem 1). The paper does not reconcile them. The implementation must choose one, and the theorem applies to the appendix version.
• Top-k ties. The definition keeps all entries with |x| >= thr §III-A. With ties at the threshold it keeps more than k entries, while the bit count assumes exactly k. Select exactly k indices if the budget matters.
• Position coding. The bit count uses enumerative coding of positions, log2 of binom(d, k) §IV-A. The paper does not discuss practical index encodings, which will cost more bits.
• Batch growth. b^(t) = floor(rho^t * b0) has no cap stated alg. 1. With larger rho, batches can exceed the local dataset size. Sampling with or without replacement is not stated.
• Constant per-iteration energy. The energy model assumes that computing energy per iteration does not change as batches grow, justified only while the batch stays under a threshold that the paper does not give §IV-A.
• Sync interval in the control loop. Algorithm 2 writes H^(i+1) = H^(i) alg. 2, so as written the master-problem solution is never fed back into the next primal problem. Treat the master's H as the next H, which the surrounding text implies. Algorithm 2 also writes "Update the upper bound value with LBD", which should be the lower bound.
• Cut naming and sign. The master constraints are called feasibility cuts §IV-B, MasProb, but the paper says the primal is always feasible. The terms are Lagrangian multiplier terms for the box constraints, so they act as optimality cuts. Check the sign convention of lambda_1 (delta_lb minus delta) and lambda_2 (delta minus delta_ub) against your own derivation.
• Closed form is unproven here. The proof of Corollary 3 is omitted "due to space limitation" §IV-B. The paper does not say which real branch of the Lambert W function is used. Check the closed form numerically against a generic solver before relying on it.
• Strong convexity needs delta >= e^(3/2). The surrogate is strongly convex only for delta >= e^(3/2), about 4.48 §IV-B. The paper asserts that "extensive empirical evidence" supports this but shows none. The paper sets delta_lb to this value without checking the solution against it.
• Horizon-dependent step size. The theoretical learning rate depends on the total number of iterations T §III-B, Theorem 1. The paper does not say how to pick it when T is not known in advance.
• Corollary substitutions. The round count in Corollary 1 uses substitutions that need sigma, J and L, and the constants alpha and beta only approximate the big-O terms §III-B §IV-A. The paper gives no values for them.
• Weights. The global objective averages device losses with weight 1/M, not by local dataset size §III-A.
• Experiments cannot be reproduced from the paper alone. The learning rate, batch sizes, epochs, H, sparsity levels and energy constants for the CIFAR-10 and MNIST runs are not stated §V. The speed-up figure ("1.5x to 100x") is stated only in the text §V.
• Simulation only. The paper's experiments are simulations with modelled energy. The introduction's comparison (about 0.16 J to send a 100 MB model over 1 Gbps, against about 0.2 J for one training step on one V100 GPU) is cited to other works §I and is not checked here.
• Typos in the appendix. The sums in Lemmas 1 and 2 are indexed from "m-1" rather than m=1 § Appendix (Proof of Theorem 1).

Lineage

• Builds on: FedAvg (local SGD in FL), cited as googleFL §II; sparsified SGD with error compensation, cited as stich2018sparsified §III-A; the perturbed iterate analysis, cited as basu2019qsparse § Appendix; dynamic batch sizes, cited as yu2019computation §III-A; generalised Benders decomposition, cited as geoffrion1972generalized §IV-B; inner convex approximation, cited as scutari2016parallel §IV-B. Keys are the paper's own bibliography keys; their arXiv ids were not checked.
• Compared against: SynSgdSpars (sparsification with a sync after every local step), GreedySpars (per-round greedy energy), UnifiedSpar (one sparsity for all devices), and ordinary distributed SGD as the baseline §V.
• Superseded by: not known from the wiki.