When Foresight Pruning Meets Zeroth-Order Optimization: Efficient Federated Learning for Low-Memory Devices

Though the sampling method in Eq. 9 is friendly to low-memory devices, the computation overhead is extremely heavy since $K$ should be large for accurate estimation. Following what was proven in (Feng et al., 2023), the estimation error is formally defined by the following theorem:

Neglecting the trivial term $o(\widehat{\boldsymbol{\delta}})$, the estimation error measured by Mean Square Error (MSE) is fully controlled by the term $||\hat{\boldsymbol{\Sigma}}-\mathbf{I}||_{2}^{2}$. Based on the proof in (Adamczak et al., 2011), we have $||\hat{\boldsymbol{\Sigma}}-\mathbf{I}||_{2}\leq\sqrt{\frac{n}{K}}$. In conclusion, we can expect a more accurate estimation if we either increase the number of Monte Carlo steps or decrease the dimensionality of the model parameter $\boldsymbol{W}$, or both. In practice, $K$ is related to the total FLOPs consumed to estimate the true gradient, and $n$ reflects the memory requirements for a device to afford the model. Developing more advanced zeroth-order algorithms to lower the value of $K$ required to achieve the same performance is a promising way to obtain more accurate estimations. Another way is to put effort into decreasing the value of $n$. To achieve the latter, we introduce model pruning to reduce computation overhead and increase the estimation precision as it aims to find the spare structure of the model, and can still maintain the performance.

Current federated pruning methods are highly integrated with the training process. Consequently, the noisy gradients negatively affect the pruning step, decreasing training performance. To mitigate the impact of noisy gradients on the pruning step, we propose decoupling the pruning processes from the training processes, which is one of the fundamental motivations of foresight pruning. Further, we dedicate efforts to developing a federated foresight-pruning method based on the NTK spectrum, ensuring full compatibility with the BP-Free training framework.

Since the symmetric property of the $\boldsymbol{\theta}_{0}$ is only valid for centralized training, we cannot directly apply it for federated training. Motivated by (Huang et al., 2021), we prove that the NTK in federated learning can be approximated by local NTK matrices.

The asymmetric property of FL-NTK prevents us from computing the nuclear norm by the Frobenius norm of gradients based on each local dataset. However, we prove that we can utilize the summation of local symmetric NTK matrices to approximate the FL-NTK and efficiently conduct the Frobenius norm computation.

Since the nuclear norm satisfies the triangle inequality, we further formulate the computation of $\|\boldsymbol{\theta}_{0}^{fl}\|_{*}$ as

Therefore, the complex FL-NTK is upper-bounded by the summation of individual local NTK matrices. Since the sparse model preserves the same training dynamic as the dense model and estimated gradients based on Stein’s Identity are unbiased, the overall convergence rate based on the training round $T$ is $O(1/\sqrt{T})$, the same as the standard FedAvg (Li and Lyu, 2023) algorithm under non-convex settings.

Unlike centralized pruning, the entire dataset $\mathcal{D}$ is distributed to multiple devices in FL. According to the form of Eq. 3, the computation of saliency score demands at least one operation of gradient computing, i.e., $\partial\mathcal{I}/\partial\boldsymbol{W}_{0}^{j}$ if not consider the function $\mathcal{I}$. We can wisely let the local devices compute the function $\mathcal{I}$ and leave the derivative and multiplication operations in Eq. 15 to the server to complete. In this way, local datasets equally contribute to the generation of the binary mask, yielding an unbiased global sparse model, and the memory required to support the auto-differential framework is shifted to the server. Things might be different when considering the specific form of the function $\mathcal{I}$. As mentioned above, the function $\mathcal{I}$ defined in Eq. 4 can be approximated by either Eq. 6 or the loss gradient w.r.t. weights $\left\|\nabla_{\boldsymbol{W}_{0}}\mathcal{L}\right\|_{2}^{2}$. The latter introduces extra gradient computation, which inevitably relies on the auto-differential framework. In contrast, Eq. 6 benefits from being independent of the auto-differential framework. Regardless, the computation of the saliency score in Eq. 3 still intertwines with it (i.e., $\partial\mathcal{I}/\partial\boldsymbol{W}_{0}^{j}$). We might introduce another zeroth-order optimization method to estimate the derivatives, but it introduces more estimation errors. As a result, we choose Eq. 6 as the function $\mathcal{I}$ to make the local devices free of the memory-inefficient backpropagation framework. The federated NTK foresight pruning is organized as

where $i$ represents the $i$-th dataset owned by the corresponding device. $\boldsymbol{W}_{0}^{j}=(\boldsymbol{W}_{0,i}^{j})_{i=1}^{m}$ denotes the $j$-th element of the initial global model. $\Delta\boldsymbol{W}_{i}$ is sampled from a Gaussian distribution with zero mean and $\epsilon\mathbf{I}$ variance. The $\mathbb{E}$ symbol denotes the expectation over $\Delta\boldsymbol{W}_{i}$, which can be approximated by the Monte Carlo sampling method. To avoid layer-collapse, we use an exponential decay schedule $\tau=(1-d^{t/T_{p}})$ to compute the pruning threshold $\tau$, where $t$ denotes the current pruning round, $d$ is the target density and $T_{p}$ is the maximum pruning round. The parameters whose saliency scores are under the threshold $\tau$ will be pruned. It may be noted that the proposed NTK pruning method introduces $T_{p}$ extra computations to the devices. In addition, the multiple-round pruning schedule requires communication between the server and all devices. To resolve the two problems in one shot, we fully explore the intrinsic property of NTK pruning. Motivated by (Shu et al., 2022), we show that under data heterogeneity scenarios, the difference between any twp local NTK matrices is $\Big{|}\left\|\boldsymbol{\theta}_{0}^{i}(P^{i})\right\|_{*}-\left\|\boldsymbol{\theta}_{0}^{j}(P^{j})\right\|_{*}\Big{|}\leq n_{0}^{-1}Z$, where $P^{i}$ and $P^{j}$ represent the local data distributions, $Z$ represent $\int\left\|P^{i}(\mathbf{X})-P^{j}(\mathbf{X})\right\|\mathrm{d}\mathbf{X}$. $n_{0}$ is the input dimension. In our cases, $n_{0}$ is $1024$ and $0<Z<2$. Therefore, the approximation error defined by the summation operation in Eq. 12 can be represented by the difference and increases as the degree of heterogeneity grows. To alleviate the approximation error, we constraint the difference between $P^{i}$ and $P^{j}$ by sampling the data $\mathbf{x}_{i}$ in Eq. 13 from the same standard Gaussian distribution, which diminishes the difference between two NTK matrices and makes our method resilient to data heterogeneity. Moreover, the computation of Eq. 13 can be fully conducted by the server, and our proposed federated NTK pruning method degrades to the centralized version, entirely saving the communication overhead.

We further integrate the proposed NTK foresight pruning with Stein’s Identity method to collaboratively build a memory-friendly federated training framework. To leverage Stein’s Identity in the context of FL, we rewrite the estimation of the aggregated gradient in the form of FedAvg (McMahan et al., 2017b) as

Note that $\boldsymbol{\delta}\in\mathbb{R}^{n}$ and $n$ is the dimensionality of gradients we want to estimate. To further decrease the communication burden, we use the Random Seed Trick technique. The Random Seed Trick smartly leverages the inherent property of the developing environment: i) $\boldsymbol{\delta}$ has the same dimensionality as the gradients we want to estimate and is generated by sampling from a Gaussian distribution; ii) the sampling from Gaussian distribution is driven by the Random Seed in software level. Therefore, the server can produce the same sampling results as the local devices by maintaining the same Random Seed and running environments between them. In detail, Eq. 16 is performed by the server and devices as:

In practice, we let the devices compute $\Delta\mathcal{L}(\boldsymbol{W}_{i};\mathcal{D}_{i})$ and leave the $\boldsymbol{\delta}$ to the server. Consequently, the communication cost is significantly reduced from large dimensional vectors to scalars level since we do not need to transmit the vector $\boldsymbol{\delta}$. The overall training processes are shown in Algorithm 1.

For the Fedavg method, the accuracy after convergence is $53.66\%$. Since it is conducted on a dense model, the FLOPs are not reduced, making it inefficient for resource-constrained devices. For the FedDST method, the FLOPs are significantly reduced to $20\%$ of the dense model. However, the accuracy drop is not negligible. In addition, the convergence speed is slower than the FedAvg method due to the consistent neuron growth and pruning steps. For the PruneFL approach, the accuracy drop after convergence is less than that of FedDST. Note that PruneFL in our experiment ends up with an all-one mask, which makes the model dense again when convergences. Therefore, the FLOPs of PruneFL remain the same compared to FedAvg. Since our NTK method collects pruning information from all local devices, the accuracy drop is largely reduced to $0.42\%$. Moreover, the NTK method avoids model sparsity readjustment during the training. As a result, the training stability shown in Figure 1 is consolidated compared with FedDST and PruneFL. Though the FLOPs for a single forward pass of NTK are not as small as FedDST, the sparsity readjustment in FedDST inevitably brings more computation expenses during the training. For the CIFAR-100 dataset, we can observe that the FedDST outperforms other methods in terms of accuracy. Yet, FedDST implements multiple local training epochs (5 epochs in the original paper) to realize the sparsity readjustment, resulting in a heavy computation burden. We further display the experimental results of the NTK pruning given $1$ local training epoch. The accuracy outperforms $5$ epoch FedDST by $0.02\%$, and the computation expense in FLOPs is $5\times$ cheaper than FedDST. Since the advantages in accuracy brought by the $1$ epoch align with the observations that fewer local steps enhance the global consistency (Liu et al., 2023), we conclude that our proposed pruning method enjoys both computation efficiency and better performance.

When utilizing ResNet-20, the improvements in the maximum accuracy are $1.00\%$, $1.93\%$, and $1.46\%$, respectively. These phenomena are mainly due to the reason that the estimation over a more complex is not as stable as that over a relatively simple model. Although the maximum accuracy improvements are not as prominent as that when using ResNet-20, the learning curves against FLOPs from Figure 3 convincingly demonstrate that NTK-BAFFLE achieves higher accuracy when consuming the same FLOPs.

Table 4 presents the maximum accuracy comparison between NTK-BAFFLE and Vanilla-BAFFLE on the CIFAR-100 dataset. Unlike the $90\%$ sparsity setting for the CIFAR-10 dataset, we uniformly set the $80\%$ sparsity level for the CIFAR-100 dataset to ensure the Figure 5 exhibits the learning curves regarding accuracy versus FLOPs. Same as Figure 3, the first row in Figure 5 presents the results on LeNet, and the second row shows that on ResNet-20. The maximum accuracy improvements based on the NTK-BAFFLE against Vanilla-BAFFLE are consistently observed through all kinds of settings. For experiments conducted on LeNet, the maximum accuracy is increased by $1.19\%$, $1.86\%$, and $2.4\%$, respectively. For experiments on ResNet-20, the accuracy is improved by $0.50\%$, $0.43\%$, and $1.77\%$ when setting $K$ to 50, 100, and 200, respectively. From Figure 3 and Figure 5, we can notice that our proposed pruning method consistently enhances the performance of the original BAFFLE method by achieving higher accuracy while consuming lower FLOPs.

The additional memory required for storing the perturbation in each forward pass and estimating the local gradient for the BAFFLE-based method is $O(n)$. It is worth noting that modern frameworks (e.g., (Elsen et al., 2020)) can efficiently store and process sparse matrices, making the memory needed to store the binary mask associated with NTK-BAFFLE negligible. We perform layer-by-layer computations, partitioning the forward computation graph into smaller segments and calculating the peak GPU memory usage during the forward pass when the input is $32$ RGB images with a size of $32\times 32$ from the CIFAR-10 dataset. In LeNet and ResNet-20, the peak memory usage of Vanilla-BAFFLE is significantly smaller compared to backpropagation-based methods, with values of $166$ MB vs. $1617$ MB and $330$ MB vs. $1671$ MB, respectively. Additionally, the proposed NTK-BAFFLE has fewer parameters. When considering an $80\%$ pruning rate, the peak memory usage of NTK-BAFFLE due to parameters is only $7.44\%$ (LeNet) and $3.88\%$ (ResNet-20) of the usage in Vanilla-BAFFLE, respectively.