The explosive growth of mobile edge devices and their local storage of data has challenged traditional centralized machine learning paradigms. Centralized approaches necessitate frequent data transmission to the central server, which results in significant privacy risks and high communication overhead (Xia et al., 2020). Federated learning (FL) has emerged as a promising distributed learning paradigm, which enables model training by aggregating locally trained models at the central server without directly transmitting raw data, thereby preserving data privacy (Konečnỳ et al., 2015). However, FL still relies on frequent communication between devices and the server, which becomes inefficient in wireless environments with heterogeneous data distributions (Lu et al., 2024).

During the evolution of FL, extensive efforts have been devoted to mitigating communication overhead from various perspectives. Hierarchical FL reduces communication frequency through staged aggregation, thereby lowering communication costs (Briggs et al., 2020; Tran et al., 2025); however, its rigid hierarchical design often limits adaptability in dynamic environments. Prototype-based clustered FL is a representative branch of personalized FL, which aligns client features through global or local prototypes, thereby effectively reducing communication costs by minimizing the need for frequent model exchanges (Yang et al., 2024; Tan et al., 2022). However, such methods often compromise global consistency, thereby limiting scalability when tasks are only partially shared. Moreover, resource-aware FL introduces client selection under bandwidth and computation constraints to improve communication efficiency but typically decouples resource optimization from model aggregation, neglecting their intrinsic interplay (Nishio and Yonetani, 2019). Against this backdrop, clustered FL (CFL) has become a more direct and effective solution. CFL reduces communication costs and improves convergence stability by performing intra- and inter-cluster aggregation after grouping clients based on their statistical or geographical characteristics (Ghosh et al., 2021; Zeng et al., 2023). For example, Yan et al. (2024) proposed adaptive clustering strategies that enable flexible client participation and dynamic cluster formation, thereby reducing redundant communication under non-independent and identically distributed (Non-IID) conditions. Similarly, Gao et al. (2023) clustered clients based on the similarity of their local data distributions and used acceleration algorithms to shorten training time and lower communication overhead. In addition, Zhang et al. (2024) developed an adaptive CFL framework that adjusts cluster size and communication intervals through online similarity measurement, thereby improving both robustness and communication efficiency. Despite these advances, existing CFL approaches still suffer from several limitations: 1) existing studies rarely explore how the clustering strategy influences both data richness and convergence dynamics. Over-reliance on data similarity for clustering may reduce intra-cluster diversity, thereby weakening model generalization. 2) Prior CFL methods primarily rely on gradient similarity or geographic proximity for clustering but often ignore joint optimization of learning performance and resource efficiency.

To bridge this research gap, we propose a data- and distance-aware clustering scheme. The proposed scheme exploits data distribution characteristics and geographical information to optimize cluster head selection and device association scheduling prior to the training process. Based on this clustering result, the CFL training procedure is subsequently carried out. The main contributions of this study are summarized as follows.

We propose a CFL framework that collectively considers learning performance and communication cost. On the learning side, a convergence analysis is conducted to theoretically demonstrate that enhancing the diversity and representativeness of intra-cluster data effectively improves the convergence behavior of CFL under data heterogeneity. On the communication side, the communication cost is modeled in terms of the transmission distance. Based on this model, a combined optimization problem for cluster head selection and device association scheduling is formulated, which simultaneously accounts for learning performance and communication cost.

We consider a wireless CFL system, as illustrated in Figure 1, which consists of a central server and a set of devices U = { 1,2 , … , U } . Each device u ∈ U has a local dataset D u = ( x u i , y u i ) i = 1 D u , where D u = | D u | . The communication distances are denoted by d u , u ′ U × U for the links between devices and d u , cs U × 1 for the links between devices and the server. Before training, a subset of high-performance devices N = { 1,2 , … , N } can be selected from the set of devices as candidates for cluster heads, and all devices are grouped through two scheduling strategies.

A binary variable a n ∈ { 0,1 } indicates whether the device n is selected as a cluster head. The set of cluster heads is N c h , satisfying a n = 1 for all n ∈ N c h . The number of clusters is N c h . We also introduce a binary variable b u , n ∈ { 0,1 } that indicates whether the device u is associated with the cluster head n . Each cluster is defined as C n = { u ∈ U ∣ b u , n = 1 } , and the overall cluster set is C = { C 1 , C 2 , … , C N } , where C n ∩ C n ′ = ∅ for n ≠ n ′ and | C n | ≥ 1 . The total number of samples in cluster n is D n = ∑ u ∈ U b u , n D u .

The global loss function is defined in Equation 1: F w = ∑ n ∈ N a n D n D ∑ u ∈ U b u , n D u D n F u w . (1) Here, F u ( w ) = 1 D u ∑ i = 1 D u ℓ ( x u i , y u i , w ) is the local loss at device u , and D = ∑ u ∈ U D u is the total number of training samples. The intra-cluster aggregation loss at the cluster head n is defined as F n ( w ) = ∑ u ∈ U b u , n D u D n F u ( w ) .

To evaluate how clustering strategies influence learning performance, we conduct a convergence analysis to understand how they influence the convergence performance of CFL. To obtain the expected convergence rate of CFL, we first make the following assumptions (Wang et al., 2020; Wan et al., 2021).

Assumption 1: Assume that the global loss function is differentiable, its gradient is uniformly Lipschitz continuous, and there exists a positive constant L that satisfies Equation 8:

Based on the aforementioned assumptions and the description of the global model, we present the following convergence results. Given the optimal global model w * , the learning rate is set to satisfy γ ≤ 1 2 6 G L . Then, it should satisfy 0 < 4 γ 2 L 2 τ 2 2 1 − 12 γ 2 L 2 τ 2 2 ≤ 1 3 and 1 < 1 1 − 12 γ 2 L 2 τ 1 2 < 1 1 − 12 γ 2 L 2 τ 2 2 ≤ 2 . The optimality gap between the expected global loss function value and the optimal global loss function value is presented in Equation 13: 1 T ∑ t = 1 T E ∇ f w t 2 ≤ 2 f w 1 − f w * γ T + 36 γ 2 L 2 τ 2 2 ϵ 2 + 3 ∑ n ∈ N a n D n D 2 ∑ u ∈ U b u , n D u D n 2 σ 2 + 48 γ 2 L 2 τ 1 2 ∑ n ∈ N a n D n D ϵ n 2 + 24 γ 2 L 2 τ 2 ∑ n ∈ N a n D n D 1 − D n D ∑ u ∈ U b u , n D u D n 2 σ 2 + 16 γ 2 L 2 τ 1 ∑ n ∈ N a n D n D ∑ u ∈ U b u , n D u D n 1 − b u , n D u D n 2 σ 2 . (13)

From Equation 13, we can analyze how each key parameter would affect the convergence of the proposed CFL algorithm. The learning rate γ determines the gradient step size, while the smoothness constant L characterizes the curvature of the loss function; increasing either parameter would enlarge the convergence upper bound, whereas decreasing them would tighten it, thereby improving convergence. Participation weights a n and b u , n , along with data sizes D n and D u , determine each device’s contribution to the aggregated gradient; uneven participation or disproportionate data sizes would further increase the bound. The total number of communication rounds T reduces the contribution of the initial optimality gap, so more rounds would tighten the bound. The parameter τ 1 specifies how many times each device updates its model before a single intra-cluster aggregation; a larger τ 1 would amplify the effects of ϵ n 2 and σ 2 , thus increasing the bound. Similarly, the parameter τ 2 specifies how many times each device updates its model before a single inter-cluster aggregation; a larger τ 2 value would amplify the effects of ϵ 2 and σ 2 , thus also increasing the bound.

Notably, by combining the convergence bound in Equation 13 with assumptions 2 and 3, it can be observed that client drift—defined as the deviation of local gradient updates from the global gradient direction—increases the convergence upper bound, thereby degrading learning performance. In particular, assumptions 2 and 3 characterize such drift through the deviation measures ϵ and ϵ n , which quantify the discrepancy between cluster-level gradients and the global gradient, along with between-individual client gradients and their corresponding cluster-level gradients. Since the derived convergence upper bound increases monotonically with both ϵ and ϵ n , stronger client drift results in a looser convergence bound and, consequently, slower convergence. We further observe that reducing ϵ 2 , which measures the discrepancy between the aggregated data distribution within each cluster and the global data distribution, yields a more significant improvement in learning performance than reducing ϵ n 2 . This is because the period of cluster-level updates, τ 2 , is typically larger than the period of device-level updates, τ 1 , causing the effect of intra-cluster heterogeneity to accumulate more strongly over τ 2 steps. This indicates that enhancing intra-cluster data representativeness plays a dominant role in mitigating client drift and motivates our focus on cluster-level optimization. Consequently, we focus on minimizing the following objective min ∑ n ∈ N a n D n D ∇ F n ( w ) − ∇ F ( w ) 2 . For analytical convenience, we assume uniform data sizes across devices and simple averaging during global aggregation. Under these assumptions, the above divergence objective can be reformulated using norm properties and further bounded via the Cauchy–Schwarz inequality. This results in an upper bound that relates the discrepancy between local and global gradients to the pairwise similarity among local gradients. As a result, minimizing the sum of intra-cluster gradient cosine similarity serves as a tractable surrogate for reducing client drift induced by data heterogeneity. Accordingly, based on the definition of cosine similarity in Equation 2, the objective is transformed into the maximization problem shown in Equation 14: min 1 N c h ∑ n ∈ N c h 1 C n 2 ∑ u , u ′ ∈ C n c u , u ′ w . (14)

Based on the above convergence analysis, the cluster-level data representativeness metric is defined as follows (Equation 15): c C = 1 N c h ∑ n ∈ N c h 1 C n 2 ∑ u , u ′ ∈ C n c u , u ′ w . (15)

By substituting the original global loss function with c ( C ) , the optimization problem in Equation 4 can be reformulated as presented in Equation 16: min A , B β d C + α c C s.t. C 1 –C 5 , (16) where Equation 14 is NP-hard and is addressed by decomposing it into device association and cluster-head selection strategies based on cosine similarity and communication distance.

Given an initial cluster head set N c h ⊆ N , the device association problem is formulated as shown in Equation 17: min B β ∑ n ∈ N c h ∑ u ∈ U b u , n d u , n + α 1 N c h ∑ n ∈ N c h 1 C n 2 ∑ u , u ′ ∈ C n c u , u ′ w s.t. C 1 , C 2 , C 5 , (17) which is an integer nonlinear programming problem. We adopt the Gurobi solver for optimal device association under a fixed cluster head set. The overall clustering utility is defined in Equation 18: Ψ N = ψ N + ∑ n ∈ N a n d n , cs , (18) where ψ ( N ) is the device association utility obtained by solving Equation 17 with the given cluster head set. The cluster head selection problem can be formulated as shown in Equation 19: min A Ψ N s.t. C 4 . (19)

To minimize Ψ ( N c h ) , a greedy iterative strategy is used to adjust the cluster head set, which includes the following three main update operations:

Cluster head addition: For any device n ∈ N \ N c h , compute the utility Ψ ( N c h + { n } ) ; if it improves the objective, add n to N c h .

Therefore, in the greedy iterative strategy for cluster head selection, the number of cluster heads varies dynamically until convergence. We summarize the alternating optimization process of device association and cluster head selection in Algorithm 1. The overall complexity is mainly determined by two components: (i) solving the device association problem using the Gurobi solver for a fixed cluster head set and (ii) the greedy iterative updates for cluster head selection, including addition, exchange, and removal operations. In particular, let N = | N | be the total number of candidate nodes and U = | U | be the number of devices. Computing the device association utility ψ ( N c h ) via Gurobi involves all devices in U , with complexity O ( U ) per evaluation. During each iteration of the greedy cluster head updates, the addition step requires up to N − | N c h | evaluations of Ψ ( N c h ) , the exchange step requires up to ( N − | N c h | ) | N c h | evaluations, and the removal step requires up to | N c h | evaluations. Hence, the per-iteration complexity can be expressed as O ( N − | N c h | ) + ( N − | N c h | ) | N c h | + | N c h | U . Since | N c h | ≤ N , this can be simplified and upper-bounded by O ( N 2 U ) per iteration. Therefore, for a total of T iterations until convergence, the overall computational complexity of the alternating optimization algorithm is O ( T N 2 U ) , which clearly highlights its scalability with respect to the number of candidate nodes and the number of devices.

Algorithm 1

Input: Initial cluster head set N c h , candidate cluster head set N , and device set U .

We simulate a wireless FL system consisting of a central server and 100 devices uniformly distributed within a circular area of 100 m radius. To model long-range communication, the server is positioned 1 km away from the device cluster center. The total number of global training rounds is fixed at 200, and the intra- and inter-cluster model aggregation periods are set to τ 1 = 1 and τ 2 = 2 , respectively. Each device is allocated to a communication bandwidth of 1 Mbps and transmits at a power of 10 mW. The noise power spectral density is set to − 173 dBm/Hz, and the path loss model follows 128.1 + 37.6 ⁡ log 10 ( d km ) . The model size is 28 kB, and each local training batch consists of 32 samples, with a learning rate of 0.01. Communication energy consumption is calculated using the Shannon formula, considering both device-to-cluster-head and cluster-head-to-server distances (Chen et al., 2020; Taïk et al., 2022). For simplicity, we assume identical transmission power and bandwidth across all devices so that both communication energy and delay depend solely on transmission distance and exhibit the same trend. In addition, we consider the single-hop communication model between devices and their associated cluster heads, which is commonly adopted in cluster-based wireless FL architectures. Under this assumption, routing overhead is negligible and does not affect the relative performance comparison. Experiments are conducted on the Fashion-MNIST and CIFAR-10 datasets. To emulate a statistically heterogeneous non-IID environment, data are partitioned across devices using a Dirichlet distribution, with a concentration parameter set to 0.5, which corresponds to a moderate-to-high level of data heterogeneity and has been widely adopted in prior federated learning studies to evaluate robustness under non-IID settings (Meng et al., 2023). We set the weights to α = β = 0.5 .

To evaluate the proposed method, three baseline algorithms are considered for comparison: (i) random clustering, where devices are grouped based on the geographical proximity and cluster heads are selected randomly within each cluster; (ii) similarity-based clustering, which groups devices with similar local data distributions using statistical distance metrics, with heads randomly assigned; and (iii) FedAvg, the conventional FL scheme without clustering, where all devices communicate directly with the central server in each round.

To verify the effectiveness of the proposed clustering algorithm in enhancing learning performance in FL, we conduct test accuracy comparison experiments on the Fashion-MNIST and CIFAR-10 datasets. The experiments evaluate the impact of different device clustering algorithms, along with the classical FedAvg, on the model’s training accuracy. Figures 2, 3 illustrate the evolution of test accuracy during training, and the corresponding test accuracies at convergence are summarized in Table 1. As shown, the proposed algorithms consistently achieve the highest test accuracies for a given number of training rounds and maintain significant advantages throughout the training process. The similarity-based clustering algorithm ranks second, suggesting that adjusting data within clusters to a more balanced distribution—i.e., aligning the data distribution across clusters with the global distribution—can lead to better convergence performance than clustering purely based on intra-cluster data similarity, which is in line with expectations. The randomized clustering algorithm ranks third because it does not account for data distribution or similarity within clusters, leading to less balanced clusters and consequently slower convergence. The FedAvg algorithm exhibits the worst performance, primarily due to the heterogeneity of local data distributions across devices under α = 0.5 , which leads to inconsistent update directions in local models and partially counteracts gradients during global model aggregation, thereby reducing convergence efficiency. Overall, these results further highlight the effectiveness of the proposed clustering algorithm in improving the training of devices in heterogeneous environments.

To verify the effectiveness of the proposed algorithm in reducing global communication energy consumption, energy simulations are conducted on the Fashion-MNIST dataset under different numbers of devices, as shown in Figure 4. FedAvg incurs the highest energy consumption due to frequent long-distance communication with the server. In contrast, the other three algorithms use intra-cluster aggregation, which shortens communication distances and reduces upload frequency, thereby lowering energy usage. The proposed method performs best by jointly optimizing inter-device distances and balancing intra-cluster data, further reducing global communication energy.

As shown in Figures 2, 4, the proposed algorithm consistently achieves high training accuracy with low communication energy consumption, highlighting its advantage in maintaining model performance while reducing communication cost. In comparison, similarity-based and random clustering exhibit slower convergence and higher energy consumption. The results demonstrate that balancing intra-cluster data enhances cluster representativeness, mitigates aggregation conflicts, and, when combined with geographical proximity, contributes to reducing overall energy consumption.

Table 2 presents the learning performance and communication energy under different values of the weighting factor α . As α increases, giving more importance to model accuracy in the optimization problem, the learning performance improves, reaching up to 85.1% for α = 1 . However, this comes at the cost of higher communication energy, which increases from 101 J to 168 J over the same range. These results clearly illustrate the trade-off controlled by the weighting factor: larger α prioritizes learning accuracy at the expense of communication efficiency, while smaller α reduces energy consumption but yields lower model performance. This demonstrates that α serves as an effective tuning parameter to balance learning quality and communication cost in the proposed CFL framework.

This study investigates the trade-off between learning performance and communication energy consumption in CFL, focusing on how cluster head selection and device association affect model training and energy overhead. To collectively optimize model performance and communication efficiency in wireless FL, we first conduct a convergence analysis linking global loss to inter-cluster data imbalance and use cosine similarity to quantify distributional dissimilarity. An optimization model of training loss is then constructed based on cosine gradient similarity, while communication energy is modeled as a function of transmission distance. Finally, a clustering algorithm is proposed to jointly schedule cosine similarity and communication distance for solving the reformulated combined optimization problem. The simulation results show that the proposed method markedly reduces communication energy while improving model accuracy.

In future work, we aim to introduce data-size–aware weighting mechanisms to further optimize client selection and matching, along with adaptive channel allocation strategies to extend the applicability of the method to heterogeneous devices and non-uniform wireless environments. These directions aim to improve both the scalability and robustness of the CFL framework, providing a more comprehensive solution for real-world federated learning scenarios.