Data-free adaptive structured pruning for federated learning

Abstract

Federated learning faces challenges in real-world deployment scenarios due to limited client resources and the problem of stragglers caused by high heterogeneity. Despite efforts to reduce the training and communication overhead of federated learning through model pruning, a uniform pruning ratio fundamentally fails to address the efficiency impact of stragglers in heterogeneous systems. Therefore, adapting the pruned sub-models to individual device capabilities is crucial yet remains under-researched. In this work, we propose AdaPruneFL, a data-free adaptive structured pruning algorithm, which formulates the adaptive pruning problem in federated learning as an optimization problem constrained by aligning the response latency of the client’s local training, to identify an adaptive fine-grained model compression ratio. Combining sequential structured pruning, we extract heterogeneous but aggregable sub-model structures based on the capabilities of client devices, achieving training acceleration in a hardware-friendly manner while mitigating the straggler effect. Our extensive experiments demonstrate that, compared to FedAvg, AdaPruneFL achieves 1.38–3.88x faster training on general-purpose hardware platforms while maintaining comparable convergence accuracy.

Appendices

Appendix

Experiment setup details

Data partition We use a Dirichlet ($\gamma$) distribution to generate the non-IID data partition among clients, where $\gamma$ controls the skewness of the data distribution. In our experiments, we partitioned the Cifra-10/100 dataset using $\gamma =0.5$. Due to the larger number of classes in the TinyImagenet dataset, we used a smaller value of $\gamma =0.2$ to partition the heterogeneous data distribution. Specifically, the results of data partitioning in different training scenarios are illustrated in Fig. 13:

Fig. 13

Data partitioning results for different datasets

Model architecture details Details of the model architecture used in our experiments are shown in Table 9.

Bandwidth allocation In our experiments, we simulate the heterogeneity of device communication capabilities by randomly allocating communication bandwidth within the 4 G range. The detailed client bandwidth allocation is shown in Table 10. We allow random perturbations within 2MB/s to simulate network fluctuations.

Table 9 Model architectures

Table 10 Bandwidth allocation of clients

Solution of the optimization problem

When using a linear form to approximate the relationship between computational complexity and computational latency, combined with equation (4) and equation (5), we have the following optimization problem:

$$\begin{aligned} \begin{aligned} \min _{p_l}-\sum ^N_{l=1}\log {p^c_l} \ \ \textrm{subject to}\, \,&|\mathcal {D}^c|\cdot \left( k \sum ^N_{l=1}\beta _l\cdot p_l^c + b\right) +\frac{\sum ^N_{l=1}\alpha _l\cdot p_l^c}{B^c} \le t_{target}, \\ 0<p^c_l\le 1,\,&\forall 1\le l\le N \end{aligned} \end{aligned}$$

(8)

According to the Lagrange multiplier method, the Lagrangian function for (8) is constructed as:

$$\begin{aligned} \begin{aligned} L(p^c_l, \mu ) = -\sum ^N_{l=1}\log {p^c_l} + \mu&\left( |\mathcal {D}^c|\cdot \left( k \sum ^N_{l=1}\beta _l\cdot p_l^c + b\right) +\frac{\sum ^N_{l=1}\alpha _l\cdot p_l^c}{B^c} - t_{target}\right) \\ 0<p^c_l\le 1,\,&\forall 1\le l\le N \end{aligned} \end{aligned}$$

(9)

where $\mu \ge 0$ is the Lagrange multiplier. According to the Karush–Kuhn–Tucker (KKT) conditions, we have:

$$\begin{aligned} \begin{aligned}&\frac{\partial L}{\partial p^c_l} = -\frac{1}{p^c_l} + \mu \left( |\mathcal {D}^c| k \beta _l + \frac{\alpha _l}{B^c}\right) = 0\\&\frac{\partial L}{\mu } = |\mathcal {D}^c|\left( k \sum ^N_{l=1}\beta _l\cdot p_l^c + b\right) + \frac{\sum ^N_{l=1}\alpha _l\cdot p_l^c}{B^c} - t_{target} = 0\\&\mu \left( |\mathcal {D}^c|\left( k \sum ^N_{l=1}\beta _l\cdot p_l^c + b\right) + \frac{\sum ^N_{l=1}\alpha _l\cdot p_l^c}{B^c} - t_{target}\right) = 0\\&0<p^c_l\le 1,\ \forall 1\le l\le N \end{aligned} \end{aligned}$$

(10)

By rearranging (10), we can conclude:

$$\begin{aligned} \left\{ \begin{aligned}&p_l=min\left( \frac{B^c}{\mu (|\mathcal {D}^c|\cdot k \cdot \beta _l \cdot B^c+\alpha _l)},1\right) \\&|\mathcal {D}^c|\left( k \sum ^N_{l=1}\beta _lp_l+b\right) +\frac{\sum ^N_{l=1}\alpha _lp_l}{B^c}=t_{target}\\&\mu >0 \end{aligned} \right. \end{aligned}$$

(11)