PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Model Pruning and Personalization for Federated Learning over Wireless Networks

Xiaonan Liu, Tharmalingam Ratnarajah, Mathini Sellathurai, Yonina C. Eldar

arXiv:2309.01816v3cs.LGcs.NI

TL;DR

Federated learning can lose accuracy when device data are heterogeneous, while model updates strain devices' computation and wireless resources. This paper combines partial pruning of a shared model with device-specific personalization and optimizes pruning and bandwidth allocation. It reports about 50% lower computation and communication latency than FL with model personalization, with comparable learning accuracy.

  • Problem

    Heterogeneous device data can reduce FL learning accuracy, while limited device computation and wireless resources increase model-update latency.

    Source: §I · p. 1

  • Method

    The framework prunes a shared global model part, keeps a device-specific personalized part, and jointly optimizes pruning ratio and bandwidth allocation using convergence analysis and KKT conditions.

    Source: §I · p. 2

  • Results

    The proposed framework reduces computation and communication latency by approximately 50% compared with FL only with model personalization, while achieving comparable learning accuracy.

    Source: §I · p. 2

  • Takeaways & Limitations

    The simulations show that partial pruning and personalization can achieve similar accuracy to partial model personalization while reducing computation and communication overheads in wireless FL.

    Source: §VI · p. 10

Abstract

from arXiv · show

Federated learning (FL) enables distributed learning across edge devices while protecting data privacy. However, the learning accuracy decreases due to the heterogeneity of devices' data, and the computation and communication latency increase when updating large-scale learning models on devices with limited computational capability and wireless resources. We consider a FL framework with partial model pruning and personalization to overcome these challenges. This framework splits the learning model into a global part with model pruning shared with all devices to learn data representations and a personalized part to be fine-tuned for a specific device, which adapts the model size during FL to reduce both computation and communication latency and increases the learning accuracy for devices with non-independent and identically distributed data. The computation and communication latency and convergence of the proposed FL framework are mathematically analyzed. To maximize the convergence rate and guarantee learning accuracy, Karush Kuhn Tucker (KKT) conditions are deployed to jointly optimize the pruning ratio and bandwidth allocation. Finally, experimental results demonstrate that the proposed FL framework achieves a remarkable reduction of approximately 50 percent computation and communication latency compared with FL with partial model personalization.

I. INTRODUCTION

The introduction motivates wireless federated learning that can handle heterogeneous device data while reducing the computation and communication burden of model updates. It presents a framework combining partial model pruning and personalization, alongside analysis and joint optimization of pruning and bandwidth.

I. INTRODUCTION

  • Non-IID device data can destabilize FL and lower learning accuracy, while limited computation and wireless resources increase local and transmission latency.Prior partial-personalization work did not explore its computation and communication overheads over wireless networks, while cited pruning methods were described as unsuitable for non-IID data.Source: §I · pp. 1–2
  • The proposed framework combines partial model pruning and personalization to adapt to data heterogeneity and changing wireless environments.The paper models computation and communication latency, analyzes convergence, and jointly optimizes the global-part pruning ratio and bandwidth allocation under constraints.Source: §I · p. 2

II. SYSTEM MODEL AND PROBLEM FORMULATION

  • The system models multiple mobile devices served wirelessly by an edge server, with device-specific datasets and no overlap between datasets.This setup frames the paper's distributed-learning and wireless-resource problem.Source: A. FL with Partial Model… · p. 2, §II · p. 2

A. FL with Partial Model Personalization Architecture

  • The model is split into a global part shared across devices and a personalized part specific to each device, addressing the limits of a single global model under heterogeneous data.The global part is used to learn shared representations, while local personalization adapts to device data.Source: §I · p. 1, A. FL with Partial Model… · p. 3

B. Learning Process

  • Each communication round broadcasts the global part; devices update both parts locally, upload the updated global part, and retain the personalized part locally.Local updating uses stochastic gradient descent on randomly selected subsets of each device's dataset.Source: B. Learning Process · p. 3, 2) Local Model Updating · p. 3

3) Uplink Transmission of Global Part:

The system sends updated global parts from devices to the edge server over an OFDMA uplink and averages participating devices' updates. Its latency model focuses on local updating and uplink transmission, with synchronous round time set by the slowest device.

3) Uplink Transmission of Global Part:

  • Devices upload updated global parts over an OFDMA uplink, with each device assigned a bandwidth fraction used to determine its transmission rate and latency.The passage identifies channel gain, transmit power, and noise power as transmission-rate parameters.Source: 3) Uplink Transmission of Global Part · p. 3

4) Global Part Aggregation:

  • The edge server averages the global parts received from participating devices and broadcasts the aggregate for the next round.The server's aggregation latency is ignored because it is assumed to have high computation capability.Source: 4) Global Part Aggregation · pp. 3–4

C. Local Computation and Uplink Latency

D. Optimization Problem Formulation

The optimization formulation minimizes global loss subject to latency and bandwidth limits. The proposed pruning procedure ranks global-part weights, removes a ratio of low-importance weights, and models the resulting local-update computation.

D. Optimization Problem Formulation

  • The optimization minimizes global loss under a computation-and-communication latency limit and constraints on device bandwidth fractions.The problem is transformed into optimizing the convergence bound after considering pruning to preserve learning accuracy while reducing latency.Source: D. Optimization Problem Formulation · p. 4

A. Model Pruning

  • Pruning is introduced because updating even the shared global part can be costly on devices with limited computation and poor wireless channels.The stated aim is to reduce model size and computation and communication latency.Source: A. Model Pruning · p. 4
  • The unstructured pruning method scores a global-part weight by the squared difference between prediction errors with and without that weight.A larger error difference means greater weight importance.Source: A. Model Pruning · p. 4
  • To reduce the cost of scoring every weight, the method can use the difference between updated and non-updated weights, which is available through backpropagation.The passage presents this as a simpler alternative to calculating prediction-error differences for all weights.Source: A. Model Pruning · p. 4

1) Local Model Updating:

  • For each device, a pruning mask removes the weights ranked in the lowest fraction set by its global-part pruning ratio.The weights are ranked by importance in descending order before masking.Source: 1) Local Model Updating · p. 5
  • Local-update computation latency is calculated from the number of cycles needed to update the personalized and global parts and the device's allocated CPU frequency.The formulation distinguishes the sizes of the personalized and global parts and the iteration count for global-part updates.Source: 1) Local Model Updating · p. 5

2) Uplink Transmission of Global Part:

This section models global-part uplink and aggregation alongside the personalized updates, then bounds convergence and identifies pruning ratio as a main determinant of that bound.

2) Uplink Transmission of Global Part:

3) Global Part Aggregation:

  • In the CNN example, convolution and pooling layers form personalized feature-extraction parts, while fully connected layers form the global pattern-recognition part targeted for pruning.The paper says pruning convolution layers would decrease CNN robustness, so it mainly considers pruning fully connected layers.Source: 3) Global Part Aggregation · p. 5, Figure 2 · p. 5

B. Proposed FL Convergence Analysis

  • Theorem 1 derives an upper bound on the proposed framework’s convergence rate, which is mainly determined by the pruning ratio under fixed iteration counts, communication rounds, and learning rate.The bound uses A1 and A2; κ* denotes the minimum occurrence of global-part weights across communication rounds.Source: B. Proposed FL Convergence Analysis · p. 6

C. Optimization Problem Transformation

The optimization is framed around maximizing convergence subject to pruning, latency, and bandwidth constraints. The pruning-ratio result relates each device’s pruning choice to its computation capability and uplink rate.

C. Optimization Problem Transformation

  • The optimization is reformulated to maximize the convergence rate, with pruning-ratio and bandwidth choices selected under the specified constraints.The pruning-ratio constraint is one of the conditions used to choose the two quantities.Source: C. Optimization Problem Transformation · p. 6

IV. PRUNING RATIO AND BANDWIDTH FRACTION OPTIMIZATION

  • The pruning-ratio and bandwidth-fraction problem is divided into two sub-problems for optimization..Source: §IV · p. 6

A. Pruning Ratio Optimization

B. Bandwidth Fraction Optimization

The bandwidth-fraction problem is shown to be convex and yields a closed-form allocation. The stated allocation gives more bandwidth to devices with poor channels and to devices with high computation capacity.

B. Bandwidth Fraction Optimization

  • Theorem 3 gives a closed-form solution for each device’s bandwidth fraction, obtained using the KKT conditions.The optimization is shown to be convex with respect to the bandwidth fraction.Source: B. Bandwidth Fraction Optimization · p. 7
  • Devices with poor channel conditions receive more bandwidth to satisfy transmission-latency constraints, while high-computation-capacity devices receive more bandwidth.The paper also states that the latter allocation improves convergence rate and decreases computation latency.Source: B. Bandwidth Fraction Optimization · p. 7

V. EXPERIMENT RESULTS

Experiments evaluate the framework on non-IID MNIST and Fashion MNIST using a CNN, with each device holding data from several labels. The algorithm updates personalized and global parts locally and transmits the pruned global part.

V. EXPERIMENT RESULTS

  • The evaluation uses one edge server and ten devices, with a CNN for image classification on MNIST and Fashion MNIST.The devices have data from several labels; convolutional layers are personalized and fully connected layers are global.Source: §V · p. 7
  • Each global round updates the personalized part, then the global part, while determining bandwidth allocation and a pruning mask.The algorithm specifies local iterations for both model parts and calculates the bandwidth fraction and pruning mask before global-part updates.Source: §V · p. 8
  • Devices transmit the pruned global part to the edge server for aggregation.The algorithm bases transmission on the allocated bandwidth fraction.Source: §V · p. 8

A. FL with Partial Model Pruning and Personalization

The experiments examine partial pruning with personalization and jointly optimized pruning and bandwidth allocation over wireless networks. The framework reports similar learning accuracy to partial model personalization with lower overhead, while accuracy and convergence vary with pruning and latency thresholds.

A. FL with Partial Model Pruning and Personalization

B. Proposed FL over Wireless Networks

  • The wireless-network evaluation examines the joint design of pruning ratio and bandwidth fraction, as well as the effect of latency thresholds.These are the two performance questions addressed in this section.Source: B. Proposed FL over Wireless Networks · p. 8

1) Pruning Ratio and Bandwidth Fraction Optimization:

  • The proposed framework has lower computation and communication latency than FL only with pruning, FL only with personalization, and equal resource pruning.The paper attributes the improvement over equal resource pruning to selecting the global-part pruning ratio based on allocated bandwidth.Source: 1) Pruning Ratio and Bandwidth… · p. 9
  • On non-IID Fashion MNIST, the proposed framework has better loss and testing accuracy than equal resource pruning and FL only with model pruning, and its learning performance is close to partial model personalization.Its communication cost is about 50% lower than FL only with model personalization.Source: 1) Pruning Ratio and Bandwidth… · p. 9, Figure 5 · p. 9

2) Effect of Latency Threshold:

  • As latency thresholds increase, loss decreases, testing accuracy increases, and fewer global rounds are needed for convergence.The reported explanation is that larger thresholds select a smaller pruning ratio, retaining more important global-part weights.Source: 2) Effect of Latency Threshold · p. 9, Figure 6 · p. 9

VI. CONCLUSIONS

APPENDIX

The appendix develops convergence bounds for the global and personalized model parts. It uses lemmas and smoothness inequalities to bound changes in the objective across updates and communication rounds.

APPENDIX

  • The convergence proof minimizes an objective over the global model and all devices’ personalized parts, using an L-Lipschitz assumption.The proof separates the personalized-part update from the more challenging global-part update.Source: APPENDIX · p. 10
  • Lemma 2 bounds the global-part update by accounting for both local stochastic-gradient updates and variation due to pruning.The proof invokes bounds on stochastic-gradient moments and model error caused by the pruning ratio.Source: APPENDIX · pp. 10–11
  • Lemmas 3 and 4 bound gradient differences using assumptions on shared weights and bounded variance.The derivation relaxes per-weight gradient subvectors to full gradients and uses the bounded-variance assumption.Source: APPENDIX · p. 11
  • Lemma 5 bounds the squared change in the global part across communication rounds, supporting the smoothness bound for the global update.The proof combines earlier bounds before extending the analysis over all communication rounds.Source: APPENDIX · pp. 11–13

Related papers

Loading 2309.01816v3…