1
AC2P2SL: Adaptive Communication-Computation Pipeline Parallel Split Learning over Edge Networks
arXiv:2606.31276v1 [cs.DC] 30 Jun 2026
Chenyu Liu, Zhaoyang Zhang, Zirui Chen, Zhaohui Yang, Chunhui Feng, and Tony Q. S. Quek Abstract—In wireless edge networks, split learning (SL) enables base station (BS) to utilize the distributed data and computing power across user equipments (UEs) to achieve collaborative model training while protecting local data privacy. However, the inherent sequential execution of computation and communication processes in conventional SL usually leads to long training times. To overcome this limitation, this paper proposes an adaptive communication-computation pipeline parallel split learning (AC2 P2 SL) framework. By conceptualizing the communication and computation processes of UEs and the BS as a unified pipeline, AC2 P2 SL achieves fine-grained pipeline parallelism across multiple micro-batches. Through this approach, effective overlapping of communication and computation is achieved which results in significant reduction of the overall training latency. Moreover, by considering the system constraints in the communication, computation, and storage dimensions as well as the heterogeneity of UEs, we formulate a joint optimization problem to minimize the training time and propose a corresponding split and pre-allocation algorithm to further enhance the pipeline efficiency. Additionally, accounting for the practical dynamic environments for the UEs, we design an adaptive re-allocation strategy to enhance the system resilience. Extensive experimental results demonstrate the effectiveness and robustness of AC2 P2 SL in reducing training time while ensuring data privacy preservation. Index Terms—Split learning, communication-computation pipeline parallelism, wireless edge network, resource allocation.
✦
1
I NTRODUCTION
In the sixth-generation (6G) era, deep integration of communication networks and artificial intelligence (AI) has become an inevitable trend [2], [3]. In edge network scenarios, such as anomaly detection monitoring [4], [5], unmanned aerial vehicle (UAV) image capturing [6]–[8], and vehicle autonomous driving [9], [10], raw data for neural network training are often stored dispersedly on local user equipments (UEs). Currently, systems typically collaborate with multiple UEs to train a common model, thereby fully leveraging their individual data to enrich the overall training dataset. With the rapid development of deep learning (DL) technologies, the parameter scale and computational load of AI models are consistently increasing [11], [12]. However, UEs in edge networks are typically resource-constrained personal devices, and relying solely on them to collaboratively train complex models presents multiple deficiencies in terms of computing power, storage, and communication capabilities [13], [14]. Meanwhile, wireless systems are also evolving, transitioning from a pure communication-oriented approach to an integrated communication and computing paradigm [15]. For instance, AI-enabled radio access network (AI-RAN) architecture [16] has transformed traditional base stations • •
•
•
A preliminary part of this work was presented at IEEE Globecom 2025 Workshop on A4E: AI/ML for Edge/Fog Networks [1]. This work was supported in part by National Natural Science Foundation of China under Grants 62394292 and 624B2129, Zhejiang Provincial Key R&D Program under Grant 2023C01021, and the Fundamental Research Funds for the Central Universities under Grant 226-202400069. (Corresponding author: Zhaoyang Zhang.) C. Liu, Z. Zhang, Z. Chen, Z. Yang, and C. Feng are with the College of Information Science and Electronic Engineering, Zhejiang University, Hangzhou 310027, China, and also with the Zhejiang Provincial Laboratory of Multi-Modal Communication Networks and Intelligent Information Processing, Hangzhou 310027, China (e-mail: [email protected]; ning [email protected]; [email protected]; yang [email protected]; [email protected]). Tony Q. S. Quek is with Singapore University of Technology and Design, Singapore 487372 (e-mail: [email protected]).
(BS) into intelligent AI processing nodes, providing additionally computational resources for edge AI tasks. Consequently, a straightforward approach is to centralize data from UEs at BS for model training, yet this raises numerous privacy-related concerns [3], [17]. By contrast, split learning (SL) [18]–[20], has emerged as an effective solution, where the model is split into two parts: the first few layers are distributed to the UEs, while the main body of the model is retained on the BS. For the split point, system utilizes wireless transmission to exchange forward propagation (FP) activations and backward propagation (BP) gradients between computing devices [21]. By splitting the model, the computational load is apportioned to the UEs and the BS, ensuring that computations directly involving user data are performed locally on the UEs without being exposed to the BS. In many supervised learning tasks, such as image recognition [22] and wireless localization [23], in addition to user data, the privacy of label information often also requires protection. To meet this demand, [18], [22], [24] introduces U-shaped SL (USL) characterized by a more ingenious twolayer cut strategy, which splits model into three sub-models: head, body, and tail models. UEs locally train head and tail models, while the BS trains the body model. This necessitates the BS serving solely as a computational unit for the intermediate model layers, without accessing either the input data or the final outputs/labels, thereby ensuring that both data, model outputs and ground-truth labels are strictly retained locally. However, while protecting privacy, SL and USL also introduce additional latency caused by wireless transmission [25], making the improvement of the overall temporal efficiency of the architecture an important research topic. Furthermore, to facilitate simultaneous multi-user participation in SL, the split federated learning approach in [26]–[28] extends the aforementioned framework by introducing the parallel model training mechanism among users. However, these methods typically require an additional
2
𝐿𝐿
𝑙𝑙1
Head Model
𝑙𝑙2
Tail Model
① Head model FP UE 2
Middle Model ①
③ Body model FP
⑧
④ Body model outputs DL
⑨
⋯
UE 1
② ⑩
⑦ ④
⑤
② Head model outputs UL
③
UE n
⑥
⑤ Tail model FP ⑥ Tail model BP ⑦ Body model output gradients UL ⑧ Body model BP
BS
⑨ Head model output gradients DL ⑩ Head model BP
Fig. 1. U-shaped parallel split learning over wireless edge networks.
aggregation server [28], thereby introducing extra synchronization overhead and posing potential privacy risks. From another perspective, [29] partitions UEs into multiple clusters to implement intra-cluster parallel training and intercluster sequential training, partly reducing training latency. Furthermore, [24] introduced parallel split learning for multiple UEs, while [30], [31] reduces the dimensionality of BP by aggregating the gradients of the BS-side layer. These approaches eliminate the need for UE-side model synchronization, thus further decreasing communication time [32]. Moreover, in edge network SL, the selection of the model split layer significantly affects training efficiency. Additionally, intrinsic heterogeneity among terminal devices in computation and communication capabilities induce substantial synchronization latency during the training process. Building upon the proposed SL schemes, [29] designed a joint optimization strategy for layer split selection and resource allocation to minimize training costs. [27] and [24] jointly optimized the split layer selection and bandwidth allocation problems within the frameworks, respectively, to minimize training latency. Furthermore, [31] additionally introduced power control to balance the trade-off with energy consumption. However, in practical heterogeneous edge networks, not only do intrinsic performance disparities exist among devices, but the communication capabilities of UEs also fluctuate due to mobility and other factors, while their computational resources remain unstable due to concurrent local computing tasks. The synchronization latency induced by these varying factors significantly degrades overall training efficiency. For the communication overhead during the training process, above approaches have primarily focused on reducing data transmission time via operations such as UEside parallelism and data aggregation. However, as shown in Figure 1, for a single data batch, the computation and communication processes remain serially executed. From head model FP to head model BP, subsequent stages needs to wait for the completion of the previous stages, inevitably introducing idle time across the various processing stages. In scenarios involving models with high-dimensional intermediate outputs or poor channel conditions, the prohibitive communication overhead results in low training efficiency. To address these limitations, we notice that the serial sit-
uation can be significantly optimized through fine-grained parallelism, similar to pipeline parallelism mechanism [33] in distributed training. By partitioning data batches into multiple micro-batches, each node propagates the results to the subsequent stage immediately after completing its sub-model computation on the current micro-batch, while concurrently processing intermediate outputs from the preceding stage for subsequent micro-batches. Stages in the computation pipeline can be parallelized through the overlap between micro-batches. Although recent works have introduced pipeline parallelism into SL, their scopes remain limited. Specifically, [34] focuses exclusively on intra-server computational pipelining, while [35] implements pipeline parallelism for the fixed stages within federated split learning framework. Critically, neither approach accounts for the inherent parallelizability between the communication and computation processes in wireless SL. This neglect of communication-computation parallelism restricts the overall training efficiency in dynamic wireless edge environments. In this paper, building upon USL, we integrate the communication pipeline consisting of uplink and downlink data transmission with the computation pipeline comprising sub-models’ FP and BP into a unified training pipeline. Treating UEs’ computation, uplink transmission, BS’s computation, and downlink transmission as distinct processing stages, we also achieve fine-grained micro-batch parallelism by partitioning data batches. This approach enables wireless transmission to occur concurrently with local model computations on both UE and BS sides to achieve parallel training. Moreover, in response to heterogeneous capability, time-varying communication quality, and realtime availability of computational resources in practical edge networks, we formulate the pipeline training time optimization problem to design the split and pre-allocation (SPA) algorithm as well as the adaptive re-allocation (ARA) strategy. Consequently, the above content collectively constitutes an adaptive communication-computation pipeline parallel split learning (AC2 P2 SL) framework in this paper. By effective parallelism between training stages, the proposed framework significantly reduces pipeline training time and enhances overall efficiency. The main contributions of this paper are summarized as follows:
3 •
•
•
•
We propose the AC2 P2 SL framework, which combines data transmission with computation tasks to achieve communication-computation pipeline parallelism between micro-batches by splitting data batches. We formulate a joint optimization problem to minimize the pipeline training time based on the computation, communication, and memory constraints of UEs and design the SPA algorithm for its solution. We implement an adaptive ARA strategy before each pipeline training round to adjust the microbatch quantity, allocated batch size and time slot in response to significant changes in UE performance. We comprehensively evaluate the effectiveness and robustness of AC2 P2 SL through extensive experiments under different models and system parameters. We also conduct ablation studies to further validate the efficacy and principle of our resource allocation strategies.
The remainder of this paper is organized as follows. Section 2 introduces the system communication, computation, and storage model. In Section 3, we present the overview of AC2 P2 SL and its training workflow. Then, problem formulation and solution approach are presented in Section 4. Next, Section 5 provides performance evaluations of our proposed scheme. Finally, we conclude our work in Section 6.
2
S YSTEM M ODEL
We consider an edge cellular network consisting of a central BS and a set of n UEs distributed within the coverage area. Table 1 summarizes the main notations of this paper. •
•
2.1
BS: By Integrating AI processing units, the BS is equipped with sufficient computational resources to train large models. Beyond performing wireless signal transmission, the BS is responsible for collecting computational and communication information from UEs to make optimizations during the training process. UE: As resource-constrained edge devices, each UE has weak computational capability to train tiny neural networks. Let N = {1, . . . , N } denote the set of all devices. The UEs locally store private datasets, where Di,j represents the local data samples of the i-th UE during the j -th training batch. Computation Model
The target model has a total of L layers, which are partitioned into head model, body model, and tail model before training. The indices of two cut layers are respectively denoted as l1 , l2 ∈ {1, 2, . . . , L}. In this way BS retains the sub-model containing the majority of parameters, while broadcasting the smaller head and tail sub-models to individual UEs. During one training round, the batch size of i-th UE’s input data is denoted as bi , and the total batch size of all UEs is given as N X B= bi . (1) i=1
In contrast to previous studies that ideally treat device computational capability as a static constant [1], we use
TABLE 1: Summary of Main Notations Notation
Description
l1 , l 2 al B , bi fi , f 0 Ii ,I0 Fi , F0 βi , β0 cfl , cbl
Cutting layers of head, body, and tail sub-models Data size of the l-th layer’s activation output Total and UE i’s batch size of one training round Attainable computing performance of UE i and BS Operational intensity of UE i and BS Peak FLOPS of UE i and BS Maximum memory bandwidth of UE i and BS FP and BP computing workload per sample of the l-th layer Fixed memory access for FP and BP of the l-th layer Activation memory access for FP and BP per sample of layer l Fixed and per sample activation memory of the lth layer Maximum memory limit of UE i Length of time slot and time frame Number of time slots for UE i and ratio of uplink to downlink time slots System bandwidth Effective antenna gain of UE i and BS Transmit power of UE i and BS Channel gain between UE i and BS Power spectral density of noise Number of micro-batches in a single data batch The j -th micro-batch data samples and labels of UE i Head/tail model activation output/gradient of the j -th micro-batch of UE i Body model activation output and gradient of the j -th micro-batch of BS Head, and tail model parameter of UE i Body model parameter of BS Forward and backward propagation mapping Loss function and j -th micro-batch loss of UE i
mfl , mbl ∆mfl , ∆mbl ul , ∆ul Ui τ, T si , ρ BW Gi , G0 pi , p 0 hi N0 k x(i, j), y(i, j) a/g h/t (i, j) ab (j), gb (j) Wh (i), Wt (i) Wb F (·), B(·) L(·), li,j
the Roofline model [36] to provide a more realistic characterization. This model formulates the attainable computing performance as a piecewise function determined by the hardware’s peak floating point operations per second (FLOPS), memory bandwidth, and the data operational intensity. Here, operational intensity is defined as the ratio of the floating point operations (FLOPs) to the amount of memory access required during the computation process. For layer l, the total memory access is composed of fixed f /b parameter memory access ml and per sample memory f /b access ∆ml , a variable component proportional to the number of input data samples. Depending on the context of f /b f /b FP or BP, the parameter ml and ∆ml takes the value f f b b of ml , ml or ∆ml , ∆ml , respectively. Consequently, for the varying lower layer bound lL , upper layer bound lU of sub-model and the number of input data samples D, the operational intensity of the i-th device for FP or BP is expressed as
D I f /b (D, lL , lU ) =
lU P l=lL
lU P l=lL
f /b
ml
+D
f /b
cl
lU P l=lL
, f /b
∆ml
(2)
4
where the superscript f /b of operational intensity indicates PlU f /b c whether the current computing task is FP or BP. l=l L l denote the FLOPs required for a single data sample. When inputting a single data sample into the neural network, the f FLOPs of l-th layer for FP and BP are denoted by cl and cbl . Furthermore, we define the borderline operational inteni sity as Iimax = F βi , where Fi represents the peak FLOPS and βi denotes the memory bandwidth. When the operational intensity falls below this threshold, the task is classified as memory-bound, which implies that the performance is constrained by the memory bandwidth, i.e. f /b fi = βi I f /b (bi , lL , lU ). Conversely, when the intensity exceeds the threshold, the task is compute-bound, limited f /b by the peak FLOPS, i.e. fi = Fi . Therefore, the attainable computational capability of the i-th device for FP or BP can be summarized as
βi bi f /b
fi
(bi , lL , lU ) = minFi , l U P l=lL
2.2
lU P l=lL
f /b
ml
+ bi
f /b cl lU P
l=lL
. (3) f /b
∆ml
Note that the activation gradients generated during BP have the same dimension as the activation outputs produced during FP for the same network layer, we uniformly denote the activation output data size of l-th layer as al . Consequently, regarding the communication interface required between the UEs and the BS, the transfer loads at the two cut layers l1 , l2 are expressed as al1 and al2 , respectively. To address the instability of UE performance in practical training scenarios, we consider a time division duplexing (TDD) communication system with time division multiple access (TDMA)1 , which facilitates dynamic resource allocation to accommodate discontinuous, time-varying, and asynchronous transmission of UEs. Specifically, TDMA partitions time into periodic frames of length T , with each frame further subdivided into multiple time slots of length τ . Each slot is assigned to a specific UE for data transmission, while the number of time slots allocated to the i-th UE is denoted by si . By dynamically scheduling resources in the time domain, this mechanism effectively reduces interuser interference, thereby ensuring parallel and reliable data transmission. Consequently, the constraint relating the time frame and the allocated time slots is expressed as
τ si ≤ T.
where BW represents the total system bandwidth shared by all UEs, pi denotes the uplink transmit power of the i-th UE, and p0 indicates the downlink transmit power of the BS. Gi and G0 stands for the antenna gain of the i-th UE and BS while N0 denotes the power spectral density (PSD) of the noise. Finally, hi represents the channel gain for the i-th UE, which includes path loss and shadow fading, as well as multipath fading. 2.3
Communication Model
N X
achievable uplink and downlink transmission rates of the i-th UE and BS are formulated as τ si ρ Gi G0 pi hi , (5) riu = BW log2 1 + T (1 + ρ) BW N0 τ si G0 Gi p0 hi rid = , (6) BW log2 1 + T (1 + ρ) BW N0
(4)
i=1
It is noteworthy that ρ is defined as the ratio of the number of uplink time slots to that of downlink time slots. We adopt the quasi-static block fading channels, holding that the UE channels remain approximately stable over the short duration of a single data batch training. This implies that the channel gain remains constant within each communication round, but may vary across different data batches. Accordingly, based on the Shannon’s theorem, the 1. Both TDMA and FDMA are merely methods for modeling communication rates and do not affect overall AC2 P2 SL framework.
Storage Model
For UEs in the wireless edge network, the constrained storage resources necessitate imposing limits on the sizes of the partitioned head and tail models. Similarly to the memory access analysis in Section 2.1, we categorize the memory usage of each layer into two components: the fixed model parameter usage ul , and the variable intermediate activation usage per sample ∆ul . Consequently, taking into account both the head and tail models, the total memory footprint of UE i is constrained by l1 X
(ul + bi ∆ul ) +
l=1
L X
(ul + bi ∆ul ) ≤ Ui ,
(7)
l=l2 +1
where Ui represents the maximum memory capacity of UE i. The storage constraint imposes a strict upper bound on both the split layer l1 , l2 and the allowable batch size bi . Conversely, for the BS, we assume it possesses sufficient storage resources for entire model training.
3
AC2 P2 SL F RAMEWORK AND T RAINING W ORK -
FLOW
In this section, we present the system overview of the AC2 P2 SL framework across its various phases, along with the parallel training workflow. 3.1
System Overview
As illustrated in Fig. 2, The overall AC2 P2 SL framework comprises four primary phases: initial profiling, adaptive allocation, parallel training and U-shaped inference phase. In the initial profiling phase, the local data on UEs is loaded while BS initializes the entire model. Subsequently, each UE uploads its communication and computation performance to the BS. Based on the parameters, the BS makes the SPA optimization which is presented in Section 4.2 to determine the near-optimal model split layer, number of micro-batches, batch size set, and allocated time slots set. Upon completing this optimization, the BS distributes the partitioned head and tail sub-model, along with the optimization results, to the UEs. The adaptive allocation phase occurs before each training round to dynamically allocate computing workloads and communication resources. Specifically, the BS monitors the performance parameters uploaded by the UEs. If
5
UEs
BS Initial Profiling Phase Body Model
Upload Performance
Split & Preallocation Optimization
Model Split & Distribution
Threshold Detection
Adaptive Allocation Phase
Parallel Training Phase
Model Initialization
Number of micro-batch
Adaptive Re-allocation
Batch size set
U-shaped Inference Phase
Tail Model
Re-allocation
Time slot set
FP
Activations
BP
Gradients Gradients
DL
Head Model
Upload Performance
Yes
UL
Data Loading
Activations
FP
Input Output Label
Input
UL
Activations Output
DL
Fig. 2. System overview of AC2 P2 SL.
variations exceed a predefined threshold δ , the ARA optimization presented in Section 4.3 concerning the number of micro-batches, batch size, and time slot is triggered to maximize operational efficiency. Furthermore, in scenarios where specific system nodes encounter failures and cease participation, this adaptive allocation functions as an elastic fault-tolerance mechanism, mitigating the adverse effects of UE heterogeneity and temporal variations. In the training workflow, serial execution of computation and communication for a single batch inevitably introduces idle time across the various processing stages. Motivated by pipeline parallelism, we adopt a similar strategy. Specifically, UEs split a single batch uniformly into k microbatches, where each micro-batch consists of bi /k data samples. This partitioning enables UE to train a sequence of micro-batches continuously without increasing the memory footprint required for each training round. By utilizing the idle time of a single micro-batch, different stages concurrently process other micro-batch. Thus, AC2 P2 SL achieves micro-batch level parallelism that effectively overlaps computation with communication. The U-shaped inference phases depicted in Fig. 2 illustrate the USL process for a single UE-BS pair. In the subsequent inference phase, UE can execute inference tasks while keeping both input data and output results local without model parameter transmission, thus achieving comprehensive privacy protection. In practice, the proposed AC2 P2 SL framework supports parallel split inference across multiple UEs. By leveraging data batch partitioning and communication-computation pipeline parallelism in the
same way, the framework achieves efficient edge inference. 3.2 Parallel Training Workflow Following the Initial profiling and adaptive resource allocation phases, the workflow of the U-shaped training phase is illustrated in Fig. 3. For convenience, we describe the training round for a single data batch, which includes the parallel training of k micro-batches. In the figure, the vertical axis sequentially represents four distinct physical stages: multi-UE computation, multi-UE uplink transmission, BS computation, and multi-UE downlink transmission, while the horizontal axis depicts the timeline as each data microbatch goes through these stages. Next, we introduce the micro-stages that occur in the different physical stages: 3.2.1 Head Model FP and Output Uplink Transmission The process begins with the parallel computation of UEs. Each UE concurrently inputs its k local data micro-batches, denoted as xi,j , into the head model for FP sequentially. These tasks are enqueued in the computation buffer. Through FP, the activation output of head model is generated as ah (i, j) = F (Wh (i), x(i, j)). (8) Upon completing FP of the current micro-batch, the activation output ah (i, j) is immediately added to the transmission queue, awaiting transmission during the allocated idle time slot in the uplink. Noting that due to heterogeneity among UEs, variations in computational and communication capabilities lead to distinct FP and uplink transmission durations, thereby inducing synchronization delays.
6
𝑘𝑘 Micro-batches
⋯
UE 1
UE n
⋯
UL 1 UL n
⋯k 1 2⋯k 1
2
BS
2
DL 1
⋯
DL n
1
1
⋯ 1
2
⋯
2
1
2
1
2
1
2
1
1
1 1
1
1
FP of head model on UE-side
2
⋯ 2 2
⋯ 2
2
1
⋯
⋯ 2
2
⋯
k
Batch 1
⋯ ⋯
2
1
2
⋯
k
1
2
k
2
1
k
1
k
k
1
k
⋯ ⋯
k
k
1
k k
⋯
2
Batch 2
⋯
2
FP of body model on BS-side
k
FP of tail model on UE-side
k
BP of head model on UE-side BP of body model on BS-side BP of tail model on UE-side
k
⋯
⋯
Batch 3
k
Uplink transmission of activations
k
Uplink transmission of gradients
Time
Downlink transmission of activations Downlink transmission of gradients
⋯
Fig. 3. Training workflow of AC2 P2 SL. The uplink and downlink transmission block of different UEs consist of multiple non-overlapping time slots.
3.2.2 Body Model FP and Output Downlink Transmission Once receiving the activation outputs of the j -th micro-batch from all UEs, the BS aggregates and concatenates them along the batch dimension, as
ah (j) = [ah (1, j); ah (2, j); . . . ; ah (n, j)].
(9)
Subsequently, the BS inputs the aggregated output ah (j) to obtain the activation output of the body model by performing FP, denoted as
ab (j) = F(Wb , ah (j)).
(10)
Later, following the logic of aggregation, the BS splits the output to derive ab (i, j) = ab (j)[i]. These partitioned outputs are then enqueued for downlink transmission to their respective UEs during available allocated time slots. 3.2.3 Tail Model FP, BP and Gradient Uplink Transmission After receiving the corresponding output of body model, the i-th UE inputs it into the tail model for FP whenever the computation queue is idle, yielding the activation output
at (i, j) = F (Wt (i), ab (i, j)).
(11)
The local loss is then calculated using the corresponding data label y(i, j), expressed as
l(i, j) = L(at (i, j), y(i, j)).
(12)
Crucially, computation of the tail model adheres to the one forward step by one backward (1F1B) principle. Specifically, immediately after completing FP of the j -th micro-batch, the UE performs the homologous BP to derive the gradients for the tail model parameters ∆Wt (i, j) and the input data gb (i, j) (i.e., activations from the body model), denoted as
∆Wt (i, j), gb (i, j) = B(l(i, j), Wt (i), ab (i, j)).
(13)
Theoretically, besides the variables explicitly denoted in the equations, BP needs to input the FP intermediate activation outputs of each layer in the sub-model. These activations constitute the computational graph retained in the device memory and are released only on completion of BP. For the sake of notational brevity, we treat these retained activations as implicit default inputs and omit them from
the mathematical formulations. Once completed, the input data gradients are added to the uplink transmission queue, pending transmission during the idle slots. 3.2.4 Body Model BP and Gradient Downlink Transmission Following receipt of the activation gradients for micro-batch j from all UEs, the BS aggregates them along the batch dimension, analogously to the forward output as
gb (j) = [gb (1, j); gb (2, j); . . . ; gb (n, j)].
(14)
Then BP is executed for the body model to obtain the gradients for the model parameters and the aggregated head model activations, represented as
∆Wb (j), gh (j) = B(gb (j), Wb , ah (j)).
(15)
Subsequently, the BS splits the activation gradients to obtain gh (i, j) = gh (j)[i], enqueuing them for downlink transmission to the original UEs. 3.2.5 Head Model BP and Parameters Update Finally, having received the corresponding activation gradients for the j -th micro-batch, the UE performs BP for the head model whenever the computation queue permits, yielding ∆Wh (i, j) = B(gh (i, j), Wh (i)). (16) During the m + 1-th round of batch training, after sequentially completing BP of all k micro-batches, the head and tail model parameters of each UE are first updated using the mini-batch gradient descent (MBGD) method [37], and then aggregated by FedAvg algorithm [38], as m+1 m Wh/t = Wh/t −
N k η X bi X m ∆Wh/t (i, j), B i=1 k j=1
(17)
where η is the learning rate. Simultaneously, once the BS completes the BP for k micro-batches and the computation queue is idle, it updates the body model parameters by using the MBGD method via k
Wb m+1 = Wbm −
ηX ∆Wbm (j). k j=1
(18)
7 2 2
Algorithm 1 The AC P SL Framework m
m
Input: x (i, j), y (i, j), η , M Output: WhM (i), WtM (i), WbM . 1: UEs upload capability parameters to BS 2: BS get {l1 , l2 , k 0 , b0i , s0i }i∈N by Algorithm 2 3: Initialize model parameters and split model 4: BS distributes split model and allocated (k 0 , b0i , s0i ) to UEs 5: for m = 1, 2, . . . , M do 6: if (tm − tm−1 )/tm−1 > δ then m 7: BS get {km , bm i , si }i∈N by Algorithm 3 8: else m m−1 m−1 m−1 9: {km , bm , bi , si }i∈N i , si }i∈N ← {k 10: end if // Runs on UEs 11: for each i ∈ N in parallel do 12: for j = 1, 2, . . . , km do m m 13: am h (i, j) ← F(Wh (i), x (i, j)) 14: Transmit am h (i, j) to BS 15: end for 16: end for 17: if Receive am b (i, j) of j -th micro-batch from BS then m 18: lm (i, j) ← L(F(Wtm (i), am b (i, j)), y (i, j)) 19: ∆Wtm (i, j), gbm (i, j) ← B(lm (i, j), Wtm (i), am b (i, j)) 20: Transmit gbm (i, j) to BS 21: end if 22: if Receive ghm (i, j) of j -th micro-batch from BS then 23: ∆Whm (i, j) ← B(ghm (i, j), Whm (i)) 24: end if 25: Update head and tail model as (17) // Runs on BS 26: if Receive am h (i, j) of j -th micro-batch from all UEs then m m m 27: am h (j) ← [ah (1, j); ah (2, j); . . . ; ah (n, j)] m m 28: am b (j) ← F (Wb , ah (j)) 29: Transmit split am b (i, j) to UEs 30: end if 31: if Receive gbm (i, j) of j -th micro-batch from all UEs then 32: gbm (j) ← [gbm (1, j); gbm (2, j); . . . ; gbm (n, j)] 33: ∆Wbm (j), ghm (j) ← B(gbm (j), Wbm , am h (j)) 34: Transmit split ghm (i, j) to UEs 35: end if 36: Update body model as (18) 37: end for
3.3 Convergence Analysis After all of the above processes, this concludes the training process for the current data batch, and the overall training framework is shown in Algorithm 1. Notably, while the synchronization of activations and gradients occurs strictly per-micro-batch across all users, the model parameters are updated synchronously only at the global batch boundary. Therefore, our pipeline mechanism does not introduce asynchronous gradient staleness. We have Theorem 1, which theoretically guarantees that the convergence speed and final accuracy of Algorithm 1 remain completely unaffected. Theorem 1. Given the same global batch size, the parameter update process of AC2 P2 SL is mathematically identical to regular model training without batch splitting. Proof. For the shared body sub-model of the BS, the cumulative parameter gradients of the B k samples in the j -th microbatch is represented as B
∆Wbm (j) = B(j−1)
k k X B(j − 1) ∆wbm ( + t), B t=1 k
(19)
where ∆wbm ( k + t) denotes the parameter gradient of the t-th sample in the j -th micro-batch. Consequently, the
average parameter gradient of k micro-batches is exactly equal to the total parameter gradient of all samples in the undivided global batch, simply as k B 1X 1 X ∆Wbm (j) = ∆wbm (s). k j=1 B s=1
(20)
Similarly, for the head and tail sub-models distributed across UEs, the parameter gradient generated by samples in the j -th micro-batch of UE i can be expressed as bi k bi (j − 1) k X m m ∆wh/t (i, ∆Wh/t (i, j) = + t), bi t=1 k
(21)
b (j−1)
m where ∆wh/t (i, i k + t) denotes the parameter gradient of the t-th sample in the i-th UE’s j -th micro-batch. After completing BP of k micro-batches, UE i aggregates k parameter gradients by FedAvg to obtain the total parameter gradient for its batch of samples, presented as bi k 1X 1 X m m (i, j) = (i, s). ∆Wh/t ∆wh/t k j=1 bi s=1
(22)
Then, by weighted averaging the parameter gradients of all UEs, the aggregated gradients of head and tail sub-models to be updated can be denoted as N k B bi X 1 X 1 X m m ∆Wh/t (i, j) = ∆wh/t (s), B i=1 k j=1 B s=1
(23)
which is also equal to the total parameter gradient of all samples in the undivided global batch. Obviously, this parallel mechanism operates entirely independently of the parameter update procedures for the head, body, and tail submodels. Consequently, the AC2 P2 SL framework is highly compatible with various federated learning model aggregation algorithms, although we default to the widely adopted FedAvg to aggregate model parameters on UE-side. As a result, Theorem 1 is proved.
4
O PTIMIZATION FOR S PLIT AND A LLOCATION
In this section, we establish a dynamic programming model to analyze training time and make optimizations on split pre-allocation and adaptive re-allocation. 4.1 Time Analysis Considering that the whole training process for a single data batch involves a complex pipelining of k micro-batch tasks across n parallel UEs, spanning 4 distinct physical queues and 9 logical micro-stages (where FP and BP of tail model are treated as a unified entity), the system is characterized by intra-stage resource re-entry, pipeline blocking, and interstage synchronization barriers. Consequently, the total training time t cannot be accurately formulated using a simple closed-form equation. In particular, since the computational overhead associated with loss calculation and parameter updating is negligible, these processes are omitted from the calculation of the total computation time. To address this, we formulate a dynamic programming model to precisely calculate the total time of the parallel process. This model tracks the progression of k micro-batches
8
for n UEs across the 9 micro-stages while simultaneously managing resource availability across the 4 physical queues. Let C(i, j, s) denote the completion timestamp (measured from the start of the training phase) of the s-th micro-stage for the j -th micro-batch of UE i, where i ∈ {1, . . . , n}, j ∈ {1, . . . , k}, and s ∈ {1, . . . , 9}. Firstly, the process commences with head model FP for the j -th micro-batch of UE i. As each micro-batch is sequentially input the head model, and in accordance with the analysis presented in Section 2.2, the computational capability of UE i is denoted as fi (bi /k, 1, l1 ) and the Pl1 f total computational workload is calculated as bki l=1 cl . Consequently, the completion time for the local computation queue is formulated as
bi C(i, j, 1) = j ·
l1 P l=1
cfl
kfif (bi /k, 1, l1 )
.
(24)
Due to the queuing constraints at the UE, the computation for a subsequent micro-batch is blocked until the preceding one completes. Building upon this, the uplink transmission for the j -th micro-batch is contingent upon two synchronization constraints: the completion of its own FP calculation and the completion of the transmission for the preceding micro-batch. With a transmission payload of bi k al1 , the completion time for the uplink transmission queue is expressed as
C(i, j, 2) = max {C(i, j, 1), C(i, j − 1, 2)} +
bi al1 . kriu
(25)
Subsequently, BS aggregates the outputs of the same micro-batch from all UEs, necessitating a synchronization barrier. Thus, body model FP start time of the j -th microbatch is determined by two values: the last uplink transmission across all UEs and the FP the previous (j − 1)-th micro-batch. Given that the aggregated input volume is the sum of batch sizes from n UEs, the BS computation queue is formulated as
B
l2 P l=l +1
cfl
1 . kf0f ( B , k l1 +1, l2 ) (26) The BS then distributes the partitioned outputs to each UE via the downlink. Since the preceding FP stage is a synchronized process, all downlink transmissions begin with the completion of that stage. With a transmission load of bi k al2 , the downlink transmission queue is defined as
C(j, 3) = max {C(i, j, 2), C(j − 1, 3)}+ i
However, distinguishing the computational and memory access differences between FP and BP, we calculate their durations sequentially. The UE computation queue is represented as
C(i, j, 5) = max{C(i, j, 4), C(i, j − 1, 5)}+ P P f b bi L bi L (28) l=l2 +1 cl l=l2 +1 cl + , f b kfi (bi /k, l2 +1, L) kfi (bi /k, l2 +1, L) where the initial value C(i, 0, 5) = C(i, k, 1). Following the tail model BP, the UE transmits the gradients via the uplink. With a payload size of al2 , the uplink transmission queue is formulated as bi al2 , (29) C(i, j, 6) = max {C(i, j, 5), C(i, j − 1, 6)} + kriu with initialization C(i, 0, 6) = C(i, k, 2). BS then receives gradients from all UEs and performs batch aggregation. Similarly to body model FP, this involves synchronization across all UEs. The BS computation queue is expressed as l2 P
cbl l=l1 +1 C(j, 7) = max {C(i, j, 6), C(j − 1, 7)}+ b B , i kf0 ( k , l1 +1, l2 ) B
(30) with initialization C(0, 7) = C(k, 3). Subsequently, BS transmits the partitioned gradients to the UEs. With a payload of al1 , the downlink transmission queue is defined as
C(i, j, 8) = max {C(j, 7), C(i, j − 1, 8)} +
bi al1 , krid
(31)
with the initial value C(i, 0, 8) = C(i, k, 4). Finally, upon receiving the downlink gradients, the UE executes BP of head model. The computation queue is given by l1 P
cbl l=1 C(i, j, 9) = max{C(i, j, 8),C(i,j − 1,9)}+ b bi , (32) kfi ( k , 1, l1 ) bi
with initialization C(i, 0, 9) = C(i, k, 5). Consequently, the total training time of a single data batch under the AC2 P2 SL framework is derived as
t = max C(i, k, 9). i
(33)
(27)
4.2 Split and Pre-allocation Optimization In the proposed framework, total pipeline training time is fundamentally governed by a complex interplay of layer split, pipeline depth, and heterogeneous resource allocation. To minimize this overall latency, it is imperative to jointly optimize these key variables.
Since the start times of the first four micro-stages correspond to the idle states of the respective physical queues, their initial values are set to zero as C(i, 0, 1) = C(i, 0, 2) = C(i, 0, 3) = C(i, 0, 4) = 0. However, the subsequent microstages involve resource re-entry into these physical queues. So, their initial values are determined by the completion times of the last micro-batch k from the preceding logical micro-stage that occupied the same physical resource queue. Upon receiving the downlink output, each UE performs FP and BP for the tail model on micro-batch j under the 1F1B principle, which we treat as a unified micro-stage.
4.2.1 Problem Formulation In USL, split layer l1 and l2 directly dictate the number of layers assigned to the head, body, and tail sub-models, satisfying the feasibility boundary 1 < l1 < l2 < L. This partition fundamentally governs the computational workload at both UEs and BS sides, as well as the data volume transmitted over uplink and downlink channels, thereby determining the training time. Crucially, because the computational intensity and output tensor dimensions of each layer vary dramatically across different neural network architectures, the selection of these split points must
C(i, j, 4) = max {C(j, 3), C(i, j − 1, 4)} +
bi al2 . krid
9
be adaptively tailored to the fine-grained layer-profiling parameters of the specific model. Simultaneously, the number of micro-batches k per training round determines the attainable hardware computing performance of devices processing each micro-batch and the pipeline depth of parallelism among micro-stages. Specifically, a larger k enhances the concurrency of the pipeline, which helps compress idle time while simultaneously reducing the sample workload contained within each individual micro-batch. This lower operational intensity degrades hardware compute utilization, thereby increasing the computational duration. Furthermore, k is naturally upperbounded by the minimum batch size allocated across the devices, requiring 1 ≤ k ≤ mini bi , k ∈ N. Navigating this trade-off to select appropriate micro-batch number is therefore vital for maximizing pipeline efficiency. Given a total batch size, the batch size allocated to each UE bi determines the number of training samples in its single round. Because both the computational workload and the communication volume are positively correlated with bi , optimizing this allocation is essential given the highly heterogeneous hardware capabilities of the UEs. Additionally, disparities in transmit power and path loss during wireless propagation introduce significant variations in communication durations among UEs. Within a TDMA system, the number of time slots allocated to each UE si maintains a positive correlation with its transmission speed. By optimizing bi , si to tune computation and communication time of UEs, the BS can balance the execution timelines of parallel UE micro-stages, thereby minimizing the synchronization waiting time caused by straggler effects at the global aggregation barrier, subject to the total batch bound in (1) and the bounded frame duration in (4). Therefore, in the initial profiling phase, the BS collects specific parameters uploaded by the UEs, including transmit power, location information, FLOPS, etc. Leveraging these heterogeneous parameters, the BS conducts joint optimization over model split layer l1 , l2 , number of micro-batches k , batch size set b, and time slot set s. We perform quantitative analysis and modeling based on the constraints of these variables with the objective of minimizing the total pipeline training time. The joint optimization problem is formulated as follows.
P1 : s.t.
min
l1 ,l2 ,k,b,s
t(l1 , l2 , k, b, s)
(34)
C1: 1 ≤ l1 < l2 ≤ L − 1, C2:
l1 X
(ul + bi ∆ul ) +
l=1 i
C4: C5:
i=1 n X
(ul + bi ∆ul ) ≤ Ui ,
l=l2 +1
C3: 1 ≤ k ≤ min bi , n X
l1 , l2 ∈ N, L X
k ∈ N,
bi = B,
bi ∈ N,
τ si ≤ T,
si ∈ N,
i=1
The optimization objective is to minimize the training time for a single data batch shown in Section 4.1. Due to the complexity of 9-micro-stage pipeline, sharing of physical resources among certain micro-stages, and presence of interstage blocking, the objective function exhibits high nonlinearity. Furthermore, constraint C1 defines the boundary
Algorithm 2 SPA Algorithm Based on AO Input: FLOPS, memory access, and storage parameters of each model layer, convergence tolerance ϵ. Output: l1∗ , l2∗ , k∗ , b∗ , s∗ . 1: Initialization: k (0) , b(0) , s(0) , tmin . 2: for l1 = 1, 2, . . . , L − 2 do 3: for l2 = L − 1, L − 2, . . . , l1 do 4: Select (l1 , l2 ) by the outer loop pruning search as Theorem 2. 5: Obtain (k, b, s) by the inner loop AO as Algorithm 3 inputting ϵ. 6: if t(l1 , l2 , k, b, s) < tmin then 7: tmin ← (l1 , l2 , k, b, s), l1∗ ← l1 , l2∗ ← l2 , k∗ ← k, b∗ ← b, s∗ ← s. 8: end if 9: end for 10: end for
conditions for the cut layers within the two-layer splitting strategy. Constraint C2 addresses the varying storage capabilities among UEs. It ensures that the aggregate memory footprint, including the head and tail model parameters as well as intermediate activations, is strictly bounded by each device’s maximum memory usage. Constraint C3 ensures that the number of pipeline micro-batches does not exceed the local batch size allocated to the UE. Constraint C4 dictates that the sum of data batches allocated across all users exactly matches the total global batch size. Lastly, constraint C5 imposes an upper bound on the total number of communication time slots allocated to the UEs. Since the decision variables (l1 , l2 , k, b, s) are all integers and the objective function and constraints are nonlinear and coupled, P 1 constitutes a mixed-integer nonlinear programming (MINLP) problem. Furthermore, the non-convexity and discontinuity of the objective function preclude the direct derivation of a global optimal solution. Therefore, we propose a SPA algorithm based on hierarchical decomposition and alternating optimization (AO) [39]. This approach decomposes P 1 into two loops: the outer loop is responsible for the feasibility pruning and search of the split layer l1 and l2 , while the inner loop solves the joint resource allocation sub-problem regarding (k, b, s) via AO. The overall algorithm is summarized in Algorithm 2. 4.2.2
Outer Loop Pruning Search
Due to the variance in computational and communication overhead across different network layers, the objective function is discrete and highly non-convex with respect to l. Given the finite layer number, we employ a double reverse loop traversal over all available l1 , l2 under constraint C1. This global search in the outer loop guarantees global optimality for the discrete structure. To reduce the computational complexity of the inner loop, we have the following simple observation: Theorem 2. For the given l1 , l2 in (34), the maximum batch size allocated to the i-th UE is as follows: (U −M ix (l1 ,l2 )) blimit (l1 , l2 ) = ⌊ iMvarf (l ⌋, where Mf ix (l1 , l2 ) = i 1 ,l2 ) Pl1 PL Pl1 u + u and M var (l1 , l2 ) = l=1 l l=l2 +1 l l=1 ∆ul + PL l=l2 +1 ∆ul . Proof. Starting from the constraint C2 in (34), after some
10
Algorithm 3 Inner Loop AO Algorithm Input: Convergence tolerance ϵ, iteration index r = 0. Output: k∗ , b∗ , s∗ . 1: Initialization: k (0) , b(0) , s(0) , tmin . 2: repeat 3: r ←r+1 (r−1) 4: for k = 1, 2, . . . , min bi with fixed b(r−1) , s(r−1) do i 5: if t(k) < tmin then 6: tmin ← t(k), k(m) ← k . 7: end if 8: end for (r) 9: Update the relaxing continuous variable b̃ by PGD r−1 with fixed s ; (r) 10: Rounding b̃ into b(r) by Algorithm 4. 11: Update the relaxing continuous variable s̃(r) by CVXPY with fixed br ; 12: Rounding s̃(r) into s(r) by Algorithm 4. 13: until |t(k (r) , b(r) , s(r) ) − t(k (r−1) , b(r−1) , s(r−1) )| ≤ ϵ. 14: if t(k, b, s) < tmin then 15: tmin ← (k, b, s), k∗ ← k, b∗ ← b, s∗ ← s. 16: end if
inequality manipulation, it is easy to arrive that P P Ui − ( l1 ul + L ul ) bi ≤ Pl1 l=1 PL l2 +1 . l=1 ∆ul + l2 +1 ∆ul
(35)
Considering that bi needs to be an integer, we thus obtain the above maximum blimit . i Based on Theorem 2, we design a two-level pruning strategy. First, a single-node feasibility check is performed: If the memory limit of any specific UE satisfies
Ui < Mf ix (l1 , l2 ),
(36)
the current split strategy is infeasible. If passed, a global capacity check is conducted then: If the sum of the maximum batch size of all UEs is less than the total volume, as N X
blimit (l1 , l2 ) < B, i
(37)
i=1
it indicates that the current split pair (l1 , l2 ) cannot fulfill the task requirements, and thus the combination is directly eliminated from the search space. Only the combinations that pass these pruning tests are passed to the inner loop to calculate the corresponding minimum time.
Theorem 3. During the inner-loop AO with fixed k , the continuous relaxations of the isolated sub-problems for b̃ ∈ Rn and s̃ ∈ Rn are strictly convex. Proof. We first analyze the nonlinearity of the computing (f /b) performance fi (b̃i /k) as a function of b̃i . Although the Roofline model introduces a step-function, the objective function maintains smoothness within the practical constraint range. According to the algebraic transformation of the computation model in (3), the computation time function Tcomp (b̃i ) transforms into the pointwise maximum of linear functions, strictly preserving convexity as ! C kM + b̃i ∆M ) Tcomp (b̃i ) = max b̃i , , (38) kFi kβi where M and ∆M represent fixed and variable memory access, C denotes the computational workload. Meanwhile, the communication time function can be simplified as Tcomm (b̃i ) ∝ b̃i . Obviously, it is a strictly convex function with respect to b̃i . Similarly, considering that the transmission rate is proportional to s̃i , the communication time is simplified as Tcomm (s̃i ) ∝ s̃−1 yields a positive second i derivative, ensuring strict convexity. The objective function (33) composed of these two time function terms is constructed recursively using only two mathematical operators: the non-negative addition of stage delays and the pointwise maximum across UEs and previous stages. According to convex analysis, if the base time functions are convex, their non-negative sums and pointwise maximums strictly preserve convexity. As a result, Theorem 3 is proved. Next, we fix (k, s), transforming P 1 into a sub-problem of the batch size set. To handle the integer constraints, we strictly relax bi into continuous variables b̃i . By simplifying C2 and C3, we derive the sub-problem P 2 of b̃i as follows. According to Theorem 3, P 2 is the minimization of a convex max function with box and equality constraints. We use projected gradient descent (PGD) method to solve it, ensuring constraint C6 is satisfied in each iteration. Upon obtaining the optimized results b̃i , we apply the largest remainder rounding algorithm for integer rounding, as detailed in Algorithm 4.
P2 :
b̃
s.t. 4.2.3 Inner Loop AO Given (l1 , l2 ), P1 is transformed into a resource allocation sub-problem of pipeline depth k , time slot set s = {si }, and batch size set b = {bi }. Due to the high coupling among variables in the objective function, we employ AO to solve them iteratively. First, we fix (b, s). An excessively small k results in low parallelism, while an excessively large k degrades computational capability, both leading to increased total time. This causes the pipeline time to exhibit a convex-like trend with k . However, the discontinuous nature of the objective function prevents a direct explicit solution. Therefore, we perform discrete search for k within [1, mini bi ] until a local optimum is identified. After fixing k , we can obtain Theorem 3.
min max C(i, k, 9|b̃) C4:
i
n X
(39)
b̃i = B,
i=1
C6: k ≤ b̃i ≤ blimit . i Finally, fixing (k, b) and relaxing si into continuous variables s̃i , P 1 becomes a sub-problem P 3 of the time slot set, which is also a min-max problem under inequality constraints. As in Theorem 3, P 3 is strictly convex with respect to s. Consequently, efficient convex optimization toolkits such as CVXPY [40] can be employed. Similarly, we apply Algorithm 4 to round s̃i into si .
P3 :
min max C(i, k, 9|s̃)
s.t.
C5:
s̃
i
n X i=1
s̃i ≤
T . τ
(40)
11
Algorithm 4 Largest Remainder Rounding Algorithm Input: Continuous variable set x̃ = {x̃i }, total number N . Output: Rounded variable set x = xi . 1: for i = 1, 2, . . . , N do 2: ∆xi = x̃i − ⌊x̃i ⌋ 3: end for PN PN 4: ∆x = i=1 x̃i − i=1 ⌊x̃i ⌋ 5: Arrange ∆xi in descending order 6: for i = 1, 2, . . . , N do 7: if ∆xi is the top ∆x in {∆xi } then 8: ∆xi ← 1 9: else 10: ∆xi ← 0 11: end if 12: xi = ⌊x̃i ⌋ + ∆xi 13: end for
4.2.4 Convergence and Complexity Analysis Bounded by the whole model layer L and the pruning strategy, the outer architectural variables (l1 , l2 ) reside in a finite discrete combinatorial set so that the outer search strictly terminates in at most γ outer steps, where γ is the number of layers under Theorem 2. Thus, the global convergence solely depends on the termination of the inner AO loop. We have Theorem 4 to solve this. Theorem 4. Under fixed (l1 , l2 ), Algorithm 3 globally converges to a stable stationary configuration {k ∗ , b∗ , s∗ } within a finite number of iterations. Proof. For any fixed (l1 , l2 ), the inner loop variables (k, b, s) are strictly constrained by C3, C4 and C5. Therefore, the global feasible state space Ω = {k, b, s} is a strictly bounded and finite discrete set. At iteration m, the continuous AO guarantees descent in Theorem 3, but Algorithm 4 introduces bounded localized perturbation, which breaks the guarantee of strict monotonic descent across successive evaluations t(r) . To mathematically neutralize this limitcycle oscillation, Algorithm 3 maintains a global historical minimum, represented as (r) (r−1) (r−1) tmin = min tmin , t(k (r) , b(r) , s(r) ) ≤ tmin , (41) which is strictly monotonically non-increasing. The physical pipeline execution time is strictly bounded below by the physical boundaries and hardware compute ceilings as (m) tmin ≥ Tbound > 0. According to the monotone convergence theorem, any sequence of real numbers that is monotonically non-increasing and bounded from below must strictly (m) converge to a finite limit. Due to the convergence of {tmin }, and the algorithm navigates exclusively through the finite discrete state space Ω, the sequence cannot take on infinitely many distinct values. The inner loop enforces a truncation tolerance |t(r) − t(r−1) | ≤ ϵ. Thus, the iterative sequence must terminate after a finite number of state transitions, stably collapsing onto the discrete stationary configuration yielding tmin . As a result, Theorem 4 is proved. In the inner AO process, we initialize the variables and sequentially solve the sub-problems until the convergence criterion ϵ is met. The complexity of Algorithm 2 comprises three main parts. Theoretically, the complexity of the outer loop is O(L2 ). But the pruning strategy in Theorem 2 significantly reduces the computational load. By leveraging the UE memory constraints, the number of layers requiring
search is compressed to γ , resulting in a complexity of O(γ 2 ). For the inner AO, the computational complexity is dominated by solving P2 and P3. For convex optimization problems with n-dimensional variables, the time complexity typically scales exponentially with the number of UEs. Therefore, the total algorithm complexity can be expressed as O(γ 2 log( 1ϵ ) · N 3 ). 4.3 Adaptive Re-allocation Optimization Following the initial profiling phase, where the model split layers (l1 , l2 ) are determined via the aforementioned SPA optimization, the BS broadcasts the initial parameters of the head and tail models to all UEs. Simultaneously, it distributes the allocated batch size bi and the number of time slots si to each respective UE. While the system is expected to operate according to these optimized settings, the inherent time-varying performance of edge devices poses significant challenges during subsequent training rounds. Specifically, the computational capabilities of UEs may fluctuate due to sudden local processes or node failures, and communication rates may vary due to device mobility. These factors can cause the actual parallel training process to deviate from the initial design, necessitating the faulttolerance mechanism for adaptive allocation. Considering the limitations on UE performance variations, coupled with the strict requirements for model parameter privacy and transmission overhead, dynamically adjusting model split layers during training is prohibitively expensive. Consequently, we only target the tunable parameters on the UE side, considering k, bi , si as optimization variables for this phase. To this end, a threshold detection mechanism is established before each training round. If the deviation ratio of training time compared to the previous round exceeds a predefined fault-tolerance threshold as
tm − tm−1 > δ, (42) tm−1 the system deems that a significant performance fluctuation of specific UE has occurred. This triggers ARA optimization at BS before next training round begins, based on the latest parameters information uploaded by UEs and the re-optimized variables are then distributed to them. This optimization problem is formulated as P 4. P4 : s.t.
min t(k, b, s)
(43)
k,b,s
C2:
l1 X
(ul + bi ∆ul ) +
l=1 i
C5:
n X i=1 n X
(ul + bi ∆ul ) ≤ Ui ,
l=l2
C3: 1 ≤ k ≤ min bi , C4:
L X
k ∈ N,
bi = B,
bi ∈ N,
τ si ≤ T,
si ∈ N,
i=1
Since P 4 constitutes a sub-problem of the original problem P 1 with l1 , l2 fixed, the ARA algorithm is easily obtained based on AO. Drawing upon the inner loop AO described in Section 4.2, we decompose P 4 into three subproblems and solve them iteratively by fixing the other variables in the same way. Consequently, we can also use Algorithm 3 to solve this optimization.
12
TABLE 3: Average per-round training time (s)
TABLE 2: System parameter settings Parameter B T BW pi Gi N0 Fi βi
Parameter N τ ρ p0 G0 Ui F0 β0
Value 8 0.125 ms 2 46 dBm 18 dBi [1,2] GB 14 TFLOPS 900 GB/s
SL
USL
E XPERIMENTAL R ESULTS
In this section, we introduce the experimental parameter settings and present a variety of experimental results to comprehensively evaluate the performance of AC2 P2 SL.
5.1
Schemes
Experiment Settings
We consider a cellular network with a 500 m radius, where the BS is centrally located and UEs are randomly distributed. We adopt the line-of-sight (LoS) channel model described in [41], characterized by an average path loss exponent of 2.1 and a shadow fading standard deviation of 3.6 dB. To emulate the heterogeneity performance, the peak FLOPS of each UE is uniformly selected from [100, 200] GFLOPs. Similarly, the memory bandwidth and maximum storage capacity follow random distributions within [5, 10] GB/s and [1, 2] GB. The remaining system parameters are specified in Table 2. To perform reasonable evaluations, we conduct the common image classification task using the ImageNet-100 dataset, which consists of 100 varied categories selected from the ImageNet dataset [42]. It contains 129,395 training samples and 5,000 test samples, and each data sample has dimensions of 224 × 224 × 3. Experiments are conducted under both independent and identically distributed (IID) and non-IID conditions. For IID data setting, we randomly shuffle the training dataset and evenly distribute it among all UEs. Under non-IID setting, we use the Dirichlet distribution to distribute the training dataset, with the Dirichlet parameter set to 0.5 to represent the heterogeneity of local terminal data. Meanwhile, we randomly shuffle the test dataset and distribute it evenly among all UEs. The models adopted for training include ResNet [43] series models and Vision Transformer (ViT) [44]. To determine specific model metrics prior to training, we calculate the FP and BP computational workloads for each network layer on a per-sample basis. Moreover, the fixed memory access volume of FP is represented by the model parameters, while for BP, it corresponds to the sum of model parameters and their gradients. The variable memory access volume is characterized by the data flow: for FP, it involves reading input data and writing output data; for BP, it involves reading the input data and output gradients and writing the input gradients. Regarding memory usage, the fixed component comprises model parameters, parameter gradients, and optimizer states, whereas the variable component consists of the size of intermediate activations generated during FP.