# Filling the Missing: Exploring Generative AI for Enhanced Federated Learning
over Heterogeneous Mobile Edge Devices

Peichun Li, Hanwen Zhang, Yuan Wu, , Liping Qian, , Rong Yu, , Dusit Niyato, ,
and Xuemin (Sherman) Shen  P. Li, H. Zhang, and Y. Wu are with the State Key
Laboratory of Internet of Things for Smart City, University of Macau, Macau,
China, and also with Department of Computer and Information Science,
University of Macau, Macau, China. (e-mail<EMAIL_ADDRESS><EMAIL_ADDRESS>yuanwu@um.edu.mo). L. Qian is with College of
Information Engineering, Zhejiang University of Technology, Hangzhou 310023,
China (email:lpqian@zjut.edu.cn). R. Yu is with School of Automation,
Guangdong University of Technology, Guangzhou, China, and also with Guangdong
Key Laboratory of IoT Information Technology, Guangzhou, China. (e-mail:
yurong@ieee.org). D. Niyato is with the School of Computer Science and
Engineering, Nanyang Technological University, Singapore, Block N4-02a-32,
Nanyang Avenue, Singapore. (e-mail: dniyato@ntu.edu.sg). X. Shen is with the
Department of Electrical and Computer Engineering, University of Waterloo,
Waterloo, ON N2L 3G1, Canada (e-mail: sshen@uwaterloo.ca). Yuan Wu is the
corresponding author.

###### Abstract

Distributed Artificial Intelligence (AI) model training over mobile edge
networks encounters significant challenges due to the data and resource
heterogeneity of edge devices. The former hampers the convergence rate of the
global model, while the latter diminishes the devices’ resource utilization
efficiency. In this paper, we propose a generative AI-empowered federated
learning to address these challenges by leveraging the idea of FIlling the
MIssing (FIMI) portion of local data. Specifically, FIMI can be considered as
a resource-aware data augmentation method that effectively mitigates the data
heterogeneity while ensuring efficient FL training. We first quantify the
relationship between the training data amount and the learning performance. We
then study the FIMI optimization problem with the objective of minimizing the
device-side overall energy consumption subject to required learning
performance constraints. The decomposition-based analysis and the cross-
entropy searching method are leveraged to derive the solution, where each
device is assigned suitable AI-synthesized data and resource utilization
policy. Experiment results demonstrate that FIMI can save up to 50% of the
device-side energy to achieve the target global test accuracy in comparison
with the existing methods. Meanwhile, FIMI can significantly enhance the
converged global accuracy under the non-independently-and-identically
distribution (non-IID) data.

###### Index Terms:

Federated learning, generative AI, data compensation, resource management

##  1 Introduction

Generative artificial intelligence (AI) has advanced significantly in content
creation. Notably, vision models such as stable diffusion and Imagen have
demonstrated their impressive capability in generating photo-realistic images
[1, 2, 3]. These advancements have sparked new possibilities that generative
models can play a significant role in enhancing the centralized training
process. Specifically, the synthesized images produced by generative models
can be effectively utilized to improve the performance of classification
models [4]. In the realm of federated learning (FL), where isolated devices
with limited storage collaborate in training a global model, data and resource
heterogeneity have emerged as crucial challenges [5, 6, 7, 8, 9]. These
factors lead to a slower convergence rate during the global training process.
Consequently, an intriguing and pertinent question arises, i.e., can
generative AI be harnessed to address these issues and improve the performance
of FL ultimately?

Figure 1: The impacts of data distribution and data amount on the training
performance.

In this paper, we first investigate how the distribution of federated local
data impacts the training performance of FL. Specifically, we focus on the
VGG9 model [10] with 20 Nvidia Jetson Nano devices on CIFAR100 dataset [11].
We regulate the local dataset into different degrees of Dirichlet
distribution, leading to varying levels of non-independently-and-identically
distribution (non-IID) data [12, 13]. As shown in the top subplot of Fig. 1,
we observe that a well-distributed Dir(0.9) significantly improves convergence
accuracy compared to Dir(0.3). Here, Dir($z$) means the Dirichlet distribution
with parameter $z$, and a larger $z$ means a more homogeneous distribution of
the samples (i.e., less non-IID distribution). The observation indicates that
augmenting the data to eliminate the non-IID feature of local data holds
promise for enhancing the learning performance of FL. However, the augmented
training data introduces extra computation costs for edge devices. As shown in
the bottom subplot of Fig. 1, doubling the amount of data results in about a
fourfold increase in energy consumption under a fixed training latency.

The observations above offer insights for developing an efficient and
effective generative AI-enhanced FL system. To address data heterogeneity, it
is recommended to leverage generative AI to compensate for the missing portion
of local data. This strategy can bridge the gaps in data distribution across
devices and improve convergence accuracy. However, it is crucial to take into
account the locally available resources and adapt the amount of synthesized
data on different devices accordingly. The majority of current literature on
generative AI-enhanced training methods tends to prioritize improving the
model accuracy in a centralized training manner, disregarding the computation
cost induced by the augmented training dataset [14, 4]. These approaches may
shorten the battery lifetime and potentially degrade the user experience when
integrated into the FL system directly. Therefore, we propose a resource-aware
data augmentation framework, called FIlling the Missing (FIMI), which is
designed for generative AI-enhanced FL systems.

The key idea behind FIMI is to leverage the concept of filling the missing
portion of data for each device. Our goal is to improve the learning
performance of FL by utilizing a pre-trained generative AI model to synthesize
customized local data while ensuring the training efficiency of edge devices.
To this end, we first outline the generic workflow of our generative AI-
empowered FL system. We then empirically measure the relationship between the
amount of training data and model learning error to quantitatively analyze the
learning performance of FL. Furthermore, we propose an optimization problem
for minimizing the device-side energy subject to the learning performance
constraints. To solve this formulated problem, we introduce auxiliary
variables and decompose the main problem into two sub-problems along with a
top-layer searching problem. To find solutions for the sub-problems, the
convex optimization methods are leveraged to obtain the optimal solution based
on the theoretical analysis. For the top-layer problem, the continuous cross-
entropy algorithm is employed to efficiently search the solution. After that,
a data entropy maximization problem for mitigating the non-IID feature of
local data is conducted to obtain the category-wise training sample amount on
each device. In this way, each wireless edge device is assigned a personalized
solution, where the data synthesis strategy and resource utilization policy
are optimized to minimize the device-side energy consumption while improving
the learning performance.

Our main contributions are summarized as follows.

  * •

We propose a generative AI-enhanced FL framework called FIMI to leverage the
pre-trained generative models for filling the missing portion of the local
data. The calibrated distribution of local data improves the convergence rate
of FL training.

  * •

By combining the theoretical and empirical analysis, we quantitatively
establish the relationship between the training data amount and the learning
performance of FL.

  * •

We formulate a joint optimization of data synthesis strategy and resource
utilization policy, with the objective of minimizing the energy consumption of
FIMI over wireless networks. We also propose an efficient algorithm to solve
the formulated problem and determine the corresponding solution for each
device.

  * •

Extensive experiments demonstrate that our proposed FIMI can outperform the
existing data augmentation FL methods in terms of the device-side energy
utilization and the accuracy of the global model after convergence.

The remainder of this paper is organized as follows. Section 2 describes
related studies. In Section 3, we illustrate the system model of FIMI. The
theoretical analysis and the corresponding solution are provided in Section 4.
The experiment evaluations are presented in Section 5. We finally conclude the
paper and discuss the future directions in Section 6.

Figure 2: left: FIMI over heterogeneous edge devices. middle: the amounts of
synthesized data are personalized for different devices to match their locally
available resources. right: each device optimizes its category-wise
synthesized data amount to mitigate the locally non-IID feature of data.

##  2 Related Work

 _1) Generative AI over edge networks_. Generative AI on edge networks has
attracted growing interest for its potential to provide low-latency AI
generative content (AIGC) services [3]. Some studies focus on achieving
efficient AIGC service by optimizing resource utilization [15] or by improving
the freshness of the age of information [16]. In addition, diffusion model-
integrated reinforcement learning presents a novel approach for solving
complex optimization problems [17]. Recently, the study in [18] incorporates
the FL framework to enhance the training and inference process of the AIGC
model. However, the potential benefits of generative AI for FL training remain
unexplored.

_2) Data augmentation methods for FL_. Data augmentation methods provide an
effective approach to address the issue of non-IID data distribution in FL.
The early study in [19] shows that exchanging a small proportion of local data
can alleviate the non-IID issue. However, this approach may raise privacy
concerns. In addition, semi-supervised FL methods utilize the unlabelled data
to enhance learning performance [20, 21]. Recently, generative adversarial
networks (GAN) have emerged as a privacy-friendly technique for data
compensation in FL [22, 23]. However, the GAN-based methods may fail to
produce photo-realistic data, and most of the existing methods overlook the
extra training cost induced by the augmented dataset.

_3) Efficient training methods for FL_. Efficient training methods for FL have
been widely studied in recent years. One popular approach is the gradient
compression, which reduces the amount of data by compressing the gradients
[24, 25]. Another approach is topology-aware methods, such as hierarchical
aggregation [26, 27] and device-to-device communication [28], which consider
the network topology to optimize the communication pattern and reduce the
communication overhead. Furthermore, energy-aware methods, such as wireless
power transfer [29] and resource management [30, 31, 32], offer a green and
sustainable framework for FL training. However, most of the existing studies
do not consider the heterogeneity of data distribution among different
participating devices, which can result in biased model updates.

##  3 System Model

In this section, we first introduce the structure of FIMI, and then present
the mathematical modeling and optimization problems for the data augmentation-
based FL training.

###  3.1 Structure of FIMI

We consider an FL scenario comprising a set of devices labeled as
$\\{1,2,\ldots,I\\}$. Each device, denoted as $i$, holds an original local
dataset ${\cal D}_{i}^{\text{loc}}$, where $|{\cal
D}_{i}^{\text{loc}}|=D_{i}^{\text{loc}}$ represents the number of local
training samples. Our goal is to leverage a pre-trained generative model to
create additional synthesized data ${\cal D}_{i}^{\text{gen}}$ for each device
$i$, with $|{\cal D}_{i}^{\text{gen}}|=D_{i}^{\text{gen}}$. In traditional FL
systems, device $i$ solely relies on its local dataset ${\cal
D}_{i}^{\text{loc}}$ for the local model update. Alternatively, the proposed
FIMI adopts a mixed local dataset ${\cal D}_{i}^{\text{mix}}={\cal
D}_{i}^{\text{loc}}\cup{\cal D}_{i}^{\text{gen}}$ to accelerate the
convergence.

We focus on the FL task for image classification across $C$ categories (or
label types). The data synthesis strategy of device $i$ can be represented as
$\\{d_{i,c}^{\text{gen}}\\}_{\forall c},(c=1,\ldots,C)$. This strategy should
meet the constraint
$\sum\nolimits_{c}d_{i,c}^{\text{gen}}=D_{i}^{\text{gen}}$, where
$d_{i,c}^{\text{gen}}$ denotes the amount of synthesized data for the $c$-th
category on device $i$. Before the beginning of FL training, a collaborative
optimization is undertaken between the server and edge devices to obtain the
data synthesis strategy. Following this, the server will employ the
corresponding label as input for the diffusion model to generate the required
data. After the synthesized data is distributed to each device, the FL
training with the mixed dataset is initiated. As shown in Fig. 2, the FIMI
framework operates as follows.

  * •

_(S1) Strategy optimization_. The server conducts a centralized optimization,
as shown in the middle subplot of Fig. 2, to determine the synthesized data
amount for each device. Each device then calculates the _category-wise_
synthesized data amount $\\{d_{i,c}^{\text{gen}}\\}_{\forall
c},(c=1,\ldots,C)$, as shown in the right subplot of Fig. 2, to rebalance the
data distribution and mitigate the non-IID feature of local data.

  * •

_(S2) Data synthesis_. Each device sends its data synthesis request,
specifying the category-wise amount of synthesized data, to the server. The
server infers the generative model and distributes the required synthesized
data to the corresponding devices in parallel.

  * •

_(S3) Training with mixed data_. The server distributes the global model to
each device. Then, the local device performs the local training with the mixed
local dataset, and then uploads its model update back to the server.

  * •

_(S4) Parameter aggregation_. The server collects and aggregates the local
updates from all devices to renew the global model. The training process jumps
to (S3) until the predetermined maximum global iterations are reached.

Since the servers are usually equipped with sufficient computational resources
and resilient power supply, we do not account for the latency and energy
consumption during the data synthesis process. Empirical computation and
communication overheads for data synthesis and distribution will be
illustrated in Section 5.1.

###  3.2 Learning Performance Analysis

In this subsection, we quantitatively analyze the relationship between the
number of local samples and the overall training performance. We first
establish the mathematical formulation that links the amount of local data to
the individual learning performance. We then determine the unknown parameters
in the formulation through the empirical fitting. Finally, we investigate how
the isolated local data from multiple devices can influence the global
learning performance.

####  3.2.1 Local Data vs. Local Model Performance

Let $\delta_{i}$ denote the local learning error of device $i$ when training
on mixed data ${\cal D}_{i}^{\text{mix}}$. According to the previous studies
in [33, 34, 35], the lower bound of the achievable local error can be modeled
as the power-law function with respect to the training sample amount of
$|{\cal D}_{i}^{\text{mix}}|$. The relationship between
$D_{i}^{\text{loc}}+D_{i}^{\text{gen}}$ and $\delta_{i}$ can be expressed as

$\delta_{i}=\alpha(D_{i}^{\text{loc}}+D_{i}^{\text{gen}})^{-\beta}-\gamma,$
(1)

where $\alpha,\beta$, and $\gamma$ are three positive hyper-parameters that
can be determined through experiments.

####  3.2.2 Parameter Fitting on Proxy Task

We next focus on determining the values of $\\{\alpha,\beta,\gamma\\}$.
However, since the local dataset is not visible to the server, direct
parameter fitting on the target learning task by the server can raise privacy
concerns. We propose an alternative approach by determining these parameters
using a public dataset. Specifically, we employ the classification task on the
CINIC10 dataset as the _proxy_ task for parameter fitting [36]. By assessing
the learning error of well-trained models on a variety of data amounts, we
obtain the numerical results as shown in Fig. 3. It is noticed that this _one-
time_ parameter fitting process can be conducted offline by the server, and
the results of fitting can be generalized to _many_ learning tasks with
different datasets.

####  3.2.3 Global Data vs. Global Model Performance

After estimating the local learning performance $\\{\delta_{i}\\}_{\forall
i}$, the average local training error $\overline{\delta}$ can be calculated as

$\overline{\delta}={\frac{1}{I}}\sum\limits_{i=1}^{I}{\delta_{i}}.$ (2)

Based on the theoretical analysis in [37, 38], when training with the average
local training error $\overline{\delta}$, the number of global iterations $N$
to reach a pre-determined global learning error $\Delta$ is bounded by

$N=\frac{{\zeta\ln(1/\Delta)}}{{1-\overline{\delta}}},$ (3)

where $\zeta$ is a positive constant parameter. Furthermore, the relationship
between the average local learning error $\overline{\delta}$ and the global
learning error $\Delta$ can be captured by

$\Delta=\exp\Big{(}\frac{{N(\overline{\delta}-1)}}{\zeta}\Big{)}.$ (4)

By combing Eqns. (1), (2) and (4), we bridge the relationship between the
amount of local mixed dataset and the global learning performance.

Figure 3: Local learning error $\delta_{i}$ with respect to local training
data amount.

###  3.3 FIMI over Wireless Networks

####  3.3.1 Computation Model

Let $\omega$ denote the computation workload required for a single training
sample and $\tau$ represent the local epoch. According to [38, 39, 30], the
computation-related configurations of device $i$ can be profiled by two
parameters, including the computing frequency $f_{i}$ and the hardware energy
coefficient $\epsilon_{i}$. Thus, the single-round energy consumption for the
local model training at device $i$ can be estimated by

$E_{i}^{\text{cmp}}=\tau\epsilon_{i}{\omega}(D_{i}^{\text{loc}}+D_{i}^{\text{gen}})f^{2}_{i}.$
(5)

Furthermore, the required latency for a single-round local model training of
device $i$ is

$T^{\text{cmp}}_{i}=\frac{\tau{\omega}(D_{i}^{\text{loc}}+D_{i}^{\text{gen}})}{f_{i}}.$
(6)

####  3.3.2 Communication Model

We next focus on the device-side uplink parameter transmission. We consider
the frequency division multiple access (FDMA) scheme with a preset total
bandwidth of $B$. For the device $i$ allocated with a sub-bandwidth of
$b_{i}$, its achievable transmitting rate for sending the local update to the
server can be given by

$r_{i}=b_{i}\log_{2}\big{(}1+\frac{g_{i}P_{i}}{N_{0}b_{i}}\big{)},$ (7)

where $g_{i}$ is the device-to-server channel gain, $P_{i}$ is the
transmitting power, and $N_{0}$ is the power spectral density of background
Gaussian noise. Here, we utilize $S$ to denote the data size of the model
update. The single-round latency required for device $i$ to send its model
update to the server is estimated by

$T^{\text{com}}_{i}=\frac{S}{r_{i}}.$ (8)

Moreover, the single-round energy consumption for device $i$ to send its model
update to the server can be calculated by

$E^{\text{com}}_{i}=\frac{SP_{i}}{r_{i}}.$ (9)

####  3.3.3 Problem Formulation

Based on the above modeling, the energy consumption of all the devices for
one-round training can be expressed as

$E^{\text{round}}=\sum\limits_{i=1}^{I}E^{\text{cmp}}_{i}+\sum\limits_{i=1}^{I}E^{\text{com}}_{i}.$
(10)

Meanwhile, the system-wise learning latency for one-round training is
calculated by

$T^{\text{round}}=\max\limits_{i}\big{\\{}T^{\text{cmp}}_{i}+T^{\text{com}}_{i}\big{\\}}.$
(11)

To optimize the proposed FIMI, we investigate an energy minimization problem
with a pre-determined learning performance constraint. Specifically, given
multiple devices with heterogeneous resources in terms of computation,
communication, and data distribution, our objective is to optimize the optimal
data synthesis strategy (i.e., $\\{D_{i}^{\text{gen}}\\},i=1,\ldots,I$) and
resource utilization policy (i.e., $\\{b_{i},f_{i},P_{i}\\},i=1,\ldots,I$) for
each device. Mathematically, we aim at studying the following problem for each
global round.

$\displaystyle({\text{P1}})$
$\displaystyle~{}~{}~{}~{}~{}\min~{}~{}E^{\text{round}}$ (12) subject to:
$\displaystyle~{}~{}~{}~{}~{}~{}\Delta\leq\Delta^{\max},$ (12a)
$\displaystyle~{}~{}~{}~{}T^{\text{round}}\leq T^{\max},$ (12b)
$\displaystyle~{}~{}0\leq{D_{i}^{\text{gen}}}\leq D^{\max},\forall i,$ (12c)
$\displaystyle~{}~{}~{}0\leq{f_{i}}\leq f_{i}^{\max},\forall i,$ (12d)
$\displaystyle~{}~{}~{}~{}\sum\nolimits_{i=1}^{I}b_{i}\leq B,$ (12e)
$\displaystyle~{}~{}~{}0\leq P_{i}\leq P^{\max}_{i},\forall i,$ (12f)
variables:
$\displaystyle~{}~{}~{}~{}\\{D_{i}^{\text{gen}},f_{i},b_{i},P_{i}\\}_{\forall
i},$

where $\Delta^{\max},T^{\max}$, and $D^{\max}$ are the pre-determined maximum
allowable global error, the maximum latency for single round learning, and the
maximum amount of synthesized data for each device, respectively. Note that
the data synthesis procedure is conducted by the server with powerful
computing capabilities. Therefore, we solely focus on the device-side energy
consumption in Problem (P1).

###  3.4 Analysis and Decomposition

Before proposing the algorithms to solve Problem (P1) in Section 4, we present
the analysis and decomposition in this subsection. We first leverage two
lemmas to simplify the main problem. After that, we decompose the simplified
problem into three sub-problems.

To facilitate the analysis, we present the following lemmas.

######  Lemma 1.

The equality in Constraint (12a) always holds when the optimal data synthesis
strategy $\\{D_{i}^{\text{gen},\ast}\\}_{\forall i}$ is achieved.

  * Proof.

The lemma can be proved by showing contradiction. Given the optimal data
synthesis strategy $\\{D_{i}^{\text{gen},\ast}\\}_{\forall i}$, suppose that
the achieved global learning error is less than the maximum allowable error,
i.e., $\Delta(\\{D_{i}^{\text{gen},\ast}\\}_{\forall i})<\Delta^{\max}$. Based
on Eqns. (1) and (4), the global learning error $\Delta$ decreases with
$D_{i}^{\text{gen}}$. Therefore, a feasible solution can always be constructed
by reducing the synthesized data of device $i_{0}$ from
$D_{i_{0}}^{\text{gen},\ast}$ to $D_{i_{0}}^{\text{gen},\prime}$, such that
$\Delta(\\{D_{i}^{\text{gen},\ast}\\}_{\forall i,i\neq i_{0}}\cup
D_{i_{0}}^{\text{gen},\prime})=\Delta^{\max}$. As $E^{\text{round}}$ increases
with $D_{i}^{\text{gen}}$, the energy consumption of the new solution
$\\{D_{i}^{\text{gen},\ast}\\}_{\forall i,i\neq i_{0}}\cup
D_{i_{0}}^{\text{gen},\prime}$ will be lower than that of the optimal one.
This completes the proof. ∎

######  Lemma 2.

The equality in Constraint (12b) always holds under the optimal resource
utilization policy $\\{b_{i}^{\ast},f_{i}^{\ast},P_{i}^{\ast}\\}_{\forall i}$.

  * Proof.

The lemma can be proved by showing contradiction. Suppose that under the
optimal resource utilization policy
$\\{b_{i}^{\ast},f_{i}^{\ast},P_{i}^{\ast}\\}_{\forall i}$, the single-round
latency required is smaller than the maximum latency, i.e.,
$T^{\text{round}}<T^{\max}$. According to Eqn. (6), the training latency
increases with $f_{i}$. Therefore, a feasible solution can always be
constructed by reducing the computing frequency of device $i_{0}$ from
$f_{i_{0}}^{\ast}$ to $f_{i_{0}}^{\prime}$. Similar to the proof of Lemma 1,
the new solution will consume less energy than the optimal one. This completes
the proof. ∎

By combing Eqn. (4) and the above two lemmas, Problem (P1) can be simplified
as

$\displaystyle({\text{P2}})$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}\min~{}~{}E^{\text{round}}$ (13) subject
to: $\displaystyle\sum\limits_{i=1}^{I}\delta_{i}=I+\frac{\zeta
I}{N}\ln\big{(}{\Delta^{\max}}\big{)},$ (13a)
$\displaystyle~{}~{}~{}T^{\text{cmp}}_{i}+T^{\text{com}}_{i}=T^{\max},$ (13b)
$\displaystyle~{}~{}\text{Constraints
(\ref{eqn:p1-ctr-4})-(\ref{eqn:p1-ctr-6})},$ (13c) variables:
$\displaystyle~{}~{}~{}~{}~{}\\{D_{i}^{\text{gen}},f_{i},b_{i},P_{i}\\}_{\forall
i}.$

Moreover, by determining the latency for both local model training and local
update transmission for each device (i.e.,
$\\{T^{\text{cmp}}_{i},T^{\text{com}}_{i}\\}_{\forall i}$), Problem (P2) can
be decomposed into two distinct sub-problems as described in the following.

The first sub-problem, denoted as Problem (P3), aims at minimizing all
devices’ energy consumption for their local model training, which is expressed
as follows:

$\displaystyle({\text{P3}})$
$\displaystyle~{}~{}~{}\min\limits_{\\{D_{i}^{\text{gen}},f_{i}\\}_{\forall
i}}\sum\limits_{i=1}^{I}\tau\epsilon_{i}{\omega}(D_{i}^{\text{loc}}+D_{i}^{\text{gen}})f^{2}_{i}$
(14) subject to:
$\displaystyle~{}~{}~{}\frac{\tau{\omega}(D_{i}^{\text{loc}}+D_{i}^{\text{gen}})}{f_{i}}=T^{\text{cmp}}_{i},\forall
i,$ (14a) $\displaystyle\text{Constraints (\ref{eqn:p1-ctr-4}),
(\ref{eqn:p1-ctr-5}), and (\ref{eqn:p1-ctr-1-1})}.$ (14b)

The second sub-problem, denoted as Problem (P4), aims at minimizing all
devices’ energy consumption for uploading their locally trained models as
follows:

$\displaystyle({\text{P4}})$
$\displaystyle~{}~{}~{}~{}~{}~{}\min\limits_{\\{b_{i},P_{i}\\}_{\forall
i}}~{}\sum\limits_{i=1}^{I}P_{i}T_{i}^{\text{com}}$ (15) subject to:
$\displaystyle\frac{S}{{b_{i}\log_{2}\big{(}1+\frac{g_{i}P_{i}}{N_{0}b_{i}}\big{)}}}=T_{i}^{\text{com}},\forall
i,$ (15a) $\displaystyle~{}~{}\text{Constraints (\ref{eqn:p1-ctr-3}) and
(\ref{eqn:p1-ctr-6})}.$ (15b)

After solving the above two sub-problems, we can obtain
$\\{D_{i}^{\text{gen}},f_{i},b_{i},P_{i}\\}_{\forall i}$ with given
$\\{T_{i}^{\text{cmp}},T_{i}^{\text{com}}\\}$. Then, we continue to determine
the optimal $\\{T_{i}^{\text{cmp},\ast},T_{i}^{\text{com},\ast}\\}_{\forall
i}$ for each device. Instead of directly optimizing
$\\{T_{i}^{\text{cmp}},T_{i}^{\text{com}}\\}$, we introduce an auxiliary
variable for each device to reduce the dimension of optimization variables.
Let $\eta_{i}\in(0,1)$ denote the time splitting factor for device $i$. Then,
we have $T_{i}^{\text{cmp}}=\eta_{i}T^{\max}$ and
$T_{i}^{\text{com}}=(1-\eta_{i})T^{\max}$. Thus, our goal can be transformed
into a top-layer searching problem that optimizes the splitting factor
$\\{\eta_{i}\\}_{\forall i}$ for each device as follows:

$\displaystyle({\text{P5}})$
$\displaystyle~{}~{}~{}~{}~{}~{}\min~{}E^{\text{round}}\big{(}\\{\eta_{i}\\}_{\forall
i}\big{)}$ (16) subject to:
$\displaystyle~{}~{}~{}~{}~{}\eta_{i}^{\min}\leq\eta_{i}\leq\eta_{i}^{\max},\forall
i,$ (16a) variables:
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\\{\eta_{i}\\}_{\forall
i},$

where $\eta_{i}^{\min}$ and $\eta_{i}^{\max}$ are the lower and the upper
bounds of the time splitting factor, respectively. Specifically, by combing
Eqns. (6), (7) and (8), the value of $\eta_{i}^{\min}$ can be obtained by

$\eta_{i}^{\min}=\frac{{\tau\omega{D_{i}^{\text{loc}}}}}{{T^{\max}f_{i}^{\max}}},$
(17)

and $\eta_{i}^{\max}$ can be calculated by

$\eta_{i}^{\max}=1-\frac{S}{{T^{\max}B{{\log}_{2}}(1+\frac{{{g_{i}}P_{i}^{\max}}}{{{N_{0}}B}})}}.$
(18)

After decomposing the main problem into two sub-problems, we will proceed to
present the solution for each problem in Section 4.

##  4 Proposed Algorithm for FIMI

###  4.1 Solution for Problem (P3)

The procedure for solving Problem (P3) is outlined as follows. We first
simplify the problem by variable substitutions. After that, we leverage the
convex optimization technique to obtain the optimal solution.

Based on Eqn. (1), $D_{i}^{\text{gen}}$ can be expressed as a function of
$\delta_{i}$ as follows:

$D_{i}^{\text{gen}}=\Big{(}\frac{\gamma+\delta_{i}}{\alpha}\Big{)}^{-\frac{1}{\beta}}-D_{i}^{\text{loc}}.$
(19)

Moreover, by combing Constraint (14a) and Eqn. (19), $f_{i}$ can be further
represented by $\delta_{i}$ as

${f_{i}}=\frac{{\tau\omega{{(\frac{{\gamma+{\delta_{i}}}}{\alpha})}^{-\frac{1}{\beta}}}}}{{T_{i}^{{\rm{cmp}}}}}.$
(20)

Next, by substituting Eqns. (19) and (20) into Problem (P3), the problem is
simplified as follows:

$\displaystyle({\text{P6}})$
$\displaystyle~{}~{}~{}\min\limits_{\\{\delta_{i}\\}_{\forall
i}}\sum\limits_{i=1}^{I}\rho_{i}(\gamma+\delta_{i})^{-\frac{3}{\beta}}$ (21)
subject to: $\displaystyle\sum\limits_{i=1}^{I}\delta_{i}=I+\frac{\zeta
I}{N}\ln\big{(}{\Delta^{\max}}\big{)},$ (21a)
$\displaystyle~{}~{}~{}~{}~{}{\delta_{i}^{\min}}\leq{\delta_{i}}\leq{\delta_{i}^{\max}},\forall
i,$ (21b)

where $\rho_{i}$ is an intermediate constant calculated by

$\rho_{i}=\frac{\epsilon_{i}{(\tau\omega)^{3}}}{\big{(}T^{\text{cmp}}_{i}\big{)}^{2}\alpha^{-\frac{3}{\beta}}}>0.$
(22)

Here, the lower bound of $\delta_{i}$ can be obtained by

$\delta_{i}^{\min}=\alpha\Big{(}\min\Big{\\{}\frac{f_{i}^{\max}T_{i}^{\text{cmp}}}{\tau\omega},D_{i}^{\text{loc}}+D^{\max}\Big{\\}}\Big{)}^{-\beta}-\gamma,$
(23)

and the upper bound of $\delta_{i}$ is calculated as

$\delta_{i}^{\max}=\alpha(D_{i}^{\text{loc}})^{-\beta}-\gamma.$ (24)

Note that the solution to Problem (P6) relies on Constraint (21a). If
$\sum\nolimits_{i=1}^{I}\delta_{i}^{\min}>I+\frac{\zeta
I}{N}\ln\big{(}{\Delta^{\max}}\big{)}$ or
$\sum\nolimits_{i=1}^{I}\delta_{i}^{\max}<I+\frac{\zeta
I}{N}\ln\big{(}{\Delta^{\max}}\big{)}$, then Problem (P6) is infeasible. In
the following, we focus on the practical case that
$\sum\nolimits_{i=1}^{I}\delta_{i}^{\min}\leq I+\frac{\zeta
I}{N}\ln\big{(}{\Delta^{\max}}\big{)}\leq\sum\nolimits_{i=1}^{I}\delta_{i}^{\max}$.

######  Theorem 1 (Solution for Problem (P6)).

Suppose that a feasible scenario where inequality
$\sum\nolimits_{i=1}^{I}\delta_{i}^{\min}\leq I+\frac{\zeta
I}{N}\ln\big{(}{\Delta^{\max}}\big{)}\leq\sum\nolimits_{i=1}^{I}\delta_{i}^{\max}$
holds. The optimal solution to Problem (P6) can be expressed as 111The
operation of $[x]_{a}^{b}=\min\\{b,\max\\{x,a\\}\\}$.

$\displaystyle\begin{split}{\delta_{i}^{\ast}}=\bigg{[}{\Big{(}{\frac{{3{\rho_{i}}}}{{\beta\nu^{\ast}}}}\Big{)}^{\frac{\beta}{{\beta+3}}}}\bigg{]}^{\delta_{i}^{\max}}_{\delta_{i}^{\min}},\;\forall
i,\end{split}$ (25)

where the variable $\nu^{\ast}$ is determined through a search process in
order to find the optimal solution $\\{\delta_{i}^{\ast}\\}_{\forall i}$ that
satisfies Constraint (21a).

  * Proof.

It can be verified that Problem (P6) is a convex optimization problem with
respect to $\delta_{i}$. We next employ Karush–Kuhn–Tucker (KKT) the necessary
condition for achieving the optimality. By introducing
$\\{\lambda_{i}\\}_{\forall i}$ and $\\{\psi_{i}\\}_{\forall i}$ as the
Lagrange multipliers for inequality constraints, and $\nu$ as the multiplier
for the equality constraint, we have

$\displaystyle\lambda_{i}\geq
0,\;\lambda_{i}(\delta_{i}^{\min}-\delta_{i})=0,\forall i,$ (26a)
$\displaystyle\psi_{i}\geq
0,\;\psi_{i}(\delta_{i}-\delta_{i}^{\max})=0,\forall i,$ (26b)
$\displaystyle-\frac{3{\rho_{i}}}{{\beta}(\delta_{i}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}-{\lambda_{i}}+{\psi_{i}}+\nu=0,\forall
i,$ (26c) $\displaystyle\text{Constraints (\ref{eqn:p3-ctr-1}) and
(\ref{eqn:p3-ctr-4})}.$ (26d)

By putting Eqn. (26c) into Eqns. (26a) and (26c), we have

$\displaystyle\lambda_{i}\geq 0,\;\psi_{i}\geq 0,\forall i,$ (27a)
$\displaystyle\bigg{(}\underbrace{-\frac{3{\rho_{i}}}{{\beta}(\delta_{i}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}+{\psi_{i}}+\nu}_{A_{1}}\bigg{)}(\delta_{i}^{\min}-\delta_{i})=0,\forall
i,$ (27b)
$\displaystyle\bigg{(}\underbrace{\frac{3{\rho_{i}}}{{\beta}(\delta_{i}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}+{\lambda_{i}}-\nu}_{A_{2}}\bigg{)}\left(\delta_{i}-\delta_{i}^{\max}\right)=0,\forall
i.$ (27c)

Furthermore, the solution to Problem (P6) can be categorized into the
following three cases, depending on the value of $\nu$.

    * –

When
$\nu>\frac{3{\rho_{i}}}{{\beta}(\delta_{i}^{\min}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}$,
we have $A_{1}>0$. Based on Eqn. (27a), we obtain
$\delta_{i}=\delta_{i}^{\min}$.

    * –

When
$\nu<\frac{3{\rho_{i}}}{{\beta}(\delta_{i}^{\max}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}$,
we have $A_{2}>0$. Thus, we get $\delta_{i}=\delta_{i}^{\max}$.

    * –

When
$\frac{3{\rho_{i}}}{{\beta}(\delta_{i}^{\max}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}\leq\nu\leq\frac{3{\rho_{i}}}{{\beta}(\delta_{i}^{\min}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}$,
we have $\lambda_{i}=0$ and $\psi_{i}=0$. Thus
$\nu=\frac{3{\rho_{i}}}{{\beta}(\delta_{i}+{\gamma})^{\frac{{\beta}+3}{{\beta}}}}$.

By summarizing the above three cases into one equation, we obtain the solution
in Eqn. (25). Thus, we complete the proof. ∎

We next aim to find the value of $\nu^{\ast}$ that satisfies Constraint (21a)
according to Theorem 1. Based on Eqn. (25), the optimal value $\nu^{\ast}$ can
be expressed as a function of $\delta_{i}^{\ast}$ as follows:

${\nu^{\ast}}=\frac{{3{\rho_{i}}}}{\beta}{\big{(}{\delta_{i}^{\ast}}\big{)}^{-\frac{{\beta+3}}{\beta}}}.$
(28)

Note that $\nu^{\ast}$ decreases as $\delta_{i}^{\ast}$ increases for
$\delta_{i}^{\ast}>0$. Thus, the search range $[\nu^{\min},\nu^{\max}]$ for
$\nu^{\ast}$ can be given by

$\displaystyle{\nu^{\min}}=\min\limits_{i}{\Big{\\{}\frac{{3{\rho_{i}}}}{\beta}{\big{(}{\delta_{i}^{\max}}\big{)}^{-\frac{{\beta+3}}{\beta}}}\Big{\\}}},$
(29a)
$\displaystyle{\nu^{\max}}=\max\limits_{i}{\Big{\\{}\frac{{3{\rho_{i}}}}{\beta}{\big{(}{\delta_{i}^{\min}}\big{)}^{-\frac{{\beta+3}}{\beta}}}\Big{\\}}}.$
(29b)

According to Eqn. (25), the value of $\sum\nolimits_{i=1}^{I}\delta_{i}$ is
non-increasing with respect to $\nu$. Therefore, the value of $\nu^{\ast}$ can
be efficiently obtained by the bi-section search. Furthermore, the algorithm
for solving Problem (P3) is shown in Algorithm 1. Specifically, the
computational complexity for Algorithm 1 can be approximated by ${\cal
O}\big{(}\log_{2}(\nu^{\max}-\nu^{\min})\big{)}$.

Input: $\nu^{\min},\nu^{\max}$, and a small positive number $\varepsilon$.

1 repeat

2 Update $\nu=(\nu^{\max}+\nu^{\min})/2$;

3 Compute $\\{\delta_{i}\\}_{\forall i}$ based on Eqn. (25);

4 if $\sum\nolimits_{i=1}^{I}\delta_{i}>I+\frac{\zeta
I}{N}\ln\big{(}{\Delta^{\max}}\big{)}$ then update $\nu^{\min}=\nu$ else
update $\nu^{\max}=\nu$;

5

6until  _$|\nu^{\max}-\nu^{\min}|\leq\varepsilon$_ ;

7Set $\nu^{\ast}=\nu$, and acquire $\\{\delta_{i}^{\ast}\\}_{\forall i}$ based
on Eqn. (25);

8 Compute $\\{D^{\text{gen},\ast}_{i}\\}_{\forall i}$ based on Eqn. (19), and
obtain $\\{f_{i}^{\ast}\\}_{\forall i}$ based on Eqn. (20);

9 return $\\{D^{\text{gen},\ast}_{i},f_{i}^{\ast}\\}_{\forall i}$

Algorithm 1 Bisection search for Problem (P3)

###  4.2 Solution for Problem (P4)

To solve Problem (P4), we begin by transforming the problem into a simplified
form that optimizes the bandwidth allocation.

Based on Constraint (15a), $P_{i}$ can be represented by $b_{i}$ as

$P_{i}=\frac{N_{0}b_{i}}{g_{i}}\big{(}2^{\frac{S}{b_{i}T_{i}^{\text{com}}}}-1\big{)}.$
(30)

Note that $P_{i}$ is decreasing with respect to $b_{i}$. Thus, by
incorporating Constraint (12f), the lower bound of $b_{i}$ can be expressed as

$b_{i}\geq b_{i}^{\min}=-\frac{S\ln
2}{T_{i}^{\text{com}}W\Big{(}-\frac{\kappa_{i}}{T^{\text{com}}_{i}}\exp\big{(}-\frac{\kappa_{i}}{T^{\text{com}}_{i}}\big{)}\Big{)}+\kappa_{i}},$
(31)

where $\kappa_{i}=\frac{N_{0}S\ln 2}{g_{i}P^{\max}_{i}}$ is an intermediate
variable and $W(\cdot)$ is the Lambert function.

By putting Eqn. (30) into Eqn. (15) and replacing Constraint (12f) to Eqn.
(31), Problem (P4) can be simplified as follows:

$\displaystyle({\text{P7}})$
$\displaystyle~{}\min\limits_{\\{b_{i}\\}_{\forall
i}}~{}\sum\limits_{i=1}^{I}\frac{N_{0}b_{i}T_{i}^{\text{com}}}{g_{i}}\big{(}2^{\frac{S}{b_{i}T_{i}^{\text{com}}}}-1\big{)}$
(32) subject to:
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\sum\nolimits_{i=1}^{I}b_{i}=B,$
(32a) $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}b_{i}\geq
b^{\min}_{i},\forall i,$ . (32b)

######  Theorem 2 (Solution for Problem (P7)).

The optimal solution to Problem (P7) can be expressed as

$\displaystyle\begin{split}{b_{i}^{\ast}}=\max\big{\\{}b_{i}^{\min},b_{i}(\varpi^{\ast})\big{\\}},\;\forall
i,\end{split}$ (33)

where the value of $\varpi^{\ast}$ is determined through a search process, and
the value of $b_{i}(\varpi^{\ast})$ is obtained by solving the equation
${Q(b_{i})+\varpi^{\ast}=0}$. Here, the function $Q(b_{i})$ is defined as

$\displaystyle\begin{split}Q(b_{i})=\frac{N_{0}T^{\text{com}}_{i}\Big{(}2^{\frac{S}{T^{\text{com}}_{i}b_{i}}}-1\Big{)}}{g_{i}}-\frac{\ln(2)N_{0}S{\cdot}2^{\frac{S}{T^{\text{com}}_{i}b_{i}}}}{g_{i}b_{i}}.\end{split}$
(34)

The optimal solution $\\{b_{i}^{\ast}\\}_{\forall i}$ under $\varpi^{\ast}$
should satisfy Constraint (32a).

  * Proof.

It can be verified that Problem (P7) is a convex optimization problem. We
further utilize KKT conditions to obtain the necessary condition for achieving
the optimality. By introducing $\\{\varphi_{i}\\}_{\forall i}$ as Lagrange
multipliers for inequality constraints and $\varpi$ as the multiplier for the
equality constraint, we have

$\displaystyle\varphi_{i}\geq 0,\;\varphi_{i}(b_{i}^{\min}-b_{i})=0,\forall
i,$ (35a)
$\displaystyle\frac{N_{0}T^{\text{com}}_{i}\Big{(}2^{\frac{S}{T^{\text{com}}_{i}b_{i}}}-1\Big{)}}{g_{i}}-\frac{\ln(2)N_{0}S{\cdot}2^{\frac{S}{T^{\text{com}}_{i}b_{i}}}}{g_{i}b_{i}}$
$\displaystyle\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad-{\varphi_{i}}+\varpi=0,\forall
i,$ (35b) $\displaystyle\text{Constraints (\ref{eqn:p5-ctr-3}) and
(\ref{eqn:p5-ctr-6})}.$ (35c)

By putting Eqn. (35b) into Eqn. (35a), we have

$\displaystyle\begin{split}\big{(}Q(b_{i})+\varpi\big{)}\cdot(b_{i}^{\min}-b_{i})=0,\end{split}$
(36)

where $Q(b_{i})$ is a function of $b_{i}$, and it is expressed as

$Q(b_{i})=\frac{N_{0}T^{\text{com}}_{i}\Big{(}2^{\frac{S}{T^{\text{com}}_{i}b_{i}}}-1\Big{)}}{g_{i}}-\frac{\ln(2)N_{0}S{\cdot}2^{\frac{S}{T^{\text{com}}_{i}b_{i}}}}{g_{i}b_{i}}.$
(37)

Let $b_{i}(\varpi)$ denote the solution of ${Q(b_{i})+\varpi=0}$. Based on
Eqn. (36), the optimal solution to Problem (P6) is given by

${b_{i}^{\ast}}=\max\big{\\{}b_{i}^{\min},b_{i}(\varpi^{\ast})\big{\\}},\;\forall
i.$ (38)

Thus, we complete the proof. ∎

We next investigate how the solution of ${Q(b_{i})+\varpi=0}$ is affected by
$\varpi$. Specifically, the first-order derivative of ${Q(b_{i})}$ with
respect to $b_{i}$ is expressed as

$\frac{{\partial{H}({b_{i}})}}{{\partial{b_{i}}}}=\frac{{{{\ln}^{2}}\left(2\right){S^{2}}{N_{0}}\cdot{2^{\frac{S}{{T_{i}^{{\text{com}}}{b_{i}}}}}}}}{{T_{i}^{{\text{com}}}{g_{i}}b_{i}^{3}}}>0.$
(39)

Since ${Q(b_{i})}$ is an increasing function with respect to $b_{i}$, the
solution for ${Q(b_{i})+\varpi=0}$, denoted as $b_{i}(\varpi)$, will decrease
with the increase of $\varpi$. Therefore, the value of
$b_{i}(\varpi)\in[b_{i}^{\min},B]$ can be obtained through the bisection
search.

After that, our goal turns to search for $\varpi^{\ast}$ that satisfies
Constraint (32a). Since $\varpi$ decreases with $b_{i}(\varpi)$, the search
range $[\varpi^{\min},\varpi^{\max}]$ for $\varpi$ can be calculated by

$\varpi^{\min}=Q(B),~{}\text{and
}\varpi^{\max}=\max\limits_{i}\big{\\{}Q(b_{i}^{\min})\big{\\}}.$ (40)

By combing Eqn. (33) and the monotonicity of $b_{i}(\varpi)$, the value of
$\sum\nolimits_{i=1}^{I}b_{i}$ is non-increasing with respect to $\varpi$.
Thus, the value of $\varpi^{\ast}$ can be obtained by an outer bisection
search.

Input: $\varpi^{\min},\varpi^{\max}$, and error bound $\varepsilon$.

1 repeat

2 Update $\varpi=(\varpi^{\max}+\varpi^{\min})/2$;

3 Invoke BandWidSearch$(\varpi)$ to obtain $\\{b_{i}(\varpi)\\}_{\forall i}$;

4 if $\sum\nolimits_{i=1}^{I}b_{i}(\varpi)>B$ then update
$\varpi^{\min}=\varpi$ else update $\varpi^{\max}=\varpi$;

5

6until  _$|\varpi^{\max}-\varpi^{\min}|\leq\varepsilon$_ ;

7Set $\varpi^{\ast}=\varpi$, and recall BandWidSearch$(\varpi^{\ast})$ to
obtain $\\{b_{i}^{\ast}\\}_{\forall i}$;

8 Compute $\\{P_{i}^{\ast}\\}_{\forall i}$ based on Eqn. (30);

9 return $\\{b_{i}^{\ast},P_{i}^{\ast}\\}_{\forall i}$

/* Function for searching $\\{b_{i}(\varpi)\\}_{\forall i}$. */

10 Function BandWidSearch$(\varpi)$.

11 for  _each device $i=1,2,\ldots,I$ in parallel_ do

12 Set $b_{i}^{\max}=B,\forall i$;

13 repeat

14 Update $b_{i}=(b_{i}^{\max}+b_{i}^{\min})/2$;

15 Compute $Q(b_{i})$ based on Eqn. (34);

16 if $Q(b_{i})+\varpi>0$ then update $b_{i}^{\max}=b_{i}$ else update
$b_{i}^{\min}=b_{i}$;

17

18 until  _$|b_{i}^{\max}-b_{i}^{\min}|\leq\varepsilon$_ ;

19 Set $b_{i}(\varpi)=b_{i}$;

20

21 end for

22return $\\{b_{i}(\varpi)\\}_{\forall i}$

Algorithm 2 Hierarchical bisection search

Finally, by first solving Problem (P7) and then putting the optimal
$\\{b_{i}^{\ast}\\}$ into Eqn. (30), we acquire the optimal solution for
Problem (P4). The workflow for solving Problem (P4) is presented in Algorithm
2, which involves a hierarchical bi-section search to obtain the optimal
solution. For the inner search, we aim to solve ${Q(b_{i})+\varpi=0},\forall
i$ and acquire $\\{b_{i}(\varpi)\\}_{\forall i}$. For the outer search, we aim
to search for $\varpi^{\ast}$ that satisfies Constraint (32a). Regarding the
computational complexity, the number of iteration steps required for inner and
outer search are dominated by ${\cal
O}\big{(}\log_{2}(\max\\{B-b_{i}^{\min},\forall i\\})\big{)}$ and ${\cal
O}\big{(}\log_{2}(\varpi^{\max}-\varpi^{\min})\big{)}$, respectively. Thus,
the overall computational complexity of Algorithm 2 is measured as ${\cal
O}\big{(}\log_{2}(\varpi^{\max}-\varpi^{\min})\cdot\log_{2}(\max\limits_{i}\\{B-b_{i}^{\min}\\})\big{)}$.

###  4.3 Cross-Entropy Searching for Problem (P5)

In this subsection, we aim at solving the top-layer problem that optimizes the
time split factor (i.e., $\\{\eta_{i}^{\ast}\\}_{\forall i}$) for each device.
However, directly obtaining the closed-form expression for function
$E^{\text{round}}\big{(}\\{\eta_{i}\\}_{\forall i}\big{)}$ is very difficult,
hindering the subsequent theoretical analysis. Inspired by [40, 41], we employ
the learning-based cross-entropy (CE) algorithm as a promising method to
address this issue.

To leverage the CE method for solving Problem (P5), we incorporate the
optimization variables into a vector denoted as
$\bm{\eta}=(\eta_{1},\ldots,\eta_{I})$. Additionally, we consider
$\bm{\eta}\in\mathbb{R}^{I}$ as a random vector, which follows the normal
distribution $\bm{\eta}\sim{\cal N}(\bm{\mu},\bm{\sigma}^{2})$. Instead of
directly tuning $\bm{\eta}$, we focus on the joint optimization of the vector
for mean value $\bm{\mu}=(\mu_{1},\ldots,\mu_{I})$ and the vector for standard
deviation $\bm{\sigma}=(\sigma_{1},\ldots,\sigma_{I})$.

Let $j$ denote the iteration index and $J$ denote the pre-determined maximum
iteration step for the CE algorithm. At the $j$-th iteration, the algorithm
samples $M$ feasible solutions $\\{\bm{\eta}_{j,m}\\}_{\forall m}$ from the
distribution profile of $\\{\bm{\mu}_{j},\bm{\sigma}_{j}\\}$, where
$\bm{\eta}_{j,m}\in\mathbb{R}^{I}$ represents the $m$-th generated solution at
the $j$-th iteration. The fundamental concept behind the CE algorithm is to
select the top $K$ solutions with the best performance from
$\\{\bm{\eta}_{j,m}\\}_{\forall m}$ for evolving the distribution profile. By
iteratively converging towards the optimal distribution with the expectation
value of $\bm{\mu}^{\ast}$, we obtain the numerical result for Problem (P5).
The detailed procedure of the CE-based method is presented in Algorithm 3,
which is explained as follows.

  * •

(Line 3: Initialization). The algorithm initializes
$\bm{\mu}_{0}=(0.5,\ldots,0.5)$ and $\bm{\sigma}_{0}=(1,\ldots,1)$.

  * •

(Line 3: Sampling). During the $j$-th iteration, the algorithm generates $M$
samples from distribution ${\cal N}(\bm{\mu}_{j},\bm{\sigma}_{j}^{2})$ to form
the solution set denoted as $\\{\bm{\eta}_{j,m}\\}_{\forall m}$.

  * •

(Line 3: Selection). After invoking Algorithms 1 and 2, we obtain the
corresponding objective values $\\{E^{\text{round}}_{j,m}\\}_{\forall m}$ for
all the generated solutions $\\{\bm{\eta}_{j,m}\\}_{\forall m}$, where
$E^{\text{round}}_{j,m}=E^{\text{round}}(\bm{\eta}_{j,m})$. By sorting
$\\{E^{\text{round}}_{j,m}\\}_{\forall m}$ in an ascending order, we record
the first $K$ indices to form a set ${\cal K}\subset\\{1,2,\ldots,M\\}$.

  * •

(Line 3: Updating). The newly estimated distribution profile parameterized by
$\tilde{\bm{\mu}}_{j+1}$ and $\tilde{\bm{\sigma}}_{j+1}$ are respectively
updated by

$\tilde{\bm{\mu}}_{j+1}=\sum\limits_{m\in{\cal
K}}\frac{\bm{\mu}_{j,m}}{K},~{}\text{and
}\tilde{\bm{\sigma}}_{j+1}=\sum\limits_{m\in{\cal
K}}\frac{\bm{\sigma}_{j,m}}{K}.$ (41)

  * •

(Line 3: Smoothness). To reduce fluctuations during iteration, the two
parameters of the distribution profile for the next iteration step are
respectively smoothed as

$\displaystyle\bm{\mu}_{j+1}=\varrho\bm{\mu}_{j+1}+(1-\varrho)\tilde{\bm{\mu}}_{j+1},$
(42a)
$\displaystyle\bm{\sigma}_{j+1}=\varrho\bm{\sigma}_{j+1}+(1-\varrho)\tilde{\bm{\sigma}}_{j+1},$
(42b)

where $\varrho\in(0,1)$ is a pre-determined smoothing factor.

1 Initialization. Mean value $\bm{\mu}_{0}=(0.5,\ldots,0.5)$, standard
deviation $\bm{\sigma}_{0}=(1,\ldots,1)$, iteration counter $j=0$, and a small
positive number $\varepsilon$.

2 while  _$j <J$ and $\max\\{\sigma|\sigma\in\bm{\sigma}_{j}\\}>\varepsilon$_
do

3 Generate $M$ samples from ${\cal N}(\bm{\mu}_{j},\bm{\sigma}^{2}_{j})$ to
form the solutions $\\{\bm{\eta}_{j,m}\\}_{\forall m}$, and
$m\in\\{1,2,\ldots,M\\}$ ;

4 Invoke Algorithms 1 and 2 to compute $\\{E^{\text{round}}_{j,m}\\}_{\forall
m}$, where $E^{\text{round}}_{j,m}=E^{\text{round}}(\bm{\eta}_{j,m})$;

5 Select the top $K$ best-performing solutions from
$\\{\bm{\eta}_{j,m}\\}_{\forall m}$, and record the corresponding indices to
form ${\cal K}\subset(1,2,\ldots,M)$ ;

6 Compute $\tilde{\bm{\mu}}_{j+1}$ and $\tilde{\bm{\sigma}}_{j+1}$ according
to Eqn. (41);

7 Smooth the distribution profile to obtain ${\bm{\mu}}_{j+1}$ and
${\bm{\sigma}}_{j+1}$ based on Eqns. (42a) and (42b);

8 Update $j=j+1$;

9

10 end while

Output: The converged solution $\bm{\eta}^{\ast}=\bm{\mu}_{j}$.

Algorithm 3 Learning-based cross entropy searching

We now analyze the overall computational complexity for solving Problem (P1).
Specifically, there are at most $J$ iteration steps in Algorithm 3. At each
iteration in Algorithm 3, it invokes $M$ executions for both Algorithms 1 and
2. We use $\chi$ to denote the maximum search range among the three bi-section
searches in Algorithms 1 and 2, where $\chi$ is computed by

$\displaystyle\small\begin{split}\chi=\max\big{\\{}\nu^{\max}-\nu^{\min},\varpi^{\max}-\varpi^{\min},\max\limits_{i}\\{B-b_{i}^{\min}\\}\big{\\}}.\end{split}$
(43)

The overall computational complexity for obtaining the solution to Problem
(P1) can be expressed as

${\cal O}\Big{(}JM\big{(}\underbrace{\log_{2}(\chi)}_{\text{Alg. \ref{alg:bin-
search}}}+\underbrace{\log_{2}^{2}(\chi)}_{\text{Alg. \ref{alg:hier-bin-
search}}}\big{)}\Big{)}={\cal O}\big{(}JM\log_{2}^{2}(\chi)\big{)}.$ (44)

Note that the proposed search algorithm can be effectively executed by the
server in practical scenarios. The empirical runtime can be disregarded when
compared to the latency consumed for model training and parameter
transmission.

###  4.4 Optimal Augmentation Policy

Upon obtaining the optimal quantity of synthesized data for each device
$\\{D_{i}^{\text{gen},\ast}\\}_{\forall i}$, each device will determine the
category-wise data amount locally. Let $C$ denote the total number of label
types (i.e., data categories) for the classification task. Each device aims at
optimizing the distribution of synthesized data at a category-wise level. This
involves determining $\\{d_{i,c}^{\text{gen}}\\}_{\forall c},(c=1,2,\ldots,C)$
for each individual device.

According to [42, 43], we adopt the concept of _data entropy_ to quantify the
distribution of local data. For device $i$, the data entropy of the local
mixed dataset ${\cal D}_{i}^{\text{loc}}\cup{\cal D}_{i}^{\text{gen}}$ is
defined as

$H_{i}=-\sum\limits_{c=1}^{C}\frac{d_{i,c}^{\text{loc}}+d_{i,c}^{\text{gen}}}{D_{i}^{\text{loc}}+D_{i}^{\text{gen}}}\log_{2}\Big{(}\frac{d_{i,c}^{\text{loc}}+d_{i,c}^{\text{gen}}}{D_{i}^{\text{loc}}+D_{i}^{\text{gen}}}\Big{)}.$
(45)

where $d^{\text{loc}}_{i}$ is the data amount of the $c$-th category before
the data augmentation at device $i$, and
$\sum\nolimits_{c=1}^{C}d^{\text{loc}}_{i}=D_{i}^{\text{loc}}$. Intuitively,
when the data distribution is uniform across all categories (i.e., $\forall
c,d_{i,c}^{\text{loc}}+d_{i,c}^{\text{gen}}=\frac{D_{i}^{\text{loc}}+D_{i}^{\text{gen}}}{C}$),
the data entropy attains its maximum value. Conversely, if data is skewed
significantly toward a single category (i.e., $\exists
c,d_{i,c}^{\text{loc}}+d_{i,c}^{\text{gen}}={D_{i}^{\text{loc}}+D_{i}^{\text{gen}}}$),
the data entropy reaches its minimum value, indicating a strong non-IID
distribution.

The augmentation policy optimization at device $i$ is to maximize its local
data entropy by regulating the category-wise data amount
$\\{d_{i,c}^{\text{gen}}\\}_{\forall c}$. Formally, the objective of device
$i$ is to solve the following optimization problem:

$\displaystyle({\text{P8}})$
$\displaystyle~{}~{}~{}~{}~{}\max\limits_{\\{d_{i,c}^{\text{gen}}\\}_{\forall
c}}H_{i}$ (46) subject to:
$\displaystyle~{}\sum\nolimits_{c=1}^{C}d_{i,c}^{\text{gen}}=D_{i}^{\text{gen}},$
(46a) $\displaystyle 0\leq d_{i,c}^{\text{gen}}\leq D_{i}^{\text{gen}},\forall
c.$ (46b)

######  Theorem 3 (Optimal local data augmentation).

For a given device $i$, characterized by its category-wise local dataset
distribution $\\{d_{i,c}^{\text{loc}}\\}_{\forall c}$ and a predetermined data
amount of synthesized data $D_{i}^{\text{gen}}$, the optimal local data
augmentation strategy to maximize the data entropy can be expressed as

$\displaystyle\begin{split}d_{i,c}^{\text{gen}}=\big{[}\pi_{i}-d_{i,c}^{\text{loc}}\big{]}_{0}^{D^{\text{gen}}_{i}}.\end{split}$
(47)

Here, $\pi_{i}$ represents the variable that needs to be sought in order to
satisfy $\sum\nolimits_{c=1}^{C}d_{i,c}^{\text{gen}}=D_{i}^{\text{gen}}$.

  * Proof.

We first introduce a set of auxiliary variables $\\{p_{i,c}\\}_{\forall c}$ to
denote the category-wise data proportion at device $i$, where
$p_{i,c}=\frac{d_{i,c}^{\text{loc}}+d_{i,c}^{\text{gen}}}{D_{i}^{\text{loc}}+D_{i}^{\text{gen}}}$.
We can further transform Problem (P8) into the following problem:

$\displaystyle({\text{P9}})$ $\displaystyle\min\limits_{\\{p_{i,c}\\}_{\forall
c}}\sum\limits_{c=1}^{C}p_{i,c}\log_{2}(p_{i,c})$ (48) subject to:
$\displaystyle~{}~{}~{}\sum\nolimits_{c=1}^{C}p_{i,c}=1,$ (48a) $\displaystyle
p_{i,c}^{\min}\leq p_{i,c}\leq p_{i,c}^{\max},\forall c,$ . (48b)

where
$p_{i,c}^{\min}=\frac{d_{i,c}^{\text{loc}}}{D_{i}^{\text{loc}}+D_{i}^{\text{gen}}}$
and
$p_{i,c}^{\max}=\frac{d_{i,c}^{\text{loc}}+D_{i}^{\text{gen}}}{D_{i}^{\text{loc}}+D_{i}^{\text{gen}}}$.
It can be verified that Problem (P9) is a convex optimization problem. By
introducing $\\{\vartheta_{i,c}\\}_{\forall c}$ and
$\\{\theta_{i,c}\\}_{\forall c}$ as the Lagrange multipliers for inequality
constraints, and $\varsigma_{i}$ as the multiplier for the equality
constraint.

$\displaystyle\vartheta_{i,c}\geq
0,\;\vartheta_{i,c}(p_{i,c}^{\min}-p_{i,c})=0,\forall c,$ (49a)
$\displaystyle\theta_{i,c}\geq
0,\;\theta_{i,c}(p_{i,c}-p_{i,c}^{\max})=0,\forall c,$ (49b) $\displaystyle
G(p_{i,c})-\vartheta_{i,c}+\theta_{i,c}+\varsigma_{i}=0,\forall c,$ (49c)
$\displaystyle\text{Constraints (\ref{eqn:p9-ctr-1}) and
(\ref{eqn:p9-ctr-2})},$ (49d)

where $G(p_{i,c})={\log_{2}}(p_{i,c})+\frac{1}{{\ln(2)}}$ is an intermediate
function. Furthermore, we have

$\displaystyle\vartheta_{i,c}\geq 0,\;\theta_{i,c}\geq 0,\forall c,$ (50a)
$\displaystyle\big{(}\underbrace{G(p_{i,c})+\theta_{i,c}+\varsigma_{i}}_{B_{1}}\big{)}(p_{i,c}^{\min}-p_{i,c})=0,\forall
c,$ (50b)
$\displaystyle\big{(}\underbrace{-G(p_{i,c})+\vartheta_{i,c}-\varsigma_{i}}_{B_{2}}\big{)}(p_{i,c}-p_{i,c}^{\max})=0,\forall
c,$ (50c)

Next, the solution to Problem (P9) can be divided into the following three
cases according to the value of Lagrange multiplier $\varsigma_{i}$.

    * –

When $\varsigma_{i}>-{G}(p_{i,c}^{\min})$, we have $B_{1}>0$. Based on Eqn.
(50a), we obtain $p_{i,c}=p_{i,c}^{\min}$.

    * –

When $\varsigma_{i}<-{G}(p_{i,c}^{\max})$, we have $B_{2}>0$. Based on Eqn.
(50b), we obtain $p_{i,c}=p_{i,c}^{\max}$.

    * –

When $-{G}(p_{i,c}^{\max})\leq\varsigma_{i}\leq-{G}(p_{i,c}^{\min})$, we have
$\vartheta_{i,c}=0$ and $\theta_{i,c}=0$. Thus
$\varsigma_{i}=-{\log_{2}}(p_{i,c})-\frac{1}{\ln(2)}$.

As a result, the optimal solution to the Problem (P9) is summarized as

${p_{i,c}^{\ast}}=\Big{[}{2^{-(\frac{1}{{\ln(2)}}+\varsigma_{i}^{\ast})}}\Big{]}_{{p_{i,c}^{\min}}}^{{p_{i,c}^{\max}}},\;\forall
c.$ (51)

By putting $p_{i,c}$ into
$d_{i,c}^{{\text{gen}}}={p_{i,c}}(D_{i}^{{\text{loc}}}+D_{i}^{{\text{gen}}})-d_{i,c}^{{\text{loc}}}$,
the value of $d_{i,c}^{\text{gen}}$ can be expressed as

    * –

When ${p_{i,c}^{\ast}}={p_{i,c}^{\min}}$, we have
$d_{i,c}^{\text{gen},\ast}=0$.

    * –

When ${p_{i,c}^{\ast}}={p_{i,c}^{\max}}$, we have
$d_{i,c}^{\text{gen},\ast}=D_{i}^{\text{gen}}$.

    * –

When ${p_{i,c}^{\ast}}=2^{-(\frac{1}{{\ln(2)}}+\varsigma_{i}^{\ast})}$, we
have
$d_{i,c}^{\text{gen},\ast}=-d_{i,c}^{\text{loc}}+2^{-(\frac{1}{{\ln(2)}}+\varsigma_{i}^{\ast})}\cdot(D_{i}^{\text{loc}}+D_{i}^{\text{gen}})$.

It can be verified that the above solution is equivalent to Eqn. (47). Thus,
we complete the proof. ∎

Note that the numerical value of $\pi_{i}^{\ast}$ can be efficiently obtained
through linear search or bi-section search. Thus, we complete the device-side
augmentation strategy optimization.

##  5 Experiment Evaluations

###  5.1 Experiment Settings

####  5.1.1 Settings for Edge Devices

We consider $I=20$ edge devices for FL training within a 400-meter radius
cell, with the base station positioned at the center of the cell. We utilize
the path loss model from [30] to model the communication link, which is
represented as $128.1+37.6\log_{10}(R_{i})$ with $R_{i}$ denoting the distance
between the server and device $i$ in kilometers. The power spectral density of
additive Gaussian noise is $N_{0}=-174$ dBm/MHz and the total bandwidth is
$B=20$ MHz. The maximum transmitting power for each device follows the uniform
distribution $P_{i}^{\max}\sim U(20,23)$ dBm. Regarding the model training,
the training workload for a single sample is measured as $5\times 10^{6}$
cycles. The maximum computing frequency and the hardware energy coefficient
for each device follow the uniform distributions of $f_{i}^{\max}\sim U(1,2)$
GHz, and $\epsilon_{i}\sim U(4\times 10^{-27},6\times 10^{-27})$,
respectively.

Figure 4: Performance comparison of FIMI and baseline methods on CIFAR10 and
GTSRB datasets with different levels of non-IID distribution. (The upper row:
energy consumption vs. test accuracy; the lower row: time consumption vs. test
accuracy.)

####  5.1.2 Settings for FL Training

Our study is targeted for the implementation of image classification using the
CIFAR10 and GTSRB datasets [11, 44]. We employ the VGG-9 model with the
parameter size of 111.7Mb during uplink transmission [10]. In terms of data
distribution, we evaluate the proposed framework using two Dirichlet
distribution levels: Dir(0.4) for skewed non-IID and Dir(0.6) for slight non-
IID scenarios [12]. For the local dataset, each device is assigned 1250
training samples, i.e., $d_{i}^{\text{loc}}=1250,\forall i$. For the other
hyper-parameters, we set the learning rate as 0.02, the batch size as $B=64$,
the local epoch as $\tau=1$, the maximum allowable global error as
$\Delta^{\max}=0.2$, the maximum latency per round $T^{\max}=60$ seconds, and
the number of the global rounds as $N=200$.

####  5.1.3 Settings for Generative Model

The pre-trained generative model is trained on the public dataset (i.e.,
CINIC10) with the classifier-free guidance technique [45, 46]. The number of
denoising steps is set as 300 and the output image resolution is $32\times
32\times 3$. Specifically, the total amount of synthesized data is 6659 under
the default setting, and it takes about 7.17 minutes for a workstation
equipped with eight RTX A5000 GPUs to generate the synthesized data. During
the synthesized data distribution, each device receives about 333 samples on
average. This yields about 8.1 Mb of downlink traffic per device on average,
which is much smaller than 111.7 Mb incurred by the one-time model upload.
Therefore, the latency associated with data synthesis and distribution is
negligible compared to the numerous rounds of FL training.

###  5.2 Performance Comparisons with Existing Methods

We conduct performance comparisons between our proposed FIMI approach and
several existing methods, including traditional approaches and data
augmentation-based FL methods. Their details are as follows.

  * •

Traditional FL (TFL). TFL adheres to the conventional FL paradigm, wherein
each device solely uses its local dataset for model training [5].

  * •

Semi-supervised FL (SEMI). SEMI incorporates unlabelled samples to enhance
model updates [21].

  * •

Heuristic data compensation (HDC). Each device adopts a heuristic scheme,
namely, generating synthetic data solely for the category with the least
available local data.

  * •

Server-side training (SST). SST preserves the synthesized data at the server,
and conducts model training to generate a complementary update every round.
The complementary update is aggregated with local model updates to renew the
global model.

  * •

GAN-based data synthesis (GAN). The server uses a pre-trained conditional GAN
model to generate additional training samples [22].

  * •

Centralized training with synthesized data (CLSD). This method only utilizes
the generated synthesized data provided by the server to perform centralized
training.

To provide a reasonable comparison, we adopt the identical optimization
algorithm as FIMI for SEMI, HDC, and GAN. In the cases of TFL and SST, the
synthesized data of each device is set to zero, and the resource utilization
policy is optimized.

Figure 4 shows the performance comparison between FIMI and baseline methods
under different levels of non-IID data distribution. As shown in Figures
4(a-d), FIMI outperforms the baseline methods in terms of the test accuracy to
reduce energy consumption. Furthermore, Figures 4(e-h) show that the proposed
FIMI can achieve a faster convergence rate to reduce the training latency. To
evaluate the learning performance, we measure the system cost in terms of
energy consumption, training latency, and uplink transmission traffic required
to achieve the desired test accuracy for each method. The results are provided
in Table I. The advantages of our proposed FIMI can be summarized into the
following three key aspects, i.e., reducing the energy consumption, reducing
the time consumption, and improving the converged accuracy.

####  5.2.1 Reducing Energy Consumption

As shown in Table I, our FIMI consumes the energy consumption smaller than the
other four methods. Specifically, given the test accuracy of 70% on the
CIFAR10 dataset under Dir(0.4), our FIMI achieves 50% reduction in the energy
consumption compared to TFL. As shown in the results, although TFL requires
the lowest energy consumption per global iteration, it consumes more
iterations than FIMI to achieve the desired accuracy, which consequently
degrades the performance of TFL in energy consumption. In contrast, FIMI can
efficiently achieve the target accuracy with less iterations than TFL.
Meanwhile, FIMI reduces the required data size of uplink traffic by
approximately 60% and 50% compared to TFL under the data distributions of
Dir(0.4) and Dir(0.6), respectively. Therefore, our FIMI can effectively
reduce the energy consumption than TFL.

####  5.2.2 Reducing Time Consumption

As shown in Table I, our FIMI can effectively reduce the latency for achieving
the desired accuracy compared to the other baseline methods. Specifically,
when evaluated on the GTSRB dataset with a target accuracy of 80% under
Dir(0.6), FIMI can reduce the latency by 80% and 70% compared to HDC and GAN,
respectively. It is noticed that the advantage of reduced latency is mainly
attributed to the faster convergence rate of FIMI. This is shown in Figures
4(e-h). Specifically, the convergence rates of the methods ranked according to
the descending order are as follows: FIMI, GAN, SST, HDC, SEMI, and TFL.

TABLE I: Performance comparison between FIMI and other baseline methods. |  | Dir=0.4 | Dir=0.6
---|---|---|---
Dataset | Method | Energy@70 (J)1 | Latency@70 (h)1 | Uplink@70 (GB)1 | Best Acc.(%) | Energy@78 (J)1 | Latency@78 (h)1 | Uplink@78 (GB)1 | Best Acc.(%)
CIFAR10 | TFL | 4858.64 (1.0$\times$) | 3.22 (1.0$\times$) | 56.45 (1.0$\times$) | 70.86 | 4506.20 (1.0$\times$) | 2.98 (1.0$\times$) | 52.36 (1.0$\times$) | 78.95
SEMI | 4913.18 (1.0$\times$) | 2.85 (0.9$\times$) | 50.02 (0.9$\times$) | 73.32 | 4625.86 (1.0$\times$) | 2.68 (0.9$\times$) | 47.09 (0.9$\times$) | 80.54
HDC | 3507.94 (0.7$\times$) | 1.68 (0.5$\times$) | 29.54 (0.5$\times$) | 79.92 | 4098.38 (0.9$\times$) | 1.97 (0.7$\times$) | 34.52 (0.7$\times$) | 82.05
SST | 2844.70 (0.6$\times$) | 1.88 (0.6$\times$) | 33.05 (0.6$\times$) | 77.92 | 3725.80 (0.8$\times$) | 2.47 (0.8$\times$) | 43.29 (0.8$\times$) | 80.23
GAN | 3403.74 (0.7$\times$) | 1.63 (0.5$\times$) | 28.66 (0.5$\times$) | 79.43 | 4723.56 (1.0$\times$) | 2.27 (0.8$\times$) | 39.78 (0.8$\times$) | 81.34
FIMI |  2361.78 (0.5$\times$) |  1.13 (0.4$\times$) |  19.89 (0.4$\times$) | 82.21 |  3264.81 (0.7$\times$) |  1.57 (0.5$\times$) |  27.50 (0.5$\times$) | 83.41
Dataset | Method | Energy@72 (J)1 | Latency@72 (h)1 | Uplink@72 (GB)1 | Best Acc.(%) | Energy@80 (J)1 | Latency@80 (h)1 | Uplink(GB)1 | Best Acc.(%)
GTSRB | TFL | N/A2 | N/A2 | N/A2 | 43.81 | N/A2 | N/A2 | N/A2 | 54.84
SEMI | N/A2 | N/A2 | N/A2 | 57.33 | N/A2 | N/A2 | N/A2 | 67.93
HDC | 6807.48 (1.0$\times$) | 3.27 (1.0$\times$) | 57.33 (1.0$\times$) | 73.40 | 6494.89 (1.0$\times$) | 3.12 (1.0$\times$) | 54.70 (1.0$\times$) | 81.59
SST | 2467.08 (0.4$\times$) | 1.63 (0.5$\times$) | 28.66 (0.5$\times$) | 86.51 | 2819.52 (0.4$\times$) | 1.87 (0.6$\times$) | 32.76 (0.6$\times$) | 89.04
GAN | 1806.07 (0.3$\times$) | 0.87 (0.3$\times$) | 15.21 (0.3$\times$) | 88.46 | 2049.19 (0.3$\times$) | 0.98 (0.3$\times$) | 17.26 (0.3$\times$) | 87.69
FIMI |  1146.16 (0.2$\times$) |  0.55 (0.2$\times$) |  9.65 (0.2$\times$) | 91.48 |  1319.82 (0.2$\times$) |  0.63 (0.2$\times$) |  11.12 (0.2$\times$) | 93.80
1“Energy@$x$”, “Latency@$x$” and “Uplink@$x$” denote the required energy,
latency and uplink traffic size to reach $x$% accuracy, respectively. 2“N/A”
means that the corresponding method cannot achieve the target global accuracy
within the predefined maximum global iteration.

####  5.2.3 Improving Converged Accuracy

As shown in Table I, the proposed FIMI can effectively improve the accuracy of
the converged global model, outperforming TFL by 11.35% and 4.46% on the
CIFAR10 dataset under the Dirichlet data distributions of 0.4 and 0.6,
respectively. Meanwhile, we also investigate the performance of model training
by solely using AI-synthesized data (i.e., the CLSD method) and show the
corresponding results in Figure 4. The results show that relying solely on the
synthesized data for model training will result in unsatisfactory convergence
accuracy, since the AI-generated dataset still shows a different distribution
from the local dataset.  This result highlights the importance of using the
generative AI model as a data augmentation method instead of solely using
synthesized data for model training.

Figure 5: Impact of key mechanisms and hyper-parameters on the generative AI-
based FL training system.

###  5.3 Detailed Performances of Different Modules in FIMI

####  5.3.1 Performance of CE Method

As shown in Figure 5(a), the proposed CE-based searching algorithm can
converge within about 30 iterations. The smaller value of the maximum
allowable error $\Delta^{\max}$ will lead to a higher energy consumption for
each round of the global iteration. Moreover, Figure 5(b) shows the
relationship between the amount of synthesized data $D_{i}^{\text{gen},\ast}$
with respect to the computation and communication conditions. We construct two
arithmetic sequences, one for hardware energy coefficients
$\\{\epsilon_{i}\\}_{\forall i}$ and the other for the distances
$\\{R_{i}\\}_{\forall i}$ between the edge devices and the parameter server,
and then optimize the data augmentation policy. The result shows that the
devices with smaller $\epsilon_{i}$ and $R_{i}$ will generate more synthesized
data for local training. This result is consistent with the intuition, namely,
it will be more beneficial for the devices with favorable channel conditions
and energy features to leverage a larger amount of synthesized data for local
training.

####  5.3.2 Impact of Generative Models and Hyper-parameters

Figure 5(c) shows the synthesized data samples using both the diffusion model
and the GAN-based model. The results show that the diffusion-based model can
outperform the GAN-based model by generating more realistic images with
detailed features. For example, the synthesized images of the fog produced by
the diffusion model exhibit a richer texture than the GAN-based model. The
images generated by the GAN-based model appear somewhat blurry. Figure 5(d)
shows the impact of resource management schemes on the energy consumption. The
results demonstrate that the bandwidth allocation will show a significant
influence, with a uniform bandwidth allocation policy leading to a 70% higher
energy consumption compared to FIMI. Figures 5(e-f) show the influence of
$\Delta^{\max}$ and $T^{\max}$ on the data augmentation-based FL training. We
utilize the GTSRB dataset under Dir(0.6) for the experiments. Specifically, a
smaller value for $\Delta^{\max}$ yields an improved test accuracy for the
converged global model. Our empirical findings suggest that setting the
maximum allowable error $\Delta^{\max}$ within the range of $[0.15,0.25]$ for
practical implementations should be appropriate. Furthermore, the single-round
maximum latency $T^{\max}$ regulates the training duration for the FL system.
Using a more stringent $T^{\max}$ ( i.e., a smaller value) can reduce the
training time, while sacrificing an increased energy consumption.

####  5.3.3 Explanation of FIMI

We illustrate the key factors that enable the FIMI to achieve superior
performance over existing methods. To shed light on this, we simulate a
virtual device holding a locally uniform dataset, resembling IID distribution
across all data categories. Let $\bm{g}_{0}$ denote the local gradient of the
virtual device, and let $\bm{g}_{i}$ denote the local gradients of device $i$
resulting from the mixed dataset. Our goal is to quantify the gradient
similarity between the edge device and the virtual device. Based on [47], we
define the gradient similarity between $\bm{g}_{0}$ and $\bm{g}_{i}$ for an
$L$-layer shared model as

$\texttt{Sim}(\bm{g}_{0},\bm{g}_{i})=\frac{1}{{2L}}\sum\limits_{l=1}^{L}{\bigg{(}\frac{{<\bm{g}_{0}^{[l]},\bm{g}_{i}^{[l]}>}}{{\|\bm{g}_{0}^{[l]}\|\|\bm{g}_{i}^{[l]}\|}}+1}\bigg{)}.$
(52)

Here, $\bm{g}_{i}^{[l]}$ denotes the gradient vector of $\bm{g}_{i}$ in the
$l$-th layer, $<\cdot,\cdot>$ is the dot product operation, and $\|\cdot\|$
calculates the Euclidean norm for a given vector. As presented in the box
plots of Figures 5(g-h), the gradient of our FIMI shows the highest level of
similarity to that of the IID local data. This result means that our FIMI can
effectively mitigate the issue of non-IID data distribution across different
edge devices, which thus accelerates the convergence rate of FL training and
improves the accuracy of the converged global model.

##  6 Conclusion and Future Work

In this paper, we have proposed a resource-aware data augmentation framework
called FIMI to leverage pre-trained generative AI models to synthesize
customized local data for FL training. FIMI can improve the convergence
accuracy of FL and save the device-side energy consumption by mitigating data
and device resource heterogeneity. Our research results have demonstrated the
potential of generative AI in FL and provided the insights for developing
efficient and effective FL systems. For our future work, we will explore
multi-server empowered AI content generation services in mobile edge computing
to facilitate the data synthesis procedure. We will also investigate the
incentive mechanisms for the collaborative synergy between generative AI and
FL.

## References

  * [1] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, “High-resolution image synthesis with latent diffusion models,” in _Proc. of CVPR_ , 2022, pp. 10 684–10 695.
  * [2] C. Saharia _et al._ , “Photorealistic text-to-image diffusion models with deep language understanding,” in _Proc. of NeurIPS_ , 2022, pp. 36 479–36 494.
  * [3] M. Xu _et al._ , “Unleashing the power of edge-cloud generative ai in mobile networks: A survey of aigc services,” _arXiv preprint arXiv:2303.16129_ , 2023.
  * [4] S. Azizi, S. Kornblith, C. Saharia, M. Norouzi, and D. J. Fleet, “Synthetic data from diffusion models improves imagenet classification,” _arXiv preprint arXiv:2304.08466_ , 2023.
  * [5] H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Arcas, “Communication-efficient learning of deep networks from decentralized data,” in _Proc. of AISTATS 2017_ , 20–22 Apr. 2017, pp. 1273–1282.
  * [6] W. Wu, C. Zhou, M. Li, H. Wu, H. Zhou, N. Zhang, X. S. Shen, and W. Zhuang, “Ai-native network slicing for 6g networks,” _IEEE Wirel. Commun._ , vol. 29, no. 1, pp. 96–103, 2022.
  * [7] M. Chen, H. V. Poor, W. Saad, and S. Cui, “Wireless communications for collaborative federated learning,” _IEEE Commun. Mag._ , vol. 58, no. 12, pp. 48–54, 2020.
  * [8] R. Yu and P. Li, “Toward resource-efficient federated learning in mobile edge computing,” _IEEE Netw._ , vol. 35, no. 1, pp. 148–155, 2021.
  * [9] Z. Yang, M. Chen, K.-K. Wong, H. V. Poor, and S. Cui, “Federated learning for 6g: Applications, challenges, and opportunities,” _Engineering_ , vol. 8, pp. 33–41, 2022.
  * [10] K. Simonyan and A. Zisserman, “Very deep convolutional networks for large-scale image recognition,” in _Proc. of ICLR_ , 2015.
  * [11] A. Krizhevsky, G. Hinton _et al._ , “Learning multiple layers of features from tiny images,” 2009.
  * [12] H. Wang, M. Yurochkin, Y. Sun, D. Papailiopoulos, and Y. Khazaeni, “Federated learning with matched averaging,” in _Proc. of ICLR_ , 2020.
  * [13] Z. Zhang, Z. Gao, Y. Guo, and Y. Gong, “Scalable and low-latency federated learning with cooperative mobile edge networking,” _IEEE Trans. Mob. Comput._ , 2022.
  * [14] C. Zheng, G. Wu, and C. Li, “Toward understanding generative data augmentation,” _arXiv preprint arXiv:2305.17476_ , 2023.
  * [15] H. Du _et al._ , “Exploring collaborative distributed diffusion-based ai-generated content (aigc) in wireless networks,” _IEEE Netw._ , pp. 1–8, 2023, doi:10.1109/MNET.006.2300223.
  * [16] M. Xu _et al._ , “Joint foundation model caching and inference of generative ai services for edge intelligence,” _arXiv preprint arXiv:2305.12130_ , 2023.
  * [17] H. Du, R. Zhang, Y. Liu, J. Wang, Y. Lin, Z. Li, D. Niyato, J. Kang, Z. Xiong, S. Cui _et al._ , “Beyond deep reinforcement learning: A tutorial on generative diffusion models in network optimization,” _arXiv preprint arXiv:2308.05384_ , 2023.
  * [18] X. Huang _et al._ , “Federated learning-empowered ai-generated content in wireless networks,” _arXiv preprint arXiv:2307.07146_ , 2023.
  * [19] Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra, “Federated learning with non-iid data,” _arXiv preprint arXiv:1806.00582_ , 2018.
  * [20] S. Itahara, T. Nishio, Y. Koda, M. Morikura, and K. Yamamoto, “Distillation-based semi-supervised federated learning for communication-efficient collaborative training with non-iid private data,” _IEEE Trans. Mob. Comput._ , vol. 22, no. 1, pp. 191–205, 2021.
  * [21] S. Wang, Y. Xu, Y. Yuan, and T. Q. Quek, “Towards fast personalized semi-supervised federated learning in edge networks: Algorithm design and theoretical guarantee,” _IEEE Trans. Wirel. Commun._ , 2023.
  * [22] B. Xin, W. Yang, Y. Geng, S. Chen, S. Wang, and L. Huang, “Private fl-gan: Differential privacy synthetic data generation based on federated learning,” in _Proc. of ICASSP_. IEEE, 2020, pp. 2927–2931.
  * [23] E. Jeong, S. Oh, J. Park, H. Kim, M. Bennis, and S. Kim, “Multi-hop federated private data augmentation with sample compression,” in _Proc. of IJCAI Workshop FML’19_ , 12 Aug. 2019, pp. 1–8.
  * [24] P. Li _et al._ , “Snowball: Energy efficient and accurate federated learning with coarse-to-fine compression over heterogeneous wireless edge devices,” _IEEE Trans. Wirel. Commun._ , 2023.
  * [25] Y. Xu, Y. Liao, H. Xu, Z. Ma, L. Wang, and J. Liu, “Adaptive control of local updating and model compression for efficient federated learning,” _IEEE Trans. Mob. Comput._ , 2022.
  * [26] L. Liu, J. Zhang, S. Song, and K. B. Letaief, “Client-edge-cloud hierarchical federated learning,” in _Proc. of ICC_. IEEE, 2020, pp. 1–6.
  * [27] S. Luo, X. Chen, Q. Wu, Z. Zhou, and S. Yu, “Hfel: Joint edge association and resource allocation for cost-efficient hierarchical federated edge learning,” _IEEE Trans. Wirel. Commun._ , vol. 19, no. 10, pp. 6535–6548, 2020.
  * [28] S. Hosseinalipour, S. S. Azam, C. G. Brinton, N. Michelusi, V. Aggarwal, D. J. Love, and H. Dai, “Multi-stage hybrid federated learning over large-scale d2d-enabled fog networks,” _IEEE/ACM Trans. Netw._ , vol. 30, no. 4, pp. 1569–1584, 2022.
  * [29] Y. Wu, Y. Song, T. Wang, L. Qian, and T. Q. Quek, “Non-orthogonal multiple access assisted federated learning via wireless power transfer: A cost-efficient approach,” _IEEE Trans. Commun._ , vol. 70, no. 4, pp. 2853–2869, 2022.
  * [30] Z. Yang, M. Chen, W. Saad, C. S. Hong, and M. Shikh-Bahaei, “Energy efficient federated learning over wireless communication networks,” _IEEE Trans. Wirel. Commun._ , vol. 20, no. 3, pp. 1935–1949, 2020.
  * [31] W. Wu, M. Li, K. Qu, C. Zhou, X. Shen, W. Zhuang, X. Li, and W. Shi, “Split learning over wireless networks: Parallel design and resource management,” _IEEE J. Sel. Areas Commun._ , vol. 41, no. 4, pp. 1051–1066, 2023.
  * [32] H. Ko, J. Lee, S. Seo, S. Pack, and V. C. Leung, “Joint client selection and bandwidth allocation algorithm for federated learning,” _IEEE Trans. Mob. Comput._ , 2021.
  * [33] I. Chen, F. D. Johansson, and D. Sontag, “Why is my classifier discriminatory?” _Proc. of NeurIPS_ , vol. 31, 2018.
  * [34] S. Wang, Y.-C. Wu, M. Xia, R. Wang, and H. V. Poor, “Machine intelligence at the edge with learning centric power allocation,” _IEEE Trans. Wirel. Commun._ , vol. 19, no. 11, pp. 7293–7308, 2020.
  * [35] Q. Hu, S. Wang, Z. Xiong, and X. Cheng, “Nothing wasted: Full contribution enforcement in federated edge learning,” _IEEE Trans. Mob. Comput._ , 2021.
  * [36] L. N. Darlow, E. J. Crowley, A. Antoniou, and A. J. Storkey, “Cinic-10 is not imagenet or cifar-10,” _arXiv preprint arXiv:1810.03505_ , 2018.
  * [37] C. Ma, J. Konečnỳ, M. Jaggi, V. Smith, M. I. Jordan, P. Richtárik, and M. Takáč, “Distributed optimization with arbitrary local solvers,” _Optim. Methods Softw._ , vol. 32, no. 4, pp. 813–848, 2017.
  * [38] N. H. Tran, W. Bao, A. Zomaya, M. N. Nguyen, and C. S. Hong, “Federated learning over wireless networks: Optimization model design and analysis,” in _Proc. of IEEE INFOCOM 2019_ , pp. 1387–1395.
  * [39] P. Li, G. Cheng, X. Huang, J. Kang, R. Yu, Y. Wu, and M. Pan, “Anycostfl: Efficient on-demand federated learning over heterogeneous edge devices,” in _Proc. of IEEE INFOCOM 2023_.
  * [40] S. Zhu, W. Xu, L. Fan, K. Wang, and G. K. Karagiannidis, “A novel cross entropy approach for offloading learning in mobile edge computing,” _IEEE Commun. Lett._ , vol. 9, no. 3, pp. 402–405, 2019.
  * [41] Z. Sun, Y. Xiao, L. You, L. Yin, P. Yang, and S. Li, “Cross-entropy-based antenna selection for spatial modulation,” _IEEE Commun. Lett._ , vol. 20, no. 3, pp. 622–625, 2016.
  * [42] G. Xu, Z. Zhou, J. Dong, L. Zhang, and X. Song, “A blockchain-based federated learning scheme for data sharing in industrial internet of things,” _IEEE Internet Things J._ , 2023.
  * [43] W. Jiang, M. Chen, and J. Tao, “Federated learning with blockchain for privacy-preserving data sharing in internet of vehicles,” _China Commun._ , vol. 20, no. 3, pp. 69–85, 2023.
  * [44] 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 _Proc. of IJCNN_ , no. 1288, 2013.
  * [45] J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” in _Proc. of NeurIPS_ , 2020, pp. 6840–6851.
  * [46] J. Ho and T. Salimans, “Classifier-free diffusion guidance,” in _Proc. of NeurIPS Workshop’2022_ , 2021.
  * [47] B. Zhao, K. R. Mopuri, and H. Bilen, “Dataset condensation with gradient matching,” in _Proc. of ICLR_ , 2020.