1
FedPLT: Scalable, Resource-Efficient, and Heterogeneity-Aware Federated Learning via Partial Layer Training
arXiv:2605.02337v1 [cs.DC] 4 May 2026
Ahmad Dabaja, and Rachid El-Azouzi
Abstract—Federated Learning (FL) has gained significant attention in distributed machine learning by enabling collaborative model training across decentralized system while preserving data privacy. Although extensive research has addressed statistical data heterogeneity, FL still faces several challenges, including high communication and computation overheads and severe device heterogeneity, which require further investigation. Prior work has addressed these issues through sub-model training and partial parameter training. However, such methods often suffer from inconsistent parameter distributions across clients, inaccurate global loss estimation, and increased bias and variance. Guided by our empirical analysis, we propose FedPLT (Federated Learning with Partial Layer Training), an innovative and structured partial parameter training approach that exhibits training behavior similar to full model training while assigning clientspecific portions of the model according to their communication and computational capabilities. In addition, we evaluate the performance of FedPLT when combined with optimal client sampling under communication constraints. We show that this integration improves FL performance by reducing sampling variance under the same communication budget. Through extensive experiments, we demonstrate that FedPLT achieves performance comparable to, or even surpassing, that of full-model training (i.e., FedAvg), while requiring significantly fewer trainable parameters per client. Moreover, FedPLT outperforms existing methods in highly heterogeneous environments, effectively adapts to client resource constraints, and reduces the number of straggling clients. In particular, FedPLT reduces the number of trainable parameters by 71%–82% while achieving performance on par with fullmodel training. Index Terms—Federated Learning, Internet of Things, Edge Intelligence, Partial Training, Communication Efficiency, System Heterogeneity, Straggler Mitigation, Client Sampling
I. I NTRODUCTION With the increasing demand for privacy-preserving and secure machine learning, particularly in light of regulatory frameworks such as the General Data Protection Regulation (GDPR) [1] and the California Consumer Privacy Act (CCPA) [2], federated learning has emerged as a promising paradigm for addressing these concerns. Ahmad Dabaja and Rachid El-Azouzi are with the Laboratoire Informatique d’Avignon (LIA), Avignon University, France (email: [email protected], [email protected]). This work was supported in part by the French National Research Agency (ANR) under Grant ANR-22-CE23-0024. This article substantially extends a conference version published in IEEE PIMRC 2025, doi: 10.1109/PIMRC62392.2025.11274718. This work has been submitted to the IEEE for possible publication. Copyright may be transferred without notice, after which this version may no longer be accessible.
Federated learning [3] is a decentralized machine learning approach that allows models to be trained on local edge devices, each with its private data, without the need to share the data itself. Instead, local devices share their model updates with a central server, where they are aggregated to form a global model. This process is repeated over multiple communication rounds, enabling gradual model improvement and enhanced generalization. However, despite its advantages, this approach still faces several significant challenges. One major challenge is the presence of heterogeneity in local data across clients, which leads to high variance in the trained model, commonly referred to as client drift [4], [5]. This issue arises because devices have different datasets that vary in terms of their source, size, and distribution, causing the global model to struggle to generalize across these diverse datasets [6]. Another challenge for federated learning is system heterogeneity, caused by variability in hardware (CPU, memory) and network connectivity [7]. Edge devices can range from smartphones and computers to IoT sensors, each with its own computational and communication limitations. This problem is further exacerbated by the growing trend of training large models, such as Transformers [8], which are resource-intensive, both in terms of communication and computation, and often exceed the capabilities of many edge devices. The third challenge lies in the communication process itself. Communication overhead remains a significant bottleneck, as frequent model updates between clients and the central server can saturate bandwidth and increase latency, especially in wireless or mobile networks. During each learning round, multiple clients perform local computations before aggregation, collectively contributing to a single global optimization step. In practice, several factors influence the overall communication and computational costs, including the model size, the number of participating clients, and the volume of locally processed data. To address these issues, three main strategies have emerged in the literature: sub-model training, partial parameter training, and client sampling. Submodeling methods (e.g., FedDrop [9], HeteroFL [10], and FedRolex [11]) reduce the model size assigned to each client according to their communication and computational capabilities. Partial parameter training approaches (e.g., FedPMT [12]) retain the full model architecture but update only a subset of parameters during each training round. Client sampling strategies, such as Optimal Client Sampling (OCS) [13], aim to reduce communication overhead by selecting only a subset of clients to participate in each round.
2
While these methods help mitigate device heterogeneity and reduce communication and computational costs, they also introduce new limitations, such as imbalanced parameter distribution across clients (e.g., in HeteroFL and FedPMT), increased variance (e.g., in FedDrop and HeteroFL), and unstable convergence (e.g., in some client sampling strategies). They may also suffer from limited scalability in federated systems composed mainly of resource-constrained devices, such as IoT devices. To address the limitations of existing sub-model training and partial parameter training approaches, we propose FedPLT (Federated Learning with Partial Layer Training), a novel framework that provides an effective and structured method for partial training. Unlike previous approaches that introduce noise through random dropout (e.g., FedDrop) or create imbalanced parameter distributions (e.g., HeteroFL, FedRolex, and FedPMT), FedPLT ensures balanced and coordinated training by dividing the layers into equal-sized sub-layers and assigning each device a group of sub-layers based on its communication and computational capacities. The design of FedPLT is guided by a series of experiments analyzing how training behaves when clients update assigned model partitions instead of the full model. By examining layerwise convergence through gradient magnitudes and directional changes across training rounds, we identify assignment strategies that promote stronger inter-layer alignment than alternative configurations. These assignment strategies exhibit training dynamics similar to those of full-model training, leading to faster and more stable convergence. Consequently, FedPLT achieves substantial improvements in both convergence speed and model accuracy relative to the number of communication rounds. Another important consideration is that the design of the assigned model partitions depends heavily on the underlying model architecture. For example, convolutional neural networks (CNNs) require careful selection of layers to preserve effective feature extraction, while ResNet-based architectures require partitioning strategies that respect their residual block structure. As a result, the design of the assigned model partitions must be tailored to the specific characteristics of each model type to ensure effective training, stable convergence, and behavior aligned with full-model training. Additionally, FedPLT is flexible, allowing large models to be trained even on low-bandwidth networks and low-resource devices by adapting the proportion of parameters trained per client. A preliminary version of this work appeared in [14]. This journal version substantially extends the prior paper with a more detailed formulation of FedPLT, expanded theoretical and efficiency analyses, additional experiments, and a broader discussion of resource-constrained federated systems. Our main contributions are as follows: •
Through empirical experiments, we analyze learning behavior when clients train assigned model partitions instead of the full model. By studying layer-wise convergence in terms of gradient magnitude and directional changes across training cycles, we identify assignment principles that promote inter-layer alignment comparable to that of full-model training.
We introduce FedPLT, a structured partial parameter training scheme that addresses key limitations in existing sub-model training and partial parameter training methods in federated learning. • We establish convergence guarantees on FedPLT under the standard assumptions, together with an additional positive masked gradient alignment assumption introduced. • We formulate the theoretical efficiency of FedPLT across three factors: communication cost, computation cost, and training time. • We extend Optimal Client Sampling (OCS) to the FedPLT setting by incorporating client-specific partial training ratios into the communication budget constraint. This yields a FedPLT-aware sampling rule that accounts for heterogeneous partial training costs while preserving the variance-reduction objective of OCS. This reduces overall communication overhead while maintaining, or even improving, the stability and accuracy of the global model. • We conduct extensive experiments to demonstrate that FedPLT achieves performance improvements over fullmodel training while reducing the number of trained parameters by 84%, thereby significantly lowering communication costs. Furthermore, FedPLT consistently outperforms other state-of-the-art submodel training and partial learning approaches. The performance gap is particularly notable on the CIFAR-10 dataset trained on an FCN model, where competing methods struggle to generalize. Specifically, FedPLT surpasses FedPMT by 7.49%, FedRolex by 9%, HeteroFL by 11.84%, and FedDrop by 20.05%. We also evaluate FedPLT under stringent communication constraints using our Optimal Client Sampling (OCS) strategy. Under the same communication budget, the combination of FedPLT and OCS outperforms OCS applied to full-model training, demonstrating the efficiency and effectiveness of integrating partial layer training with optimal client selection. •
II. R ELATED W ORK In this section, we review existing literature that addresses communication and computation efficiency and system heterogeneity in federated learning (FL). Broadly, these strategies fall into two main categories: 1) Sub-modeling / Sub-model Training Approach: This approach reduces the communications and computational burden by shrinking the global model with a factor that depends on the client’s capabilities. This is achieved by omitting certain training parameters, resulting in smaller sub-models that are less resource demanding. Several works follow this approach, including HeteroFL [10], FedDrop [9], and FedRolex [11]. The difference between these methods lies in the way sub-models are created. 2) Partial Parameters Training Approach: Unlike submodeling, this method keeps all the model parameters intact but updates only a subset of them during training. This reduces communication and computation without altering the overall model architecture. FedPMT [12] is a common method that follows this approach.
3
Each of the two approaches has distinct advantages and limitations. In all cases, inference is performed using the full global model, which is used to compute both the loss and the evaluation metrics on client data. However, during training, sub-model training approaches rely on reduced models that omit a subset of the trainable parameters. Consequently, the forward pass and the resulting loss are computed on an incomplete architecture that no longer reflects the true local objective associated with the full model. This mismatch leads to inaccurate gradients and, in turn, suboptimal updates being propagated to the server. Partial parameter training methods address these challenges by maintaining the full model architecture during the forward pass. As a result, the loss is computed on this full architecture, allowing each client to optimize the local objective of the full model. The resulting gradient is therefore not an inaccurate approximation but a projected gradient obtained by masking the coordinates of untrained parameters while preserving the optimization direction over the active subset. This leads to a more consistent optimization trajectory and reduces gradient error. However, these benefits come with slightly higher communication and computation overhead. Several sub-modeling methods have been proposed to operationalize the idea of training reduced architectures on heterogeneous clients. Although they share the common goal of lowering computational and communication costs, they differ in how sub-models are constructed, assigned to clients, and rotated across training rounds. These design choices influence how much of the global model is trained, how biased the resulting updates are toward certain clients, and how these updates interact during aggregation, leading to different convergence behaviors. FedDrop [9], HeteroFL [10], and FedRolex [11] are three existing sub-model training schemes. Figure 1 illustrates their differences by showing the local models sent on several round for each method.
Fig. 1. Illustration of sub-model training strategies. FedDrop applies random masks each round, HeteroFL uses a fixed mask, and FedRolex introduces a rolling mask mechanism.
FedDrop [9] adopts a stochastic neuron-dropping mechanism, where each neuron in fully connected layers is independently deactivated according to a Bernoulli distribution at each round during the training. This randomized sparsification
reduces computational and communication overhead while allowing different subsets of neurons to be trained across rounds, thereby improving overall parameter coverage over time. This randomness offers simplicity and adaptability but introduces substantial variance in the training process. The fluctuating sub-model structures introduce gradient noise and cause oscillatory loss and unstable global updates. As the random masks change each round, the accumulated variance may slow convergence, especially under non-IID data distributions. HeteroFL [10] addresses this instability by adopting a fixed sub-modeling scheme tailored to device capacity. A dropout mask is defined prior to training and remains constant throughout the entire process. Clients with greater computational resources train larger sub-models, while resourceconstrained clients train smaller ones formed by truncating neurons on one side of each layer. This eliminates perround randomness and improves training stability. However, it creates a learning imbalance: Parameters retained across all sub-models are updated by every client and thus benefit from exposure to diverse data. In contrast, parameters in the truncated region are updated only by high-resource clients. This uneven update frequency leads to biased parameter learning and can degrade overall model performance, especially when smaller sub-models lack sufficient capacity to generalize well on complex or highly heterogeneous datasets. FedRolex [11] extends HeteroFL by introducing a rolling or shifting dropout mask. Instead of fixing the truncated region, the mask is gradually shifted across the layer as training progresses, ensuring that different subsets of parameters are activated at different rounds. This mechanism reduces the parameter learning bias observed in HeteroFL, as all neurons eventually receive updates from clients with diverse data. However, a complete rotation requires as many rounds as the width of the layer. Moreover, the shifting masks reintroduce part of the variance seen in FedDrop, and parameters that are reactivated after long intervals may overwrite earlier learning, leading to a form of catastrophic forgetting. As a result, FedRolex often needs significantly more communication rounds to converge. FedPMT [12] is a partial parameter training method that generates a limited number of local training configurations by freezing the shallow dense layers of a fully connected network. In this method, high-resource clients with good bandwidth train the full model, whereas low-resource or bandwidthconstrained clients update only the deeper layers, thereby reducing computational load and communication burden without modifying the model architecture. While FedPMT improves stability compared to sub-modeling approaches, it suffers from several structural limitations. First, the flexibility of training proportions is highly limited because FedPMT relies on freezing prefixes of layers. As a result, the number of valid partial-training configurations is bounded by the number of dense layers, and the corresponding training ratios are discrete and often unevenly spaced due to the non-uniform parameter sizes of the layers. Second, FedPMT introduces an inherent imbalance in parameter exposure across clients. The deeper layers updated by all clients benefit from more diverse data distributions, whereas the shallower layers trained only by high-resource clients receive updates from a much narrower
4
subset of clients. This issue becomes more pronounced in architectures such as FCNs, where the early layers, which contain the majority of parameters, are trained by only a few clients and on less diverse data. This imbalance can bias training, degrade generalization, and slow global convergence, particularly under strong data heterogeneity. A common limitation of existing sub-model training and partial parameter training methods is their reliance on the availability of sufficiently powerful clients. In fully resourceconstrained systems, such as IoT-dominated environments, this assumption may not hold, making large-model training impractical since some model components may remain insufficiently trained or never updated. III. BACKGROUND AND P RELIMINARY A NALYSIS OF PARTIAL PARAMETER T RAINING A. Federated Learning (FedAvg) We consider a federated system composed of a set of clients K = {1, . . . , K} and a central server. Each client k holds a k private local dataset Dk = {(xi , yi )}ni=1 of size nk . The global training objective is defined as min F (W ) :=
W ∈Rd
K X nk k=1
n
Fk (W ),
(1)
s = 0, . . . , τ − 1.
where s is the current local epoch and ηk is the learning rate of client k. After local training, each client k returns its final model Wkt,τ to the server which aggregates them using a weighted average: W t+1 =
K X nk k=1
n
B. Formalizing Partial Parameter Training To formalize partial-training configurations adapted to device capacities, we represent the configuration of client k by a layer-wise allocation vector Qk = (qk,1 , qk,2 , . . . , qk,L ),
where Fk (W ) = E(x,y)∼Dk [ℓ(W ; x, y)] denotes the local expected P loss of client k, W denotes the model parameters, K and n = k=1 nk . We recall the standard local update rule of FedAvg. At each communication round t, the server broadcasts to each client k the current global model W t . Each client initializes its local model as Wkt,0 = W t and performs τ steps of local stochastic gradient descent (SGD): Wkt,s+1 = Wkt,s − ηk ∇Fk (Wkt,s ),
Fig. 2. Overview of the standard Federated Learning (FedAvg) workflow. (1) The server broadcasts the global model to participating clients; (2) each client performs local training on its private dataset; (3) clients send their local model updates back to the server; (4) the server aggregates these updates to obtain the new global model.
Wkt,τ .
This process is repeated for multiple rounds R until convergence. Figure 2 provides an overview of the four stages of a standard FedAvg round, illustrating model broadcast, local training, update transmission, and server aggregation. Although FedAvg is widely used, training and communicating the entire global model presents significant challenges in realistic heterogeneous environments. These challenges arise from device resource constraints, high communication costs, and the presence of straggler clients with limited computational capabilities or unstable connectivity. To address these issues, prior work has explored resource-aware local training strategies, including sub-model training [9]–[11] and partial parameter training [12]. However, even with such solutions, designing effective client-specific assignments of trainable model partitions remains challenging, particularly when the goal is to preserve training behavior and performance comparable to full-model training.
where qk,l ∈ [0, 1] denotes the proportion of parameters in layer l that are trainable on client k. This layer-wise allocation induces an overall training ratio for client k, defined as the fraction of model parameters that the client is allowed to train: PL l=1 qk,l hl rk = P , L l=1 hl where hl denotes the number of trainable parameters in layer l. While Qk specifies the trainable proportion within each layer, it does not directly describe how the overall trained portion is distributed across layers. To capture this distribution, we introduce the contribution vector Xk = (xk,1 , xk,2 , . . . , xk,L ), where xk,l =
qk,l hl . PL rk · j=1 hj
By construction, xk,l ∈ [0, 1] and
PL
l=1 xk,l = 1.
C. Exploring Parameter-Efficient Training In this subsection, we empirically investigate how different trainable parameter allocations affect training dynamics across model architectures (FCNs and CNNs), with the goal of identifying an effective allocation strategy. To study these dynamics, we use diagnostic metrics such as Magnitude Gradient (MG) and Effective Perturbation (EP). 1) Diagnostic Metrics for Training Dynamics:
5
TABLE I PARTIAL PARAMETER TRAINING CONFIGURATIONS FOR THE MNIST FCN MODEL . E ACH CONFIGURATION IS DEFINED BY THE VECTOR Qk , THE PROPORTION OF TRAINED PARAMETERS rk , THE PER - LAYER CONTRIBUTION VECTOR Xk , AND THE STANDARD DEVIATION OF Xk . Config Full model Deep-heavy Shallow-dominant
Qk
rk (%)
Xk
Std(Xk )
(1.00, 1.00, 1.00, 1.00) (0.08, 0.68, 0.68, 1.00) (0.30, 0.10, 0.10, 1.00)
100.0 25.6 24.4
(0.708, 0.231, 0.058, 0.002) (0.222, 0.615, 0.154, 0.009) (0.872, 0.095, 0.024, 0.009)
0.278 0.224 0.361
Fig. 3. The global update norm (left) and the effective perturbation (mid) at the model level and the model accuracy (right) for different partial training configurations.
Fig. 4. Layer-wise effective perturbation for fc1-fc4 weights under different partial training configurations.
a) Magnitude Gradient (MG): To characterize the strength of parameter updates, we track the Magnitude Gradient (MG) at each round: M Gt = ∆wt
2
,
where ∆wt = wt − wt−1 .
Here, wt denotes the parameter vector at round t, which may represent either the full model or a specific layer, and ∆wt is the corresponding update between two successive rounds. The Magnitude Gradient reflects how much the parameters change during training. Larger values correspond to stronger parameter changes, while smaller values indicate weaker or more stable updates. b) Effective Perturbation: We quantify the temporal alignment of parameter updates using Effective Perturbation (EP) [15], defined as Pt t′ t′ =t−τ +1 ∆w t EP = Pt . t′ t′ =t−τ +1 ∥∆w ∥ EP t can be computed at the model level or the layer level. It takes values in [0, 1] and measures the directional
consistency of parameter updates over time. Values close to one indicate that successive updates are well aligned and optimization proceeds coherently, whereas values close to zero reflect oscillatory or inconsistent update directions caused by noise or conflicting gradients. We note that EP t does not indicate proximity to the optimum. A learning trajectory may exhibit high directional alignment while moving away from the optimum, or low alignment while oscillating near a solution. Therefore, this metric should be interpreted together with amplitude-based measures, such as the update norm, to assess the quality of the optimization process. 2) Empirical Analysis of Parameter-Efficient Training: To better understand how partial parameter training affects optimization dynamics and to identify an effective layer-wise allocation strategy, we conduct a series of experiments on two model families: fully connected networks (FCNs) and convolutional neural networks (CNNs). This analysis allows us to examine the optimization behavior induced by different trainable parameter allocations and to compare it with that of full-model training within each model family.
6
TABLE II THE DIFFERENT PARTIAL TRAINING CONFIGURATIONS FOR CI FAR-10 R ES N ET-8 M ODEL . H ERE Qk IS THE CONTRIBUTION VECTOR , rk THE TRAINING RATIO ( IN %), Xk THE INDUCED CONTRIBUTION VECTOR , AND Std(Xk ) THE STANDARD DEVIATION OF Xk .
Config Full Model Deep-heavy Shallow-dominant
Qk
rk (%)
Xk
Std(Xk )
(1.00, 1.00, 1.00, 1.00, 1.00) (0.06, 0.06, 0.12, 0.42, 1.00) (1.00, 1.00, 0.50, 0.25, 1.00)
100.00 34.75 35.38
(0.00598, 0.06193, 0.18537, 0.73805, 0.00868) (0.00108, 0.01114, 0.06669, 0.89613, 0.02497) (0.01690, 0.17505, 0.26198, 0.52154, 0.02452)
0.27677 0.34878 0.18556
Fig. 5. The global update norm (left) and the effective perturbation (mid) at the model level and the model accuracy (right) for different partial training configurations.
Fig. 6. Effective perturbation for each ResNet-8 layer under different partial training configurations.
a) Fully Connected Networks (FCN): We consider a four-layer fully connected neural network for FashionMNIST, which takes grayscale images of size 28 × 28 as input and consists of three hidden fully connected layers of widths (512, 256, 128) followed by a 10-class output layer. We evaluate this model in a homogeneous setting, where all clients use the same partial-training configuration. Table I reports the three studied configurations obtained from different layer-wise allocations Qk : (i) Full Model, where all layers are trained; (ii) Deep-Heavy, where larger values of qk,l are assigned to deeper layers; and (iii) Shallow-Dominant, where larger values of qk,l are assigned to earlier layers. We first analyze the accuracy and model-level training dynamics induced by different partial-training configurations to determine which one remains closest to full-model training in terms of model-level behavior and performance. Figure 3 reports these metrics for the different configurations. At the model level, both the deep-heavy and full-model configurations exhibit a decreasing update norm during the
early rounds, followed by smooth convergence toward low values, whereas the shallow-dominant configuration presented a noisier pattern that diverges since early rounds. This suggests that the deep-heavy configuration preserves training dynamics closer to those of full-model training and that partial training has not adversely affected performance, whereas the shallowdominant configuration departs from the full model behavior by drifting into a suboptimal region of the optimization landscape. A similar trend appears in the perturbation metric. The deep-heavy configuration remains closely aligned with full-model training until around round 400, maintaining values above 0.7, and then diverges only slightly while still reaching a perturbation value of about 0.5 by round 1000. In contrast, the shallow-dominant configuration exhibits much lower perturbation values from the beginning of training, decreasing to around 0.2 by round 1000. This suggests weaker coordination across layers, which may lead to less stable optimization and poorer generalization. To further support this observation, we analyze the layer-level perturbation induced by these
7
configurations and examine how closely their perturbation profiles align across layers and match those of full-model training. These profiles are shown in Figure 4. We observe that, in both the full-model and deep-heavy configurations, the layers remain strongly aligned, leading to more coordinated and effective learning. In contrast, the shallow-dominant configuration exhibits greater divergence across layers, which makes training more difficult and may lead to suboptimal convergence. Overall, the shallow-dominant configuration departs from full-model training behavior and underperforms, whereas the deep-heavy configuration remains closer to fullmodel dynamics and achieves stronger performance, as also reflected in the accuracy curves in Figure 3. This makes the latter the more suitable choice for FCN training. b) Convolutional Neural Networks (CNN): CNNs differ from fully connected networks in how parameters are distributed across depth. In particular, early layers generally contain fewer parameters, while deeper layers account for a larger proportion of the trainable weights due to increasing channel dimensionality. This raises the question of whether the allocation strategy identified for FCNs also applies to CNNs, or whether it must be adapted to this structural difference. To investigate this, we train a ResNet-8 model [16] on CIFAR-10 using different layer-wise training configurations Qk , analogous to those examined in the FCN study: full model, deep-heavy, and shallow-dominant. These configurations are summarized in Table II. Note that ResNet-8 contains 8 layers, but we group each block of two convolutional layers into a single unit, yielding H = [448, 4640, 13888, 55296, 650], corresponding to (step conv, block 1, block 2, block 3, fc). We first report the model-level update norm, perturbation metric, and model accuracy for the different configurations in Figure 5. In the CNN setting, we observe the opposite trend: the shallow-dominant configuration achieves better performance, whereas the deep-heavy configuration underperforms. As shown in Figure 6, this behavior may be related to stronger inter-layer divergence in the deep-heavy setting, which can hinder representation learning and negatively affect convergence. These empirical findings suggest a simple principle: layers with higher parameter density should receive lower training ratios, whereas layers with lower parameter density should receive larger allocations. Consequently, the best-performing configurations yield a more balanced contribution vector Xk , corresponding to a more uniform distribution of trained parameters across layers. Tables I and II confirm that these configurations consistently achieve the lowest standard deviation of Xk (marked in bold). In the next section, we describe how to determine the allocation vector Qk from rk o that the resulting contribution vector Xk is balanced.
both MLP and CNN architectures. However, we emphasize that although configurations with balanced contribution vectors consistently achieve strong accuracy, this study does not guarantee optimal performance among all feasible choices of Qk . We formulate below the optimization problem to identify the most balanced feasible contribution vector: min
X∈RL
s.t.
J(X) L X
(2)
xl = 1,
(3)
l=1
0 ≤ xl ≤ x̄l ,
∀l.
(4)
where 2
J(X) :=
1 1 X− 1 2 L 2
and x̄l =
rk
hl PL
j=1 hj
denotes the maximum feasible contribution of layer l. Since the objective function is strictly convex and the feasible set is convex, a unique global minimizer exists. By applying the KKT conditions, the optimal contribution vector X ∗ is given by x∗l = min(x̄l , τ ), l = 1, . . . , L, (5) where the scalar τ is chosen such that L X
min(x̄l , τ ) = 1.
(6)
l=1
The optimal contribution vector Xk∗ can then be mapped to a corresponding layer-wise allocation vector Qk , yielding a balanced training configuration for the client k under the fixed ratio rk . This strategy provides a principled way to distribute each client’s training budget across layers within the proposed framework. IV. F ED PLT: F EDERATED L EARNING WITH PARTIAL L AYER T RAINING Motivated by the limitations of sub-model and partial parameter training strategies discussed in Section II, and guided by the empirical analysis in Section III-C2, we propose a new federated learning framework, FedPLT (Federated Learning with Partial Layer Training). Its core principle is a finegrained allocation of training ratios across layers, achieved by decomposing each layer into smaller sub-layers and assigning subsets of these to clients. In this section, we introduce the FedPLT framework and describe its main design components. A. Partial Layer Training and Sublayer Assignment
D. The Optimal Allocation Vector for Parameter Training For a fixed client training ratio rk , different layer-wise configurations Qk can induce markedly different optimization behaviors. In particular, configurations whose induced contribution vectors Xk are relatively balanced exhibit coherent layer-wise training dynamics and strong performance across
After determining the layer-wise allocation vector Qk for each client, FedPLT divides every trainable layer into multiple equal-sized sub-layers (blocks), and instantiates each allocation by assigning to client k a subset of these blocks according to Qk . This defines the client-specific assignment: Ak = {Ak,1 , . . . , Ak,L },
8
where Ak,l ⊆ {1, . . . , Hl } denotes the set of indices of sublayers in layer l assigned to client k, and Hl is the number of sub-layers in layer l. To illustrate a client assignment, figure 7 compares two partial-training configurations with identical training ratios generated by FedPMT and FedPLT. For a client with rk ≈ 30%, FedPMT applies Qk = {0, 0, 1, 1} and concentrates training in deeper layers, whereas FedPLT realizes Qk = {1/6, 1/6, 1/2, 1} by distributing updates across layers. In this FCN illustration, the assigned sub-layers in Ak are shown as brown blocks, and the corresponding trainable parameters are the incoming weights to these selected blocks, shown as brown solid lines, whereas the frozen parameters are shown in dotted lines. Unlike FedPMT, which restricts training to contiguous layers, FedPLT enables fine-grained block-wise allocation, allowing balanced participation of all layers in the learning process. Thus, it distributes the training depth and width rather than activating or freezing entire layers. This enables flexible training ratios.
Fig. 7. Comparison of two partial-parameter training configurations generated by FedPMT (left) and FedPLT (right) under the same training ratio. Brown blocks indicate the sub-layers assigned to the client, and brown solid lines indicate the corresponding trainable weights.
Moreover, these assignments are generated in a structured, rotating manner across clients, such that over all clients each sub-layer is trained by approximately the same number of clients (Figure 8). This prevents persistent under-training of any parameter group and ensures uniform exposure of all model parts to heterogeneous data. It also ensures that all parts of the global model are updated across clients, even in fully resource-constrained systems. Importantly, these assignments are computed once before training begins and remain fixed throughout the entire federated process.
sub-layers, freezing all others. More precisely, if client k is assigned sub-layer index h ∈ Ak,l in layer l, then it updates the corresponding parameter block Wl,h . B. Aggregation During each round, each client transmits only the updated parameters of the sub-layers it trained. The server aggregates these updated sub-layer weights independently using weighted averaging over the clients that trained them, with weights proportional to their local dataset sizes. Clients that did not train a given sub-layer are excluded from its aggregation. Let Wl,h denote the parameter block associated with the ht,τ th sub-layer of layer l, and let Wk,l,h denote its locally updated version returned by client k at round t. The global parameter block for sub-layer h in layer l is computed as: X nk t,τ t+1 , Wl,h = ck,l,h Wk,l,h , ck,l,h = P j∈Sl,h nj k∈Sl,h
where Sl,h is the set of clients that trained sub-layer h in layer l, and nk is the dataset size of client k. The formulation above is written at the sub-layer level to make the aggregation mechanism explicit. However, it can also be expressed compactly at the model level. Let W = {Wl,h }l=1,...,L; h=1,...,Hl denote the full model obtained by stacking all sub-layer parameter blocks. For each client k, define a binary mask ( 1, if k ∈ Sl,h , mk = {mk,l,h }l,h , mk,l,h = 0, otherwise, where each mk,l,h is understood to act on the whole block Wl,h . We also define the corresponding block-wise aggregation weight vector n P k , if k ∈ Sl,h , j∈Sl,h nj ck = {ck,l,h }l,h , ck,l,h = 0, otherwise. Then the aggregation can be written equivalently as W
t+1
=
K X
ck ⊙ mk ⊙ Wkt,τ ,
(7)
k=1
where ⊙ is applied element-wise. C. FedPLT scheme
Fig. 8. Illustration of FedPLT sub-layer allocation. Sub-layers are assigned to clients in a rotational manner so that each sub-layer is trained by approximately the same number of clients.
During training, each client receives the full global model but updates only the parameters associated with its assigned
Given each client’s training ratio rk , the server first determines a layer-wise allocation vector Qk using optimal solution of (2-4), and partitions each layer into equal-sized sub-layers. Based on Qk , sub-layers are assigned to clients in a fixed rotational manner such that, across the population, each sublayer is trained by approximately the same number of clients. During federated training, each client receives the full global model, updates only the parameters associated with its assigned sub-layers while freezing the remaining parameters, and sends back only these updated parameters to the server. The server then aggregates each sub-layer independently using
9
weighted averaging over the clients that trained it. This process is repeated for multiple rounds until convergence. The complete FedPLT procedure is summarized in Algorithm 1. Algorithm 1 FedPLT Scheme 1: Input: Initial global model W 0 , number of layers L, client set K = {1, . . . , K}, client training ratios {rk }K k=1 , rounds count R, local iterations count τ , learning rates {ηk }K k=1 2: 1. Client Initialization 3: for each client k ∈ K do 4: Compute Qk = {qk,1 , . . . , qk,L } from rk 5: end for 6: for each layer l = 1, . . . , L do 7: Partition layer l into Hl equal-sized sub-layers 8: end for 9: for each client k ∈ K do 10: Generate Ak = {Ak,1 , . . . , Ak,L } according to Qk 11: end for 12: 2. Federated Training 13: for each round t = 0, . . . , R − 1 do 14: for each client k ∈ K in parallel do 15: Download the global model W t 16: Initialize the local model Wkt,0 ← W t 17: for each local iteration s = 0, . . . , τ − 1 do 18: Sample a mini-batch ξkt,s from Dk 19: for each layer l = 1, . . . , L do 20: for each sub-layer h = 1, . . . , Hl do 21: if h ∈ Ak,l then 22: Update the assigned block using SGD: t,s+1 t,s Wk,l,h = Wk,l,h − ηk ∇Wl,h ℓ(Wkt,s ; ξkt,s )
23: 24: 25: 26: 27: 28: 29: 30: 31: 32: 33: 34:
else t,s+1 t,s Freeze the block: Wk,l,h = Wk,l,h end if end for end for end for t,τ Upload the updated blocks Wk,l,h ; h ∈ Ak,l end for for each layer l = 1, . . . , L do for each sub-layer h = 1, . . . , Hl do Define Sl,h = {k ∈ K : h ∈ Ak,l } Aggregate the updated block: X nk t,τ t+1 Wl,h = ck,l,h Wk,l,h , ck,l,h = P k∈Sl,h
j∈Sl,h nj
35: end for 36: end for 37: end for 38: Output: Final global model W R
V. C ONVERGENCE A NALYSIS In this section, we provide a theoretical analysis of the convergence behavior of FedPLT. Our goal is to establish that, despite the use of fixed binary masks that restrict each client
to updating only a subset of the model parameters, the global model converges under standard optimization assumptions. For the convergence analysis, we use the simplified client weight ck = nnk instead of the coordinate-dependent aggregation weights defined in (7). This simplification is introduced only to ease the proof and does not affect the main convergence argument. We make the following standard assumptions commonly used in the convergence analysis of federated learning. • Assumption 1 (Smoothness): Each local objective function Fk is L-smooth. That is, for all W, W ′ ∈ Rd , ∥∇Fk (W ) − ∇Fk (W ′ )∥ ≤ L∥W − W ′ ∥.
(8)
This implies that the gradient of Fk does not change too rapidly. Equivalently, by the Descent Lemma, we have: L Fk (W ′ ) ≤ Fk (W )+∇Fk (W )⊤ (W ′ −W )+ ∥W ′ −W ∥2 . 2 (9) • Assumption 2 (Strong Convexity): Each local objective function Fk is µ-strongly convex. That is, for all W, W ′ ∈ Rd , µ Fk (W ′ ) ≥ Fk (W )+∇Fk (W )⊤ (W ′ −W )+ ∥W ′ −W ∥2 . 2 (10) • Assumption 3 (Unbiased Gradient Estimates): The stochastic gradients computed from local mini-batches are unbiased estimators of the true local gradient: Eξk [∇ℓ(W ; ξk )] = ∇Fk (W ),
(11)
where the expectation is taken over the random minibatch ξk sampled at client k. • Assumption 4 (Bounded Variance): The variance of the stochastic gradients is uniformly bounded: h i 2 Eξk ∥∇ℓ(W ; ξk ) − ∇Fk (W )∥ ≤ σ 2 , (12) for all W ∈ Rd and all clients k. • Assumption 5 (Bounded Global Gradient Norm): The norm of the gradient of the global objective function F (W ) is uniformly bounded: ∥∇F (W )∥2 ≤ G2 , •
∀W ∈ Rd .
(13)
Assumption 6 (Positive Masked Gradient Alignment): We assume that for all clients k, the masked gradient retains positive alignment with the descent direction: ρk :=
⟨mk ⊙ ∇Fk (W ), W − W ∗ ⟩ > 0, ⟨∇Fk (W ), W − W ∗ ⟩
where W ∗ is the global minimizer. This ensures that masking does not reverse the gradient direction, preserving progress toward optimality. A full geometric interpretation of this assumption is provided in Appendix A. Let W t ∈ Rd denote the global model at communication round t, and let W ∗ ∈ Rd be an optimal solution minimizing the global objective F (W ). We define the expected squared distance to optimality as: Dt := E[∥W t − W ∗ ∥2 ].
10
To simplify notation, define
VI. O PTIMAL C LIENT S AMPLING WITH PARTIAL L AYER T RAINING
z := µ ρ,
A. Optimal Client Sampling
and B := 2τ 2 (G2 + σ 2 )Γ + 2LνΛ + σ 2 Γ, where ρ := mink,t ρtk , Γ captures the aggregate masking effect, Λ is the aggregate heterogeneity gap, τ is the local training steps, and ν is a constant depending on the masking and smoothness terms. All these quantities are defined explicitly and derived in Appendix B. Proposition 5.1 (One-step recursive bound): Under Assumptions 1–6, the sequence {Dt } satisfies, for every round t, Dt+1 ≤ (1 − zη t ) Dt + (η t )2 B. Proposition 5.1 shows that the convergence behavior is governed by two competing terms: a contraction term (1 − zη t )Dt , which drives the iterates toward the optimum, and a second-order error term (η t )2 B, which captures the effect of stochasticity, heterogeneity, and masking. The above recursion immediately yields the following finitehorizon convergence bound. Theorem 5.2 (General finite-horizon convergence bound): Under Assumptions 1–6, for any horizon T ≥ 1, the iterates generated by FedPLT satisfy ! TY −1 T −1 TY −1 X T t D ≤ (1 − zη ) D0 + B (η t )2 (1 − zη s ). t=0
t=0
s=t+1
Theorem 5.2 provides a unified upper bound that explicitly captures the influence of all algorithmic and problemdependent parameters through the constants z and B. In particular, z measures the effective contraction induced by strong convexity and masked gradient alignment, while B aggregates the impact of stochastic gradient noise, local update drift, smoothness, and objective heterogeneity. We next specialize this result to two standard learning-rate schedules. Corollary 5.3 (Constant step size): If η t = η for all t, with 0 < η < 1/z, then ηB . z Thus, the method converges geometrically to a neighborhood of W ∗ , with asymptotic error floor ηB z . A smaller learning rate yields a tighter neighborhood, at the cost of slower convergence. Corollary 5.4 (Decaying step size): 1 If η t = z(t+1) , then 1 Dt = O . t Dt ≤ (1 − zη)t D0 +
Hence, the method converges sublinearly to the exact optimum, with vanishing error over time. These results show that FedPLT preserves standard convergence guarantees despite relying on static partial updates and masked gradients. The complete derivation and all intermediate steps are provided in Appendix B.
Client Sampling methods are widely studied approaches where only a subset of clients participates in each FL round. This family of approaches reduces the overall bandwidth consumption and lowers system-wide resource usage, which is particularly beneficial for limiting computation and communication costs. However, client sampling also introduces gradient noise and oscillations, especially when clients’ data are heterogeneous. The randomness of client selection, combined with the variability in local data distributions, can cause the aggregated global update to drift away from the true global objective. This drift manifests over rounds as an oscillatory training trajectory, which may slow convergence and degrade model stability. A key distinction between client sampling methods lies in the criteria used to assign sampling probabilities, whether based on dataset size, computational capability, or training time. A principled strategy is to choose these probabilities so as to directly minimize the variance introduced by partial participation. One method that follows this principle is Optimal Client Sampling (OCS) [13], which selects clients to minimize the variance of the aggregated update under a fixed communication budget κ. Formally, OCS solves the optimization problem 2 K X X nk t (14) nk Ukt , min EAt t Uk − p {ptk } k t k=1
k∈A
subject to the communication constraint K X
ptk = κ,
(15)
k=1
where nk is the number of local samples on client k, Ukt denotes its model update (or gradient), ptk is its sampling probability, and At is the set of selected clients at round t. The constraint ensures that, in expectation, only κ clients participate per round. The optimal solution assigns higher sampling probabilities to clients with larger gradient norms and larger local datasets, thereby prioritizing those with stronger contributions to the global update. OCS has demonstrated significant improvements in communication efficiency and convergence speed compared to uniform sampling. However, OCS requires estimating gradient information for all clients at each round, which may be computationally expensive. More importantly, its formulation assumes that every client trains the full global model. As a result, OCS does not naturally extend to settings where clients train only fractions of the model, as is the case with sub-model training or partial parameter training methods. Once clients no longer train the full set of parameters, as in sub-modeling or partial parameter training, the communication cost assumed in OCS fundamentally changes. This is because each client now contributes updates on only a fraction of the parameters and incurs a proportionally smaller communication
11
cost. This shift naturally raises the question of how the OCS formulation should be modified to account for partial training ratios. The following subsection addresses this by extending OCS to the FedPLT setting. B. Extending OCS to Partial Layer Training We now extend the OCS formulation to the FedPLT setting, where each client trains only a fraction of the global model determined by its partial training ratio rk . Under this regime, the communication cost of selecting client k scales proportionally with rk , rather than being uniform across clients as assumed in the original OCS formulation. As a result, the standard constraint must be replaced by a FedPLT-aware communication budget that accounts for heterogeneous training fractions: K X
rk ptk = κ.
bytes. For each parameter, the forward computation requires α FLOPs and the backward computation requires β FLOPs, following the standard relation β ≈ 2α in typical neural networks. Each training round consists of a fixed number of local iterations, denoted by τ , and incurs a fixed latency term δ accounting for connection setup, synchronization delay, or other non-payload overhead. The server is assumed to act only as an aggregator with unbounded computational resources. A. Computation Cost Under full-model training, the per-round computation cost of client k is Ξfull k = (α + β)τ P. Under FedPLT, only the backward part scales with the trained fraction rk , hence Ξplt k = (α + βrk )τ P.
k=1
This constraint ensures that the expected total ratios of trained parameter per round is bounded by the communication budget κ. Solving the associated optimization problem with Lagrangian methods yields a modified rule for client selection probabilities: PK P nk ∥Ukt ∥ κ − j=1 rj + j∈O rj ·P , √ √ t pk = rk rj nj ∥Ujt ∥ j∈O 1,
if k ∈ O if k ∈ K − O (16)
where the set O is defined by: ( O=
k∈K
.√
rk nk ∥Ukt ∥ <
κ−
√ j∈O
rj nj ∥Ujt ∥
PK
P
P
j=1 rj +
j∈O rj
) . (17)
This formulation preserves the variance-minimization objective while adapting to the heterogeneous communication costs introduced by partial layer training. Clients with larger datasets, larger gradient norms, and smaller rk (i.e., cheaper partial participation) receive higher sampling probabilities. Full mathematical details are provided in Appendix C. This extension enables client sampling to explicitly account for heterogeneous partial training costs, allowing FedPLT to jointly exploit parameter-level adaptivity (through partial layer training) and client-level adaptivity (through variance-optimal sampling). While FedPLT balances layer-wise contributions across heterogeneous devices, optimal client sampling further controls which clients participate at each round under a fixed communication budget. In the experimental section, we evaluate the combined impact of these two mechanisms and study how variance-aware client selection interacts with partial layer training under heterogeneous data and system conditions. VII. E FFICIENCY A NALYSIS In this section, we analyze the efficiency of FedPLT relative to FedAvg along three axes: computation cost, communication cost, and round time, under a simple system model. We consider a system with K clients, each characterized by a computation speed γk (in FLOPs/s), an uplink bandwidth Bk↑ (in B/s), and a downlink bandwidth Bk↓ (in B/s). We consider a global model containing P parameters, each occupying s
Therefore, the relative computation reduction is ∆comp =1− k
Ξplt α + βrk k =1− . α+β Ξfull k
Thus, FedPLT reduces computation through the backward stage, with gains increasing as rk decreases. B. Communication Cost In full-model FL, each client downloads and uploads the entire model. Hence, the bidirectional communication volume for each client k per round is Θfull k = 2P s. In FedPLT, the full model is still downloaded, but only the updated fraction is uploaded: Θplt k = (1 + rk )P s. Hence, the relative uplink reduction is ∆comm,up = 1 − rk , k while the total bidirectional communication reduction is ∆comm,tot =1− k
Θplt 1 − rk k = . 2 Θfull k
Therefore, FedPLT preserves the downlink cost while significantly reducing uplink communication. C. Round Time and Stragglers In synchronous federated learning, the duration of each round is determined by the slowest participating client. Under full-model training, the round time is ατ P P full Ps Ps Tround = δ + max + βτ + B + B ↓ ↑ γ γ k k k k k |{z} |{z} |{z} |{z} forward
= δ+
backward
(α+β)τ P Ps Ps + B̂ + B̂ ↓ ↑ γ̂
downlink
,
uplink
12
where k̂ denotes the straggler client attaining the maximum, with γ̂ = γk̂ , B̂ ↓ = Bk̂↓ , and B̂ ↑ = Bk̂↑ . FedPLT can mitigate the straggler effect by selecting the fractions {rk } so that clients finish their local computation and communication in approximately the same time: plt plt ≈ Tround . T1plt ≈ T2plt ≈ . . . ≈ TK
Hence, the round time under FedPLT can be expressed as plt Ps Ps Tround = δ + ατγ̂P + B̂ + r̂ βτγ̂P + B̂ , ↓ ↑ where r̂ denotes the fraction assigned to the limiting client under the balanced FedPLT allocation. The relative round-time efficiency is Ps plt (1 − r̂) βτγ̂P + B̂ ↑ Tround . ∆time = 1 − full = (α+β)τ P Ps Ps Tround + B̂ + B̂ ↓ ↑ γ̂ Thus, since synchronous FL is governed by the slowest participating client, assigning smaller fractions rk to slower clients directly shortens the overall round duration compared with full-model training and mitigates the straggler effect. Detailed derivations and a concrete numerical example are provided in Appendix D. VIII. E XPERIMENTATION AND R ESULTS We evaluate FedPLT through three experiments covering different settings: • Homogeneous low-resource system: We consider a system of resource-constrained devices, with limited communication and computation capabilities, such as IoT devices or drone networks, where no client can realistically train the full model locally. In this setting, all clients use the same small training ratio rk . We evaluate whether FedPLT can enable effective training under such constraints and compare it with full-model training (FedAvg). • Highly heterogeneous system: We consider a system with devices of varying computation and communication capacities, leading to different feasible training ratios rk . We compare FedPLT with existing sub-model training and partial parameter training methods under this heterogeneous setting. • FedPLT with optimal client sampling: We investigate the impact of integrating FedPLT with optimal client sampling. Specifically, we compare the original OCS formulation with our FedPLT-aware extension and analyze the resulting performance gains and trade-offs. A. Experimental Setup •
Datasets and Data Distribution: We use two benchmark datasets: Fashion-MNIST [17] and CIFAR-10 [18], each split into 50,000 training samples and 10,000 validation samples. Data are distributed across clients using a Dirichlet partitioning strategy [19] to simulate non-IID settings. For Experiments 1 and 2, data are partitioned among 50 clients with Dirichlet concentration parameter α = 0.2. For Experiment 3, data are partitioned among 100 clients with a more heterogeneous split using α = 0.1
Model Architectures: For Fashion-MNIST, we use a fully connected network (FCN) with four hidden layers following the architecture, (Input → 512 → 256 → 128 → Output), with ReLU activations in the hidden layers and Softmax at the output. For CIFAR-10, we evaluate two models: the same FCN model and ResNet8 [16]. • Training Configuration: In Experiments 1 and 2, all 50 clients participate in every communication round. For (Fashion-MNIST, FCN), we train for 300 rounds, while for (CIFAR-10, FCN), we train for 1,000 rounds. In both cases, each client performs one local epoch per round with a batch size of 64. For the (CIFAR-10, ResNet8), each client performs 16 local iterations per round, each on a randomly sampled mini-batch of 64 images. In Experiment 3, we use the (CIFAR-10, FCN) configuration with client sampling, so the number of participating clients varies by strategy. These experiments run for 500 rounds with three local epochs per round. • Optimization and Loss Function: Across all experiments, we use stochastic gradient descent (SGD) with a fixed learning rate of 0.01 and optimize the categorical cross-entropy loss. Performance is measured using classification accuracy. • Federated Aggregation and Settings: Client updates are aggregated using a weighted average based on local sample counts, each time adapted to the used scheme. For our FedPLT scheme, we use the aggregation rule defined in Eq. 7. For the (CIFAR-10, ResNet-8) setting, we additionally adopt FedBN [20], keeping batch-normalization layers local to preserve client-specific feature statistics. • Reproducibility and Randomness Control: To ensure robustness and reproducibility, each experiment is repeated multiple times with different random seeds, and reported results are averaged across runs. We maintain a consistent experimental setup across all runs. •
B. Evaluating FedPLT in a homogeneous system In this experiment, we consider a homogeneous lowresource federated system in which all clients have limited computation and communication capabilities, so that fullmodel local training is impractical. We therefore evaluate FedPLT in this setting, where each client trains only a small fraction of the model. For each experimental setting, all clients share the same training ratio rk and the same layer-wise allocation vector Qk . To study the effect of parameter allocation, we repeat the experiment with several choices of Qk corresponding to different balancing levels of the induced contribution vector Xk . We conduct this evaluation on three dataset-model pairs: FashionMNIST with a custom fully connected network, CIFAR-10 with the same fully connected model, and CIFAR-10 with ResNet-8. Recall that the most balanced contribution vector Xk∗ is obtained by solving (2)–(4). To quantify how far a given
13
allocation deviates from this balanced solution, we define the relative error J(X) − J(X ∗ ) E(X) = , J(X ∗ ) which measures the imbalance of Xk relative to the optimal balanced vector Xk∗ . For each dataset-model pair, we therefore consider the default allocation Q∗k associated with Xk∗ together with alternative allocations of increasing imbalance. The corresponding configurations are reported in Table III. TABLE III E XPERIMENTAL SETTINGS USED IN E XP. 1. F OR EACH SETTING , FULL - MODEL TRAINING WITH rk = 100% IS ALSO SIMULATED AS A BENCHMARK . (rk = 29%)
Fashion-MNIST + custom FCN Qk
Config. Most-balanced (Q∗k ) Mildly-unbalanced Moderately-unbalanced Largely-unbalanced
(0.16, 0.50, 1.00, 1.00) (0.20, 0.45, 0.76, 0.80) (0.10, 0.73, 0.86, 0.90) (0.25, 0.29, 0.78, 1.00)
CIFAR-10 + custom FCN Most-balanced (Q∗k ) Mildly-unbalanced Largely-unbalanced
(0.15, 1.00, 1.00, 1.00) (0.17, 0.81, 1.00, 1.00) (0.20, 0.45, 0.76, 0.80)
CIFAR-10 + ResNet-8 Config.
(rk = 23%) Qk
Config.
E 0 30.53% 66.5% 87.43%
E 0 18.79% 92.15%
(rk = 18%) Qk
E Most-balanced (Q∗k ) (1.00, 0.88, 0.31, 0.08, 1.00) 0 Mildly-unbalanced (0.75, 0.75, 0.25, 0.10, 1.00) 25.25% Largely-unbalanced (0.50, 0.50, 0.25, 0.12, 1.00) 85.94% For ResNet-8, each entry of Qk corresponds to a block: stem convolution, block 1, block 2, block 3, and the final fully connected layer.
For each dataset-model pair, full-model training with rk = 100% is also included as a benchmark. Each experiment is repeated over multiple random seeds. We report the validation accuracy curves in Figure 9 and the final validation accuracy (mean ± standard deviation) in Table IV. Both the table and the figure show that FedPLT remains competitive with full-model training across all three homogeneous low-resource settings, despite each client training only a small fraction of the model. On Fashion-MNIST with the FCN model, all FedPLT configurations outperform FedAvg, with gains ranging from 1.02% to 3.58% while training only 29% of the parameters. On CIFAR-10 with the same FCN model, the most-balanced and mildly-unbalanced configurations improve upon FedAvg by 0.76% and 0.16%, respectively, whereas the largely-unbalanced configuration incurs a modest drop of 1.44% while reducing the trained parameters by 77%. On CIFAR-10 with ResNet-8, where only 18% of the model is trained, FedPLT remains close to full-model training, with accuracy drops of only 0.57%, 1.29%, and 1.98% for the mostbalanced, mildly-unbalanced, and largely-unbalanced configurations, respectively. These results indicate that partial-layer training does not inherently compromise accuracy and may even improve generalization in some cases. A possible explanation lies in
how FedPLT distributes parameter updates across clients. In FedAvg, all clients update all parameters, so each parameter receives heterogeneous gradients from the whole population. Under non-IID data, this can create strong inter-client interference and allow dominant client distributions to bias the entire model. In contrast, in FedPLT, each parameter block is updated by fewer clients, and different blocks receive gradients from different subsets of the population. This reduces per-parameter gradient conflict and prevents any single client group from dominating the whole model. In this sense, FedPLT distributes client influence across the parameter space, which may reduce client drift, encourage block-wise specialization, and improve convergence and generalization in heterogeneous settings. Table IV also highlights the variance in accuracy across runs. On Fashion-MNIST with FCN, FedAvg exhibits a standard deviation of 9.53 points, which is markedly higher than those of the FedPLT configurations. Similarly, on CIFAR10 with FCN, FedAvg attains a standard deviation of 4.12, compared to 0.43-2.01 for FedPLT configurations, all of which remain below half of FedAvg’s value. On CIFAR-10 with ResNet-8, FedAvg and FedPLT have similarly low variability, with standard deviations below 2, indicating stable training. Overall, these results suggest that FedPLT exhibits greater robustness and provides more consistent and reliable training behavior, especially for the fully connected models. Across FedPLT configurations, the most-balanced allocation consistently provides the best performance. As allocations become increasingly unbalanced, accuracy degrades across all three settings. Comparing the most-balanced and largely-unbalanced configurations, accuracy drops by 1.59% on Fashion-MNIST + FCN, 2.20% on CIFAR-10 + FCN, and 1.41% on CIFAR-10 + ResNet-8. The same trend is reflected in the standard deviations, which generally increase with imbalance. These results confirm that balanced parameter allocation improves both performance and training stability. In terms of convergence speed, FedPLT and FedAvg appear to exhibit similar learning curves when plotted against communication rounds. However, the horizontal axis in Fig. 9 does not reflect the actual communication or computation cost. To account for this, Table V reports target-based comparisons under three scenarios: equal accuracy, equal communication budget, and equal computation budget. As discussed in Section VII, FedPLT preserves the full forward pass and downlink cost, while the backward pass and uplink communication scale with the training ratio rk . Consequently, its lower per-round cost allows more training rounds under the same total resource budget. Table V shows that, across all datasets and models, FedPLT consistently achieves higher accuracy than FedAvg under the same total communication or computation budget. Conversely, for a fixed target accuracy, FedPLT requires substantially less total system communication and computation than FedAvg. Under the equal-computation comparison, the communication gain may become negative, since FedPLT can require more rounds to match the same computation budget despite its lower perround backward cost. Overall, these results indicate that FedPLT is well suited to homogeneous low-resource federated settings in which no
14
Fig. 9. Validation accuracy curves of FedPLT in homogeneous low-resource systems under different layer-wise allocation vectors Qk , compared with fullmodel training, for (Fashion-MNIST, FCN), (CIFAR-10, FCN), and (CIFAR-10, ResNet-8). TABLE IV F INAL VALIDATION ACCURACY ( MEAN ± STANDARD DEVIATION ) OVER MULTIPLE RUNS OF F ED PLT AND F EDAVG IN HOMOGENEOUS LOW- RESOURCE SYSTEMS ACROSS DIFFERENT DATASETS , MODELS , AND ALLOCATION BALANCE LEVELS . Methods
FashionMNIST + FCN
CIFAR-10 + FCN
Accuracy (%)
Allocation
CIFAR-10 + ResNet-8
Accuracy (%)
Allocation
Allocation
Accuracy (%)
FedAvg
Full model
74.33 ± 9.53
Full model
47.42 ± 4.12
Full model
79.61 ± 1.10
FedPLT
Most-balanced Mildly-unbalanced Moderately-unbalanced Largely-unbalanced
77.91 ± 5.13 77.65 ± 4.51 75.35 ± 5.35 76.32 ± 5.50
Most-balanced Mildly-unbalanced Largely-unbalanced
48.18 ± 0.43 47.58 ± 0.56 45.98 ± 2.01
Most-balanced Mildly-unbalanced Largely-unbalanced
79.04 ± 1.22 78.32 ± 1.87 77.63 ± 1.95
TABLE V F ED PLT GAINS OVER F EDAVG UNDER THREE TARGET BUDGETS ( OR GOALS ) IN HOMOGENEOUS LOW- RESOURCE SYSTEMS : ACCURACY, COMMUNICATION , AND COMPUTATION . FashionMNIST + FCN
CIFAR-10 + FCN
CIFAR-10 + ResNet-8
Fixed Target
Eq-Acc. 74%
Eq-Comm. 43.75 GB
Eq-Comp. 27.50 TFLOPs
Eq-Acc. 47%
Eq-Comm. 425 GB
Eq-Comp. 260 TFLOPs
Eq-Acc. 78%
Eq-Comm. 5.25 GB
Eq-Comp. 512.50 TFLOPs
Accuracy Gain (%) System Comm. Gain (GB) System Comp. Gain (TFLOPs)
– +36.08 +31.43
+5.17 – +6.13
+5.91 −8.27 –
– +208.39 +204.44
+1.75 – +68.20
+2.61 −90.04 –
– +2.23 +396.32
+2.09 – +151.15
+4.02 −1.29 –
client can train the full model locally. Although each client updates only a small portion of the network (29%, 23%, and 18% in our scenarios), the global model is still effectively optimized through complementary sub-layer updates across clients. This makes FedPLT a practical and scalable alternative to full-model training, particularly for large models, where full local optimization is often infeasible on resource-constrained edge devices such as IoT devices.
C. Benchmarking FedPLT against Existing Methods in a Heterogeneous Setting In this experiment, we evaluate FedPLT in a heterogeneous federated system where clients differ significantly in computation and communication capacity. Such heterogeneity creates both resource imbalance and straggler effects, making this setting particularly challenging for federated training. Our objective is to compare FedPLT with existing federated learning methods under these realistic system constraints.
As in the previous subsection, we conduct three sets of experiments using different dataset-model pairs: Fashion-MNIST with an FCN model, CIFAR-10 with an FCN model, and CIFAR-10 with ResNet-8. For all three cases, we consider a realistic heterogeneous system with three levels of computational and communication capability: Highly capable devices: 10% of clients, capable of training the full global model (rk = 1.00). • Moderately capable devices: 30% of clients, capable of training 29% of the model for Fashion-MNIST + FCN, 23% for CIFAR-10 + FCN, and 18% for CIFAR-10 + ResNet-8. • Low-capability devices: 60% of clients, capable of training only 6% of the model. •
For each training ratio rk , we use the default layer-wise allocation vector Q∗k whose induced contribution vector Xk is as balanced as possible. The resulting FedPLT configurations are reported in Table VI.
15
TABLE VI F ED PLT CONFIGURATIONS USED ACROSS DIFFERENT DATASET- MODEL PAIRS IN THE HETEROGENEOUS SETTING . Dataset/Model
rk
Q∗k
Fashion-MNIST + FCN
1.00 0.29 0.06
(1.0, 1.0, 1.0, 1.0) (0.16, 0.5, 1.0, 1.0) (0.03, 0.08, 0.31, 1.0)
CIFAR-10 + FCN
1.00 0.23 0.06
(1.0, 1.0, 1.0, 1.0) (0.15, 1.0, 1.0, 1.0) (0.02, 0.27, 1.0, 1.0)
CIFAR-10 + ResNet-8
1.00 0.18 0.06
(1.0, 1.0, 1.0, 1.0, 1.0) (1.0, 0.88, 0.31, 0.08, 1.0) (1.0, 0.25, 0.06, 0.03, 1.0)
We benchmark FedPLT against FedAvg, FedPMT, FedRolex, FedDrop, and HeteroFL. For fairness, each baseline is simulated under a configuration that matches, as closely as possible, the corresponding FedPLT trained parameter count. Each setting was repeated multiple times with different random seeds. We report the validation accuracy curves in Figure 10 and the final validation accuracy (mean ± standard deviation) in Table VII. TABLE VII F INAL VALIDATION ACCURACY ( MEAN ± STANDARD DEVIATION ) COMPARED WITH EXISTING METHODS IN A HETEROGENEOUS AND LOW- RESOURCE FL SYSTEM ACROSS DIFFERENT DATASETS AND MODELS .
Methods
Fashion MNIST + FCN
CIFAR-10 + FCN
CIFAR-10 + ResNet-8
FedAvg FedPLT FedPMT FedRolex FedDrop HeteroFL
74.33 ± 9.53 78.08 ± 4.65 65.81 ± 5.8 74.32 ± 3.85 52.88 ± 8.35 16.56 ± 19.70
47.42 ± 4.12 46.56 ± 0.71 39.16 ± 3.47 37.65 ± 1.32 26.60 ± 2.94 34.81 ± 8.31
79.61 ± 1.10 78.21 ± 1.57 75.57 ± 3.04 9.96 ± 1.48 14.53 ± 0.81 46.22 ± 1.58
Figure 10 and Table VII show that FedPLT consistently outperforms existing partial-training and submodeling approaches across all heterogeneous settings. On Fashion-MNIST with the FCN model, FedPLT achieves the highest final accuracy among all methods, exceeding FedAvg by 3.75% and substantially outperforming FedPMT, FedRolex, FedDrop, and HeteroFL. On CIFAR-10 with the FCN model, FedPLT remains close to FedAvg (46.56% vs. 47.42%) while outperforming FedPMT by 7.40%, FedRolex by 8.91%, HeteroFL by 11.75%, and FedDrop by 19.96%. A similar trend is observed for CIFAR-10 with ResNet-8, where FedPLT remains within 1.4% of FedAvg while improving upon FedPMT by 2.64% and HeteroFL by 31.00%. In this setting, FedRolex and FedDrop fail and remain below 15% accuracy. These results help explain why FedPLT performs better than existing heterogeneity-aware methods. FedPMT restricts low-capability clients to training the deeper layers, leading to highly imbalanced contribution vectors and leaving the shallow layers undertrained across the heterogeneous client datasets. HeteroFL also trains the full model unevenly, since clients operate on reduced-width local models, which results in biased training. FedDrop and FedRolex use temporal masking,
where the set of trained parameters changes from one round to another, leading to inconsistent training. In contrast, FedPLT assigns trainable model parts in a fixed and balanced manner, while rotating sublayers across clients so that all parameters are trained more uniformly across the heterogeneous client population. This leads to more stable and less biased training and helps explain the stronger accuracy observed in Fig. 10 and Table VII. Beyond mean accuracy, FedPLT also exhibits markedly reduced variance across runs. For example, on CIFAR-10 with the FCN model, FedPLT achieves a standard deviation of 0.71, compared to 4.12 for FedAvg and larger values for the competing methods. This improved stability reflects the structured and deterministic nature of FedPLT’s partial updates, which avoids both the rigidity of suffix-layer training and the randomness of sub-model training schemes. Table VIII complements the final-accuracy comparison by reporting equal-accuracy and equal-communication evaluations on CIFAR-10. Across both models, FedPLT reaches the target accuracy with substantially lower total system communication and computation than FedAvg, while several competing baselines fail to reach the target. For example, with ResNet8, FedPLT reaches the 75% target using roughly half the communication and computation required by FedAvg. Under the same total communication budget, FedPLT also achieves higher accuracy than FedAvg for both models and remains clearly stronger than the other heterogeneity-aware baselines in the CIFAR-10 settings. These results show that, even in heterogeneous systems, FedPLT remains considerably more resource-efficient and converges faster under a limited budget than FedAvg and the other baselines. Overall, these experiments demonstrate that, despite severe system heterogeneity, where 60% of clients are limited to training only 6% of the model, FedPLT consistently achieves accuracy comparable to FedAvg while outperforming the other baselines across all configurations. The results also highlight its stability, robustness, and resource efficiency under highly imbalanced client capabilities. Together, these findings suggest that FedPLT effectively reduces optimization bias and variance through fixed and balanced layer-wise training and more uniform parameter exposure across the client population, making it a high-performing and resource-efficient solution for heterogeneous federated systems with severe device constraints. D. FedPLT and Optimal Client Sampling This experiment evaluates how FedPLT interacts with variance-optimal client sampling under severe communication constraints. In particular, we compare the original Optimal Client Sampling (OCS) formulation with our FedPLT-aware OCS to assess whether accounting for partial training ratios improves convergence and generalization in heterogeneous systems. We consider two heterogeneous systems with 100 clients: • Low-resource system: 10% of clients train the full model, while 90% are restricted to training only 20% of the model. • Moderate-resource system: 10% of clients train the full model, while 90% train 50% of the model.
16
Fig. 10. Validation accuracy curves comparing FedPLT with existing methods in a heterogeneous and low-resource FL system, for (Fashion-MNIST, FCN), (CIFAR-10, FCN), and (CIFAR-10, ResNet-8). TABLE VIII C OMPARISON UNDER EQUAL - ACCURACY AND EQUAL - COMMUNICATION EVALUATIONS IN THE HETEROGENEOUS SETTING . F OR EQUAL - ACCURACY, WE REPORT THE TOTAL SYSTEM COMMUNICATION AND COMPUTATION REQUIRED TO REACH THE TARGET ACCURACY. F OR EQUAL - COMMUNICATION , WE REPORT THE ACCURACY ACHIEVED UNDER THE COMMON TOTAL SYSTEM COMMUNICATION BUDGET. CIFAR-10 + FCN Method
FedAvg FedPLT FedPMT FedRolex FedDrop HeteroFL
Target Acc. (45%)
CIFAR-10 + ResNet-8
Comm. Budget (139.26 GB)
Target Acc. (75%)
Comm. Budget (1.77 GB)
Comm.
Comp.
Acc.
Comm.
Comp.
Acc.
304.67 GB 260.26 GB ∞ ∞ ∞ ∞
233.96 TFLOPs 155.90 TFLOPs ∞ ∞ ∞ ∞
39.37% 40.52% 33.45% 37.65% 26.60% 34.81%
4.05 GB 2.72 GB 3.82 GB ∞ ∞ ∞
511.31 TFLOPs 264.82 TFLOPs 353.00 TFLOPs ∞ ∞ ∞
64.98% 71.26% 70.20% 9.96% 14.53% 46.22%
In both settings, we fix the communication budget to κ = 5 out of a maximum of 100, corresponding to extremely limited resource availability. We compare the original OCS formulation, where the budget constrains only the number of selected clients, against our FedPLT-aware OCS, where the budget accounts for the partial training ratios rk .
Fig. 11. Validation accuracy curves comparing OCS with and without FedPLT. Left: original OCS formulation. Right: FedPLT-aware OCS.
With the original OCS formulation, where the communication budget counts only the number of clients but ignores their layer ratios rk , we observe mixed behavior. In the moderateresource case (50% training), FedPLT combined with OCS improves validation accuracy by +3.51% relative to full-model training while reducing communication by 36% (6,375 vs. 10,000 units). However, in the low-resource case (20% training), accuracy drops slightly by -1.11%, despite achieving 43% lower communication (5,700 vs. 10,000 units). This occurs
because the original constraint limits the expected number of selected clients to κ, regardless of whether they train the full model or only a fraction, thereby restricting data coverage per round and biasing sampling toward less informative updates under strong heterogeneity. By contrast, the FedPLT-aware OCS consistently improves performance under the same normalized communication budget of 10,000. Accuracy gains of +3.84% in the moderateresource system and +1.29% in the low-resource system confirm that incorporating the rk constraint yields a more efficient allocation of the communication budget. Since the budget now accounts for trained fractions rather than client counts, the expected number of sampled clients can exceed κ when rk < 1, allowing multiple partial trainers to be included instead of a single full trainer. Importantly, this improvement is not merely due to selecting more clients per round, but also to the interaction between communication-aware sampling and FedPLT’s balanced partial training. By allowing multiple low-rk clients to participate simultaneously, FedPLT-aware OCS exposes a larger portion of the global model to more diverse data, leading to improved generalization and more stable convergence. Overall, these experiments highlight that FedPLT benefits significantly from being integrated with communicationaware sampling schemes. While standard OCS already reduces communication cost, our FedPLT-aware extension achieves both higher accuracy and fairer client participation under heterogeneous system and data constraints, further reinforcing
17
FedPLT’s scalability and robustness. IX. C ONCLUSION In this paper, we proposed FedPLT, a federated learning framework that enables flexible and fine-grained partial parameter training in heterogeneous systems. Unlike existing submodeling and partial-training approaches that rely on coarse layer-wise freezing or stochastic parameter selection, FedPLT decomposes each layer into sub-layers and assigns them deterministically across clients, ensuring balanced parameter exposure and more consistent optimization of the global model despite highly unequal device capabilities. We introduced a principled formulation of partial training through the training ratio rk , layer-wise allocation vector Qk , and contribution vector Xk , and showed empirically that balanced contribution configurations yield stable optimization and strong performance across both FCN and CNN architectures. Building on this insight, FedPLT distributes training effort across layers in a structured manner while preserving the full forward pass, allowing all model components to be optimized collaboratively even when individual clients can train only small fractions of the model. Extensive experiments demonstrate that FedPLT achieves accuracy comparable to or exceeding FedAvg while substantially outperforming existing heterogeneity-aware baselines under severe system constraints, including scenarios where up to 60% of clients train only 6% of the model. Moreover, FedPLT consistently exhibits reduced variance across runs, indicating improved stability and robustness in non-IID environments. We further extended Optimal Client Sampling to account for partial training ratios and showed that FedPLTaware sampling yields additional gains in communication efficiency and generalization. Collectively, these results establish FedPLT as a scalable and resource-efficient alternative to full-model federated learning, capable of maintaining near full-model federated performance while accommodating extreme device heterogeneity. By enabling flexible partial participation without sacrificing model coherence, FedPLT is well suited to IoT environments, where edge devices such as sensors, wearables, smart cameras, drones, and gateways often operate under strict computation and communication constraints. It thus opens new opportunities for deploying large-scale federated learning systems on resource-constrained edge devices and heterogeneous realworld infrastructures.
18
A PPENDIX A G EOMETRIC I NTERPRETATION OF THE M ASKED G RADIENT A LIGNMENT A SSUMPTION In the convergence analysis, we introduce the following assumption specific to static partial training schemes. For each client k and model iterate W , define the masking attenuation ratio ρk (W ) :=
⟨mk ⊙ ∇Fk (W ), W − W ∗ ⟩ , ⟨∇Fk (W ), W − W ∗ ⟩
whenever the denominator is nonzero, where W ∗ denotes a global minimizer. The assumption requires that ρk (W ) > 0. This ratio measures how masking affects the alignment between the local gradient and the direction W − W ∗ . To interpret it geometrically, let ∗ • θ1 be the angle between ∇Fk (W ) and W − W , ∗ • θ2 be the angle between mk ⊙ ∇Fk (W ) and W − W , • θ3 be the angle between mk ⊙ ∇Fk (W ) and ∇Fk (W ). Then ⟨∇Fk (W ), W − W ∗ ⟩ = ∥∇Fk (W )∥ ∥W − W ∗ ∥ cos(θ1 ), ⟨mk ⊙ ∇Fk (W ), W − W ∗ ⟩ = ∥mk ⊙ ∇Fk (W )∥ ∥W − W ∗ ∥ cos(θ2 ), and, since masking corresponds to a coordinate-wise projection onto the active coordinates, ∥mk ⊙ ∇Fk (W )∥ = ∥∇Fk (W )∥ cos(θ3 ). Substituting into the definition of ρk (W ) gives ρk (W ) =
cos(θ3 ) cos(θ2 ) . cos(θ1 )
Interpretation: Since cos(θ3 ) ≥ 0, the condition ρk (W ) > 0 requires cos(θ1 ) and cos(θ2 ) to have the same sign. In other words, masking must preserve the overall descent orientation of the local gradient relative to the direction W − W ∗ . Thus, if the full gradient is aligned with progress toward W ∗ , then the masked gradient must remain aligned in the same general direction. Moreover: ∗ • If masking improves alignment with the direction W − W , that is, cos(θ2 ) > cos(θ1 ), then ρk (W ) > cos(θ3 ). • If masking degrades alignment, that is, cos(θ2 ) < cos(θ1 ), then ρk (W ) < cos(θ3 ). Therefore, ρk (W ) quantifies how much masking preserves or degrades the directional usefulness of the local gradient. Why the assumption is reasonable: Under balanced participation and mild heterogeneity, local gradients are typically expected to contribute to global progress. The assumption ρk (W ) > 0 rules out degenerate masking patterns that would reverse this contribution. It is therefore a natural condition ensuring that static partial training preserves a meaningful descent direction at the client level.
19
A PPENDIX B C ONVERGENCE P ROOF OF S TATIC PARTIAL PARAMETER U PDATE S TRATEGIES A. Introduction 1) Objective: In this subsection, we present a detailed convergence analysis of any static partial model training scheme, including our proposed method FedPLT. The objective is to show that, under standard assumptions, such methods converge. We consider the case where each client’s objective is smooth and strongly convex, and the global objective is defined as a weighted sum of local losses. Our analysis takes into account the key conditions: • Static Partial Layer Training: Clients update only a fixed subset of model parameters using predefined binary masks over all training rounds. • Stochastic Gradient Descent (SGD): Clients train locally using mini-batch SGD, leading to stochastic updates. The convergence analysis aims to: • Quantify how these factors (partial updates and stochastic gradients) influence the global model update and convergence. • Derive a bound on the expected squared distance between the global model at any round and its optimal value over t rounds. • Establish a convergence rate to stationary points in the non-convex setting. The overall goal is to demonstrate that despite the constraint imposed by restricted update masks, FedPLT, as well as other static partial model update methods, remains an effective and theoretically grounded algorithm with provable convergence guarantees. 2) Problem Formulation: We consider the standard federated learning optimization problem, where the goal is to minimize a global objective function defined as the weighted average of local losses across K clients: min F (W ) :=
W ∈Rd
K X nk k=1
n
Fk (W ),
with n =
K X
nk ,
(18)
k=1
where Fk (W ) is the empirical risk on client k given by: n
k 1 X (k) (k) Fk (W ) := ℓ(W ; xi , yi ), nk i=1
(k)
(19)
(k)
and ℓ(·; x, y) is the per-sample loss function evaluated on a data point (xi , yi ) from client k. 3) Partial Update Rule: Let W t ∈ Rd denote the global model at communication round t. Each client k is associated with a binary mask vector mk ∈ {0, 1}d , which specifies the subset of model parameters that client k is responsible for updating. Specifically, for each coordinate i ∈ {1, . . . , d}: ( 1 if client k is assigned to update coordinate i, (20) (mk )i = 0 otherwise. The mask mk reflects the client’s computational budget and its assigned sub-layers in the model. Let rk = ∥mk ∥0 /d denote the fraction of coordinates that client k is allowed to update. At the beginning of round t, each selected client k receives the current global model W t and initializes its local model: Wkt,0 := W t . Each client performs τ local update steps, only modifying the coordinates specified by mk . We distinguish between two possible settings: 1) Local Stochastic Gradient Descent (SGD): In the default and practical setting, each client uses mini-batch SGD. At (s) each local step s, a random mini-batch ξk is sampled from client k’s local dataset, and the update rule is: (s) Wkt,s+1 = Wkt,s − ηk mk ⊙ ∇ℓ(Wkt,s ; ξk ) , s = 0, . . . , τ − 1, (21) where ηk > 0 is the local learning rate, and ∇ℓ(W ; ξ) is the stochastic gradient over the mini-batch ξ. 2) Local Full Gradient Descent (FGD): In the idealized setting, each client uses full-batch gradient descent on its local data. The update rule becomes: Wkt,s+1 = Wkt,s − ηk mk ⊙ ∇Fk (Wkt,s ) , s = 0, . . . , τ − 1, (22) where ∇Fk (W ) is the full local gradient: n
∇Fk (W ) :=
k 1 X (k) (k) ∇ℓ(W ; xi , yi ). nk i=1
20
After completing local training, the client returns only its local updates: mk ⊙ Ukt in case of SGD or mk ⊙ Ūkt for FGD, where: Ukt = Wkt,τ − Wkt,0 = −ηk
τ −1 X
(s)
(23)
∇Fk (Wkt,s )
(24)
∇ℓ(Wkt,s ; ξk )
s=0
and Ūkt = Wkt,τ − Wkt,0 = −ηk
τ −1 X s=0
The server aggregates the received partial local updates using coordinate-wise weighted averaging. We define the normalized weight of client k as: nk ck := PK
j=1 nj
,
so that
K X
ck = 1.
k=1
Then the global model update at round t can be written as: t
∆ :=
K X
! ck · (mk ⊙ Ukt )
,
k=1
B. Convergence Analysis In this subsection, we present the convergence analysis of static partial training schemes, including our proposed method, FedPLT. We begin by introducing a set of standard assumptions commonly adopted in the federated learning literature to ensure theoretical rigor. Additionally, we propose a new assumption specifically tailored to the partial update setting, which captures the alignment properties between masked and full gradients. Our proof follows the general structure of the analysis in [21], which we extend to accommodate the static partial update mechanisms used in schemes like FedPLT. We detail each step of the convergence proof, highlighting where our contributions diverge from or generalize existing results. In the cited paper, virtual aggregation is applied at each local step s, whereas in a real setting, aggregation occurs only once every τ steps. We will be using their proving flow by using this notation and then, when necessary in the analysis, we revert to the per-step formulation by explicitly indexing local iterations s. Thus, we use the simplified update notation: Ukt = −η t ∇ℓ(Wkt ; ξk ),
and Ūkt = −η t ∇Fk (Wkt ),
1) Finding the Expected Distance Toward the Optimal Global Model: We aim to analyze the expected squared distance to the optimal global model W ∗ after round t + 1 of federated training with partial updates. Specifically, our goal is to bound the following quantity: E ∥W t+1 − W ∗ ∥2 , where the expectation is taken over the randomness in local mini-batch selection. This quantity captures the expected gap between the global model and the optimum across rounds. Step 1: Expanding the Distance to the Optimum: We begin by expanding the global model’s parameters using the global update as follow: W t+1 = W t + ∆t , where the global update is defined as: t
∆ :=
K X
ck mk ⊙ Ukt ,
k=1
and each client’s stochastic local update Ukt is given by: Ukt = −η t ∇ℓ(Wkt ; ξk ), thus we have: ∥W t+1 − W ∗ ∥2 = ∥W t + ∆t − W ∗ ∥2 .
21
We then introduce the full (non-stochastic) update: ¯ t := ∆
K X
ck mk ⊙ Ūkt ,
where Ūkt = −η t ∇Fk (Wkt ).
k=1
by add and subtract it inside the norm: ¯ t ) + (∆t − ∆ ¯ t) ∥W t+1 − W ∗ ∥2 = (W t − W ∗ + ∆
2
.
Applying the Euclidean norm expansion (a.k.a. the cosine or parallelogram identity), we obtain: ¯ t ∥2 + 2 ∆t − ∆ ¯ t, W t − W ∗ + ∆ ¯t . ¯ t ∥2 + ∥W t − W ∗ + ∆ ∥W t+1 − W ∗ ∥2 = ∥∆t − ∆ {z } | {z } | | {z } A1
A2
(25)
A3
Step 2: Expectation of the Cross-Term A3 : We now analyze the expectation of A3 . The key idea is that the stochastic ¯ t arises from the randomness due to mini-batch sampling at round t, while the other term in the inner error term ∆t − ∆ product is deterministic. Specifically, we compute: ¯ t , ∆t − ∆ ¯ t⟩ . Eζ t [A3 ] = Eζ t 2⟨W t − W ∗ + ∆ Note that: t ∗ ¯ t are deterministic with respect to the mini-batch sampling at round t, • W , W , and ∆ t t ¯ is a zero-mean random variable: • ∆ −∆ ¯ t ] = 0. Eζ t [∆t − ∆ Therefore, the expectation of the inner product vanishes: ¯ t , ∆t − ∆ ¯ t ⟩ = 0, Eζ t ⟨W t − W ∗ + ∆ which implies that: E[A3 ] = 0. The term A3 contributes nothing in expectation. Hence, the expected squared distance simplifies to: E ∥W t+1 − W ∗ ∥2 = E [A1 ] + E [A2 ] . Step 3: Bounding the Stochastic Gradient Variance Term (A1 ): We now provide a concise upper bound for the variance of the global update due to stochastic gradient noise. Recall that the actual update is: ∆t := W t+1 − W t =
K X
ck (mk ⊙ Ukt ),
k=1
and the full gradient update is: ¯ t := ∆
K X
ck (mk ⊙ Ūkt ),
k=1
where Ukt and Ūkt are the stochastic and full local updates, respectively. Letting δkt := Ukt − Ūkt , the update difference becomes: ¯t = ∆t − ∆
K X
ck (mk ⊙ δkt ).
k=1
We now bound the squared norm: ¯ t ∥2 = ∥∆ − ∆ t
K X
2
ck (mk ⊙ δkt )
.
k=1
We take the expectation of the squared norm: 2 K K X X t ¯ t ∥2 = E E ∥∆ − ∆ ck (mk ⊙ δkt ) ≤ · c2k E ∥mk ⊙ δkt ∥2 . k=1
Now applying the masking inequality ∥mk ⊙ δkt ∥ ≤ αrk ∥δkt ∥, we get:
k=1
22
K X ¯ t ∥2 ≤ · E ∥∆t − ∆ c2k αr2k · E[∥δkt ∥2 ]. k=1
Finally, applying Assumption 3 (bounded variance): E[∥δkt ∥2 ] = E[∥Ukt − Ūkt ∥2 ] ≤ (η t )2 σ 2 , We conclude: ¯ t ∥2 ] ≤ (η t )2 σ 2 · E[∥∆t − ∆
K X
c2k αr2k .
k=1
Step 4: Decomposition of A2 and Finding its expectation: Recall the second term in our main inequality: ¯ t ∥2 , A2 := ∥W t − W ∗ + ∆ We expand this term as: ¯ t ∥2 + 2⟨W t − W ∗ , ∆ ¯ t⟩ . A2 = ∥W t − W ∗ ∥2 + ∥∆ {z } | {z } | A2.1
A2.2
We define the following components for further analysis: ¯ t ∥2 , A2.1 := ∥∆ ¯ t ⟩. A2.2 := 2⟨W t − W ∗ , ∆ Step 4.1: Bounding A2.1 (Ideal Update Norm): Recall the definition of the ideal global update: ¯ t := ∆
K X
ck mk ⊙ Ūkt ,
k=1
where Ūkt = −η t ∇Fk (Wkt ) is the full local gradient descent. Using the masking inequality ∥mk ⊙ v∥ ≤ αrk ∥v∥, and applying the convexity of the squared norm ∥ · ∥2 , we get: ¯t 2
∥∆ ∥ =
K X
2
ck (mk ⊙ Ūkt )
≤
K X
2
ck (mk ⊙ Ūkt )
≤
ck ∥mk ⊙ Ūkt ∥2 .
k=1
k=1
k=1
K X
Applying the masking bound: ∥mk ⊙ Ūkt ∥ ≤ αrk ∥Ūkt ∥,
∥mk ⊙ Ūkt ∥2 ≤ αr2k ∥Ūkt ∥2 ,
so
we obtain: ¯ t ∥2 ≤ ∥∆
K X
ck αr2k ∥Ūkt ∥2 .
k=1
We now bound ∥Ūkt ∥2 : ∥Ūkt ∥2 = (η t )2 ∥∇Fk (Wkt )∥2 . Using the L-smoothness of Fk and the local optimal value Fk∗ := minW Fk (W ), we apply: ∥∇Fk (Wkt )∥2 ≤ 2L Fk (Wkt ) − Fk∗ . Hence, ∥Ūkt ∥2 ≤ 2(η t )2 L Fk (Wkt ) − Fk∗ . ¯ t ∥2 , we obtain: Substituting into the bound on ∥∆ ¯ t ∥2 ≤ 2L(η t )2 A2.1 = ∥∆
K X k=1
αr2k ck Fk (Wkt ) − Fk∗ .
23
Step 4.2: Bounding the Inner Product Term A2.2 : We now focus on bounding the term: ¯t . A2.2 := 2 W t − W ∗ , ∆ ¯ t using the aggregation rule: We start by expanding ∆ ¯ t := ∆
K X
ck (mk ⊙ Ūkt ),
where Ūkt := −η∇Fk (Wkt ),
k=1 ∗
t
We decompose W − W as: W t − W ∗ = (W t − Wkt ) + (Wkt − W ∗ ). Substituting this into the inner product and splitting the sum: * + K X t ∗ t A2.2 = 2 W − W , ck (mk ⊙ Ūk ) k=1
=2
K X
W t − Wkt , ck mk ⊙ Ūkt
+ Wkt − W ∗ , ck mk ⊙ Ūkt
.
k=1
Part 1: Applying the Cauchy–Schwarz inequality, we get: W t − Wkt , ck mk ⊙ Ūkt ≤ ∥ ck mk ⊙ Ūkt ∥ · ∥W t − Wkt ∥ ≤ ·ck · ∥mk ⊙ Ūkt ∥ · ∥W t − Wkt ∥, = ·η t · ck · ∥mk ⊙ ∇Fk (Wkt )∥ · ∥W t − Wkt ∥, √ Then, applying the masking bound ∥mk ⊙ v∥ ≤ αrk ∥v∥, where αrk := min(1, drk ), we obtain: W t − Wkt , ck mk ⊙ Ūkt ≤ η t · ck · αrk · ∥∇Fk (Wkt )∥ · ∥W t − Wkt ∥. Now applying the AM–GM inequality with parameter η t > 0, we get: 1 ∥W t − Wkt ∥ · ·αrk · ∥∇Fk (Wkt )∥ ≤ t ∥W t − Wkt ∥2 + η t ( αrk )2 ∥∇Fk (Wkt )∥2 . η Multiplying both sides by the positive scalar ck · η t , we obtain: W t − Wkt , ck mk ⊙ Ūkt ≤ ck ∥W t − Wkt ∥2 + (η t αrk )2 ∥∇Fk (Wkt )∥2 . Recall the L-smoothness of Fk , which implies ∥∇Fk (Wkt )∥2 ≤ 2L Fk (Wkt ) − Fk∗ , Substituting it into the previous inequality: W t − Wkt , ck mk ⊙ Ūkt
≤ ck ∥W t − Wkt ∥2 + (η t αrk )2 L Fk (Wkt ) − Fk∗
.
Part 2: We aim to bound the inner product: Wkt − W ∗ , ck mk ⊙ Ūkt ,
where Ūkt := −η t ∇Fk (Wkt ).
Expanding, we obtain: Wkt − W ∗ , ck mk ⊙ Ūkt
= −η t Wkt − W ∗ , ck mk ⊙ ∇Fk (Wkt ) .
Then we apply the client-wise masked alignment ratio: ρtk :=
⟨mk ⊙ ∇Fk (Wkt ), Wkt − W ∗ ⟩ , ⟨∇Fk (Wkt ), Wkt − W ∗ ⟩
which measures how well the masked gradient aligns with the true gradient. Under Assumption 5, we assume the alignment is uniformly lower bounded across all rounds and clients: ρ :=
min k∈[K], t∈N
ρtk > 0.
Using the strong convexity of Fk , we have: ∇Fk (Wkt ), Wkt − W ∗ ≥ Fk (Wkt ) − Fk (W ∗ ) +
µ ∥Wkt − W ∗ ∥2 . 2
24
Combining the above, we obtain: Wkt − W ∗ , (ck mk ⊙ Ūkt ) = −η t ck Wkt − W ∗ , mk ⊙ ∇Fk (Wkt ) ≤ −η t ck Wkt − W ∗ , mk ⊙ ∇Fk (Wkt ) = −η t ρtk ck ∇Fk (Wkt ), Wkt − W ∗ h i µ ≤ −η t ρtk ck Fk (Wkt ) − Fk (W ∗ ) + ∥Wkt − W ∗ ∥2 . 2 Finally, using ρtk ≥ ρ, we get the uniform lower bound: i h µ Wkt − W ∗ , (ck mk ⊙ Ūkt ) ≤ −η t ρ ck Fk (Wkt ) − Fk (W ∗ ) + ∥Wkt − W ∗ ∥2 . 2 Part 3: Final Bound on A2.2 Combining the two parts, we obtain: K X ck ∥W t − Wkt ∥2
¯t ≤ A2.2 = 2 W t − W ∗ , ∆
k=1
+ 2(η t αrk )2 L Fk (Wkt ) − Fk∗ µ t ∗ 2 t t ∗ , − 2η ρ ck Fk (Wk ) − Fk (W ) + ∥Wk − W ∥ 2 Step 4.3: Grouping and Rearranging the Terms of A2 : Recall the decomposition: ¯ t ∥2 = ∥W t − W ∗ ∥2 + ∥∆ ¯ t ∥2 + 2⟨W t − W ∗ , ∆ ¯ t⟩ . A2 := ∥W t − W ∗ + ∆ | {z } | {z } | {z } Initial distance
A2.1
A2.2
Substituting the bounds from previous steps, we obtain: A2 ≤ ∥W t − W ∗ ∥2 | {z }
Initial distance K X t 2 2 + 2(η ) αrk L ck (Fk (Wkt ) − Fk∗ ) k=1
|
{z
}
Ideal update norm (A2.1 )
K X + ck ∥W t − Wkt ∥2 − µη t ρ ck ∥Wkt − W ∗ ∥2 k=1
| +
{z
}
Model difference terms from A2.2
K X t 2(η αrk )2 L ck (Fk (Wkt ) − Fk∗ ) − 2η t ρ ck (Fk (Wkt ) − Fk (W ∗ )) . k=1
|
{z
}
Loss terms from A2.2
Combining the result and regrouping all components, we obtain: K X A2 ≤ ∥W t − W ∗ ∥2 + ck ∥W t − Wkt ∥2 − µη t ρ ck ∥Wkt − W ∗ ∥2 k=1 t
+ 4(η αrk )
2
L ck (Fk (Wkt ) − Fk∗ ) − 2η t ρ ck (Fk (Wkt ) − Fk (W ∗ ))
We recall that the global model at round t can be written as a weighted average of local models: Wt =
K X
ck Wkt ,
with
k=1
nk ck := PK
j=1 nj
Using the convexity of the squared norm, we apply: ∥W t − W ∗ ∥2 ≤
K X k=1
ck ∥Wkt − W ∗ ∥2 .
.
.
25
Substituting into the earlier bound, we obtain: A2 ≤(1 − µη t ρ) · ∥W t − W ∗ ∥2 +
+
K X
ck ∥W t − Wkt ∥2
k=1 K X
4(η
t
αrk ) L ck (Fk (Wkt ) − Fk∗ ) − 2η t ρ ck (Fk (Wkt ) − Fk (W ∗ )) 2
.
k=1
|
{z
}
A2.4
Step 4.4: Bounding A2.4 : We define A2.4 to capture all the loss-related contributions arising in the expansion of A2 : K X A2.4 := 4(η t αrk )2 L · ck (Fk (Wkt ) − Fk∗ ) − 2η t ρ · ck (Fk (Wkt ) − Fk (W ∗ )) . k=1
We regroup the above term by adding and subtracting Fk∗ in the expression (Fk (Wkt ) − Fk (W ∗ )): K X t 2 t t ∗ t ∗ ∗ A2.4 := 4(η αrk ) L − 2η ρ ·ck (Fk (Wk ) − Fk ) + 2η ρ ·ck (Fk (W ) − Fk ) . | {z } | {z } k=1 := −γkt
:= β t
∗
We insert and subtract F inside both gap terms accordingly: K X t t ∗ ∗ ∗ t ∗ ∗ ∗ ∗ − γk · ck Fk (Wk ) − F + F − Fk + β · ck (Fk (W ) − F + F − Fk ) . A2.4 = k=1
We regroup and simplify to highlight the three key components: K X A2.4 = − γkt · ck Fk (Wkt ) − F ∗ k=1
+ (β t − γkt ) · ck (F ∗ − Fk∗ ) + β t · ck (Fk (W ∗ ) − F ∗ ) . Noting that F ∗ =
PK
k=1 ck Fk (W
∗
), this implies: K X
ck (Fk (W ∗ ) − F ∗ ) = 0.
k=1
We thus eliminate the final term and rewrite the expression as: K X t t ∗ t t ∗ ∗ A2.4 = − γk · ck Fk (Wk ) − F + (β − γk ) · ck (F − Fk ) . k=1
Bounding (Fk (Wkt ) − F ∗ ) gap term By adding and subtracting Fk (W t ), we write: Fk (Wkt ) − Fk∗ = Fk (Wkt ) − Fk (W t ) + Fk (W t ) − F ∗ . Using the convexity of Fk , we have: Fk (Wkt ) − Fk (W t ) ≥ ∇Fk (W t ), Wkt − W t . Applying the AM–GM inequality with parameter η t , we get: ∇Fk (W t ), Wkt − W t ≥ − Using the L-smoothness of Fk :
1 ηt ∥∇Fk (W t )∥2 − t ∥Wkt − W t ∥2 . 2 2η
∥∇Fk (W t )∥2 ≤ 2L Fk (W t ) − Fk∗ ,
we obtain the following bound: 1 Fk (Wkt ) − F ∗ ≥ Fk (W t ) − F ∗ − η t L Fk (W t ) − Fk∗ − t ∥Wkt − W t ∥2 . 2η
26
Now we decompose the bound by adding and subtracting F ∗ in (Fk (W t ) − Fk∗ ): Fk (Wkt ) − Fk∗ ≥ 1 − η t L (Fk (W t ) − F ∗ ) 1 − η t L(F ∗ − Fk∗ ) − t ∥Wkt − W t ∥2 . 2η We substitute the lower bound on Fk (Wkt ) − F ∗ into the original upper bound for A2.4 . This gives: K X 1 A2.4 ≤ γkt · ck (η t L − 1)(Fk (W t ) − F ∗ ) + η t L(F ∗ − Fk∗ ) + t ∥Wkt − W t ∥2 2η k=1 + (β t − γkt ) · ck (F ∗ − Fk∗ ) . We now regroup the terms in the corrected upper bound of A2.4 : K X A2.4 ≤ γkt (η t L − 1) · ck (Fk (W t ) − F ∗ ) k=1
+ γkt (η t L − 1) + β t · ck (F ∗ − Fk∗ ) 1 t t t 2 + γk · t · ck ∥Wk − W ∥ . 2η We now simplify the bound on A2.4 under the following conditions: 1 t t • η < L ⇒ (η L − 1) < 0, t t t • γk > 0, so γk (η L − 1) < 0, γkt 1 t t t 2 t t 2 • 2η t < 1, so γk · 2η t · ck ∥Wk − W ∥ < ck ∥Wk − W ∥ , • The combined coefficient: γkt (η t L − 1) + β t = 2(η t )2 L · 2αr2k (1 − η t L) + ρ ≤ 2(η t )2 Lν, where ν := maxt 2αr2k (1 − η t L) + ρ > 0. Applying these observations to the previous bound: K X A2.4 ≤ 2(η t )2 Lν · ck (F ∗ − Fk∗ ) k=1
+ ck ∥Wkt − W t ∥2
.
We define the aggregate heterogeneity gap: Λ :=
K X
ck (F ∗ − Fk∗ ),
k=1
thus, Applying these observations to the previous bound: K X A2.4 ≤ 2(η t )2 Lν · ck Λ k=1
+ ck ∥Wkt − W t ∥2 . Step 4.5: Regrouping the Final Bound on A2 : We now regroup and simplify the final upper bound on A2 , using the decompositions derived in previous steps. Recall that: A2 ≤(1 − µη t ρ) · ∥W t − W ∗ ∥2 +
+
K X
ck ∥W t − Wkt ∥2
k=1 K X
4(η
t
αrk ) L ck (Fk (Wkt ) − Fk∗ ) − 2η t ρ ck (Fk (Wkt ) − Fk (W ∗ )) 2
.
k=1
|
{z
A2.4
}
27
We substitute the simplified bound for the loss-related component A2.4 derived in step 6.4, yielding: A2 ≤(1 − µη t ρ)∥W t − W ∗ ∥2 +
K X
2ck ∥Wkt − W t ∥2
k=1
+ 2(η t )2 Lν · Λ. Step 4.6: Taking the Expectation of A2 : To compute E[A2 ], we note that: "K # X t t ∗ 2 t t 2 E[A2 ] ≤ (1 − µη ρ)∥W − W ∥ + 2 E ck ∥Wk − W ∥ + 2(η t )2 Lν · Λ. k=1
We now bound the expected local divergence term: "K # X t t 2 E ck ∥Wk − W ∥ . k=1
Recall that: t,τ t • Wk = Wk is the local model of client k after τ local steps, t,0 t • W = Wk is the global model at round t, identical to the local initialization, • Local updates are computed with coordinate-wise masks: (s) Wkt,s+1 = Wkt,s − η t mk ⊙ ∇ℓ(Wkt,s ; ξk ) . By unrolling the update over τ local steps, we obtain: Wkt − W t = Wkt,τ − Wkt,0 = −η t
τ −1 X
(s) mk ⊙ ∇ℓ(Wkt,s ; ξk ) .
s=0
Taking the squared norm and applying Jensen’s inequality: 2 τ −1 X (s) E ∥Wkt − W t ∥2 = (η t )2 · E mk ⊙ ∇ℓ(Wkt,s ; ξk ) s=0
≤ (η t )2 · τ
τ −1 X
h i (s) E ∥mk ⊙ ∇ℓ(Wkt,s ; ξk )∥2 .
s=0
Applying the masking bound ∥mk ⊙ v∥ ≤ αrk · ∥v∥ and the variance assumption: h i (s) E ∥∇ℓ(Wkt,s ; ξk )∥2 ≤ G2 + σ 2 , We obtain: E ∥Wkt − W t ∥2 ≤ (η t )2 · τ 2 · αr2k · (G2 + σ 2 ). Finally, summing over clients with weights ck , we get: "K # K X X t t 2 E ck ∥Wk − W ∥ ≤ (η t )2 · τ 2 · (G2 + σ 2 ) ck αr2k . k=1
k=1
We denote this term by: Γ :=
K X
ck αr2k ,
k=1
and thus conclude: E
"K X
# ck ∥Wkt − W t ∥2
≤ (η t )2 · τ 2 · (G2 + σ 2 ) · Γ.
k=1
Combining this with the earlier result for A2 , we obtain the following upper bound in expectation: E[A2 ] ≤ (1 − µη t ρ) · ∥W t − W ∗ ∥2 + 2 (η t )2 · τ 2 · (G2 + σ 2 ) · Γ + 2 (η t )2 Lν · Λ,
28
Step 5: Recursive bound over the expected optimal gap: We recall the Euclidean expansion: ¯ t ∥2 + ∥W t − W ∗ + ∆ ¯ t ∥2 + 2 ∆t − ∆ ¯ t, W t − W ∗ + ∆ ¯t . ∥W t+1 − W ∗ ∥2 = ∥∆t − ∆ | {z } | {z } | {z } A1
A2
A3
Taking expectations over the stochasticity at round t, and applying the derived bounds, we get: E ∥W t+1 − W ∗ ∥2 = E[A1 ] + E[A2 ] + E[A3 ] ≤ (η t )2 σ 2 · Γ + (1 − µη t ρ) · E ∥W t − W ∗ ∥2 + 2 (η t )2 τ 2 (G2 + σ 2 ) · Γ + 2(η t )2 Lν · Λ. We define the sequence: Dt := E ∥W t − W ∗ ∥2 , and obtain the recursive inequality: Dt+1 ≤ 1 − z η t · Dt + (η t )2 B, where: z := µ ρ,
B := 2 τ 2 (G2 + σ 2 ) Γ + 2Lν Λ + σ 2 Γ.
Convergence Analysis of the Recursive Inequality: We recall the recursive inequality derived in Step 5: Dt+1 ≤ (1 − zη t ) · Dt + (η t )2 B, We now analyze the convergence behavior under two scenarios: (1) constant step size, and (2) decaying step size. Case 1: Constant Learning Rate: Assume a constant step size η t = η ∈ (0, z1 ). Then the recursion becomes: Dt+1 ≤ (1 − zη) Dt + η 2 B. Let ω := 1 − zη < 1. Then by recursion: Dt+1 ≤ ω t+1 D0 + η 2 B
t X
ω j = ω t+1 D0 + η 2 B ·
j=0 t+1
Since 1−ω 1−ω
1 − ω t+1 . 1−ω
1 1 ≤ 1−ω = zη , we obtain the final bound:
Dt+1 ≤ ω t+1 D0 +
ηB . z
Interpretation: The first term decays geometrically, while the second term is a residual error that depends on η. Therefore, the convergence ∗ is fast, but with a gap of size: ηB z around the optimal point W . Implication: Choosing a smaller η reduces the residual error, but slows down convergence. This is a speed-error trade-off: a smaller learning rate implies lower asymptotic error but slower descent. Case 2: Decaying Learning Rate: We now analyze the convergence behavior under a decaying learning rate of the form: ηt =
1 , z(t + 1)
Recall the recursive inequality: Dt+1 ≤ (1 − zη t )Dt + (η t )2 B. Substituting the decaying step size into the recursion, we get: 1 B t B Dt+1 ≤ 1 − Dt + 2 = Dt + 2 . t+1 z (t + 1)2 t+1 z (t + 1)2 We now prove by mathematical induction that for all t ≥ 0, there exists a constant C ≥ zB2 such that: Dt ≤ Base Case:
C . t+1
29
At t = 0, D1 ≤ D0 +
B C ≤ , z2 0+1
provided that: C ≥ D0 + Inductive Step: Assume the hypothesis holds at step t, i.e., Dt ≤
B B ≥ 2. z2 z
C . t+1
Then: t C B · + 2 t + 1 t + 1 z (t + 1)2 Ct + zB2 = . (t + 1)2
Dt+1 ≤
C We now compare this to t+2 by defining the function:
f (t) :=
Ct + zB2 C − . 2 (t + 1) t+2
It can be proved through typical methods of calculus that f (t) ≤ 0 if C ≥ zB2 . C holds. This completes the inductive step. Thus, for any C ≥ zB2 , the inequality Dt+1 ≤ t+2 Conclusion: By induction, for all t ≥ 0, we have: C B , for some constant C ≥ 2 . t+1 z Therefore, the expected squared distance to the optimum decays at the rate: 1 t D =O . t Dt ≤
Interpretation: This gives asymptotic convergence to the optimum W ∗ , but at a sublinear rate. The error shrinks slowly over time, especially during early rounds. Summary and Comparison: ηB ∗ • With a constant step size, we converge fast to a small neighborhood around W , with asymptotic error z . ∗ • With a decaying step size, we converge exactly to W , but at a slower O(1/t) rate. Practical Guideline: In practice, we can use a hybrid strategy, starting with a large constant step size to ensure fast initial progress, then decaying it slowly to improve final accuracy.
30
A PPENDIX C S TATIC PARTIAL PARAMETER T RAINING WITH OCS A. Problem Formulation We now consider the optimization problem: min
EAt
{ptk }
X nk k∈At
0 ≤ ptk ≤ 1,
such that
K X
Ukt −
ptk
K X
2
nk Ukt
k=1
∀ k = 1, . . . , N,
rk ptk = m.
k=1
where: ptk is the probability of sampling client k at round t, t • A is the set of clients selected at round t, • nk is the number of data samples for client k, t • Uk is the gradient transferred from client k ar round t, • rk is a ratio in [0, 1] representing the training proportion of client’s k model.
•
B. Problem Simplification Observe that "
X nk
#
Ut ptk k t k∈A
E
At
=
K X
nk Ukt .
k=1
Thus, the objective function X nk
EAt
k∈At
ptk
Ukt −
K X
2
nk Ukt
k=1
is the variance of the first term. By the variance identity, 2 2 2 K K X nk X nk X X t t t t EAt = EAt U − nk Uk U − nk Uk . ptk k ptk k t t k=1
k∈A
k=1
k∈A
By expanding the squared norms into inner products and distributing the transpose over the sum, we obtain !T 2 T ! X X nk X nk X nj X nk nj t t t t t = = U . t Uk t Uk t Uj t Uk p p p p ptj j j k k k t t t t t k∈A
k∈A
j∈A
k∈A
j∈A
Thus, the squared norm can be written as X nk
t t Uk p k∈At k
2
=
X X nk k∈At j∈At
Ut ptk k
T
nj t U ptj j
! .
Next, we split the double summation into two terms: one for the case k = j and one for k ̸= j: ! T T X X T X X nk X nk nj t nk t nk t t t U U = U U + U ptk k ptj j ptk k ptk k ptk k k∈At k∈At j∈At k∈At j∈At | {z } j̸=k | {z terms with k=j
terms with k̸=j
nj t U ptj j
! . }
31
Taking the expectation over At and noting that each client k is selected with probability ptk , we have " !# T X X nk nj t t EAt U U ptk k ptj j k∈At j∈At ! T X T K K X K X nk t nk t nk t nj t = U U + U U E[Ik ] E[Ik Ij ] ptk k ptk k ptk k ptj j j=1 k=1
k=1
j̸=k
=
K X
ptk
k=1
=
nk t U ptk k
T
nk t U ptk k
+
K X K X
ptk ptj
k=1 j=1 j̸=k
K X n2
K X K X t 2 k ∥U ∥ + nk K j k pt k=1 k k=1 j=1 j̸=k
Ukt
T
nk t U ptk k
T
nj t U ptj j
!
Ujt .
Similarly, for the second term we have K X
2
nk Ukt
K X
=
!T K K X K X X T nj Ujt = nk Ukt nk Ukt nj Ujt . j=1
k=1
k=1
k=1 j=1
Dividing the double summation into the diagonal and off-diagonal parts, we obtain K X K X
nk Ukt
K X K K T X T X nk Ukt nj Ujt . nj Ujt = nk Ukt nk Ukt +
T
k=1 j=1
k=1 j=1 j̸=k
k=1
{z
|
}
k=j
|
{z
}
k̸=j
Thus, from the previous reduction, we have 2 K K X nk X X T n2k t t 2 EAt U ∥U nk Kj Ukt Ujt , = ∥ + k k t t pk pk t k=1
k∈A
K X
2
nk Ukt
=
k=1
K X
k,j=1 i̸=j
n2k ∥Ukt ∥2 +
k=1
K X
nk Kj Ukt
T
Ujt .
k,j=1 i̸=j
By rewriting the variance in terms of the separated summations, the total variance is given by 2 2 K X X nk t t U − nk Uk Var = EAt pt k k=1 k∈At k =
K X
k=1
=
=
n2k ∥Ukt ∥2 + ptk
K 2 X n k=1 K X k=1
k − n2k ptk
n2k ∥Ukt ∥2
K X
nk Kj Ukt
T
k,j=1 i̸=j
Ujt −
K X
n2k ∥Ukt ∥2 +
k=1
k,j=1 i̸=j
∥Ukt ∥2
1 − 1 . ptk
Thus, the original optimization problem can be rewritten as min
{ptk }
K X
n2k ∥Ukt ∥2
k=1
1 −1 ptk
subject to 0 ≤ ptk ≤ 1,
∀ i,
and
K X k=1
K X
rk ptk = m.
nk Kj Ukt
T
Ujt
32
C. Variable Substitution We perform the substitutions: qk =
1 −1 ptk
ptk =
=⇒
1 , qk + 1
and define xk = n2k ∥Ukt ∥2 . Then, the objective becomes K X
min {qk }
xk q k ,
k=1
with the constraints: qk ≥ 0,
∀ i,
and the modified equality constraint is now K X k=1
rk = m. qk + 1
D. Lagrangian Formulation and KKT Conditions Define the Lagrangian with Lagrange multiplier λ for the equality constraint and multipliers µk ≥ 0 for the non-negativity constraints on qk : ! K K K X X X rk µk qk . L({qk }, λ, {µk }) = xk q k + λ −m − qk + 1 k=1
k=1
k=1
The KKT conditions are: 1) Stationarity: For each k = 1, . . . , K, ∂L rk − µk = 0. = xk − λ ∂qk (qk + 1)2 2) Primal Feasibility: qk ≥ 0,
∀ i,
and
K X k=1
rk = m. qk + 1
3) Dual Feasibility: µk ≥ 0,
∀ i.
µk qk = 0,
∀ i.
4) Complementary Slackness: These conditions will be used to solve for the optimal {qk }, from which the original probabilities are recovered via ptk =
1 . qk + 1
E. Solving the KKT Conditions The stationarity condition is xk −
λ rk − µk = 0. (qk + 1)2
Isolating µk yields xk −
λ rk = µk . (qk + 1)2
We now consider two cases: Case 1: qk > 0.: Then complementary slackness implies µk = 0, so that xk = Solving for qk gives
λ rk . (qk + 1)2
r qk =
λ rk − 1. xk
33
Case 2: qk = 0.: In this case, the stationarity condition becomes xk − λ rk = µk ≥ 0, which implies xk ≥ λ r k . Thus, we have: q For qk > 0: qk = λxrkk − 1. • For qk = 0: xk ≥ λ rk .
•
F. Determining λ via the Equality Constraint The modified equality constraint is K X k=1
rk = m. qk + 1
Define the index set O = {k ∈ {1, . . . , K} : xk < λ rk }, so that for k ∈ O (by Case 1) we have 1 = qk + 1
r
xk , λ rk
and for indices k ∈ Oc = K − O (where qk = 0 by Case 2) we have 1 = 1. qk + 1 Thus, the equality constraint becomes X
r rk
k∈O
X xk + rk = m. λ rk c k∈O
note that X
rk =
k∈O c
K X
X
rk −
k=1
rk .
k∈O
Thus, the equality constraint becomes K X
!
1 X√ √ rk xk + λ k∈O
k=1
P
√
Isolating λ yields λ=
rk −
X
rk
= m.
k∈O
rk xk k∈O P PK m − k=1 rk + k∈O rk
!2 .
G. Finding the Explicit Solution for qk For k ∈ O (i.e. where qk > 0), substituting the expression for λ into r λ rk − 1, qk = xk yields
√ qk =
m−
PK
j=1 rj +
√
P
rk P
j∈O rj
·
j∈O
√
rj xj
xk
For k ∈ Oc , we set qk = 0. Thus, the piecewise definition is: P √ √ rj xj rk j∈O · − 1, √ PK P x qk = m − j=1 rj + j∈O rj k 0,
− 1.
if k ∈ O, if k ∈ Oc .
34
H. Returning to the Original Variable ptk Recall that ptk =
1 . qk + 1
Thus, for k ∈ O (where qk > 0) we have ptk =
m−
PK
j=1 rj +
√
√
P
j∈O rj
rk
xk , √ rj xj j∈O
·P
c
and for k ∈ O (where qk = 0) we obtain ptk = 1. Substituting the expression and recalling that xk = n2k ∥Ukt ∥2 , the final solution is PK P t m − j=1 rj + j∈O rj P nk ∥Uk ∥ , · √ √ t pk = rk rj Kj ∥Ujt ∥ j∈O 1, Where
( O=
k ∈ {1, . . . , K} :
√
if k ∈ Oc .
√ rj Kj ∥Ujt ∥ j∈O PK P m − j=1 rj + j∈O rj P
rk Kk ∥Ukt ∥ <
if k ∈ O,
) .
I. Determining the Index Set O Since for indices in Oc we have ptk = 1, the equality constraint K X
rk ptk = m
k=1
becomes
X
rk ptk +
X
rk = m.
k∈O c
k∈O
and by splitting the second term, we got: X
rk ptk +
K X
rk −
k=1
k∈O
X
rk = m.
k∈O
equivalent to X
K X rk − m. rk 1 − ptk =
k∈O
k=1
Thus, the indices in O must satisfy X k∈O
rk >
K X
rk − m.
k=1
PK That is, the total r-mass in O must exceed k=1 rk − m. The procedure to determine O and its complement Oc is as follows: √ 1) Initialization: Order the indices in increasing order of rk Kkt ∥Ukt ∥. Begin by selecting the indices corresponding to P PK the smallest values until the cumulative sum k∈O rk exceeds k=1 rk − m. √ 2) Iteration: Let Oc = {1, . . . , N } \ O. For the next candidate index k in Oc (with the next smallest rk ntk ∥Ukt ∥), check whether adding k to O maintains the inequality P √ t t √ j∈O∪{k} rj Kj ∥Uj ∥ t t rk nk ∥Uk ∥ < PK P m − j=1 rj + j∈O rj 3) Update: c c • If the inequality holds, update O ← O ∪ {k} and O ← O \ {k}, and then repeat the iteration. c • If the inequality fails for the candidate k, NO further indices can be added to O; then O and O are finalized.
35
A PPENDIX D D ETAILED D ERIVATIONS AND N UMERICAL A NALYSIS OF E FFICIENCY G AIN A NALYSIS A. Derivations of Efficiency Gains In this appendix, we provide the detailed derivations underlying the efficiency analysis presented in Section VII. 1) Detailed Derivation of Computation Cost: Under full-model training, each client performs both forward and backward computation over all P parameters during each of the τ local iterations. The forward and the backward workloads are, respectively: Ξfull k,fwd = ατ P,
Ξfull k,bwd = βτ P
Hence, the total per-round computation cost under full-model training is full full Ξfull k = Ξk,fwd + Ξk,bwd = (α + β)τ P.
Under FedPLT, the forward pass still involves the entire model. However, only the assigned fraction rk of the model participates in backward computation. Therefore, the forward backward remains unchanged, but backward workload scales with rk . Ξplt Ξplt k,fwd = ατ P, k,bwd = βrk τ P. The total per-round computation cost under FedPLT becomes then plt plt Ξplt k = Ξk,fwd + Ξk,bwd = ατ P + βrk τ P = (α + βrk )τ P.
The absolute computation reduction is plt ∆Ξk = Ξfull k − Ξk = (α + β)τ P − (α + βrk )τ P = β(1 − rk )τ P.
Dividing by Ξfull yields the relative computation reduction: k ∆comp =1− k
Ξplt (α + βrk )τ P α + βrk k =1− =1− . (α + β)τ P α+β Ξfull k
Therefore, the computation gain comes entirely from the backward stage, and it increases as rk decreases. 2) Derivation of Communication Cost: Under full-model training, each client downloads the full global model and uploads the full local update. Since the model size is P s bytes, the per-round communication volume is Θfull k = P s + P s = 2P s. Under FedPLT, the full model is still downloaded, but only the updated fraction rk is uploaded. Thus, Θplt k = P s + rk P s = (1 + rk )P s. The absolute communication reduction is plt ∆Θk = Θfull k − Θk = 2P s − (1 + rk )P s = (1 − rk )P s.
and the elative communication reduction is ∆comm =1− k
Θplt 1 − rk (1 + rk )P s k =1− = . full 2P s 2 Θk
It is also useful to separate downlink and uplink: plt Θfull k,↓ = Θk,↓ = P s,
Θfull k,↑ = P s,
Θplt k,↑ = rk P s.
Hence, the relative uplink reduction is ∆comm,up =1− k
Θplt k,↑ Θfull k,↑
=1−
rk P s = 1 − rk . Ps
Therefore, FedPLT preserves the downlink cost while reducing the uplink proportionally to the assigned fraction rk .
36
3) Derivation of Round Time and Straggler Mitigation: In synchronous FL, the duration of a round is determined by the slowest participating client. For a client k, under full-model training, the forward and backward computation time are respectively, ατ P , γk
βτ P . γk
Ps
Ps
and the download and upload times are respectively Bk↓
,
Bk↑
.
Including the fixed overhead δ, the total per-round time of client k is Tkfull = δ +
βτ P Ps Ps ατ P + + ↓+ ↑ γk γk Bk Bk
=δ+
(α + β)τ P Ps Ps + ↓ + ↑. γk Bk Bk
Hence, the full-model round time is full Tround = max Tkfull = δ + max k k
(α + β)τ P Ps Ps + ↓+ ↑ γk Bk Bk
! .
Let k̂ the straggler client be the one attaining this maximum. Defining γ̂ = γk̂ ,
B̂ ↓ = Bk̂↓ ,
B̂ ↑ = Bk̂↑ ,
full Tround =δ+
(α + β)τ P Ps Ps + + . γ̂ B̂ ↓ B̂ ↑
we obtain
Under FedPLT, the forward computation and model download remain unchanged, whereas the backward computation and upload scale with rk . Therefore, the per-round time of client k becomes ! ! ατ P Ps βτ P Ps plt Tk = δ + + ↓ +rk + ↑ . γk γk Bk Bk | | {z } {z } fixed part
scaled part
FedPLT mitigates the straggler effect by adapting rk to clients’ capacity so that client completion times become approximately balanced: plt plt T1plt ≈ T2plt ≈ · · · ≈ TK ≈ Tround . Choosing a target round duration T imposes Ps ατ P + ↓ + rk Tkplt = δ + γk Bk
βτ P Ps + ↑ γk Bk
! ≈ T.
Hence, to mitigate stragglers, the training ration to be applied by each client k is P Ps T − δ − ατ + ↓ γk B k . rk ≈ βτ P Ps γk + B ↑ k
In practice, this expression is meaningful when the resulting value lies in [0, 1]; otherwise, the target T must be relaxed or the value clipped to the feasible range. Under approximate equalization, the FedPLT round time can be written in terms of the limiting client as ατ P Ps βτ P Ps plt Tround = δ + + + r̂ + , γ̂ γ̂ B̂ ↓ B̂ ↑ where r̂ denotes the fraction assigned to the limiting client.
37
The absolute round-time reduction is then plt full ∆T = Tround − Tround (α + β)τ P Ps Ps ατ P Ps βτ P Ps = δ+ + + − δ+ + + r̂ + ↓ γ̂ γ̂ γ̂ B̂ ↑ B̂ ↓ B̂ ↑ B̂ βτ P βτ P Ps Ps − r̂ = + + ↑ γ̂ γ̂ B̂ ↑ B̂ βτ P Ps = (1 − r̂) . + γ̂ B̂ ↑ full Dividing by Tround gives the relative round-time efficiency:
Ps plt (1 − r̂) βτγ̂P + B̂ ↑ Tround ∆time = 1 − full = (α+β)iP . Ps Ps Tround + B̂ ↓ + B̂ ↑ γ̂
Therefore, the round-time gain comes from reducing the variable part of the slowest client’s workload, namely backward computation and uplink transmission, which directly mitigates the straggler effect. B. Numerical Example We consider a system with K = 5 heterogeneous clients. The global model contains P = 5 × 106 parameters, stored using s = 4 bytes per parameter, so that P s = 20 × 106 bytes ≈ 20 MB. Each client performs τ = 150 local iterations per round. The per-parameter forward and backward costs are α = 2 and β = 4, respectively, with β ≈ 2α. The fixed latency is δ = 0.2 s. Client k has computation speed γk (in GFLOPs/s), downlink bandwidth Bk↓ , and uplink bandwidth Bk↑ (in Mb/s).1 1) Full-Training Per-Client Round Time and FedPLT Equalization: Under full-model training, the per-round time of client k is ! ! P s βτ P P s ατ P + . Tkfull = δ + + ↓ + ↑ γk γk Bk /8 Bk /8 We first compute the workload terms: ατ P = 2 × 150 × 5 × 106 = 1.5 × 109 FLOPs = 1.5 GFLOPs, βτ P = 4 × 150 × 5 × 106 = 3 × 109 FLOPs = 3 GFLOPs. Thus, under full-model training, each client performs (α + β)τ P = 4.5 GFLOPs per round. Let T ∗ := min Tkfull k
denote the target round duration used for equalization. In this example, the fastest full-training client determines the target. FedPLT selects rk so that ! ! Ps βτ P Ps ατ P plt Tk = δ + + ↓ + rk + ↑ ≈ T ∗. γk γk Bk /8 Bk /8 Solving for rk gives T∗ − δ − rk =
ατ P Ps γk + B ↓ /8
βτ P Ps γk + B ↑ /8
k
.
k
Table IX reports the client parameters, the full-training round times, and the resulting FedPLT fractions. The value of T ∗ is highlighted in bold. 1 If bandwidth is expressed in Mb/s, the corresponding transmission rate in MB/s is B/8. Hence, the transmission time of a payload of size P s is
Ps . B/8
38
TABLE IX P ER - CLIENT PARAMETERS , FULL - TRAINING ROUND TIME Tkfull , AND F ED PLT ALLOCATION rk . U NITS : γk IN GFLOP S / S , Bk↓ , Bk↑ IN M B / S , AND Tkfull IN SECONDS . B OLD INDICATES T ∗ = mink Tkfull . Client
γk
Bk↓
Bk↑
Tkfull
rk
C1 (smartphone, WiFi) C2 (smart TV, Ethernet) C3 (drone, 5G) C4 (smart AC, WiFi) C5 (IoT device, LTE/4G)
80 30 18 10 8
120 200 60 50 45
40 50 25 16 12
5.59 4.35 9.52 13.85 17.65
0.69 1.00 0.21 0.078 0.03
2) Computation Efficiency: Under full-model training, Ξfull k = (α + β)τ P = 4.5 GFLOPs. Under FedPLT, Ξplt k = ατ P + βrk τ P = 1.5 + 3rk Therefore, ∆k,comp = 1 −
(GFLOPs).
Ξplt α + βrk k =1− . α+β Ξfull k
For interpretability, we also report the derived computation time (·)
(·)
Tk,comp = with γk in GFLOPs/s. For example, for client C1 ,
Ξk , γk
Ξplt 1 = 1.5 + 3(0.69) = 3.57 GFLOPs,
4.5 = 0.05625 s = 56.3 ms, 80 Table X summarizes the resulting values for all clients.
plt T1,comp =
full T1,comp =
3.57 = 0.0446 s = 44.6 ms. 80
TABLE X P ER - CLIENT COMPUTATION WORKLOAD AND DERIVED COMPUTATION TIME . Ξ IS REPORTED IN GFLOP S AND Tk,comp = Ξ/γk IN MILLISECONDS .
Ξfull (GFLOPs) k full Tk,comp (ms) Ξplt (GFLOPs) k plt Tk,comp (ms) ∆k,comp (%)
C1
C2
C3
C4
C5
4.50 56.3
4.50 150.0
4.50 250.0
4.50 450.0
4.50 562.5
3.57 44.6 20.7
4.50 150.0 0.0
2.13 118.3 52.7
1.73 173.4 61.5
1.59 198.8 64.7
As expected, clients with weaker computational capabilities (C3 –C5 ) receive smaller fractions and thus achieve the largest computation savings. 3) Communication Efficiency: In full-model training, each client exchanges the full model: Θfull k = 2P s = 40 MB. Under FedPLT, Θplt k = (1 + rk )P s = 20(1 + rk ) MB. The relative uplink and total communication savings are ∆comm,up = 1 − rk , k
∆comm,tot = k
1 − rk . 2
For example, for client C4 with r4 = 0.078, Θplt 4 = 20(1 + 0.078) = 21.56 MB, ∆comm,up = 1 − 0.078 = 0.922 = 92.2%, 4 ∆comm,tot = 4
1 − 0.078 = 0.461 = 46.1%. 2
39
Table XI reports the per-client communication costs and the corresponding efficiency gains. TABLE XI P ER - CLIENT COMMUNICATION COST AND EFFICIENCY GAINS . VALUES ARE REPORTED IN MB, ASSUMING P s = 20 MB.
Θfull (MB) k Θplt (MB) k comm,up
∆k (%) ∆comm,tot (%) k
C1
C2
C3
C4
C5
40.0 33.8 31.0 15.5
40.0 40.0 0.0 0.0
40.0 24.2 79.0 39.5
40.0 21.56 92.2 46.1
40.0 20.6 97.0 48.5
These results highlight that the communication gain comes entirely from the reduced uplink, while the model download remains unchanged. C. Round Time and Straggler Mitigation Under full-model synchronous training, the global round time is determined by the slowest client: full Tround = max Tkfull = 17.65 s. k
Under the FedPLT equalization target, all clients are adjusted to finish in approximately plt Tround = T ∗ = 4.35 s.
Therefore, the relative round-time efficiency is ∆time = 1 −
plt Tround 4.35 = 75.35%. =1− full 17.65 Tround
Table XII summarizes the reduction in overall round duration. TABLE XII G LOBAL ROUND TIME BEFORE AND AFTER F ED PLT ( SECONDS ). full Tround
plt Tround
∆time
17.65
4.35
75.35%
Under full training, each client may finish early and then remain idle while waiting for the slowest client. The idle time of client k is full Ikfull = Tround − Tkfull ≥ 0. Under ideal FedPLT equalization, this idle time is removed. For example, for client C1 , I1full = 17.65 − 5.59 = 12.06 s, which corresponds to
12.06 = 68.33%. 17.65 Table XIII reports the idle time avoided for all clients. TABLE XIII full . I DLE TIME AVOIDED BY F ED PLT. A BSOLUTE VALUES ARE IN SECONDS AND RELATIVE VALUES ARE WITH RESPECT TO Tround
Tkfull (s)
Idle avoided (s) Idle avoided (%)
C1
C2
C3
C4
C5
5.59 12.06 68.33
4.35 13.30 75.35
9.52 8.13 46.06
13.85 3.80 21.53
17.65 0.00 0.00
Overall, FedPLT reduces the global round time from 17.65 s to 4.35 s, while also yielding substantial savings in backward computation and uplink communication. This numerical example supports the analytical results by showing how partial-layer training improves efficiency while mitigating the straggler effect in heterogeneous systems.
40
R EFERENCES [1] E. Union, “General data protection regulation (gdpr),” Official Journal of the European Union, L119, pp. 1-88, Apr. 2016. [Online]. Available: https://eur-lex.europa.eu/eli/reg/2016/679/oj. [Accessed: Sep. 9, 2024], 2016. [2] S. of California, “California consumer privacy act (ccpa),” California Legislative Information, AB-375, Jun. 2018. [Online]. Available: https://leginfo.legislature.ca.gov/faces/billTextClient.xhtml?bill id= 201720180AB375. [Accessed: Sep. 9, 2024], 2018. [3] H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. Agüera y Arcas, “Communication-efficient learning of deep networks from decentralized data,” Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS), vol. 54, pp. 1273–1282, 2017. [4] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the convergence of fedavg on non-iid data,” in International Conference on Learning Representations, 2020. [Online]. Available: https://openreview.net/forum?id=HJxNAnVtDS [5] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, “SCAFFOLD: Stochastic controlled averaging for federated learning,” in Proceedings of the 37th International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, H. D. III and A. Singh, Eds., vol. 119. PMLR, 13–18 Jul 2020, pp. 5132–5143. [Online]. Available: https://proceedings.mlr.press/v119/karimireddy20a.html [6] M. Ye, X. Fang, B. Du, P. C. Yuen, and D. Tao, “Heterogeneous federated learning: State-of-the-art and research challenges,” ACM Comput. Surv., vol. 56, no. 3, oct 2023. [Online]. Available: https://doi.org/10.1145/3625558 [7] A. K. Sahu, T. Li, M. Sanjabi, M. Zaheer, A. Talwalkar, and V. Smith, “Federated optimization in heterogeneous networks,” arXiv: Learning, 2018. [Online]. Available: https://api.semanticscholar.org/CorpusID: 59316566 [8] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, “Attention is all you need,” arXiv preprint arXiv:1706.03762, 2017. [9] D. Wen, K.-J. Jeon, and K. Huang, “Federated dropout—a simple approach for enabling federated learning on resource constrained devices,” IEEE wireless communications letters, vol. 11, no. 5, pp. 923–927, 2022. [10] E. Diao, J. Ding, and V. Tarokh, “Heterofl: Computation and communication efficient federated learning for heterogeneous clients,” in International Conference on Learning Representations (ICLR), 2021. [Online]. Available: https://arxiv.org/abs/2010.01264 [11] S. Alam, L. Liu, M. Yan, and M. Zhang, “Fedrolex: Modelheterogeneous federated learning with rolling sub-model extraction,” Advances in neural information processing systems, vol. 35, pp. 29 677– 29 690, 2022. [12] H. Wu, P. Wang, and C. V. A. Narayana, “Straggler-resilient federated learning: Tackling computation heterogeneity with layer-wise partial model training in mobile edge network,” arXiv preprint arXiv:2311.10002, 2023. [Online]. Available: https://arxiv.org/abs/2311. 10002 [13] H. Zhang, Z. Li, Z. Gong, M. Siew, C. Joe-Wong, and R. El-Azouzi, “Poster: Optimal variance-reduced client sampling for multiple models federated learning,” in 2024 IEEE 44th International Conference on Distributed Computing Systems (ICDCS), 2024, pp. 1446–1447. [14] A. Dabaja and R. El-Azouzi, “Fedplt: Scalable, resource-efficient, and heterogeneity-aware federated learning via partial layer training,” in Proc. IEEE International Symposium on Personal, Indoor and Mobile Radio Communications (PIMRC), 2025. [15] C. Chen, H. Xu, W. Wang, B. Li, B. Li, L. Chen, and G. Zhang, “Synchronize only the immature parameters: Communication-efficient federated learning by freezing parameters adaptively,” IEEE Transactions on Parallel and Distributed Systems, vol. 35, no. 7, pp. 1155–1173, 2023. [16] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 770–778. [17] H. Xiao, K. Rasul, and R. Vollgraf, “Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms,” arXiv preprint arXiv:1708.07747, 2017. [18] A. Krizhevsky and G. Hinton, “Cifar-10 and cifar-100 datasets,” 2009. [Online]. Available: https://www.cs.toronto.edu/∼kriz/cifar.html [19] J. Scott and Á. Cahill, “Improved modelling of federated datasets using mixtures-of-dirichlet-multinomials,” arXiv preprint arXiv:2406.02416, 2024.
[20] X. Li, M. Jiang, X. Zhang, M. Kamp, and Q. Dou, “Fedbn: Federated learning on non-iid features via local batch normalization,” arXiv preprint arXiv:2102.07623, 2021. [21] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the convergence of fedavg on non-iid data,” arXiv preprint arXiv:1907.02189, 2019.