LBI-FL: Low-Bit Integerized Federated Learning with Temporally Dynamic Bit-Width Allocation

Thank you for your valuable suggestions and below are our detailed responses to the raised weaknesses and questions.

W1: Experiments on different scenarios.

Response: We have made comprehensive evaluations on our LBI-FL, including image classification on CIFAR-10 using diverse architectures of ResNet-18/50/101, MobileNet-V2, and ViT-S and on Tiny-ImageNet using ResNet-18, and image segmentation on DSB2018 using U-Net. In all the evaluations, our LBI-FL achieves acceptable performance (tolerable 1.5% loss) with over 45% reduction of BitOPs, compared with UI8 training.

W2: Experiments on more tasks and datasets and comparison with advanced compression methods.

Response: i) We train U-Net for Image segmentation on DSB2018. We adopt the image size of 96$\times$96. The table below shows that, compared with UI8 training, we achieve acceptable performance of about 0.6% loss in Dice similarity coefficient (DSC) with over 50% reduction of BitOPs.

ii) We train ResNet-18 for image classification on Tiny-ImageNet. The table below shows that our LBI-FL achieves over 50% reduction of BitOPs with less than 0.1% accuracy loss

iii) We further employ centralized low-bit training method AMPA [D1] in FL for comparison. AMPA uses layer-wise bit-width allocation based on sensitivity measurement during training. Our LBI-FL with RL agents obtains superior performance with evidently reduced BitOPs in training ResNet-18, MobileNet-V2, and ViT-S on CIFAR-10.

[D1] Li Ding et al. AMPA: Adaptive mixed precision allocation for low-bit integer training. ICML 2024.

W3: Discussion on combining pruning and quantification.

Response: This paper focuses on quantization for low-bit training in FL. Our LBI-FL introduces a temporally dynamic bit-width allocation scheme for weights, activations, and gradients, which evolves along the training trajectory. It achieves higher efficiency than UI8 training with great flexibility in diverse FL scenarios.

We will discuss the meaningful topic of combining pruning and quantization such as [D2] in the final version. However, pruning is more difficult to deploy on hardware than quantization and relies on manually tuned hyperparameters such as pruning ratios and threshold values. [D2] only offers results of LeNet and ResNet-20 on CIFAR-10. We will explore to combine pruning with low-bit training (simultaneously quantizing weights, activations and gradients) in future.

[D2] Pavana Prakash et al. IoT device friendly and communication-efficient federated learning via joint model pruning and quantization. IEEE IoT Journal 2022.

W4: Figure 1 not referenced.

Response: Thank you. We will include the reference to Figure 1 in the final version.

Q1: Can the source code be made public?

Response: We will release the source code upon acceptance.

Q2: Difference with other methods like in deep learning.

Response: While low-bit training has been widely studied in conventional deep learning, its direct application to federated learning (FL) is challenging due to the decentralized and heterogeneous nature of FL. In particular, applying a uniform low-bit strategy across all clients can lead to training instability caused by heterogeneity in data distributions and model dynamics. To address this, we propose a aware bit width adaptation framework based on reinforcement learning. Specifically, we introduce an agent that dynamically allocates bit-widths for each client by considering local states, including the current bit-width, training phase, and quantization-induced loss.

Thank you for your constructive suggestions.

W1: Claims may raise concerns.

Response: Claim i): We use this sentence to emphasize that our LBI-FL can reduce both the overhead in the uplink and downlink communication with low-bit training. In fact, the paper "Prune at the Clients, Not the Server" cannot reduce the downlink communication only using client-side pruning and has to resort to accelerated server pruning, which could completely fail in accuracy and loss as mentioned in the paper. For rigidity, we will change the claim to "existing methods based on quantization" in the final version.

Claim ii): In lines 19-21 and 55-59, "low-bit training" refers to simultaneously quantizing the weights, activations, and gradients. The paper "IntSGD: Adaptive floatless compression of stochastic gradients" develops adaptive integer compression operators for distributed SGD. It only quantizes the gradients and does not quantize weights and activations. To our best knowledge, we are the first to achieve "low-bit training" in the FL scenario.

Claim iii): We have improved the theoretical result by relaxing Assumption 4.2 without affecting the final convergence rate. Instead of assuming the quantization noise to follow a Gaussian distribution with identical variance across all the clients, we now assume only that (a) The gradient quantization noise on each client has a well-defined expectation and variance ($\mu_N^i$ and $(\sigma_N^i)^2$ for the $i$-th client), and (b) The expectation $\mu_N^i$ can be viewed as zero.

Based on this assumption, we provide an updated derivation of the theorem. The noise term does not influence the $A_1$ term in Equation (12). Neglecting the differences across epochs, we reformulate Equation (13) in the appendix as:

$$\begin{aligned} A_2 & \le \frac{\eta^2}{m^2} \mathbb{E}\_t\left[\left\|\sum\_{i=1}^m \sum_{k=0}^{K-1} \mathbf{g}\_{t, k}^i\right\|^2\right] + \frac{\eta^2}{m^2} K \cdot \mathbb{E}\_t \left[ \left\| \sum\_{i=1}^m n_t^i \right\|^2\right] \end{aligned}$$

Considering the quantization noise on each client can be supposed independent, the expectation in the second term is calculated as:

$$\mathbb{E}\_t \left[ \left\| \sum\_{i=1}^m n_t^i \right\|^2\right] = \sum\_{i=1}^m \mathbb{E}\_t \left[ \left\| n_t^i \right\|^2\right] + 2\sum_{i<j} \mathbb{E}_t\left[n_t^i n_t^j\right] = \sum\_{i=1}^m ((\mu\_N^i)^2 + (\sigma^i\_N)^2) + 2\sum\_{i<j} \mu_N^i \mu_N^j$$

Therefore, $\Phi$ in Equation (10) changes as

$$\Phi = \frac{1}{c}\left[\frac{L \eta}{2 m} \left(\sigma_L^2 + \frac{1}{m} \left( \sum\_{i=1}^m ((\mu\_N^i)^2 + (\sigma^i\_N)^2) \right)\right) +\frac{5 K \eta^2 L^2}{2}\left(\sigma_L^2+6 K \sigma_G^2\right) \right]$$

However, the convergence rate does not change, since the quantization noise expectation is usually at 0.

Note that previous convergence result is a special case of this new result.

W2: Assumption for theoretical claims.

Response: We perform experiments with ResNet-18 on CIFAR-10 to support the assumption (b) in W1 that the expectation $\mu_N^i$ can be viewed as zero. The distribution of quantization noise is shown to concentrate around zero.

W3: Evaluations on more complex architectures.

Response: We have evaluated on larger-scale ResNet50 and ResNet101 on CIFAR-10. The table below shows that our LBI-FL yields less than 1.5% performance loss with over 45% BitOPs reduction.

We further evaluate our LBI-FL by training U-Net for image segmentation on DSB2018. The table below shows that, compared with UI8 training, our LBI-FL obtains about 0.6% loss in Dice similarity coefficient (DSC) with over 50% reduction of BitOPs.

W4: References Not Discussed.

Response: We carefully read the two papers and will definitely cite them in the final version. However, they are different from our paper in the contents, as elaborated in W1.

Q1: Convergence results without Assumption 4.2 or with a relaxed version.

Response: We have provided a relaxed version of Assumption 4.2. We only assume the expectation and variance of the distribution of the gradient quantization noise on each client exist and the expectation can be viewed as zero, as elaborated in W1. The relaxed assumption does not change the convergence rate in the theorem.

Thank you for your valuable suggestions and below are our responses to the weaknesses and questions.

W1: Proof of the theorem.

Response: According to your comment, we have improved the theoretical analysis by relaxing Assumption 4.2 without affecting the final convergence rate. Specifically, instead of assuming the quantization noise to follow a Gaussian distribution with identical variance across all the clients, we now assume only that (a) The gradient quantization noise on each client has a well-defined expectation and variance (denoted as $\mu_N^i$ and $(\sigma_N^i)^2$ for the $i$-th client), and (b) The quantization noise expectation $\mu_N^i$ can be viewed as zero.

Based on this assumption, we provide an updated derivation of the theorem. The noise term does not influence the $A_1$ term in Equation (12). Neglecting the differences across epochs, we reformulate Equation (13) in the appendix as:

$$\begin{aligned} A_2 & \le \frac{\eta^2}{m^2} \mathbb{E}\_t\left[\left\|\sum\_{i=1}^m \sum\_{k=0}^{K-1} \mathbf{g}\_{t, k}^i\right\|^2\right] + \frac{\eta^2}{m^2} \mathbb{E}\_t \left[ \left\| \sum_{i=1}^m \sum\_{k=0}^{K-1} n_{t, k}^i \right\|^2\right] \\\\ & = \frac{\eta^2}{m^2} \mathbb{E}\_t\left[\left\|\sum\_{i=1}^m \sum_{k=0}^{K-1} \mathbf{g}\_{t, k}^i\right\|^2\right] + \frac{\eta^2}{m^2} K \cdot \mathbb{E}\_t \left[ \left\| \sum\_{i=1}^m n_t^i \right\|^2\right] \end{aligned}$$

Considering the quantization noise on each client can be supposed independent, the expectation in the second term is calculated as:

$$\begin{aligned} \mathbb{E}\_t \left[ \left\| \sum\_{i=1}^m n_t^i \right\|^2\right] = \sum\_{i=1}^m \mathbb{E}\_t \left[ \left\| n_t^i \right\|^2\right] + 2\sum_{i<j} \mathbb{E}_t\left[n_t^i n_t^j\right] = \sum\_{i=1}^m ((\mu\_N^i)^2 + (\sigma^i\_N)^2) + 2\sum\_{i<j} \mu_N^i \mu_N^j \end{aligned}$$

Therefore, $\Phi$ in Equation (10) changes as

$$\begin{aligned} \Phi =& \frac{1}{c}\left[\frac{L \eta}{2 m} \left(\sigma_L^2 + \frac{1}{m} \left( \sum\_{i=1}^m ((\mu\_N^i)^2 + (\sigma^i\_N)^2) \right)\right) +\frac{5 K \eta^2 L^2}{2}\left(\sigma_L^2+6 K \sigma_G^2\right) \right] \end{aligned}$$

However, the convergence rate does not change, since the quantization noise expectation is usually at 0.

Note that previous convergence result is a special case of this new result.

W2: References Not Discussed.

Response: Thanks for the reminder. We will include FedAvg for the standard FL setting in the final version.

W3: Some typos.

Response: We will correct these typos in the final version.

Q1: Bit-width change for other models on other datasets.

Response: Yes, there are no significant differences among clients but there is a difference between activations and gradients for other models on other datasets. In the table below, we provide the bit-width change process for training ViT-S on CIFAR-10 (200 epochs) and LeNet on CIFAR-100 (2000 epochs) as further evidence. The numbers in brackets indicate the epochs of bit-width change.

Q2: Regarding quantization error as Gaussian noise.

Response: First of all, we emphasize that we no longer assume Gaussian noise for quantization error to achieve the convergence result. We assume that the expectation and variance of the noise distribution on each client exist are well defined and the expectation can be viewed as zero. We further perform experiments with ResNet-18 on CIFAR-10 to support the assumption, where the distribution of quantization noise is concentrated around zero.

In the previous result, we made the assumption as in [B1]-[B4] where the Gaussian distribution is assumed to find the optimal clipping value to reduce the quantization loss.

We will revise the assumption and clarify claims for notations in the final version.

[B1] Ron Banner et al. Scalable methods for 8-bit training of neural networks. NeurIPS 2018.

[B2] Ruizhou Ding et al. Regularizing activation distribution for training binarized deep networks. CVPR 2019.

[B3] Zhezhi He and Deliang Fan. Simultaneously optimizing weight and quantizer of ternary neural network using truncated gaussian approximation. CVPR 2019.

[B4] Xishan Zhang et al. Fixed-point back-propagation training. CVPR 2020.

Thank you for your valuable comments and below are our detailed responses to the raised weaknesses and questions.

W1: Experiments on diverse tasks and datasets.

Response: We have performed extensive experiments on diverse tasks and datasets to demonstrate the effectiveness of our LBI-FL. We use the RL agent trained for image classification on CIFAR-10 without retraining, which proves the generality of our RL agent.

i) Image segmentation with U-Net on DSB2018. We adopt the image size of 96 $\times$ 96. The table below shows that, compared with UI8 training, we achieve acceptable performance of about 0.6% loss in Dice similarity coefficient (DSC) with over 50% reduction of BitOPs.

ii) Image classification with ResNet-18 on Tiny-ImageNet. The image size is 64$\times$64. Training on Tiny-ImageNet is much more difficult. The table below shows that our LBI-FL achieves over 50% reduction of BitOPs with less than 0.1% accuracy loss.

W2: Baselines using layer-wise quantization.

Response: Most layer-wise quantization based methods (e.g., HAWQ v1-v3) are designed for inference only, and cannot be employed for FL. We compare our LBI-FL with AMPA [A1] that achieves layer-wise quantization based on the sensitivity measurement for centralized low-bit training. We train ResNet-18, MobileNet-V2, and ViT-S on CIFAR-10 using the two methods. The results below show that our LBI-FL achieves higher compression ratios and better performance.

[A1] Li Ding et al. AMPA: Adaptive mixed precision allocation for low-bit integer training. ICML 2024.

W3: Discussion about works on energy-latency tradeoffs with quantization in FL.

Response: Thanks for the suggestion. We will include relevant works [A2][A3] in the Related Work section. Optimizing energy consumption is indeed essential for practical deployment of FL systems. For instance, [A2] derives the time and energy consumption models for FL and proposes a iterative algorithm to allocate resources. [A3] proposes an optimization framework that minimizes the total energy consumption of local computation and wireless transmission by adaptively selecting the quantization level. These methods are complementary to our LBI-FL that simultaneously reduces computational and communication overheads, and thereby improves energy efficiency with reduced latency in federated learning.

[A2] Yang, Zhaohui, et al. "Energy efficient federated learning over wireless communication networks." IEEE TWC 2020.

[A3] Ouiame Marnissi, Hajar El Hammouti, and El Houcine Bergou. Adaptive sparsification and quantization for enhanced energy efficiency in federated learning.IEEE OJCOMS 2024.

Q1: Overhead of RL agent's decision-making.

Response: The overhead of RL agent is very small compared to the computational savings from low-bit training.

i) The RL agent consists of two linear layers with only 1.92K parameters, and requires 1.97G BitOPs for making one decision (i.e., 1.92K Mac$\times$32$\times$32).

ii) The RL agent makes decision for every 5 epochs rather than each epoch.

We provide results on CIFAR-10 as examples. When training LeNet using 100 clients, each client has 500 training images and the RL agent makes a decision after training on 250 images on average (participation rate is 0.1). Its inference cost is only 1.1% that for UI4 training and 0.275% for UI8 training. When training ResNet-18 using 10 clients, each client has 5000 training images. The inference cost of the RL agent is 0.0044% and 0.0011% those for UI4 and UI8 training.