On-demand Quantization for Green Federated Generative Diffusion in Mobile Edge Networks
Abstract
Generative Artificial Intelligence (GAI) shows remarkable productivity and creativity in Mobile Edge Networks, such as the metaverse and the Industrial Internet of Things. Federated learning is a promising technique for effectively training GAI models in mobile edge networks due to its data distribution. However, there is a notable issue with communication consumption when training large GAI models like generative diffusion models in mobile edge networks. Additionally, the substantial energy consumption associated with training diffusion-based models, along with the limited resources of edge devices and complexities of network environments, pose challenges for improving the training efficiency of GAI models. To address this challenge, we propose an on-demand quantized energy-efficient federated diffusion approach for mobile edge networks. Specifically, we first design a dynamic quantized federated diffusion training scheme considering various demands from the edge devices. Then, we study an energy efficiency problem based on specific quantization requirements. Numerical results show that our proposed method significantly reduces system energy consumption and transmitted model size compared to both baseline federated diffusion and fixed quantized federated diffusion methods while effectively maintaining reasonable quality and diversity of generated data.
Index Terms:
Federated Diffusion, Energy Efficient, Generative AI, Generative Diffusion, On-demand Quantization.I Introduction
As the carrier of content flow, mobile edge networks become essential fundamentals of next-generation applications like Metaverse [1] and the Industrial Internet of Things. The generative models like GAN [2] have demonstrated excellent performance in trajectory prediction [3], education [4], healthcare [5], and other scenarios involving the Internet of Things and the Internet of Vehicles. Therefore, more creative generative diffusion models are expected to be deployed in mobile edge networks for next-generation application scenarios such as the 6G communication networks [6] and vehicular metaverses [7], [8]. Towards deploying the generative diffusion models in mobile edge networks, distributed training schemes based on federated learning called federated diffusions [9, 10]are proposed. These innovative models enable mobile edge networks to achieve higher productivity and efficiency in next-generation application scenarios.
During the training phase of federated diffusions, the model needs to be transferred between the server and edge devices at each training step to update the global model [11]. This is not a problem when training traditional AI models using federated learning due to their small number of model parameters. However, generative diffusion models are usually large, leading to significant energy expenditure during the federated training process [12]. Therefore, reducing training energy costs is crucial for improving overall operational efficiency in mobile edge networks [13].
Recent research on federated diffusions has primarily focused on improving their task performance [9]. This includes endeavors to elevate the quality and diversity of generated content. However, there is still limited depth and scope in studies that aim to optimize the overall training cost of these models. The authors in [14], [15] studied the problem of energy-efficient resource allocation of FL over wireless communication networks. They derived the energy consumption models for FL based on the convergence rate analysis. The authors in [16] explored post-training quantization techniques for diffusion models, allowing direct quantization into 8 bits without significant performance degradation, with no substantial decline in performance. However, existing works fail to take into account the substantial training costs of diffusion models or the trade-off between performance and efficiency in the context of complex generative diffusion models [17].
To explore the deployment of a green generative diffusion model in mobile edge networks, we propose a dynamic quantization scheme for transmitting models during federated diffusion training. Firstly, we compress the diffusion models using a quantization scheme before transmission. We then study an energy consumption optimization problem and its solution. The performance of our proposed scheme is evaluated through simulations on the DDPM [17] model. Our main contributions can be summarized as follows:
-
•
We design a new and environmentally friendly federated generative diffusion framework that utilizes a dynamic method for parameter quantization and training in mobile edge networks.
-
•
We formulate an optimization problem for resource allocation in dynamic quantized federated diffusion, aiming to minimize total energy consumption while maintaining commendable performance.
-
•
Numerical results demonstrate the effectiveness of our proposed method compared to other baseline methods, particularly in terms of energy efficiency and sample quality.
The structure of the paper is organized as follows. The system model and the proposed on-demand quantized federated diffusion framework are introduced in section II. Next, we study the energy efficiency optimization problem in section III. Finally, We show the simulation results in section IV and discuss the conclusion and future work in section V.
II System Model
As shown in Fig. 1, we consider a mobile edge network scenario where a central server and edge devices collaborate to train a diffusion model using federated learning. Given the inherent characteristics of large diffusion models, training them in federated learning scenarios can be exceptionally energy-intensive. A promising model compression method named stochastic quantization [18] is implemented prior to the transmission of model parameters from edge devices to the edge server for aggregation, this is done to mitigate the transmission costs of the training process. Additionally, we take into account the variable quantization level needs of edge devices, ensuring the flexibility of quantization to accommodate different device requirements. Furthermore, considering the heterogeneous nature of edge devices and their varying resource capacities, an energy optimization problem is formulated to further minimize the energy consumption during federated diffusion training. After the efficient training is done, the server could utilize the final global diffusion model for efficient and high-quality content generation. The learning process for each round of iteration is as follows:
-
•
Step 1: Given different quantization requirements, the central server determines the optimal strategy for each edge device to balance computing and communication resources based on the resource status of different devices.
-
•
Step 2: The edge devices then perform a local diffusion computation and transmission according to the optimal strategy.
-
•
Step 3: After receiving all local diffusions from edge devices, the central server uses an aggregation scheme (such as Fedavg [19]) to unite the local diffusion into a new global diffusion and send it back to the edge devices for next round of training.
Hereinafter, we introduce the concept of quantization, which is a promising method for compressing neural networks. Stochastic quantization can be efficiently used in the federated learning process to significantly reduce energy consumption while maintaining minimal impact on model performance. To minimize the cost of transmitting a comparatively large model, as in Fig. 1, we propose quantizing the local diffusion model before uploading it to the server, since the resources of edge devices are often limited. To train the federated diffusion model with a quantization scheme, we first define the stochastic quantization function as . Given the local diffusion weight , the quantized weight can be expressed as . Let denotes the absolute value of element in , The stochastic quantization function is defined as
(1) |
Here, is the scale factor and denotes the sign function which represents the sign of . Moreover, is the quantization interval such that for any there exists . With the given quantization level , can be calculated as
(2) |
Where and represent the maximum and minimum value of non-zero element , respectively.
Next, we formulate the computation and communication models of the proposed scheme. Let represent the computation frequency of local client , and denotes the data size of the local dataset. The computation time for training diffusion model is expressed by
(3) |
where and denote the local iteration times in each communication round and the workload of local diffusion training, respectively. Following that, given the energy coefficient , the energy consumption of client is estimated by
(4) |
In the distributed diffusion setting, each edge device uploads the local diffusion model in order to generate a better global model. Moreover, the local diffusion model is quantized as for efficiency improvement. To this end, we adopt a frequency division multiple access (FDMA) transmission scheme [20] for quantized local diffusion transmission. Therefore, with the transmission power , the uplink transmission rate of client is deduced by
(5) |
Here, and denote the bandwidth and noise power-spectral-density each while corresponds to the distance between the client and server. Meanwhile, and represent the Rayleigh channel coefficient and pathloss exponent, separately. Subsequently, given the updated model size , and quantization level , the time spent by client to transmit the local model to the server is
(6) |
Thus, the corresponding energy consumption is calculated by
(7) |
III Energy Efficiency Optimization
Before the discussion of the energy efficiency problem of the proposed methods started, similar to the work in [15], we first present the following assumption and theorem. Where represents the unique bound demand of various edge devices, a smaller indicates that the resource of edge device is relatively insufficient, which leads to a lower quantization level strategy:
Assumption 1.
the expectation of the square norm of the local weight uploaded by edge devices is bounded: for any uploaded weight, .
Theorem 1.
Based on Assumptions 1, the square of local weight quantization error is bounded by:
(8) |
Through this theorem, we can easily obtain the corresponding quantization levels for each heterogeneous edge device at different demands, thereby further constructing our energy consumption optimization model, which we will discuss in detail in the next section.
III-A Problem Formulation
As a consequence of variations in resource capabilities among edge devices, variability exists in the required quantization level demands. In simpler terms, each device has its distinct upper bound for quantization error. Leveraging Theorem 1, we establish the energy minimization problem within the confines of this quantization error constraint as follows:
(III-A) | |||||
subject to: | (9a) | ||||
(9b) | |||||
(9c) | |||||
(9d) |
With the unique given by the different bound requirements of local edge devices, we can always obtain the optimal solution for the optimization problem. As a result, we simplify P1 as
(III-A) | ||||||
subject to: | (10a) | |||||
(10b) | ||||||
(10c) |
Following that, we transform P2 into a more tractable form by introducing two intermediate variables and . Moreover, we let and represent the weight factors of maximum time budget for client such that
(11) |
Here, the lower bound of and can be easily acquired given the optimal
(12) |
Furthermore, the total energy consumption of client during the fine-tuning process can be rewritten in the following form
(13) |
Thus, we convert problem P3 into the following form
(III-A) | ||||
subject to: | (14a) | |||
(14b) | ||||
(14c) |
Through this basic form, we can readily acquire the numerical solution for the original problem. In the next subsection, we will present the solution to address the current matter.
III-B Solution
It can be easily proved that problem P3 is a convex problem, which can be effectively solved by applying the Karush-Kuhn-Tucker (KKT) conditions [21]. With the optimal energy optimization solution, we can decide the final resource allocation scheme for the federated diffusion. The Lagrange function of P3 is as follows:
(15) |
Here, is the equality Lagrange multiplier associated with equality constraint (14a), while and denote the inequality Lagrange multiplier for constraints (14b) and (14c), respectively. In order to accomplish optimality for problem P3, we derive the necessary equations from the Lagrange function as follows:
(16) |
Based on Eqns 16, there exist two cases that satisfy constrain (14b) concerning variable . If , the optimal solution of is obtained by
(17) |
Otherwise, we always have . Similarly to , when , the optimal is acquired the same way. Given the equality Lagrange multiplier , we have
(18) |
Where is the zero point of function . In general, the optimal solution of and can be acquired by
(19) |
It’s worth mentioning that seeking the problem’s optimal solution directly can be rather intricate, which is why we employed binary search to find the optimal strategy for the Lagrange multiplier . Utilizing the most favorable Lagrange multiplier value, the optimal approach for variables and are computed based on (19). To be specific, given the searching range of and error tolerance , the optimal is obtained with the constraint (14a). Additional and more detailed information is provided in Algorithm 1. Finally, the overall algorithm of the proposed method is shown as algorithm 2.
IV Numerical Results
IV-A Simulation Settings
To simulate the practical case of federated diffusion in mobile edge networks, we fine-tuned the pre-trained DDPM [17] on CIFAR10 [22] using the GTSRB [23] dataset. The dataset is divided into 10 subsets for 10 edge devices to perform federated learning. We fine-tuned the federated model with 1000 epochs for performance evaluations. The sample steps are set as 1000 for image generation. For computation and communication hyper-parameters {} and {}, the default settings are {MCycles} and {MHzWmdbm/MHzM}.
IV-B Performance Evaluations
Fig. 2 illustrates the performance and energy consumption of the proposed algorithm. We employed the Fréchet Inception Distance (FID) [24] as an evaluative metric for assessing the quality of the images generated by the model. A lower FID value indicates a higher degree of similarity between the distribution of the generated dataset and that of the original dataset. To enhance the precision of our evaluation concerning the quality of the generated dataset, we ensured that the number of generated datasets equaled the number of original images. Remarkably, our proposed methodology takes into consideration distinct quantization error constraints customized for heterogeneous edge devices. Subsequently, an energy minimization problem is optimized. The range of quantization levels spans from 6 bits to 8 bits, a range normally associated with a substantial reduction in energy consumption while concurrently upholding commendable performance. We conduct a comparative analysis between the baseline method, Fedavg, in addition to the fixed quantization methods employing 6-bits, 7-bits, and 8-bits quantization levels. It is evident from our results that our approach surpasses the more economical 8-bits quantization scheme in terms of both performance and cost-efficiency. It is essential to note that the compared methods did not specifically optimize for energy consumption, utilizing only time budget for computation and another for communication by default.
Fig. 3 illustrates the successful convergence achieved by our proposed binary search algorithm in addressing the energy optimization problem we have established. Specifically, provided various quantization level requirements, the optimal solution for minimizing the energy cost is determined after about 20 searching iterations. Furthermore, a study was conducted to examine the impacts of two different hyperparameters, the time budget and the distance of communication. It can be clearly seen that the proposed method can converge well under different settings. As the allotted time budget reduces, the requisite energy consumption by the system escalates. In parallel, with an augmentation in the distance of communication between edge devices and central servers, there is a concurrent amplification in energy expenditure. Fig. 4 presents a comparative analysis of our method alongside other approaches across various time budget ranges. As observed, with an increase in the allotted time budget ranging from 13s to 18s, the system’s energy costs diminish, and our solution consistently outperforms the baseline approach. This indicates that our method can adapt to parameter settings within certain ranges.
V Conclusion and Future Work
In this paper, we first design a dynamic quantized federated diffusion training considering each edge device’s demand. Subsequently, our study turns towards addressing the challenge of energy efficiency, taking into account the unique constraint imposed by quantization demand. Our simulation results demonstrate that our proposed method outperforms both the baseline federated diffusion approach and fixed quantized federated diffusion in substantially reducing system energy consumption and transmitted model size. Remarkably, this reduction is achieved without compromising the reasonable quality and diversity of the generated data, underscoring the effectiveness of our approach.
To achieve the benefits of efficient federated generative diffusion, there still exist several open and challenging issues. For distributed diffusion models, the matter of proficient sampling remains an unsolved problem, primarily due to the distinctive characteristics inherent to diffusion itself. Diverging from conventional AI models, the interference phase of diffusion entails a substantial energy outlay, particularly in the denoising sampling steps. This heightened energy consumption may be deemed unmanageable for certain edge devices. Consequently, further study is imperative to enhance the efficiency of sampling within the context of diffusion, particularly in the domain of distributed edge intelligence scenarios.
References
- [1] B. Mao, Y. Liu, J. Liu, and N. Kato, “AI-assisted edge caching for metaverse of connected and automated vehicles: Proposal, challenges, and future perspectives,” IEEE Vehicular Technology Magazine, vol. 18, no. 4, pp. 66–74, 2023.
- [2] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio, “Generative adversarial networks,” Communications of the ACM, vol. 63, no. 11, pp. 139–144, 2020.
- [3] S. Choi, J. Kim, and H. Yeo, “Trajgail: Generating urban vehicle trajectories using generative adversarial imitation learning,” Transportation Research Part C: Emerging Technologies, vol. 128, p. 103091, 2021.
- [4] G. Cooper, “Examining science education in chatgpt: An exploratory study of generative artificial intelligence,” Journal of Science Education and Technology, vol. 32, no. 3, pp. 444–452, 2023.
- [5] J. Kang, J. Wen, D. Ye, B. Lai, T. Wu, Z. Xiong, J. Nie, D. Niyato, Y. Zhang, and S. Xie, “Blockchain-empowered federated learning for healthcare metaverses: User-centric incentive mechanism with optimal data freshness,” IEEE Transactions on Cognitive Communications and Networking, 2023.
- [6] B. Mao, X. Zhou, J. Liu, and N. Kato, “Digital twin satellite networks towards 6G: Motivations, challenges, and future perspectives,” IEEE Network, pp. 1–1, 2023.
- [7] J. Kang, J. He, H. Du, Z. Xiong, Z. Yang, X. Huang, and S. Xie, “Adversarial attacks and defenses for semantic communication in vehicular metaverses,” IEEE Wireless Communications, vol. 30, no. 4, pp. 48–55, 2023.
- [8] X. Luo, J. Wen, J. Kang, J. Nie, Z. Xiong, Y. Zhang, Z. Yang, and S. Xie, “Privacy attacks and defenses for digital twin migrations in vehicular metaverses,” IEEE Network, 2023.
- [9] F. V. S. Jothiraj and A. Mashhadi, “Phoenix: A federated generative diffusion model,” arXiv preprint arXiv:2306.04098, 2023.
- [10] M. de Goede, “Training diffusion models with federated learning: A communication-efficient model for cross-silo federated image generation,” 2023.
- [11] B. Mao, J. Liu, Y. Wu, and N. Kato, “Security and privacy on 6g network edge: A survey,” IEEE Communications Surveys & Tutorials, vol. 25, no. 2, pp. 1095–1127, 2023.
- [12] X. Huang, P. Li, H. Du, J. Kang, D. Niyato, D. I. Kim, and Y. Wu, “Federated learning-empowered AI-generated content in wireless networks,” arXiv preprint arXiv:2307.07146, 2023.
- [13] J. Wen, J. Kang, M. Xu, H. Du, Z. Xiong, Y. Zhang, and D. Niyato, “Freshness-aware incentive mechanism for mobile AI-generated content (aigc) networks,” in 2023 IEEE/CIC International Conference on Communications in China (ICCC), pp. 1–6, 2023.
- [14] Z. Yang, M. Chen, W. Saad, C. S. Hong, and M. Shikh-Bahaei, “Energy efficient federated learning over wireless communication networks,” IEEE Transactions on Wireless Communications, vol. 20, no. 3, pp. 1935–1949, 2020.
- [15] P. Li, G. Cheng, X. Huang, J. Kang, R. Yu, Y. Wu, M. Pan, and D. Niyato, “Snowball: Energy efficient and accurate federated learning with coarse-to-fine compression over heterogeneous wireless edge devices,” IEEE Transactions on Wireless Communications, 2023.
- [16] X. Meng and Y. Kabashima, “Quantized compressed sensing with score-based generative models,” in The Eleventh International Conference on Learning Representations, 2022.
- [17] J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” Advances in neural information processing systems, vol. 33, pp. 6840–6851, 2020.
- [18] R. Chen, L. Li, K. Xue, C. Zhang, M. Pan, and Y. Fang, “Energy efficient federated learning over heterogeneous mobile devices via joint design of weight quantization and wireless transmission,” IEEE Transactions on Mobile Computing, 2022.
- [19] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-efficient learning of deep networks from decentralized data,” in Artificial intelligence and statistics, pp. 1273–1282, PMLR, 2017.
- [20] H. G. Myung, J. Lim, and D. J. Goodman, “Single carrier fdma for uplink wireless transmission,” IEEE vehicular technology magazine, vol. 1, no. 3, pp. 30–38, 2006.
- [21] S. P. Boyd and L. Vandenberghe, Convex optimization. Cambridge university press, 2004.
- [22] A. Krizhevsky, G. Hinton, et al., “Learning multiple layers of features from tiny images,” 2009.
- [23] S. Houben, J. Stallkamp, J. Salmen, M. Schlipsing, and C. Igel, “Detection of traffic signs in real-world images: The German Traffic Sign Detection Benchmark,” in International Joint Conference on Neural Networks, no. 1288, 2013.
- [24] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, “High-resolution image synthesis with latent diffusion models,” in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684–10695, 2022.