Source-linked AI summary
Accelerating DNN Training in Wireless Federated Edge Learning Systems
Jinke Ren, Guanding Yu, Guangyao Ding
TL;DR
Centralized DNN training can be time-consuming, resource-hungry, privacy-sensitive, and communication-intensive. The paper uses federated edge learning and jointly optimizes batchsize selection with communication resource allocation through a learning-efficiency criterion, deriving CPU and GPU solutions. Experiments report reduced training time and improved learning accuracy, including speedups over individual and model-based FEEL schemes.
Problem
Centralized DNN training requires substantial computation and memory, while transmitting distributed wireless data to the cloud creates privacy and latency concerns.
Method
The paper defines learning efficiency as global loss decay divided by end-to-end latency and derives joint batchsize-selection and communication-allocation solutions for CPU and GPU scenarios.
Results
The proposed scheme reduces training time and improves learning accuracy, reaching about 3.76 times training speed over model-based FEEL in IID data and 2.78 times in non-IID data.
Takeaways & Limitations
The CPU-derived policy and its GPU extension provide batchsize and resource-allocation guidelines for accelerating DNN training in FEEL systems.
Abstract
from arXiv · showhide
Training task in classical machine learning models, such as deep neural networks, is generally implemented at a remote cloud center for centralized learning, which is typically time-consuming and resource-hungry. It also incurs serious privacy issue and long communication latency since a large amount of data are transmitted to the centralized node. To overcome these shortcomings, we consider a newly-emerged framework, namely federated edge learning, to aggregate local learning updates at the network edge in lieu of users' raw data. Aiming at accelerating the training process, we first define a novel performance evaluation criterion, called learning efficiency. We then formulate a training acceleration optimization problem in the CPU scenario, where each user device is equipped with CPU. The closed-form expressions for joint batchsize selection and communication resource allocation are developed and some insightful results are highlighted. Further, we extend our learning framework to the GPU scenario. The optimal solution in this scenario is manifested to have the similar structure as that of the CPU scenario, recommending that our proposed algorithm is applicable in more general systems. Finally, extensive experiments validate the theoretical analysis and demonstrate that the proposed algorithm can reduce the training time and improve the learning accuracy simultaneously.
I. INTRODUCTION
The paper uses federated edge learning to address centralized DNN training’s latency, resource, and privacy challenges by jointly optimizing computation and communication. It defines learning efficiency, derives closed-form CPU solutions, extends the approach to GPUs, and evaluates the resulting schemes experimentally.
- Motivation: Centralized training is costly and problematic because wireless devices hold fragmented data and transmitting it to the cloud creates latency and resource burdens.Device-side learning also suffers from isolated data islands and long training latency.
- Federated edge learning: Federated edge learning aggregates local learning updates at the wireless edge instead of transferring users’ raw data to a remote cloud.FEEL combines federated learning with edge-based centralized aggregation to relieve core-network congestion.
- Learning efficiency: Learning efficiency is defined as global loss decay divided by end-to-end latency, providing a system-level training-performance criterion.The framework also analyzes global loss decay as a function of training batchsize.
- CPU scenario: In the CPU scenario, the paper derives closed-form joint solutions for training batchsize selection and communication resource allocation under resource budgets.The reported batchsize structure depends on local training speed, training priority ratio, and uplink data rate.
- GPU scenario: The GPU extension develops a latency–batchsize function and yields an optimal solution with a structure similar to the CPU case.This similarity supports applying the proposed algorithm to more general systems.
- Evaluation: Experiments with popular DNN models and a real dataset show better learning performance than benchmark schemes.The evaluation compares convergence speed and learning accuracy in IID and non-IID settings.
B. DNN Model
The paper models a general DNN trained collaboratively across statistically independent local datasets, with computation and wireless transmission represented in the FEEL system model. The global objective is minimizing the edge server’s global loss.
- DNN model: A general DNN with p parameters is trained using local datasets held by distributed user devices.Each device’s local loss measures prediction error between model output and ground-truth labels.
- Global objective: The global dataset size is N = Σ_k N_k, and the edge server expresses the global loss from the devices’ local data.The training objective is to optimize model parameters toward minimizing global loss.
- Communication model: The communication model uses device uplink and base-station downlink powers and channel gains, with channel gains available at the base station each round.The paper adopts TDMA and notes that its design principle also applies to other access methods.
- Communication model: Average uplink and downlink data rates replace instantaneous rates because short-slot channel gains vary while a communication round spans a longer timescale.The paper attributes the longer round duration to backpropagation complexity and limited device computation resources.
- System scope: The edge server is assumed to have a powerful GPU, while user devices may use either CPUs or GPUs for training.The subsequent analysis therefore treats CPU and GPU device-training scenarios separately.
III. FEDERATED LEARNING IN CPU SCENARIO: PROBLEM FORMULATION
The CPU formulation models training progress through global loss decay and models each communication round through computation, upload, and download latency. These quantities define the learning-efficiency optimization over batchsize and communication resources.
- A. Training Loss Decay Analysis: Mini-batch SGD is used because training batchsize strongly affects learning accuracy and is therefore optimized for acceleration.The formulation seeks faster training while maintaining a desired learning accuracy.
- A. Training Loss Decay Analysis: Global loss decay measures the difference in the global loss function across one communication round.It is introduced as an auxiliary function for evaluating training performance.
- A. Training Loss Decay Analysis: Each device computes a local gradient from one selected batch, and the global batchsize equals the sum of local batchsizes across devices.The edge server aggregates the local gradients after devices perform local computation.
- A. Training Loss Decay Analysis: Global loss decay does not increase linearly with global batchsize because the learning rate adapts dynamically to batchsize.The paper uses a learning-rate scaling law to keep gradient-expectation variance fixed for convergence and learning accuracy.
- B. End-to-End Latency Analysis: End-to-end latency includes local gradient calculation, local gradient upload, global gradient download, and local model update latency.The round latency is determined by the slowest device because aggregation waits for all local gradients.
- B. End-to-End Latency Analysis: CPU gradient-calculation latency depends on device CPU frequency and the cycles required per data sample, with serial processing across the batch.The model represents each device’s computation capability by its CPU frequency.
- B. End-to-End Latency Analysis: Each quantized local gradient has s = b×p bits, so uplink latency depends on gradient size, frame length, and the allocated uplink time slot.The paper ignores quantization latency because it is low relative to gradient-upload latency.
- B. End-to-End Latency Analysis: TDMA downlink allocation is designed to equalize devices’ end-to-end latency and support synchronized model updates.The downlink analysis assigns each device a time-slot duration for global-gradient delivery.
C. Problem Formulation
The paper reframes training acceleration as maximizing learning efficiency, defined by global loss decay relative to end-to-end latency. It formulates this objective under communication, computation, and batchsize constraints, while noting an extension to broadcast downlink models.
- Learning Efficiency: Learning efficiency is the ratio of global loss decay to end-to-end latency and reflects the global loss decrease rate during training.Maximizing it is presented as equivalent to minimizing overall training time.
- Optimization Formulation: The training-acceleration objective is transformed into a learning-efficiency maximization problem.
- Scope Extension: The analytical framework can be extended to a broadcast model by modifying global-gradient download latency according to the minimum device data rate.
- Optimization Formulation: The formulation constrains uplink and downlink resources, processed data per round, and device batchsizes, with Bmax determined by memory size and CPU configuration.Batchsizes are relaxed to continuous variables for analysis and can be rounded because Bmax is typically large.
IV. FEDERATED LEARNING IN CPU SCENARIO: OPTIMAL SOLUTION
The CPU-scenario problem is decomposed into local-gradient upload and global-gradient download subproblems. Their coupling through the global batchsize enables separate resource-allocation analyses before final optimization.
- Problem Decomposition: The CPU training-acceleration problem is decomposed into subproblems for local gradient calculation and uploading, and global gradient downloading and local model updating.The decomposition uses the independence of local upload and global download latencies.
- Uplink Subproblem: The first subperiod jointly addresses local gradient computation and upload under the listed uplink-related constraints.
- Downlink Subproblem: The second subperiod addresses global-gradient download and local model update under downlink and time-sharing constraints.
- Problem Decomposition: The global batchsize B must be identical in the two subproblems and is optimized last as a shared global variable.
B. Solution to Subproblem P2
For a fixed global batchsize, the uplink subproblem is made tractable through auxiliary efficiency indicators and convex optimization, yielding closed-form joint batchsize and resource-allocation structure.
- Problem Transformation: The uplink subproblem is a difficult min-max optimization, so the maximum reciprocal uplink learning efficiency is introduced as an auxiliary indicator.
- Convexity: With fixed global batchsize, the transformed problem is convex, allowing classical convex optimization methods and structural analysis of its solution.
- Closed-Form Solution: Theorem 1 gives the joint optimal batchsize-selection and uplink resource-allocation policy.The solution is characterized for fixed B and supports the subsequent search procedure.
- Batchsize Selection: The optimal batchsize follows a threshold-based structure determined mainly by local training speed, training priority ratio, and uplink data rate.It increases linearly with local training speed and sublinearly with the training priority ratio and uplink data rate.
- Resource Allocation: Adaptive uplink allocation equalizes end-to-end latency by accounting for uplink data rate, local training speed, and training batchsize.Devices with better channels request fewer time slots, while slower local training can require more time-slot resources.
- Search Algorithm: Algorithm 1 iteratively updates batchsizes and uplink time slots until time-sharing and global-batchsize constraints are satisfied.
C. Solution to Subproblem P3 and Global Discussion
The downlink subproblem has a closed-form resource-allocation solution, and combining uplink and downlink solutions reduces the original optimization to a one-variable search over global batchsize.
- Downlink Subproblem: The downlink subproblem is reformulated using the maximum reciprocal downlink learning efficiency among devices.
- Downlink Subproblem: Given the global batchsize, the downlink problem has characteristics similar to the uplink problem and can be solved using KKT conditions.
- Closed-Form Solution: Theorem 2 provides the optimal downlink resource-allocation policy.
- Consistent Resource Allocation: Allocating more time slots to devices with worse downlink channels enables simultaneous local-model updates and synchronous FEEL operation.
- Global Discussion: The combined closed-form policies reduce the original problem to univariate optimization over global batchsize B.Gradient descent or one-dimensional bisection can solve the remaining problem.
V. EXTENSION TO GPU SCENARIO
The paper extends its training-acceleration framework to GPU-equipped devices and shows that the GPU-case optimum has a structure similar to the CPU-case solution.
- GPU-equipped devices are incorporated by proposing a GPU training function and extending the training-acceleration optimization problem.
A. GPU Training Function
The GPU training model captures parallel execution through a data-bound region with constant latency and a compute-bound region where latency grows with batchsize. The resulting optimization retains the CPU scenario’s solution structure and is supported by experiments.
- A. GPU Training Function: GPU parallelism makes local gradient latency constant for small batchsizes because all data can be processed simultaneously.This regime is called the data bound region because computation resources are not fully exploited.
- A. GPU Training Function: Local gradient latency grows linearly after the batchsize exceeds the GPU’s processing threshold, defining the compute bound region.The threshold marks the point beyond which the computation resource cannot process all samples simultaneously.
- A. GPU Training Function: Experiments with DenseNet, GoogleNet, and PNASNet show latency first remains invariant and then increases almost linearly with batchsize, fitting the theoretical GPU model.The experiments used three NVIDIA GeForce GTX 1080 Ti GPUs.
- A. GPU Training Function: The optimal GPU batchsize places each device in the compute bound region, consistent with fully exploiting computation resources.The data bound region can consequently be neglected when reformulating the optimization problem.
- A. GPU Training Function: The GPU scenario preserves the CPU scenario’s optimization structure because the relevant local gradient latencies are proportional to batchsize in the corresponding formulations.Therefore, the algorithms developed for the CPU problem remain applicable to the GPU problem.
VI. EXPERIMENTS
Experiments evaluate the proposed algorithms across DNN models and learning rates in a simulated wireless federated edge-learning system. The reported results show fast convergence, desired accuracy, and preserved ultimate accuracy across models and learning rates.
- VI. EXPERIMENTS: The experiments use DenseNet121, ResNet18, and MobileNetV2 with CIFAR-10 to evaluate training across different DNN models.CIFAR-10 contains 50,000 training images and 10,000 validation images.
- VI. EXPERIMENTS: The simulated FEEL system contains 12 devices with CPU frequencies distributed across 0.7 GHz, 1.4 GHz, and 2.1 GHz.Four devices are configured at each frequency.
- VI. EXPERIMENTS: The evaluation tests convergence speed and learning accuracy under different learning rates using a non-IID data partition.
- VI. EXPERIMENTS: The proposed scheme attains desired learning accuracy with fast convergence speed for all tested models.
- VI. EXPERIMENTS: Ultimate learning accuracy remains well guaranteed across different learning rates, demonstrating the algorithm’s reported generalization ability.
C. CPU Scenario: Performance Comparison among Different Schemes
The paper compares individual learning, model-based FEEL, gradient-based FEEL, and the proposed joint batchsize-selection/resource-allocation scheme. Across CPU and GPU evaluations, the proposed scheme achieves faster training while maintaining or improving learning accuracy, including in non-IID settings.
- CPU Scenario: The proposed scheme jointly selects local training batchsizes and allocates communication resources, unlike gradient-based FEEL, which does not model the communication–computation trade-off.The scheme optimizes both decisions to balance communication and computation costs.
- CPU Scenario: 1.09 times speedup and 0.83% learning-accuracy improvement over individual learning occur in the IID case with K = 6 devices.In the non-IID case, the corresponding values are 1.03 times and 2.27%.
- CPU Scenario: 3.76 times training speed over model-based FEEL in IID and 2.78 times in non-IID are achieved with slight learning-accuracy improvement.The proposed scheme also reaches about 2 times and 1.5 times the training speed of gradient-based FEEL in IID and non-IID cases, respectively, with almost the same accuracy.
- CPU Scenario: With K = 12 devices, the proposed scheme remains fastest while attaining desired learning accuracy, and its speed improvement over individual learning becomes more evident.Periodic aggregation makes the IID/non-IID accuracy gaps smaller for collaborative schemes than for individual learning.
- GPU Scenario: In the GPU scenario, the proposed scheme achieves the fastest convergence speed and highest learning accuracy across the IID and non-IID comparisons.The reported explanation is that joint time-slot allocation and batchsize selection balance communication and computation costs.
- GPU Scenario: The GPU optimal solution has a structure similar to the CPU solution, supporting application of the algorithm to more general systems.The GPU formulation uses a new function relating training latency to training batchsize.
- Limitation: Optimizing batchsize for learning efficiency may bias the global gradient, especially with non-IID data.The analytical framework could incorporate unbiased-gradient constraints as an additional batchsize constraint.
APPENDIX A PROOF OF THEOREM 1
The appendix proves the CPU optimization structure using convexity and KKT conditions, then derives bounds and monotonicity properties through special cases and loss-decay behavior.
- Theorem 1: For a fixed B, problem P4 is convex, so the Lagrange multiplier method and KKT conditions provide necessary and sufficient conditions.The multipliers correspond to the stated resource and constraint conditions.
- Theorem 1: The optimal batchsize-selection policy follows from the KKT conditions and simple mathematical calculation.The appendix then combines the resulting conditions with the equality condition for EU* to obtain optimal uplink time-slot allocation.
- Corollary 1: Equivalent resource allocation supplies an upper bound on EU*, while infinite memory supplies a lower bound after relaxing the batchsize limitation.Combining the two cases yields the range stated in Corollary 1.
- Corollary 2: Online learning and full-batch learning serve as the two boundary cases for analyzing the optimal batchsize.When one device has an interior batchsize, the associated multiplier satisfies µℓ ≤ µ* ≤ µh.
- Lemma 2: In the data-bound region, global loss decay increases with batchsize while end-to-end latency remains unchanged, so learning efficiency increases.Therefore, the optimal batchsize cannot lie in the data-bound region.