FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
Authors: Srijith Nair
Michael Lin
Peizhong Ju
Amirreza Talebi
Elizabeth Serena Bentley
Jia Liu
Venue: ICML 2025 poster (2025)
Keywords: federated split learning
communication efficient split learning
client-constrained federated learning
Primary area: optimization->large_scale_parallel_and_distributed
TL;DR: We propose a new federated split learning algorithm with sublinear convergence rate and superior communication efficiency compared to state-of-the-art.

Abstract: Collaborative training methods like Federated Learning (FL) and Split Learning
(SL) enable distributed machine learning without sharing raw data. 
However, FL assumes clients can train entire models, which is infeasible for large-scale
models.
In contrast, while SL alleviates the client memory constraint in FL by offloading most training to the server, it increases network latency due to its sequential nature. 
Other methods address the conundrum by using local loss functions for parallel client-side training to improve efficiency, but they lack server feedback and potentially suffer poor accuracy. 
We propose FSL-SAGE (Federated Split Learning via Smashed Activation Gradient Estimation), a new federated split learning algorithm that estimates server-side gradient feedback via auxiliary models.
These auxiliary models periodically adapt to emulate server behavior on local
datasets. 
We show that FSL-SAGE achieves a convergence rate of $\mathcal{O}(1/\sqrt{T})$, where $T$ is the number of communication rounds.
This result matches FedAvg, while significantly reducing communication costs and
client memory requirements. 
Our empirical results also verify that it outperforms existing state-of-the-art
FSL methods, offering both communication efficiency and accuracy.

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
Srijith Nair Michael Lin
Peizhong Ju
Amirreza Talebi Elizabeth Serena Bentley Jia Liu
Abstract
vron et al., 2023) pretrained on large corpora of data have demonstrated near human-level performance variety of tasks. These models demand heavy computational re- sources and very large datasets for training (Kaplan et al., 2020; Hoffmann et al., 2022). The datasets are often owned by different organizations, making it impossible for single entity train model on all the data (Sani et al., 2024b). Distributed learning paradigms like federated learning (FI.) (McMahan al., 2016) and split learning (SL) (Vepakomma et al., 2018), which preserve data privacy while leveraging the compute power of various clients, have become increas- ingly popular in LLM pretraining (Sani et al., 2024a) and fine-tuning (Sun et al., 2024; Wang et al., 2024).
Collaborative training methods like Federated Learning (FL) and Split Learning (SL) enable distributed machine learning without sharing raw data. However, FL assumes clients can train entire models, which is infeasible for large-scale models. In contrast, while SL alleviates the client memory constraint in FL. by offloading most training to the server, it increases network latency due to its sequential nature. Other methods address the co- nundrum by using local loss functions for parallel client-side training improve efficiency, but they lack server feedback and potentially suffer poor accuracy. We propose FSL-SAGE (Federated Split Learning via Smashed Activation Gradient Estimation), a new federated split learning algo- rithm that estimates server-side gradient feedback via auxiliary models. These auxiliary models peri- odically adapt to emulate server behavior on local datasets. We show that FSL-SAGE achieves con- vergence rate of ((I/VT), where T is the num- ber of communication rounds. This result matches FedAvg, while significantly reducing communica- tion costs and client memory requirements, Our empirical results also verify that it outperforms existing thrthe FSL methods, offering both communication efficiency and accuracy.
Despite their benefits, FL and SL both have some limitations. Although FL can preserve data privacy by training the model the client and aggregating the client-side models at server, an inherent assumption of FL is that client devices have enough computational power to process the model, which is impractical for large models. SL, on the other hand, reduces compute requirements for clients by splitting the model into two parts, processing a smaller piece at the client, and offloading the other piece the server. This, however, implies that processing the complete model in each iteration involves transmitting the cut-layer features and gradients back and forth between the client and the server. Moreover, SL loses the benefits of data parallelism in FL, since the server needs to assist gradient computation for each client in a sequential fashion. Consequently, SL suffers significant increases in communication load, power, and latencies.
1. Introduction
To reduce the communication overheads incurred by SL, some recent works (Han et al., 2021; Mu & Shen, 2023) pro- pose to use auxiliary models at the clients. These auxiliary models continue the forward propagation of the client-side model to compute a local loss function, which is then used to update both models via gradient descent. This elimi- nates the need for awaiting the server's response, which significantly speeds up the training process. Meanwhile, the server-side model is updated at lower frequency by using the cut-layer features from the clients. While enjoying the benefit that the clients can conduct training in parallel, this can significantly reduce communication costs and latencies. However, the client-sidd side models do not get any feedback from the server-side model, i.e., the clients only updated
Background and Motivation: In the recent years, large foundation models (Devlin et al., 2019; Radford et al., 2021; Ramesh et al., 2022) and LLMs (OpenAl et al., 2024; Tou-
'Deparmment of Electrical and Computer Engincering, The Ohio State University, Columbus, Ohio, USA 'Department of Computer Science, University of Kentucky, Lexington, Kentucky, USA DDepartment of Industrial Engineering, The Ohio State University, Columbus, Ohio, USA Air Force Research Labo- ratory, Rome, New York, USA. Correspondence to: Srijith Nair <nair.203@osil oow.eoou.ssu.uu................. Jia Liu cetu@cccc.ou.........
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
using the auxiliary models and thus never get trained via the server-side model.
result, we further show that under a mild in-expectation PAC-Ieamabilityy assumption on the auxiliary models, FSL-SAGE can achieve an O(1/VT) convergence rate, which is the same as those of sfate-of-the-ar FL algo- rithms even with heterogeneous datasets. Our analysis sheds theoretical lights on the impacts of the class of auxiliary model functions on the final-time of conver- gence, which could be of independent interests to other FSL methods with auxiliary models.
To overcome these limitations, we propose a new approach called Federated Split Learning via Smashed Activation Gradient Estimation (FSL-SAGE). Note that the auxiliary model other FSL frameworks can be interpreted as mim- icking the server-side model at each client. Since, from the client-side model's perspective, the role of the server-side model is to provide the gradients of the loss function with re- spect to the cut-layer features, we train the auxiliary models to estimate these cut-layer gradients returned by the server- side model, which are then used to update the client-side models. This allows FSL-SAGE to continue to enjoy the benefits of FL. and SI., while using the auxiliary models to accurately play the server's role, Also, we propose to update the auxiliary models periodically in less frequent fashion, thus significantly reducing communication costs.
We further propose a "lazy version" of our FSL-SAGE method, where the auxiliary models are frozen beyond certain point of alignment. This lazy version can be useful in practice to further reduce communication costs. We rigorously establish explicit trade-off between this time and the accuracy of the final model.
We conduct extensive experiments using large ResNet- 18 and meedimm models to verify our theoretical results and demonstrate that while being much more communication efficient than existing state-of- state-oftthe-art FL/SL algorithms, the accuracy of FSL-SAGE is either on-par or even better.
However, due to the following technical challenges, establishing finte-time convergence rate performance guarantee for FSL-SAGE is highly non- trivial and necessitates new proof and analysis techniques. First, as will be shown later in Fig. 1, there are multiple coupled time-scales involved during training, which include 1) the local step, whereby the client-side model is updated over mini-batch of local data; 2) the lesser frequent server step, where the server-side model is updated given the cut- layer features; 3) a federated averaging step for the client- side models; and finally, 4) the least frequent alignment step, where the auxiliary models are updated match the server- side model. Also, one needs to account for the auxiliary alignment step, which has not been studied in conventional SL methods (e.g., SplitFed (Thapa et al., 2022) or in the precursor to this work (Mu & Shen, 2023)). Exacerbating the problems in ESL-SAGE convergence rate analysis is the fact that the auxiliary models inject highly heterogeneous estimation errors during the training process, which are further coupled with all four time-scales.
2) Technical Challenges:
2. Related Work
In this section, we provide a quick overview on several closely related areas, namely FL, SL, and FSL, thus putting our work in comparative perspectives.
1) Federated Learning: Federated learning was first pro- posed in McMahan et al. (2016) and theoretically analyzed in Yu et al. (2019), as an alternative to centralized learning, whereby several clients could collaborate on training a sin- gle ML model under the supervision of a server without the need for sharing client data. In FL., the model parameters are transmitted over the network to the client devices, where the clients can train the model on their local datasets be- fore transmitting them back to the server where the models get aggregated. This process is repeated iteratively until the model converges. Federated learning enjoys the ben- efit that the model needs to be transmitted only once per communication round, and the clients can train their local models in parallel, for several iterations per round. Several FL methods (Yang et al., 2021; Karimireddy et al., 2019; Sahu et al., 2018; Reddi et al., 2021; Li & Lyu, 2023) have been analyzed in the recent years, most of which are vari- ants of the original FedAvg algorithm (McMahan et al., 2016). As mentioned earlier, one of the key limitations of conventional FedAvg-type algorithms is that they assume each client has sufficient memory capacity to store the entire model to perform local gradient-based updates. However, this assumption becomes increasingly problematic as ma- chine learning models continue to grow in size.
3) Main Contributions:
The key contribution of this paper is that we overcome the aforementioned technical challenges and rigorously establish the finite-time convergence rate guarantees of our proposed FSL-SAGE. Our main technical results are summarized as follows:
We propose a new federated split learning algorithm called FSL-SAGE to train models that are too large to be trained on commodity client devices. By using a periodically aligned auxiliary model at each client to estimate the server-side model response, FSL-SAGE enables large model federated training, while enjoying low communication costs and the same data parallelism as in conventional FL.
We establish a general finite-time stationarity conver- gence rate bound for FSI--AGE. Based on this general
2) Split Learning: To alleviate the memory constraint in training large models with FL., an alternative strategy was
2

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
proposed in Vepakomma et al. (2018), where the idea is to split large neural network model into two parts; the smaller part that contains the input layer is trained on the clients and the remaining larger portion is trained on the server. The layer at which the model is split is often called the cut-layer. During training, the cut-layer features for each mini-batch of data are communicated from the clients to the server, and the gradients of the loss with respect to (w.r..) the cut-layer features are in-turn relayed from the server back to the clients. Moreover, the server needs to sequentially process the server-side portion of the model for each client as discussed in Gupta & Raskar (2018), which renders high latency in using the vanilla SL approach. It is thus clear that SL reduces client memory requirements at the expense of slow and highly sequential training and increased communication costs. CPSL (Wu et al., 2023) is an extension of traditional SL, where clients are divided into clusters and clients within each cluster train in parallel. However, CPSL still suffers the same latencies as in SL.
(2023) is similar to SplitFedv2, where only one server-side model is updated a firt-comm-fiist-serve basis. How- ever, our work differs from these existing methods in the following key aspect: the auxiliary models at the clients in our proposed FSL-SAGE are periodically aligned with the server-side model to ensure high training accuracy with low communication costs. In contrast, the auxiliary models in the aforementioned methods do not perform any alignments with the server-side model, thus leading to potentially low training accuracy. Notably, we circumvent all limitations of the above methods. Our method minimizes computational burden on the clients and communication costs that are in- evitable in the other algorithms, and reduces latencies by facilitating parallel training, without compromising accu- racy, Oh et al. (2022) addressed the inherent limitations of using local losses and the dual frequency of updates in FSL frameworks, and proposed to locally regularize the client- side models to maximize the mutual information between raw and cut-layer data, while also augmenting the smashed data so that the server-side model to improve latency and accuracy in parallel SL.
3) Federated Split Learning: Since FL and SL each have inherent limitations, an important research is to explore whether we can design a new algorithmic framework that combines the Sata-parallel training efficiency of FL with the model-splitting capability of SL, thus enabling both faster training and reduced computational demands on clients. To- ward this end, the SplitFed algorithm with two variants was proposed in Thapa et al. (2022), namely SplitFedv1 and SplitFedy2. SplitFedy1 maintains multiple copies of the sidver-side model at the server, one corresponding to each client, thereby allowing all clients to train in parallel, but at the expense of the non-scalable computational demand the server. By contrast, in SplitFedv2, the server performs each client's backpropagation sequentially on a first-come- first-serve basis. The key advantage to both methods is that the clients can perform forward propagation in paral- lel. AdaptSFL (Lin et al., 2024) is a recent extension of SplitFedv2, where the authors adapted the client-side model size based on the available resources making SplitFed more flexible to resource-constrainee clients. However, the com- munication overhead of these methods is the same as SL.
Group Knowledge Transfer (GKT) (He al., 2020) is an- other approach related to our work. The basic idea of GKT is to train smaller client-side auxiliary models and then trans- fer their knowledge to the server via a variant of knowledge distillation (Hinton et al., 2015) by using the logits from the auxiliary model to train the server-side model. GKT is classified as a knowledge distillation approach with mod- ified loss functions to incorporate server feedback during client-side model training, and is not exactly federated split learning method. Thus, its performance is not directly comparable to FSL methods due to its different training objective function.
Although auxiliary models have been used in centralized deep network training Marquez et al. (2018); Bhatti & Moon (2022), the training of the auxiliary models in such works indirect, and can be interpreted as training the network in parts. In our work, we explicitly train the auxiliary models to approximate the server-side models, thus ensuring their most accurate use.
4) Auxiliary Models: The most related works to our paper are Han et al. (2021) and Mu & Shen (2023), where algo- rithms using local cost functions were proposed avoid the need for the clients to wait for the server to return the cut-layer gradients. The client-side models are updated in parallel by backpropagating gradients via the local loss function through auxiliary models on the corresponding clients. The client-side and auxiliary models are then ag- gregated every round. The server-side model is updated by receiving lut-l activations once in several local it- erations. The method in Han et al. (2021) is similar to SplitFedv1 I in the sense that separate server-side models are maintained per client, while the method in Mu & Shen
3. The Proposed FSL-SAGE Algorithm
In this section, we present our system settings and the pro- posed FSL-SAGE algorithm.
1) System Setting:
As shown in Fig. 1, our federated system consists of a server and m clients indexed by i. Each client is associated with a local dataset D,. The server contains two server processes: (i) the " "F-seever" for conducting federated aggregation, and (ii) the "S-server" that processes the server-side model and updates the auxiliary models for each client. In this paper, we call one cycle of federated
3

---

FSI-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
Algorithm 1 The Code of Client in FSL-SAGE.
Require: Fa.
(1):
for f- ...,.. -T do # Inirialize weights from F-server
F
x 1.0 x
3: if 0 mod / then
1
4:
Receive o' from S-server
S
5: else
MSB
IC P.re
D2
L
(6z: pz)
6:
x t- a,t
3,P
7:
end if
for k
1, ... (K 1) do
9.2
2a.2%
9:
(g :4
1. Dj
# Sample local mini-batch
10:
1.k zi.s(a (.3
Forward pass on client
Ce
11:
if k ===00 =0 mod (K/Q) then
2
D
(c.. Mom)
12:
Send (z Mi
l.k t.k
to S-server
f
Zmnb
13:
end if
14:
L.k. Li + (2' fif)
st.k
# Compute grad estimate
(0% + 3 A. pa.)
15:
1.++1 t Fe - mr.
s.kkk Rib
# Update client
Server Processes
Clients
16:
end for
17: Send Fe to F-server
Figure 1. Schematic diagram of the FSL-SAGE algorithm. Text on the arrows indicates (sender -> receiver, message, rate), e.g., the red arrows transmit æ, between C; and the F-server at the rate of I per round (p.r.,, the updated Va.i is sent to C, once in rounds. Arrows color coded: brown arrows represent local operations, blue arrows represent the smashed data and label transmission the server, red arrows indicate federated averaging, and green arrows indicate auxiliary model transfer; circled numhers indicate the order of operations
18:
end for
Algorithm 2 FSL-SAGE F-Server
Initialize model x
2: (a. (Q] f- split (I:²) 3: Send x. to Server S
# Split model
4: Broadcast x² to all clients
5: for 1 4- 0, 1,..,... T 1 do
i.K 6: Receive o, from clients i = 1,2,.....
aggregation as one "round" and assume that FSL-SAGE runs for total o rounds. Our objective to learn a model Rdd that minimizes a global loss function f: RR R:
Σm= x 1,K x/ m
#Aggregate client model
8: Broadcast x
++1 to all clients
end for
m
min f(x) = m Σ P((a); wER
(I)
are in turn used to update the client-side model; (4) Ev- ery client performs K local iterations before sending their models, Te,is to the F-server for aggregation. Algorithms and 2 summarize the client and F-server operations in FSL-SAGE.
where F((-) = == -D, [Fi (o; E,)l corresponds to the objec- tive function of the i-th client.
2) The FSL-SAGE Algorithm: The entire process of FSL- SAGE is illustrated in Fig. 1. Similar to split learning, in FSL-SAGE, we split the model æ into the client-side model Ze,i for all clients E [m]] and the server-side model x. As shown in Fig. 1, a single local iteration at the i-th client con- sists of the following key steps: (1) The client-side model takes as input sampled batch of data &, D; and outputs where the subscript f
To update the server-side model, the clients periodically send the computed cut-layer features and labels to the S- server Q times per round. The S-server computes the true loss function Fa(:;; Zi,f ys) and the gradient of F; with respect to zi,f, which is denoted by 22,b- The server-side model is then updated a gradient descent fashion. The tuple (zt,/, y:). called the aligument dataset, is stored for training the auxiliary model. Algorithm 3 summarizes the operations of the S-server.
the cut-layer features ( E.), in z denotes the output of the "forward" pass; (2) Each client has an auxiliary model Ca,i, which takes as input Zi,f and computes the cost function estimate ( (3) In the backward pass, the auxiliary model computes the gradients of the cost estimate with respect to the cut layer features, denoted as ( zi Vi), where y is the label corresponding to E and the subscript b in 2,,6 denotes the result of the 'backward' pass. These gradients
Once every | rounds, the server initiates an alignment pro- cess, where each client's auxiliary model is updated using the respective alignment dataset. Assuming round t is a multiple of I the auxiliary model at round updated as

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
Algorithm 3 (Lazy) FSL-SAGE S-Server
Reddi al., 2021; Yang et al., 2021).
1: Receive xº. from F-server
Assumption 4.1 (Smoothness). For all i 1,2,...m. the following objective functions corresponding to the i- th client are smooth, i.e., for all Eg Ra and yı:
2: for E [T] do
3: for [K] if
mod K/Q do
4: Receive (z 36) Vi on FCFS basis
5: Compute grad g, (xf:;; (i ye) via backprop
L.k
(a)
The gradient of 9: Le- Le-Lipschitz with respect to I = i.e., there exists a constant Le > 0 such that Ilg:(ee, usi Eis (i, - 9i (e UsiGi,
VFi(a Tzi (.. y1) is (Ic: at).
6: -. = x - -gg 1.
# Update server model
7: Store for client i: (2.,,.
S
8:
end for
Le u - v|l, for all u = (Uo and = (Ve, v.) E X.
9:
if = 0 mod and < T" then
#T S
for client i t 1.2..... do Compute 2,,6(2),,, ut Vj from
11:
(b) The gradient of the local client objective Fi, Ic., V V.Fi Lasi Ei, Mi) is Lo-Liscchitz with respect xc: there exists constant The > 0 such that |1 (Ue, Da Si, Mi) - -Ve, Darti Ei, S E. Uu well for all ue. Us E Nec
{ {
12:
x², = align
13:
Send a to client i
14:
end for
15:
end if
16:
end for
Assumption 4.2 (Bounded Variance). The following quan- tities measuring variability of the local stochastic estimates of the objective function and the variability of the client loss function are bounded:
the following:
Qt Σ 2
2e, = arg min, 2a
(a f, ,.,,, vt) - 2/ 2/,0 (2)
(a) The local gradient estimate for a mini-batch sam- pled at client is unbiased for all i, i.e. V ã. Bg,,D V (T:( = vf())) and has a bounded variance, i.e., or, > 0 such that - vÅ o
where the backward gradients = are computed by passing zi and of through the server-side model. In practice, the minimization in (2) could be approximately done by using a finite number of steps of gradient descent, Adam (Kingma Ba, 2015), or other optimizers. Also, since the alignment dataset grows larger in size with the number of rounds increases, one can discard the older entries in order to maintain fixed maximum size of the dataset. For simplicity in theoretical analysis, we assume that the server has sufficiently large storage to accommodate the growing alignment dataset size.
EG,DD
(b) The local gradient for client i is unbiased for all i, i.e., V T. BE,-D, (V; &.] = (F, and has a bounded variance, i.e., 3 do > 0 such that EgWD, [[VFF((; - V=(+))] oo
(c) The global variability of the local client gradi- ent is bounded, i.e., for all æ and E [m]] [[V V[(x)]] soz og.
In Algorithm 1, if the auxiliary models are only updated for the first T' < T rounds of communication and then are frozen until the end of training, we will refer this version as the Lazy FSL-SAGE algorithm. Note that, in practice, choosing a smaller I' can further reduce communication costs at the expense of final accuracy of the learned model. On the other hand, the FSL-SAGE algorithm is non-lazy if T" = T. Section 4, we will characterize the stationarity gaps of both lazy and non-lazy versions of FSL-SAGE.
Assumptions 4.1 and 4.2 are commonly used in analyzing convergence of many optimization algorithms from the clas- sical SGD (Ghadimi & Lan, 2013) many variants of fed- erated learning (Yu et al., 2019; Yang et al., 2021). The vari- ance terms øi. of and og in Assumption 4.2 characterize the degree of intra-client and inter-client data heterogeneity.
4.1. General Finite-Time Convergence Rate Results of Non-lazy FSL-SAGE
4. Theoretical Convergence Analysis
In what follows, we provide an upper bound on the station- arity gap of the FSL algorithm, i.e., a bound on the expected squared norm of the gradient of the model parameters.
For convenient reference, we summarize the key notation used the following analysis in Appendix and Table 2. We begin by making the following standard assumptions widely adopted in the FL. literature (Karimireddy et al., 2019;
Theorem 4.3 (Convergence of Non-Lazy FSL-SAGE). Un- der Assumptions 4. and 4.2, given step-sizes T and n
5

---

FSL-SAGE: Accelcrating Federated Split Learning via Smashed Activation Gradient Estimation
satisfying v S AV2QL, and AL S 2vI0KL.' FSL-SAGE, as described in Algorithmss 1-3, has the following finite-lime stationarity convergence rate bound:
"yes" if we further impose a mild learnability assumption on the auxiliary model. Toward this end, first introduce the notion called in-expectation learnability (Mey, 2022):
Definition 4.5 (In-Expectation Learnability). The hypoth- esis class G is said to be in-expectation learnable by the empirical risk minimization (ERM) algorithm if and only if A > O, there exists a rg(() such that, for 2 rg(() training samples, the following bound holds:
[(x") -
E |1
min RE
cmin {NL. mT
mk
8' ((111)))
(3)
2cQ (,. mn)
where
E [l((0r;))) f(a=)1
(TT. = mc min mm)
eKK
[nL
Smnn E [l)(0=))- - s(x))1²] +e,
+ 40400m 12 + Ko²)
where the expectation over æ; 31 22...., the ERM hypothesis learned from {æ,)%_. that are i.i.d sampled from D:
,....,,a, Fy and 0, is
+m²1jo 32Qm² 1/ L + Qoa)]-
(4)
and E is the gradient estimation error incurred by the aux- iliary model averaged over clients i E [m]:
Or ((2:))) = arg min Age
(1g(0: - f(x:)]l (8)
e'
Vî(Ea) -FF
(5)
Next, we introduce additional notation pertaining to the loss function used to train the auxiliary model. We start by denoting the hypothesis class of auxiliary model functions for the 7th client as A; := (Z.,(ZZ,rt)) Ia, E Xa), which is parameterized by æa,;- Omitting obvious parameters for readability, we can denote the empirical loss on training samples {(22, "/)))=12 as
A full proof of Theorem 4.3 is provided in Appendix B.I. The following result directly follows from Theorem 4.3. Corollary 4.4. For the step-size choices 1) O (mVT)
and nL O (T)) Algorithms 1-3 achieves the following
finite-time stationarity convergence rate:
1 Σ
min E [[V = ( (VT) LETT
0(1)
E, (2a, x)) =
[a,((tas)) - =,o(2=))?? (9)
(6)
and the true loss as
The last term in (6) corresponds to the round-average of the auxiliary models' estimation error in estimating the cut-layer gradients returned by the server, which plays an important role in the convergence of our method. While these terms similar to terms 1 and 3 in the RHS of (3) appear in con- vergence proofs of FedAvg (McMahan et al., 2016) and its variants (Yang et al., 2021; Reddi et al., 2021), the estima- tion error term a direct result of using our auxiliary model to estimate the server gradients. We note that, although CSE-FSL (Mu& & Shen, 2023) also uses auxiliary models, their analysis only considers the client-side and server-side models separately, and thus they bypass analyzing such important estimation error term.
(
Li x)) = ESS.NDDD ( (10)
Recall that the auxiliary model is trained once every rounds on an increasing number of smashed data samples r == Lt/tj. Thus, one can define the ERM solution round from Algorithms I & 3, and the corresponding optimal solution follows:
sy arg min L, x') 2. EXX
(11)
x 1A arg min Li x²).
(12)
0,
a, EXY.
4.2. O ) Convergence Rate of Non-lazy ESL-SAGE with Agnostic PAC Learnable Auxiliary Models
Then, the auxiliary model (11) is an in-expectation learner of A, if e > 0. there exists a re(e) such that, for 2 ri(e):
With the general finite-time convergence rate results The- orem 4.3, it remains unclear whether one can achieve the same O ) convergence rate as the conventional FedAvg- type methods. In this section, we show that the answer is
E [ at)] s C, x²) +
(13)
This leads us to introduce the following learnability assump- tion on the auxiliary model.
6

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
Assumption 4.6 (Learnable Auxiliary Model), The auxil- iary model of the i-th client aa,i in (9) is in-expectation learner of A, according to Definition 4.5 with sample com- plexity r.(c) = O (1/e))
models, there are several results such as the universal ap- proximation theorem (Cybenko, 1989; Hornik et al., 1989), and bounds for sufficiently wide or deep networks (Lu et al., 2017; Hanin & Sellke, 2018; Hanin, 2019) that show that et can be made arbitrarily small. Finally, from the above discussions, we have the following complexity results:
In (Grunwald et al., 2021), in-expectation bounds under the PAC-Bayesian formulation have been derived for loss functions satisfying the so-called (B. 3))Bernsteinn condi- tion (Bartlett & Mendelson, 2006) with B > 0, E [0, 1]. The sample complexity of those bounds given by 0 (1 comp 5' where comp is some complexity mea- sure that appears in the form of KL-divergence (Audibert, 2004; Zhang. 2006) or the conditional mutual information (CMI) (Steinke & Zakynthinou, 2020) of the posterior and prior distributions in PAC-Bayesian framework. Provided the comp term does not increase drastically with r, Assump- tion 4.6 can be satisfied by such models. Lastly, we make the following Lipchitz client assumption:
Corollary 4.9. For the step-size choices TΙ = O(1/vT) and r = (1//m))))) the non-lazy FSL-SAGE with PAC- learnable auxiliary models achieves a finite-tmme conver- gence rate of OL//TT).
r
4.3. Convergence Results of Lazy FSL-SAGE
Since the alignment process is expensive in terms of commu- nication costs, in practice one may want to stop performing alignment of the auxiliary model after T" < T rounds, i.e., the Lazy FSL-SAGE apporach. Since the only difference in the convergence bound of non-lazy FSL-SAGE in The- orem 4.8 and Lazy FSL stems from the estimation error's 's contribution to the upper bound s², we immediately have the following result for the Lazy FSL-SAGE with slight modification of the bound in Theorem 4.8:
Assumption 4.7 (Lipschitz client). For all clients í = 1, 2.....o m and for all uer Ue E Xx. the client-side cut-layer activation z (a E,) is Lxg-Liisschiz with respect to the client parameters Fe, i.e., there exists constant Laf > such that || (uc: E z1, (Vci E s L ||ue- vell-
Corollary 4.10 (Convergence Rate of Lazy FSL). Under Assumptions 4.1, 4.2, 4.6 and 4.7, and the step-sizes (1) ML) satisfying the conditions in Theorem 4.3, let us additionally assume that A is in- in-expectation learnable as per Defini- tion 4.5 with (e) = O (1/2) ) for all 5, then after >> T" rounds Lazy FSL-SAGE satisfies:
With these assumptions, we are now in position to state our next main result:
Theorem 4.8 (FSL-SAGE with PAC Auxiliary Models). Under Assumptions 4.1, 4.2, and 4.7, and with step-sizes (n, mL) satisfying the conditions in Theorem 4.3, the auxil- iary model satisfies Assumption 4.6 with re(e)= O (1/2), then after T rounds, the iterates in FSL-SAGE satisfy:
E ( 11- S eminffL: mnn QT
((x)) l
min =E(1...... (T/1)
(1+
2( n)
min
E
||}]
f(xo) f
2Q min mm) VT"
T
nE
T/1
emin Emin QT
3CКЛN
& m)
3,
1
Σ
T
(15)
2cQ min (LL, mn) T
2Q min( mm VT
T
i=1
L3 1
T
where c, C and ex, are as defined in Theorem 4.8.
Σ
+
(14)
2c0 min mn)
T
1=1
Although Han et al. (2024) studied the convergence of SplitFed without any auxiliary models local losses, to add to the challenges faced in their setting, in case, the client-side model is updated through an approximation of the server-side model, which introduces approximation er- rors in the client-side model. To the best of our knowledge, our work is the first to provide an explicit relationship be- tween the convergence rate and the approximation power of the auxiliary models.
where C 0 and 0- <.< < 0.5 - 20K2 22 are some constants, and E¹, ;= J Em C, x²).
Note that, for the step-size choices nT (L/VTT) and " ((I/(m)T)) the first three terms in the upper bound in (14), decrease with rounds as O(1/VT). The last term, given by C/2 ΣT= et where C O (1), is the average of el over all rounds. The term et, the error achieved by the best hypothesis in A at round and is entirely determined by the architecture of the auxiliary model. It is trivially zero for cases when the server-side model can be obtained for some realization of the auxiliary model parameters. A more inter- esting case is when the auxiliary model is typically much smaller than the server-side model. For 2-layer auxiliary
An important limitation between FSL with sequentially pro- cessed server-side model and FI., which manifests in our convergence analysis is the lack of linear speedup (Yang et al., 2021), i.e., the convergence rate does not reduce with the increase in the number of clients m. Note that in all FSL approaches (Han et al., 2024; Mu Shen, 2023), including
7

---

FSI-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
0.5
Toat Aemmey 0 50.00 GIB
0.5
PSL-SAGEE SpiitaaaE
Tee Accuracy
Teat Accurwey
0.4
PALBAGE
Sphn
0.7
Aed SpintFak-88
0.3
ESL-
0.6
0.2
CSE-FSL
SpiitFed-MS
0.8
0.1
0.5
FedAvg
0.2
SplitFod-SS
0.0
0.4
100
200
100
200
Commmnicstion Load (CB)
Commmniiattin Load (CB)
Figure 2. Tost accuracy vs. communication load for ResNet-18 CIFAR- 10 distributed homogeneously across 10 clients. Curves are labeled from best (leftmost) to worst (rightmost) in final accuracy.
Figure 3. Test accuracy vs. communica-
Figure 4. Best accuracy vs. Dirichlet a E
tion load for ResNet-18 CIFAR-100 dis-
(0, oc] for ResNet-18/CTFAR with non- iid. data for 10 clients up to 80 GiB comm.
tributed homogeneously across 10 clients.
in our work, where a single server-side model is maintained, the server-side model must be sequentially processed, thus losing linear speedup.
66.2M parameters for the client-side model and 365.7M 7M parameters for the server-side model. For the auxiliary model, we use the first 3 attention blocks cascaded with the language decoder of the server-side model (92.4M params).
5. Experimental Results and Discussion
I-d) Hyperparameters: For image classification, we use batch-size of 256. Clients train their models for 1 epoch on their local dataset per federated averaging round. For CSE- FSL and FSL-SAGE, the cut-layer features are sent to the server-side model every 5 local steps, and for FSL-SAGE, the auxiliary models are aligned with the server every = 10 rounds. We stop training when the communication cost incurred exceeds 200 GiB. The client-side and server-side models are optimized using Adam (Kingma & Ba, 2015) with a learning rate 10-3, weight decay 10-4, and 31 = 0.9, 33 = 0.999. For alignment, we use the same optimizer settings with no weight decay. For better interpretability. we use SplitFed-MS (multi-server) lieu of SplitFedvl and SplitFed-SS (single-server) in lieu of SplitFedv2.
In this section, we conduct numerical experiments to verify the efficacy of our proposed FSL-SAGE algorithm'.
1) Experiment Settings: 1-a) Compute and Baselines: We compare FSL-SAGE with FedAvg (McMahan al., 2016), SplitFedv1 I and SplitFedv2 (Thapa et al.. 2022), and CSE- FSL (Mu& & Shen, 2023). We use PyTorch for training on single NVIDIA H100 NVL GPU with 80GB of memory.
1-b) Datasets: Although FL generally applies to a wide range of machine learning tasks, we focus on two tasks: 1) image classification on CIFAR-10 10 and CIFAR- 100 datasets (Krizhevsky et al., 2009); and 2) natural language gener- ation using the E2E (Novikova et al., 2017) dataset. The CIFAR datasets contain 60K 32 x 32 3-channel images with 10 and 100 classes respectively. To simulate the effect of data heterogeneity we use the Dirichlet distribution parame- terized by a E (0, ) to determine the proportion of class labels (Hsu et al., 2019), where smaller a indicates more heterogeneity of class label distribution.
2) Results and Discussions: 2-a) Image Classification: In Fig. 2, we plot the test accuracy against communication load on 10 clients on CIFAR-10, averaged over 4 training runs. FSL-SAGE outperforms all the other methods in final accu- racy. More importantly. for a given level of accuracy, say 81%, FSL-SAGE achieves almost 2.2x lesser communica- tion cost than CSE-FSL, and 10x lesser than SplitFed-SS. Figure 3 depicts similar plot for CIFAR-100. Table 1 summarizes the best accuracies CIFAR datasets.
1-c) Models: For image classification, we use the ResNet- 18 (He et al., 2016), which comprises of 4 ResNet blocks. We split the network in between block 2 and 3 to create client-side model with 685K parameters and a server-side model with 10.5M parameters. For the auxiliary models, we arbitrarily cascade the first server ResNet block with the final fully- connecco layer, yielding 2.1M parameters. For natural language generation, we perform LoRA finetuning (Hu al., 2021) of GPT2-medium (Radford et al., 2019). We split the model after the first attention block, yielding
In Fig. 4, we plot the best accuracy achieved by each method
CIFAR-10
CIFAR- FFFR-1100 100
Algorithm
iid
non-iid
iid
non-iid
FedAvg
82.08
12.72
45.79
35.26
SplitFedv1
84.21
81.48
50.78
51.03
SplitFedv2
81.42
46.79
45.54
35.57
CSE-FSL
83.74
80.40
50.94
49.55
FSI-SAGE
85.71
82.75
56.06
51.90
'Our source code is available https:://// / htttps:://iittu.com// ri #ijith19e 9966/2966622 //SSL-S SAGE.
Table 1. Best test accuracy (%) for ResNet-18 trained up to 200 GiB. Best and second best results are colored and bold.
8

---

FSI-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
savings of our proposed FSL-SAGE method.
CSE-FSL
PSL-SGGE
1.0
SplitFed-SS
Acknowledgements
Test Loss
0.9
FidAvg
0.8
This work is supported in part by ONR grant N00014- 24-1-2729; NSF grants CAREER 21 10259, 21 2112471, and 2324052; DARPA YFA D24AP00265 and DARPA HR0011- 25-2-0019.
0.6
0.5
0.4
0
1000
2000
3000
4003
Impact Statement
Communication Load (GB)
Figure Test loss vs. communication load LoRA finetuning of GPT2-m on Lid. E2E distributed across 3 clients.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
over the course of training on CIFAR-10, upto a commu- nication load of 80 GiB, against the Dirichlet distribution parameter a (0, sc). We vary a from 10-1 (high hetero- geneity) 10 (low heterogeneity) simulate the effect of uneven data distribution among 10 clients. While Fe- dAvg and SplitFed-SS do not converge within the specified communication budget under high heterogeneity, the other three methods are much more robust to data heterogeneity. FSL-SAGE is robust within the given range of heterogeneity, and outperforms other methods. We present some supple- mentary plots for the image classification experiments in Appendix D.
References
Audibert, J.-Y. PAC-Bayesian Statistical Learning Theory.
PhD thesis, Paris 6, Jul 2004.
Bartlett, P. L. and Mendelson, S. Empirical minimization. Probability theory and related fields, 135(3):311-334 2006.
Bhatti, H. L. and Moon, J. Locally supervised learn- ing with periodic global guidance. arXiv preprint arXiv:2208.0088 2022.
2-b) LoRA Fine-tuning GPT2 on E2E: In order to study the convergence rates of different FL/SL methods, we also perform an experiment on fine-tuning the pretrained GPT2- medium (Radford et al., 2019) on the E2E dataset (Novikova al., 2017), which is a tabular natural language genera- tion problem. In Fig. 5, we plot the masked cross-entropy loss against communication load for test data. The meth- ods are run upto 4 TB of communication or 50 rounds, whichever is earlier. We observe that FSL-SAGE converges faster, and is more accurate than its auxiliary- based counter- part CSE-FSL, and almost as fast as SplitFed-SS. FedAvg converges the fastest since, for GPT2-medium, communicat- ing the model once per round consumes much lesser bytes than more frequently transmitting the smashed data.
Cybenko, G. V. Approximation by superpositions of sigmoidal function. Mathematics of Control. Signals and Systems, 2:303-314, 1989. URL http:: //ap:::::///..// manticscholaz: wtt.:/oo......o/ooouuI.c.D: 39583699,
Devlin, Chang, M.-W,, Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. In Proceedings of the 2019 confer- ence of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pp. 4171-4186, 2019.
Ghadimi, S. and Lan, G. Stochastic first-and zeroth-orde methods for nonconvex stochastic programming. SIAM journal on optimization, 23(4):2341-236 2013.
6. Conclusion
Grunwald, P., Steinke, T., and Zakynthinou, L. Pac-bayes, mac-bayes and conditional mutual information: Fast rate bounds that handle general ve classes. In Conference Learning Theory, pp. 2217-2247. PMLR, 2021.
In this paper, we proposed new federated split learning algorithm, FSL-SAGE, which facilitates the training of large models using FL, while enjoying the benefits of data paral- lelism. Our method leverages parallel training of client-side models while incorporating server feedback via auxiliary models. FSL-SAGE has a finite-time convergence rate of ((LVVT) for I communication rounds, which matches Fe- dAvg. We conducted extensive experiments with large-size computer vision and natural language models to verify the efficacy and the significant amount of communication cost
Gupta, O. and Raskar, R. Distributed learning of deep neural network over multiple agents. Journal of Network and Computer Applications, 116:1 -8, 2018.
Han, D.-J., Bhatti, H. I., Lee, J., and Moon, J. Accelerating federated learning with split learning on locally generated losses. JCMI. 202/ workshop on federated learning

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
for user privacy and data confidentiality, ICML Board, 2021.
Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S. J., Stich, S. U., and Suresh, A. T. SCAFFOLD: stochastic con- trolled averaging for on-device federated learning, CoRR, abs/ fbss/9910.03378 2019. URL httt:://aaxi///
Han, P., Huang, C., Tian, G., Tang, M., and Liu, X. Con- vergence analysis of split federated learning hetero- gencous data. In The Thirty-eighan Annual Conference on Neural Information Processing Systems, 2024. URL net forum? idduddR
abs / 91 0 06378.
Kingma, D. P. and Ba, J. Adam: A method for stochas- tic optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference on Learning Representa- tions, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: //arxiv. org/ abtp////////// 1412. 6980.
BkdBfE.
Hamin, B. Universal function approximation by deep neural nets with bounded width and relu activations. Mathemat- ics, 7(10):992, 2019.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers features from tiny images. 2009.
Hanin, B. and Sellke, M. Approximating continuous functions by relu nets of minimal width, 2018. URL httpps:/arre
Li, Y. and Lyu, X. Convergence analysis of sequential feder- ated learning on heterogeneous data. Advances in Neural Information Processing Systems, 36:56700-567555 2023.
org/ abs/1710. 11278.
He, C., Annavaram, M., and Avestimehr, S. Group knowl- edge transfer: Federated learning of large enns the edge. Advances in Neural Information Processing Systems, 33: 14068-14088 2020.
Lin, Z., Qu, G., Wei, W., Chen, X., and Leung, K. K. Adaptsfl: Adaptive split federated learning in resource-constrained edge networks. arXiv preprint arXiv:2403.33100 2024.
He, K., Zhang. X., Ren, S., and Sun, Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference computer vision and pultern recognition, pp. 770-778, 2016.
Lu, Z., Pu, H., Wang, F, Hu, Z., and Wang, L. The ex- pressive power of neural networks: A view from the width. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL httpss//pprceedinn...eerrippppp cc /paper pppper Fiiee//ppeer/0177/iee/2222181121 4674407067117611dd3e
Hinton, G., Vinyals, O., and Dean, J. Distilling the knowledge in a neural network. arXiv preprint arXiv: 1503.02331, 2015.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., et al. An empirical analy- sis of compute-optimal large language model training. Advances in neural information processing systems, 35: 30016-30030, 2022.
Marquez, E. S., Hare, J. S., and Niranjan, M. Deep cascade learning. IEEE transactions neural networks and learning systems, 29(11):
2018
McMahan, H. B., Moore, E., Ramage, D., and y Arcas, A. Federated learning of deep networks using model averaging. CoRR, (b2./660.00/69,, 2016. URL http:
/ / www.iv..rrg/bb....
Hornik, K., Stinchcombe, M., and White, H. Multilayer feedforward networks universal approximators. Neu- ral networks, 2(5):359-366 1989.
Mey, A. A note on high-probability versus in-expectation guarantees of generalization bounds in machine learning, 2022. URL 2576.
Hsu, T.-M. H., Qi, H., and Brown, M. Measuring the effects of non-identical data distribution for federated visual clas- sification. arXiv preprini
2019.
Mu, Y. and Shen, C. Communication and storage efficient federated split learning. ICC 2023 -EEE International Conference on Communications, pp. doi: 10. 10011//CA444122211102777
Hu, E. J.. Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106. 09660, 2021.
2023.
Novikova, J., Dušek, O., and Rieser, V. The E2E dataset: New challenges for end-to-end generation. In Jokinen, K., Stede, M., DeVault, D., and Louis, A. (eds.), Proceedings of the 18th Annual SIGdial Meeting on Discourse and Dialogue, pp. 201-206, Saarbrücken, Germany, August 2017. Association for Computational Linguistics. doi:
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
10

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
10.1 18653/11//1-555225 URL htttt:t//cccaaateeeeet wwgyyorg//ii/ 7
Sun, Y., Li, Z., Li, Y., and Ding, B. Improving loRA in Privac--rreeevingg federated learning. In The Twelfth International Conference on Learning Representations, 2024. URL https: / http//pppeeeeeeeeeeeeeeeeeeeeeee net/ for 11d-NIP226HWNN
Oh, S., Park, J., Vepakomma, P., Baek, S., Raskar, R., Ben- nis, M., and Kim, S.-I. Locfedmix-sl: Localize, federate, and mix for improved scalability, convergence, and la- tency in split learning. In Proceedings of the ACM Web Conference 2022, wWW '22 pp. 3347-3357, New York, NY, USA, 2022. Association for Computing Machinery. ISBN 9781450390965. doi: 10.1 10.114453385544.3552255 URL http::://....ooeoeee/o 111.://///..11111144444414 3512153.
Thapa, C., Arachchige, P. C. M., Camtepe, S., and Sun, L. Splitfed: When federated learning meets split learn- ing. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 8485-8493, 2022.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B.. Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lam- ple, G. Llama: Open and efficient foundation language models, 2023. URL https: / /aaaiviiivv////c///aaaa/aaaa org/ abs 23 02.13971.
OpenAI, Achiam, J., et al. Gpt-4 technical report, 2024.
https:///xxix org/ abs/ 2302 200.//833.... 08774.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Vepakomma, P., Gupta, O., Swedish, T., and Raskar, R. Split learning for health: Distributed deep learn- ing without sharing raw patient data. arXiv preprint 2018.
Radford, A., Kim, J. W., Hallacy. C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision, 2021. URL wtw.thtp://// //rxiiv / abs/2103 00020.
Wang. Z., Shen, Z.,, He, Y., Sun, G., Wang, H., Lyu, L., and Li, A. FLoRA: Federated fine-tuning large language models with heterogeneous low-rank adaptations. In The eight Annual Conference on Neural Information Processing Systems, 2024. URL https:/ / openrevi ew . wwe.ee//ffee forumm id=TcCorXxNJ2.
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. Hierarchical text-conditional image generation with clip latents, 2022. URL https::/p//......... / /arxivvv.. arg s s/2204. 06125.
Reddi, S. J., Charles, Z., Zaheer, M., Garrett, Z., Rush, K., Konečný, J., Kumar, S., and McMahan, H. B. Adaptive federated optimization. In International Conference Learning Representationss 2021. URL https: //o:://///pnnnn// revie ew net. wor.moormmmm/mm 21t.221311dd.f333333333 fet.
Wu, w., Li, M., Qu, K., Zhou, C., Shen, X., Zhuang, W., Li, X., and Shi, W. Split learning over wireless networks: Parallel design and resource management. IEEE Journal on Selected Areas in Communications, 41(4):1051-1066, 2023.
Sahu, A. K., Li, T., Sanjabi, M., Zaheer, M., Talwalkar, A., and Smith, V. On the convergence of federated optimiza- tion in heterogeneous networks. CoRR, abs/1812.0611 2018. http:////arxx wwwpxiii.rrrgaaoo www.org///aaoo. /bss///////// 111./.111.0000 061 27.
Yang, H., Fang, M., and Liu, J. Achieving linear speedup with partial worker participation in non-lID federated learning. In International Conference on Learning Rep- resentations, 2021. URL https: oop.revie....nn et / ft//oorrmm/mm iddihmuhu-ddddd
Sani, L., Iacob, A., Cao, Z., Lee, R., Marino, B., Gao, Y., Cai, D., Li, 7.,, Zhao, W., Qiu, X., et al. Photon: Feder- ated I1m pre-training. arXiv preprint arXiv:2411.02900 2024a.
Yu, H., Yang. S., and Zhu, S. Parallel restarted sgd with faster convergence and less communication: Demystify- ing why model averaging works for deep learning. In Proceedings of the AAAI conference on artificial intelli- gence, volume 33, pp. 5693-5700, 2019.
Sani, L., Iacob, A., Cao, Z., Marino, B., Gao, Y., Paulik, T., Zhao, W., Shen, W. F., Aleksandrov, P., Qiu, X., and Lane, N. D. The future of large language model pre-training is federated. In International Workshop on Federated Foundation Models in Conjunction with NeurIPS 2024, 2024b. URL https: /ttttp:://ppereeeeeee wwwpeerrevew.net//Eeee
Zhang, T. From E-entropy to kl-entropy: Analysis of min- imum information complexity density estimation. The Annals of Statistics, pp. 2180-2210. 2006.
fid=hfeH5APp 9NY.
Steinke, T. and Zakynthinou, L. Reasoning about general- ization via conditional mutual information. In Conference on Learning Theory, pp. 3437-3452. PMLR, 2020.

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
A. Notation
Jo go,i(mes (ey Boi (c)
2 (o 5z Ui)
J. Ji(
Z ( E.)
Fi 2,,, Ui)
Ios Fos Basi : (i)
Zi.b (M 2i,f Us)
Da,i
( Zs We)
Figure 6. Outline and notation of gradient flows in FSL: Blue arrows denote gradients w.r.t model parameters: green arrows represent gradients w.r.t. the cut-layer: brown arrow indicates the Jacobian the cut-layer w.r.t the client-side parameters.
We use the following additional notation in our convergence analysis. We denote the gradients of the loss F E y:) associated with client i, w.r.t. the server-side model Is.; 00:0/00 and the client-side model as:
Ic.i = OF. 0z OF. dz.,f
Ji,i-Zihh-
(A.1)
dæ dz.
where we use J, to denote the transpose of the Jacobian matrix 0z.,, of the cut-layer activations žt, with respect to the client-side model weights Fe.i for the 7thh client. The returned gradient 2i,b is function of the model æ, and the client's mini-batch (E,, y:), and can be written as <x;; E., 1/).
We denote the gradients of the loss (a;;;; E,, v/) associated with client as 9; (g), where Pe, and g, denote the gradients with respect to the client-side and server-side models respectively.
Figure outlines the gradients used in our analysis. The following equations define the notation used there.
OF, 9 (Z.i v.) = (x Ri, 21)
(A.2a)
dæ.
3F;
21,0 (2,( 2k) Pz.,f
(æsi Zi,f vb)
(A.2b)
2z'sf (wei Ei)
(.( = Dx.
(A.2c)
90.6 (Mes 2:i E.. = c Ei) 2200( (1))
(A.2d)
aF;
(1, 2t,f āz,,
(Da,s) 2i,f yk)
(A.2e)
(x (Bes Da.si E,, y)) = J, J, (m.i E.)
(Ba Zi,f, Vi)
(A.2f)
Also, we denote the gradient w.r.t. the client-side model with V. and the w.r.t. the server-side model V..
Table 2 summarizes all the notation used in this paper.
B. Proofs
B.I. Proof Theorem 4.3
We first state the following useful lemma that bounds the 4tth iterate averaged across clients, which we call the client drift.
The proof is provided in Appendix B.3.
12

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
Table 2. Notations for FSI-SAGE
Quantity
Meaning
m
Number of clients
T
Total number of communication rounds
K
Number of client local updates
1
Frequency of server-side or auxiliary model updates
p
IT
Fraction of rounds for which server-side model is updated
Client-side, server-side and auxiliary model parameters
30
Concatenation of client-side and auxiliary model parameters
Concatenation of client-side and server-side model parameters
9c,a 9,
Gradients of f()) w.r.. client-, server-side and concatenated model parameters Learning rates for client- and server-side models respectively cost function values at initialization and end of T rounds Gradients w.r.t. client- and server-side models
ML. 7)
Jo, f
Ve, V.
Les Le
Lipschitz constants of V P; and F.;
Lif
Lipschitz constant of 2i,f w.r.t. client-side model
Bound on variance local stochastic gradients
oG
Bound on variance estimating global cost function gradient
Lemma B.1. (Client drift bound) For any step-size satisfying !rr. S ViuKL.' be bounded as:
the client drift for any & E O...., K can
E
x
z=l1
S 20K
(62 + Kog) + K V
1=1
K
V - -
+
(B.1)
For the following, we denote the expectation conditioned on all randomness up to the eth step as Eg [.. First, from Assumption (a), and the fact that x = (Ie, Fy), we get
E, [[(x*t]]] S f(x²) (Ve[(x'), E. a
x:()) ) +
Le
Ee
|| - ar:l?
2
&C c
L
(Vof(x'). E, [ x²)) +
E,
1|
x
(B.2)
2
Hereafter, the proof is divided into two parts: the client-side bound (C) and the server-side bound (S).
1. Client-side bound: At the th iteration of the tt round, the net update for client i takes the form:
l.k 1.k- c, x
x.
ME
ges
(B.3)
Together with the model averaging at server F, the client-side model update in one communication round is given by
K-1
K
1+1
9 0.
x
M
9 Oai
(B.4)
i=1
1=1
Ñ k=
km
13

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
We first bound the quantity C in (B.2):
C = vof(x'),
7
Ee
K
1
c,i
+
LE
Be
1
K k=
c,
(B.5)
k=0
A
where we get (B.5) from (B.4). For the sequel, we define ã; == (Ie, La,i) as the set of parameters at client i. Then, we can bound the expression A in (B.5) as follows.
K-1
Km kmvef(at))
A
U Ee Vef(x'),
KmVaf(a') +
k=0
K
KnL of (z1)
Vof(2'') M
Es
(ofitttittt - veFi(e'))
2m
k=U
Knn [ff(z')|1
(
-I
- MT
( -vc()))))()))
+ E,
(B.6)
mvF K
Using the property (a, b) 1/2 ||al12 1/2 |1611 1 l|a
b112. we can rewrite as
A
KTL
(x+()))
n Er
Σ
2
2m K
IL E,
MKV.f(o')
(. - V-F((2)))
2m
K-1
KnL
[=(("")
TL E 2m K
ΣΣR -.
2
K=
K-1
Ee
V (a
2m K
km0
KnL
Vef(=*)
K-1 E,
V. -
2
2m
k=0
AA
m K-1
TJL
E,
ΣΣ> V. (.t
(B.7)
2m
Next, we bound A₁ in (B.7) as follows:
Fl - (
A E,
V.F v.f.oa V.Pi(Et-)) -
E,
< 2E,
V. (...
v.f( 9 1²| +2E,
(B.8)
|| 2: =;'17 + V. ((AE))-v((()) ?
< 2.
(B.9)
14

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
where, n := [t/i] in (B.8) we use the property |E a:ll 2 Ei=1 |a,1l with 2, in (B.9) we use Assump- tion 4.b. Next, bounding the quantity B in (B.5):
B = ni Le 2
K-1
Es
ci.k ge,s
k=
ni Li Et
K-
T Le E.
m K
9c.s
Sh.k
V Pi(
v
(B.10)
2
2m
k=
k=0
n Ki.Koj, Le
n Le 2m2
K-1
(B.11)
2m
R=0
where (B.10) follows from E (B.9) in (B.7), and then (B.7) and (B.11) (B.5) we get the following.
||==||]]]
=E E
[x]]
+
| [=]] and (B.11) from Assumption 4.2(a). Substituting
C = A+ B
S KnL
KKKLL Le 2m
V [ef(z')
2
a
E
22 x [1
voFi(())
- 19]
k=0
K-1
71 2m K
(1 - L-)LL,KEE,
V
(B.12)
R=
KıL
ni Ko
K-1
172
xe 111)
(+)(2'))
7I,
12
E
2m
771
k=0
K
+
Σ V.
(B.13)
where we use the fact that the last term in (B.12) is strictly negative for nr Lemma B.1 and rearranging, we get:
kT. to get (B.13). In (B.13), using
C c 2
KiLL V (x.)
i KoJ
m K
Σ
(.F) - V.Pl
2m
m
K m
20 1, 12
(07. + Ko
+ K Vef
V.Fi V.FF(=>))'
m
=KKn 1/ - 17 Vef(E>)| Ko Le
+ 7 + Koo)
2m
KNL (1+ + 20 x 12) Σ v.)((a))
v.Pi(o') |*
m
ni Ko
< -cKnn [al())))) ||2
+ 20K of 12 + Koo)
2m
3K n
V. v
(B.14)
2m
We get 14) because, 3 c 0 such that (0.5- - 20 > provided on 7jr. in the last term well. Rewriting (B.14) in terms of + (nr.) = ! m2o8 Le
nr 2v10KL. and by reusing the bound
+ 20K [2 + Koğ) and the
2m
15

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
estimation enror metric şf in (5), we get:
3 e' (n))) c cKTL (-I[va[("))|++
(B.15)
2c
2. Server-side bound: The clients transmit the smashed data to the S-server once every K/Q local iterations, thus the server-side model is updated Q times per round per client. The update equation for the server-side model is then:
(B.16)
where the second superscript on g indicates the 4thh server update, and is different from the local client iteration index k. Then, the server-side bound S over one communication round in (B.2) can be written as follows:
S (Vaf(x'), [ot-1 -e'{>> +
Le
Ee
$1
= EE (v
Σ
m Q-1
Le
17
m
Q-1
93.5
Ex
go.
(B.17)
2
q=0
0=0
Equation (B.17) is identical to (B.5) with the substitutions n1/m 17, the same steps as the client-side bound we can say that under Assumptions 4. (a) and 4.2(b), and provided S 4.2QL.' S can be bounded as
Te.
ct.k
g V. V, and K Q. Following
S s cmm@n
(B.18)
with #2(()) noo Le 16Q Q02)
Finally, we can substitute (B.15) and (B.18) back in (B.2) to get the following:
E, [f(] s f(x") cKyL (-11 (-1|va((>>))++
3 ex (2(m)) cmQn (-I +
2c
which can be rearranged to get:
3K e' tcKn + (m)
cKnL [[]((=)] + cmmn "
S f(x") - E, [f(x*+1)]
2
Next, by redefining e min(c, c') and d' = max(c, C... After further rearranging we get:
+ 3kTL ex c'K (ILL) + mn)2)(n)
oQ min(nz, mn) | [[+)('''') || < f(x') - E. [f(o++1]] |/)(')|| f(x') - EE [f(x(+1)]
3Km
e¹ e'K (m.1(m)) m)) mm))2)))
cQ min m))
2cQ min min(mL mn)
cQ
min(ni,
Next, taking full expectation on both sides, and summing over t= 1,2 and recognizing that minimum is lesser than the average, we get the final bound:
f(x")-
3KTL 2cQ mn.. mn) T
$ 9) T
min E [|V
+
(B.19)
EE
emin {NL, QT
where (1I., 11) is given by:
(11I,1)) = eK
34
1L Le 2m
20 + Ko²2)
no? La 16Q
cQ
min
07 Le
20K'
12
12
02 Koo)
mofoiL
16Qmi) 3 L2 + Qoà) ((TL.))
< C.
2m
(B.20)
=Qmin(ni, mn)
Thus, substituting ((),, 11) instead of 56nL, 1)) (B.19) we get the statement of the theorem.
16

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
B.2. Proof of Theorem 4.8 and Corollary 4.10
The statement of the theorem follows by further bounding the estimator term + Σズ= et in the RHS in Theorem 4.3, There, we obtained an upper bound on the stationarity gap, i.c., min; | [=x))l which contained a term involving
E
- -
VoR(o)))
(B.21)
where the gradients V ((())) and (a:)) are written in terms of the cut-layer activations as:
v.ii(T)= = E, [J.(2(i))) Enp =l. v.f('')= = E, [J.(e( 2
(B.22)
(B.23)
where E, [:] refers expectation with respect to the randomness at iteration We can then bound et as follows:
E E v.R(() - V -f((""))
|| (z 4) (2.0 ,,, - 2,,((<t =1,, s()i [1]]
E
[|J.(a(: 4) (Bi wi) 2,,(00 .f,y vfi) ||]
(B.24)
|| v vi) 21,000 =L., w/)1
(B.25)
E || 1| zl., 200( , st)1l"]
(B.26)
m
13 m m
(z x*)
where we Jensen's inequality in (B.24) and the spectral norm of J.(-)) in (B.25). (B.26) use Assumption 4.7 and the fact that the spectral norm of the Jacobian is bounded by the Lipschitz constant.
Note that r;(c) O (1/2) implies that O (1/vF) = O (1/vt). Now, from (13), we have:
L m
C x: +
E1. S
C (a 0
(B.27)
Vi
m
1/7 st
for some C₁ 0. Note that ef, is the lowest achievable error rate by the hypothesis class A, time and reducible further. Finally, substituting (B.27) back into (6) and using 1/r Σt= 1/vi & 2/VT from Lemma B.2, we get the desired result in Theorem 4.8.
17

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
For the proof of Corollary 4.10, the only difference is that we need bound the round average of E as follows:
1 T-1
C1
B
L
T 1=0
T
VE.
T-1 Σ
C1 T-1 1 T
+
T
Vt
1-0
1=0
1
T
Σ
Vt
T t=
VT
1
+ { /T
(1 (1-T/T)
$
+
VT
VT'
Σ
(T+T/T) VT'
3 +
Ci
T
for some C1 > O. Note that when T" = T, alignment happens till the last round, and this recovers the bound in Theorem 4.8. On the other hand when T" 0, alignment occurs, and this makes the upper bound - co.
B.3. Proof of Lemma B.1
The proof is along similar lines as (Reddi et al., 2021) and proceeds as follows. We start with the iterate:
E₁
x.
= Eg
c,a
1,k-1
x -,19
kt.k-1
= Eg
x cj
1
x²¹- - 11 { 9cst
{ 1gt.--
VePi(
(T.t.- V 1) -vF(()) -
( - v.Fi(=t) + V.F((x') - vaf(z') + vaf(at))
+
1 - x 2;11 [1]]
+ Bn 1
< Eg
v.R(=t)] 17] + 5m v.F((z)) v.F((a')]'
+ 55m Es
VoFi(z') - Vef(='))| 2 + मा || (=')(=))l" 2 + 55
(B.28)
where in (B.28), we use the property that on the right hand side as follows:
1>; 2
< n Ee ||ao|| ||| Since K > 1, we can continue bounding the expression
Ee
A.k F2
1
1
E2
p2 1-k- 2!! ==|1]
+ +5n2 Ee
00
V VoP((Et 11
2K-
V (T - V.Pi(=3)| 1|
Vori(a)) VePP(ars)1
+ 5K E.
+ +5Kgi
+ 5Km
V.
- vof(xt))12 + 5K Vof(2'))
1
2 12 Es || - x: |]
1+ 2K 1
+ 5Km
+ 5 + Koo)
+ 5Kni.
v.F.(z)) - VoFi(a')
||
+ +5K Vef(z')) |
(B.29)
1
E, 12 k-k - 2 ||]
+ K --1 1
+ t + +Koz)
+ 5km2.
v.P((z))) - voFi(at))
+ +5K7 Vef(z'))* 2
(B.30)
18

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
where in (B.29), we Assumptions 4.2 and 4. (b);; and (B.30) holds when AI S and unrolling the recursion we can obtain
Next, averaging over m clients
VIūKL.
772
)²
s.k cF
5n + KKog) + 5K [a]()')
x
77A
K
p=u
5K v
m
Vo v.F((E))) /
+
(B.31)
Representing the term in square brackets as H² we can write:
Σ
t.k c,
1"]
S( -1
1+ K - - 1 H
(1+ K - 1 11¹
(K -1 1)
(B.32)
4
(B.33)
where, use the fact that & < K and that 1+1/(K -1) 2 1 in (B.32), and the fact that for K 1, (1,(1+1((K11) in (B.33). Finally, substituting the value of H' in (B.33), we obtain the inequality in the lemma.
<5
B.4. Additional Lemmas
= 1/v2 = 1/v -
1/v2
1/v3
1/P9
1/VPT1
2
P
p
Figure 7. Illustration of the inequality in (B.36)
Lemma B.2 (Series bound). Göven p N, the series:
1 P
1
5(p)
(B.34)
p
Vr
satisfies the bounds
2
2 vp
< S(p)
(B.35)
1p+I+1
In other words S(p) = (+)
Proof. The series in (B.34) be bounded using integrals as follows:
1 dx
5() <5))
1,
1 p
p+1
dx
(B.36)
p Vズ
Vz-1
The reason for this becomes clear from the illustration in Figure 7. The integrals in (B.36) simplifies to the desired lower and upper bounds.
19

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
2000
2000
2000
2000
1000
1000
1000
1000
5 Client
Client
Client
Client
2300
2000
2000
2030
1000
1000
1000
1000
Test Accuracy
5
0
5
Client
FedAvg
2000
2000
0.6
Fed
Fed-S
1000
1000
0.5
CSE-FSL
0e:000
0
FSI-SAGE
Client
Client 9
Communication Load (GB)
Figure Class label distribution Dirichlet distributed data with simulating heterogeneous data distribution across clients.
Figure Scatter plot of best accuracy vs, communica- tion load for ResNet- ResNet-18/CIFAN 1E/---//////RRRII 10 on 10 clients.
C. Notes on experimental setup
Dirichlet Sampling We use the Dirichlet distribution to simulate client data heterogeneity. Specifically, we sample the proportion of class labels in each client from a Dirichlet distribution, and then sample examples uniformly from each class per client while respecting the proportions sampled from the Dirichlet distribution. The parameter a controls the degree of heterogeneity with smaller values indicating a higher variation in the proportion of class labels distributed across the clients. In Fig. 8, we show instance of class label distribution for 1 for CIFAR-10.
Note on the cut-layer and auxiliary models: There are two important choices in the design of FSL-SAGE: the choice of the cut-layer and the choice of auxiliary models. 1) Cut-Layer: The location of the cut-layer determines the communication cost of transmitting the cut-layer features. the cut-layer features are too large compared to the size the auxiliary models, the communication gain of FSL-SAGE is not too large compared to CSE-FSL, since both methods expend resources in transmitting the cut-layer features. Although, note that FSL-SAGE does strictly better than CSE-FSL in terms of communication cost. 2) Auxiliary: For our experiments, we choose auxiliary models as small subsets of the server model for training. This arbitrary choice is able demonstrate competitive performance for us, but it is important to note that the size of the chosen auxiliary can impact the communication advantage of our method. We perform ablation experiments demonstrating the effect of auxiliary model size on final test performance in Section D.4.
D. Additional Experiments
D.1. Communition performance
In Fig. we plot the best accuracy achieved and the corresponding communication cost incurred by the five algorithms for various values of Dirichlet a. Each point represents different (algorithm, Dirichlet-a) combination. The best performing algorithms, which maximize accuracy while minimizing communication costs, would be found the north-west corner of the plot. The radius of the points in the plot, indicates the degree of heterogeneity, with larger points corresponding to lower and hence more non-i.i.d. data. We observe that FSL-SAGE is able to achieve high level of accuracy for all tested levels of heterogeneity, being the only method to occupy the north-west corner of the plot, even for certain high degrees of heterogencity.
D.2. Latency Analysis
In Fig. 10, we analyze the differences in the computation and communication latencies of the various FL methods used in our CIFAR- 1IFAAR-O/RRsNet--1 18 experiments. Note that the computation latency of FSL-SAGE is higher than CSE-FSL due to the auxiliary alignment process, which can is computationally expensive for the server. However, the communication
20

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
2h 46m 40s
Compatatton
Cocaunmication
16m 40s
Latency
1m 40s
10s
FedAvg
SplitFed-MS SplitFod-SS
CSE-FSL
FSL-SAGE
Figure 10. Computation and communication latency in various FL methods for 200 rounds. (Note that the y-axis in logarithmic scale.)
latency of FSL-SAGE is about half that of CSE-FSL, since CSE-FSL requires the server to send the auxiliary models back to the clients after aggregation, while FSL-SAGE only transmits the updated auxiliary models once from the S-server to the clients. The communication latency of the three baselines, FedAvg. SplitFed-MS and SplitFed-SS high because full model transmission for aggregation, and smashed data transmission every iteration are very communication intensive.
D.3. Effect of alignment interval (Ablation study)
In Figures 11 and we plot the final test accuracy of FSL-SAGE using ResNet-18 18 on CIFAR-10 10 and CIFAR- datasets respectively, for 200 rounds of training. We vary the alignment interval I from once in 2 rounds that, interestingly, there is clear pattern between and the final test accuracy for the CIFAR-10 dataset, likely because CIFAR-10 is not difficult dataset for ResNet-18 to learn. Hence, the server-side model is not very complex, and the auxiliary model learns to mimic it fairly quickly, thus requiring far fewer update frequency. On the other hand, for more difficult dataset like CIFAR-100, we see a clear pattern of deterioration of test accuracy with increasing The alignment interval is an important trade-off parameter in FSL-SAGE, as it trades-off increased computational load on the server with the final test performance. For our main results, we chose a value of 10, which is suitable compromise between the two.
in 10 rounds. Note
D.4. Effect of auxiliary model size (Ablation study)
In this section, we study the effect of the size of the auxiliary model on the performance of the trained model. use the CIFAR datasets, and the ResNet-18 model split into client-side model and server-side model as previously described in main experiments. The client-side model comprises the first two blocks of ResNet-18 18, and the last two blocks with the final fully-connected layer are used as the server-side model. We try out the following four variations in the auxiliary model:
1. Linear: with only the final fully-connected layer of ResNet-18. The size of the auxiliary model is only k
2. Half:
with 3rd and 4th ResNet-18 18, but with only half the original number of layers, i.e. 1 layer per block, and the final fully connected layer. this case, the size is Rs 3.5 MB.
3. Ours: with only the 3rd block and the fully connected layer. The size for this case is ss 8 MB.
4. Full: with the auxiliary model being the same architecture the server-side model. The size is 40 MB.
In Figures 13 and 14, we see the effect of using the above four auxiliary models on the CIFAR-10 and CIFAR-100 datasets.
21

---

FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation
0.90
Test Accuracy (200 rounds)
Test Accuracy (200) rounds)
0,540
0.535
0.86
0,530
0.84
0.525
0.82
0.520
0.515
10 Alignment interval (1)
10a Aligmment interval (1)
Figure 11. Effect of auxiliary model size on final test perfor-Figure 12. Effect of auxiliary model size of final test perfor- mance of CIFAR-10 after 200 rounds.
mance of CIFAR-100 after 200 rounds.
0.860
0.56
Tost Accuracy (@) 200 GB)
0.855
Test. Accuracy (0) 200 GB)
0.54
0.850
0.845
0.840
0.50
0.830
CSE-FSL
CSE-FSL
0.825
0.48
ESI-SAGE
FSL-SAGE
0.820
10-2
101
10
10"
10a
Auxiliary size (in MB)
Auxiliary size (in MB)
Figure Effect of ResNet-18 auxiliary model size on CIFAR- -iiuur Effect of ResNet- auxiliary model size on CIFAR- 10.
100.
For CIFAR-10, the increasing auxiliary model size shows a drop test performance. We attribute this to the fact that CIFAR-10 relatively simple classification problem, and perhaps small auxiliary model able to replicate the nuances in the function learned by the server-side model, while larger auxiliary model overfits the patterns on the training data. In able to learn simpler linear auxiliary model much quicker resulting in good fact, both CSE-FSL and FSL-SAGE, performance for the linear model in all cases. More importantly, note that the difference in performance of the same auxiliary model increases between the two algorithms as the model size increases. This shows that FSL-SAGE is able to learn more complex auxiliary model more efficiently than CSE-FSL. In Fig. 14, the performance of the FSL-SAGE improves with increasing auxiliary model sizes, in contrast to CSE-FSL.
22


Reviews:
[{"review_id": "i9ekD1BMkt", "reviewer": "Reviewer_bte9", "summary": "The paper proposed a new federated split learning algorithm called FSL-SAGE. It builds upon existing works on local split federated learning, where client updates are derived from local approximations using an auxiliary model attached to each client. A key challenge with previous approaches is that, due to the lack of feedback from the server, these local auxiliary modules can drift away from the optimal solution. FSL-SAGE addresses this issue by periodically aligning these auxiliary modules via an additional optimization loop conducted on the server side.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and suggestions. Due to space limitation, we could only respond to a subset of more critical comments in this rebuttal.  But we are happy to continue to complete our responses to your remaining comments in the discussion stage when new space opens up. Please refer to the abbreviations in **HNw2**'s rebuttal.\n\n> **Comment 1:** Marginal contribution; server feedback already mentioned in [1]. Periodic alignment is related to LSL-PGG[1] which should be cited.\n\nThank you for pointing us to Ref. [1].  We would like to clarify that our core contributions are distinct from Refs. [1-2] as follows:\n\n* Unlike LSL-PGG [1], which is a centralized algorithm with a goal to reduce training memory footprint, our work, is a **federated** split learning algorithm designed to train large models using FL on data distributed over commodity hardware.\n\n* In [1] the AMs are updated indirectly via a global loss update.  In contrast, we update the AMs directly to _mimic_ the SSM by minimizing the MSE loss.\n\n* To the best of our knowledge, our approach is the first AM based FSL approach to have a convergence guarantee on the joint model. Previous works [2] only provide separate convergence guarantees for the CSM and SSM, which do not imply convergence of the join model to a minimum.\n\nWe will cite [1] and clarify these differences in the next revision\n\n[1] Bhatti el al. \"Locally Supervised Learning with Periodic Global Guidance.\" arXiv preprint arXiv:2208.00821 (2022).\n[2] Mu el al. \"Communication and storage efficient federated split learning.\" ICC 2023-IEEE International Conference on Communications. IEEE, 2023.\n\n> **Comment 2:** Single local epoch setting, small number of clients $m$, and manual splitting impractical.\n\n1. **Single Epoch:** We chose the local epoch to be 1 in our experiments primarily for a fair comparison to previous literature [2], which also has the local epoch as 1. To clarify, our algorithm statement and theoretial analysis makes **no such assumption** regarding the number of local iterations or local epochs. \n\n2. **Manual Splitting:** Our framework does **not** depend on the splitting technique used, and we do **not** used the term 'manual splitting' anywhere in our manuscript. Developing an algorithm that can optimally split a model would indeed be very interesting, but this fundamental topic deserves a seperate paper that is beyond the scope of this manuscript. We will pursue this topic in our future studies and we thank the reviewer for pointing out this direction.\n\n3. **Number of Clients:** Our theoretical analysis reveals that $m$ in the FL setup does not impact our convergence rate. We have only simulated 10 clients in our experiments due to resource limitations.  We note that our shared source code is configurable to accomodate for much higher $m$ given enough GPU resources.\n\n> **Comment 3:** 1. In-Expectation learnability assumption challenging in practice. 2. AMs are larger that the CSMs adds to client computation. What is the trade-off between AM size and performance? 3. How do we compare to baselines when smaller AMs are used.\n\n* Please see our response to **Comment 6** by **iZqi**.\n* In the next revision, we will also compare our method against baselines using smaller AMs as you suggested.\n\n> **Comment 4:** Assumes honest server; Include discussion of potential privacy concerns. Cite works analyzing privacy risks. Does sending labels to server compromize privacy?\n\n* Please see our response to **Comment 1** by **HNw2**. Thank you for pointing us to these references. We will cite them in our revision as you suggested.\n* Sharing labels is the same in all FSL methods. Hence, our method achieves the same privacy performance as in other FSL methods.\n\n> **Comment 5:** 1. Server computational load in large-scale FL; Storage and computation bottleneck due to alignment. 2. How does server deal with partial client participation and client drop-out?\n\nStoring the alignment dataset at the server and alignment can indeed be a bottleneck at the server. To mitigate the need for storage, a simple solution is to align the AMs on the most recent batch of cut-layer activations on an on-demand basis, thus also solving the problems associated with client drop-out and partial client participation. Also, one can perform alignment in another process on the server so it doesn't bottleneck the SSM optimization. While these are interesting engineering extensions to our algorithm, they do not affect our theoretical claims and results, which are our main contributions.\n\n> **Comment 6:** Key hypothesis is periodic alignment; analysis of auxiliary gradient error is essential. \n\nThanks for the comment.  We analyze the auxiliary gradient error given by Eq. (5) in the manuscript, in **Section 4.2**. We can bound this term as $\\mathcal{O}(1/\\sqrt{T})$ provided the in-expectation learnability assumption holds, which shows that the gradient error decreases with the same sub-linear rate.", "claims_and_evidence": "The core contribution of this work is rather marginal. The core element of this technique is the server feedback to the client’s auxiliary network which is carried out indirectly via periodic alignment of these modules on the server-side is straightforward and has been discussed before.", "methods_and_evaluation_criteria": "The authors use decent datasets and large models for evaluation, and the baselines are generally appropriate. However, the setting of a single local epoch is limiting and not widely practiced. Additionally, the manual nature of model splitting may hinder generalization to other real-world scenarios. Moreover, the number of clients used in the experiments is very low, whereas typical federated learning setups involve hundreds of clients with partial participation.", "theoretical_claims": "1. The proof sketch appears technically sound, and most assumptions are standard. However, the in-expectation PAC learnability assumption, though theoretically mild, may be challenging in practice since it requires the auxiliary model to be sufficiently expressive (large enough). \n\n2. In previous works, the auxiliary network was kept small because it was not intended to fully replace the server model, whereas in this paper, a larger auxiliary network is used, even larger than the client model itself. \n\n3. Additionally, the paper assumes an honest server, yet in practice, the server could be prone to various attacks. A discussion on potential attacks, particularly regarding privacy concerns, is missing.", "experimental_designs_or_analyses": "1. The paper does not analyze the computational load on the server side, which is crucial for large-scale federated learning scenarios where thousands of devices participate simultaneously. Storing alignment datasets for every client and running the secondary optimization loop might create a bottleneck on the server. \n\n2. Also, the added computation due to increased auxiliary network size as compared to small aux networks in existing techniques should be discussed. \n\n3. The paper does not provide guidance for practitioners on how to select the optimal hyperparameters introduced e.g. the alignment period, neither does the paper provide an ablation study. \n\n4. Given that the key hypothesis is that periodic alignment is what makes the difference, an analysis of auxiliary gradient error over time is essential to confirm it. \n\n5. Wall-clock time comparison or an estimate would be great to understand the potential latency introduced by the alignment overhead.", "relation_to_broader_scientific_literature": "The paper is well-positioned within the federated and split learning literature. However, it would benefit from discussing privacy analyses and through evaluation of the trade-offs introduced.", "essential_references_not_discussed": "Since FSL-SAGE involves sending smashed data (activations) and labels to the server, the paper should reference works that analyze privacy risks and provide mitigation approaches.\n\nThe periodic alignment of these auxiliary modules in the centralized setting is discussed in a related paper LSL-PGG[1] which should be cited. \n\n[1] Bhatti, Hasnain Irshad, and Jaekyun Moon. \"Locally Supervised Learning with Periodic Global Guidance.\" arXiv preprint arXiv:2208.00821 (2022).", "other_strengths_and_weaknesses": "1.\tA potential bottleneck for the system is when client partially participates. In the low participation ratio, the server still has to store alignment datasets for every client. Also it is unclear how outdated data is managed when clients drop out.\n2.\tThe choice of the split cut-layer determines the size of the smashed activation, which can be substantial. In this method, the server must store these activations for each client for alignment, potentially imposing a significant computational and memory burden. A more detailed analysis of this aspect is needed.\n3.\tThe size of the auxiliary network is quite large compared to existing baselines. For ResNet-18, the auxiliary network (2.1M parameters) is nearly three times the size of the client’s original network, and it is even larger (92.4M parameters) for the language task. I would like to see how does it compare to the baselines when the aux networks are used as suggested by those works.", "other_comments_or_suggestions": "The paper has some typos, e.g. “limtations” and “sever-side”. The text “For LoRA finetuning” can be misleadin, consider clarifying the context.. A brief section mentioning limitations would be helpful. Also, readers might wonder: does sending labels to the server compromise privacy? A line addressing this could be added for reader to understand the potential risk associated with FSL techniques.", "questions_for_authors": "1.\tDoes the client-side auxiliary module also get updated during the local update step?\n2.\tWhat happens if the auxiliary network is smaller? Is there a minimum size required for the auxiliary network to effectively serve as a surrogate for the server model?\n3.\tCan you provide more insight or data on choosing the alignment interval? The theory gives a trade-off, but practically how does it affect the performance. \n4.\tHow does the server scale when the number of clients increases significantly (e.g., hundreds or thousands)?\n5.\tHow do you justify the trade-off between increased auxiliary network size and overall performance, especially considering the potential burden on client devices?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vX3KhihogB", "reviewer": "Reviewer_HNw2", "summary": "The paper addresses the computational burden faced by clients when performing local updates on whole (possibly large) models by splitting the model into server-side and client-side components, following the approach used in prior split learning (SL)-based methods. Unlike previous works, the authors propose an approach to estimate gradient feedback from the server using an auxiliary model on the client side, which significantly reduces latency and communication overhead.\nSpecifically, the true cut-layer gradient is periodically received from the server, and the auxiliary model is trained using an MSE loss to approximate the gradient. The paper provides theoretical convergence analysis and demonstrates the broad applicability of their method by conducting the experiment on language model as well.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.  For brevity we will use the abbreviations FSL for federated split learning, FL (SL) for Federated (resp. Split) Learning, CSM (SSM) for client-side (resp. server-side) model, and AM for auxiliary model.\n\n> **Comment 1:** Sending alignment dataset to server poses potential privacy risk.\n\n**Response:** We agree that there are privacy concerns about uploading cut-layer data to the server, but one can argue that this risk would be equally present in any SL algorithm which sends smashed data to the server. The study of attacks and privacy guarantees for our algorithm is indeed important, which deserves a separate paper and is beyond the scope of this manuscript. Also, we note that the FSL setting in this paper is the same as in other SL algorithms published in the literature, hence having the same privacy performance.\n\n> **Comment 2:** The AM size is larger than the CSM; imposes compute burden on clients.\n\n**Response:**  We agree that the AMs used in our experiments are large.  Please see our response to **Comment 6** of reviewer **iZqi**. We will conduct a more thorough ablation study to demonstrate the effect of AMs size in our revision.\n\n> **Comment 3:** Optimization of $x_a$ requires double derivative, causing compute burden on server. \n\n**Response:** It is true that the alignment of AMs requires a double derivative, and this could pose a computation burden, especially if we use large AMs. Our ablation experiments on the AMs, as discussed above, will help clarify if this is indeed a significant burden.\nIn our experiments with GPT2m, we used an AM of size 92.4M and didn't face very high compute burdens.\n\n> **Comment 4:** In theoretical results, does FSL-SAGE enjoy linear speedup?\n\n**Response:** Thank you for raising this interesting point, it will help clarify our manuscript. Unfortunately all FSL methods using a single server-model, including our method, which handles one client at a time, lacks linear speedup [1]. Primarily, this is because the server-side model is trained sequentially on cut-layer activations received from the clients. We will add a note about this aspect to the revision. \n\n[1] Han, Pengchao, et al. \"Convergence analysis of split federated learning on heterogeneous data.\" arXiv preprint arXiv:2402.15166 (2024).", "claims_and_evidence": "The claims in this study are supported by experimental evaluations and detailed analyses.", "methods_and_evaluation_criteria": "The proposed method is conceptually well-grounded, and the evaluation is appropriate.", "theoretical_claims": "The convergence bound is rigorously analyzed, including both FSL-SAGE and its lazy variant.", "experimental_designs_or_analyses": "The experiments and evaluations are well designed and valid.", "relation_to_broader_scientific_literature": "The paper builds on prior SL based methods for addressing key challenges in FL, e.g., computational efficiency, communication overhead, and latency. The paper introduces an auxiliary model for gradient feedback estimation, significantly reducing latency and communication costs. The authors also provide rigorous convergence analysis for both FSL-SAGE and its lazy variant, ensuring theoretical foundation of the paper.", "essential_references_not_discussed": "I have not noticed any key references that were missing from the paper.", "other_strengths_and_weaknesses": "### Strengths\n\n- The proposed method effectively addresses the challenges of existing SL-based approaches by introducing an auxiliary model to approximate gradient feedback.\n- The gradient feedback estimation mechanism significantly reduces communication overhead and latency while ensuring high performance.\n- A convergence analysis is provided, covering both FSL-SAGE and its lazy version.\n- Experimental results on language model are presented, demonstrating the broad applicability of their method.\n\n### Weaknesses\n- Sending the alignment dataset to the server poses a potential risk of privacy leakage.\n- The auxiliary model is not particularly small. it appears to be even larger than the client-side split model, which may impose additional computational burden on the client compared to traditional SL-based methods.\n- From the server-side perspective, the optimization of $x_a$ is required for 'all' the clients. As suggested by the equation (2), if I understand correctly, this optimization might require double derivative with respect to $x_a$, which could introduce a significant computational burden on the server.", "other_comments_or_suggestions": "I appreciate the careful consideration of the alignment mechanism in the convergence analysis, and the attempt to analyze the lazy version as well. These aspects add significant value to the paper.\n\nRegarding the convergence rate, while the rate with respect to $T$ matches existing results, one of the key focuses is how the rate is expressed in terms of the number of clients. In fact, a linear speedup with respect to the number of clients is often achieved in recent FL works. However, in the convergence presented in this paper, it is unclear how this aspect is handled.\n\nAdditionally, while uploading the alignment dataset itself poses a potential privacy leakage risk, this issue is not discussed in the paper. Including a discussion on this aspect would strengthen the paper.", "questions_for_authors": "Please refer to the weaknesses and comments parts.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "K1TNA0k06B", "reviewer": "Reviewer_iZqi", "summary": "The paper studies split federated learning with a focus on reducing training latency/communication. It uses an auxiliary model to facilitate the computation of cut-layer gradients. To mitigate potential accuracy drop, the paper aligns the auxiliary model with the server-side model periodically. A solid convergence analysis is provided for the proposed algorithm, and some experiments are given to showcase the effectiveness.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.  Please refer to the abbreviations in reviewer **HNw2**'s rebuttal.\n\n> **Comment 1:** In Theorem 4.3, physical meaning of $c$?\n\n**Response:** Thank you, we missed defining the constant $c$ in **Theorem 4.3**, but it is defined later in **Theorem 4.8**. We will add the definition in **Theorem 4.3** in the next revision. $c < 0.5 - 20K^2 \\eta_L^2 \\widehat{L}_c^2$ is a positive constant defined to simplify the final expression of convergence.\n\n> **Comment 2:** Theorem 4.3 is agnostic to $l$; does convergence depend on $l$?\n\n**Response:** **Theorem 4.3** is indeed agnostic to $l$, the alignment interval, not because the convergence is unaffected by $l$ or that the bound is loose, but because $l$ directly affects the _auxiliary estimation error_, $\\varepsilon^t$, given in **Theorem 4.3**, left-hand-side of Eq. (5). Our final result in **Theorem 4.8** Eq. (14) is obtained by further bounding $\\varepsilon^t$. There, we have a bound on every $l^{th}$ round, which implies that a larger $l$ would need a proportionally larger number of rounds $T$ (i.e., slower) to converge and vice-versa.\n\n> **Comment 3:** [1] though most relevant, is not included in baselines.\n\n**Response:** We did not include [1] in our baselines for the following reasons:\n* [1] uses multiple copies of the SSM, one for each client, and also uploads cut-layer activations to the server at every iteration, making it very memory and communication inefficient.\n* In CSE-FSL [2], which is a baseline for our work, the authors demonstrated that they performed much better than Han et. al [1] in terms of communication efficiency. Thus, it suffices to compare with [2].\n* Lastly, our method is not a generalization of [1]. The following are two important differences between our method and [1]:\n    1. [1] updates the AMs via local loss functions at each local iteration, while we do not.  This means that even as $l\\to\\infty$, our method will behave very differently.\n    2. Han et. al [1] used multiple SSMs, one for each client, while we use only 1 SSM.\n\n[1] Han, Dong-Jun et al. “Accelerating Federated Learning with Split Learning on Locally Generated Losses.” (2021).\n[2] Mu, Yujia, and Cong Shen. \"Communication and storage efficient federated split learning.\" ICC 2023-IEEE International Conference on Communications. IEEE, 2023.\n\n> **Comment 4:** Arbitrary choice of AMs has non-trivial impact on performance. Cut strategy should also be investigated.\n\n**Response:** The choice of an appropriate AM architecture and cut-layer are interesting problems in their own right, but somewhat out of scope of this work. Like all other SL methods [1-2], we focus on a given split and our work works with **general models**. Similarly, we use a given AM and compare to CSE-FSL [2]. As we mention in **Comment 6** below, we will include an ablation study on AM sizes in the next revision.\n\n> **Comment 5:** Training latency as a metric.\n\n**Response:** We will estimate the latency assuming a fixed communication and computation bandwidth between the clients and server. We will include this simulated wall-clock result to the revised manuscript.\n\n> **Comment 6:** Empirically justify $\\epsilon$ in learnability assumption.\n\n**Response:** This is also a response to **bte9**'s **Comment 3**.  It is intractable to check if a given AM is in-expectation PAC learnable or not.  However, in our experiments, an AM which is chosen by taking a small portion (of size $\\leq 0.1\\times$) of the respective SSM demonstrates comparable, if not better convergence to the tested baselines. We would also like to note that **Assumptions 4.6** and **4.7** are sufficient but not necessary conditions for convergence.  In our revision, we will add an ablation study of the AM size on performance. There, we will also test smaller AMs than the ones presented in the paper.\n\n> **Comment 7:** Include references for convergence of SplitFed.\n\n**Response:** We thank the reviewer for pointing out the references on convergence analyses of sequential FL and SplitFed. We will add these to the next revision.\n\n> **Comment 8:** Some typos; inclusion of source code.\n\n**Response:** We thank the reviewer for detecting these oversights.  In the next revision we will:\n* Proofread and fix typos.\n* Add numbers to **Fig. 1** to indicate the order of operations.\n\nWe will also share our source code after the review process (we have already shared our code as supplementary material in a `.zip` file for the review process).\n\n> **Comment 9:** Technical challenges faced in proofs and implementation.\n\n**Response:** While [1] originally introduced the idea of using auxiliary models to train clients in parallel, our work involves aligning the auxiliary model directly to _mimic_ the server-side model. This brings about several technical challenges in convergence analysis, all of which are detailed under the heading **2) Technical Challenges, Lines 70-90** in our manuscript.", "claims_and_evidence": "The claims are provided clearly and convincingly with theoretical and/or numerical evidence.", "methods_and_evaluation_criteria": "The choice of models and datasets is suitable and is widely adopted in other related work on (split) federated learning.", "theoretical_claims": "I have quickly checked the proofs of the theorems, and no obvious errors have been spotted.\n\nHowever, there are some issues. \n- In Theorem 4.3 Eq. (3), it seems that $c$ is a typo (originally defined for clients). Can the authors specify its physical meaning?\n- Theorem 4.3 is agnostic to $l$, i.e., how frequently the auxiliary models are aligned with the server-side model. Does this mean the convergence is not affected by $l$, or the bound is loose?", "experimental_designs_or_analyses": "The experimental results well presented and analyzed. While the reviewer appreciates the theoretical contributions, more experiments (in particular ablation studies) are needed to demonstrate the effectiveness of the approach. \n\n- Han et. al 21 seems to be the most relevant work to this submission. However, it has not been included in the benchmark. If I understand correctly, FSL-SAGE is a generalized version of that in Han et. al 21, where one could use $l=\\infty$. At least an ablation study of different choices of $l$ should be included. \n\n- In line 369, the authors ``arbitrarily'' chose the structure of the auxiliary model. This seems to be inappropriate, as one can imagine this choice will have a non-trivial impact on the model performance. Also, the model cut strategy should be investigated empirically as well. \n\n- It is suggested to include training latency (e.g., in wall clock time) as a metric. \n\n- Sec. 4.2-4.3 included more assumptions to establish the convergence results, and those assumptions should be empirically justified. For example, can the authors quantify the value of $\\epsilon$ for the proposed algorithm?", "relation_to_broader_scientific_literature": "- A convergence analysis is provided for split federated learning with client-side update using auxiliary models.", "essential_references_not_discussed": "Since the main contribution is a convergence analysis for split federated learning, the authors are suggested to include the following references on split (federated) learning convergence, and discuss the differences/improvements. \n\n- Li, Yipeng, and Xinchen Lyu. \"Convergence analysis of sequential federated learning on heterogeneous data.\" NeurIPS, 2023.\n- Han, Pengchao, et al. \"Convergence analysis of split federated learning on heterogeneous data.\" NeurIPS, 2024.", "other_strengths_and_weaknesses": "- The paper implemented LoRA fine-tuning of LLMs in a split federated learning environment, which is very relevant and could benefit future research efforts.", "other_comments_or_suggestions": "The current version of Fig. 1 is a bit confusing, and the authors may consider adding steps/ordering. \n\nThe paper could benefit from proofreading and correcting a list of typos. Here are a few examples.\n\n- In line 71, ``due'' -> ``due to''\n- In line 396, CIFAR-100 should be corrected to be CIFAR-10. \n\nIt would be better to include codes for reproducibility purposes.", "questions_for_authors": "Can the authors specify what technical challenges model alignment brings to the proofs and/or implementations compared to those in Han et. al 21?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6C5tDaHz9v", "reviewer": "Reviewer_vMHm", "summary": "The paper introduces FSL-SAGE, a framework designed to address the limitations of Federated Learning (FL) and Split Learning (SL). FL struggles with training large models due to client-side memory constraints, while SL incurs high communication latency due to sequential processing. FSL-SAGE combines the data parallelism of FL with the model-splitting efficiency of SL by employing client-side auxiliary models to estimate server-side gradients. These auxiliary models are periodically aligned with the server to minimize communication overhead while preserving accuracy.  \n\nAlgorithm: Clients use auxiliary models to estimate server gradients locally, enabling parallel training and reducing reliance on frequent server communication. A \"lazy\" variant further reduces communication by freezing auxiliary models after alignment.  \nConvergence Guarantees: Theoretically, FSL-SAGE achieves an O(1/√T) convergence rate, matching FedAvg, despite reduced communication costs.\nEmpirical Results: Experiments on ResNet-18 (CIFAR-10/100) and GPT2-medium (E2E dataset) demonstrate superior accuracy and communication efficiency compared to FedAvg, SplitFed, and CSE-FSL. FSL-SAGE reduces communication costs by up to 10× while maintaining robustness to data heterogeneity.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.  Please refer to the abbreviations in reviewer **HNw2**'s rebuttal.\n\n> **Comment 1:** Experimental validation relies on small-scale models. Weakens claim that FSL-SAGE applies to large models.\n\n**Response:** We appreciate the reviewer's point about using large-scale LLMs for our experimental evaluation. It is true that smaller models like ResNet-18 or GPT-2m are not large enough to necessitate splitting. However, our goal is to demonstrate the efficacy of our method in relation to other state-of-the-art methods for a given ML model. We make no assumptions on the model architecture or size, so we would expect our theoretical results to hold even for large models. We are very interested to test larger LLMs in our framework, the manual process required for implementing a split architecture and the AM is a challenge for us. However, we didn't have sufficient resources to implement them for larger models in this rebuttal. If this paper is accepted, we will share our source code and welcome collaborations to train larger LLMs.\n\n> **Comment 2:** AM architecture is fixed without ablation study. How does AM size affect performance?\n\n**Response:**  Please see our responses to **Comments 4** and **6** by **iZqi**.  We will add an ablation study of the AM size to the revision.\n\n> **Comment 3:** The alignment interval $l = 10$ chosen without ablation study.  Sensitivity analysis missing.\n\n**Response:** We thank both reviewers **bte9** and **vMHm** for raising this point. We will include an ablation study analyzing the impact of $l$ in the next revision.  However, we would like to note that our **Theorem 4.8** has shown the trade-off effect of the alignment interval $l$. In fact, the bound in Eq. (14) applies to every $l^{th}$ communication round. Thus, if $l$ is very large, the algorithm would take longer to convergence, which makes sense intuitively, since the AMs would get aligned infrequently thus mostly misguiding the client-model.  On the other hand, if $l$ is too small the first few iterations are spent in overtraining the AMs to mimic a near randomly initialized SSM. For our experiments, we manually tuned $l$.\n\n> **Comment 4:** Newer hybrid methods like AdaptSFL, CPSL, not included in baselines. Important references missing.\n\n**Response:** From an optimization perspective, AdaptSFL and CPSL are the same as SplitFed and vanilla SL, respectively.  These hybrid methods either change the structure of model splitting or the sequence of parallel or sequential operations in SL, but these do not impact the algorithmic convergence of the optimization of the model.  Also, the communication efficiency of AdaptSFL and CPSL would be a scaled version of vanilla SL or SplitFed, since these methods do not cut down on the communication cost due to client-server messages at every local iteration, which are both included in our baselines. For these reasons, these two works do not represent new baselines. However, we thank the reviewer for pointing out these and additional references and, we will cite them in the revision.\n\n> **Comment 5:** 1. Experiments use outdated models/datasets. 2. Baselines are limited. 3. Insufficient ablation studies.\n\n**Response:** Please see our responses to **Comment 1, 2** and **3**. We plan to strengthen our experiments by adding ablation studies for $l$ and AM size.\n\n> **Comment 6:** Second-best results in experimental figures.\n\n**Response:** Thank you, we will indicate the second-best results in our figures.\n\n> **Comment 7:** How does system/data heterogeneity affect convergence speed?\n\n**Response:** In our theoretical analysis, we allow for system and data heterogeneity via the quantities $\\sigma_l$, i.e., the s.t.d. of the mini-batch of data from its true value, and $\\sigma_g$, the s.t.d. of the client's loss function from the loss overall loss function. The expressions in **Theorem 4.3, Eqs. (3)** and **4.8, Eq. (14)** indicate the effect of these quantities on the convergence rate. Larger values of these variance terms decrease the rate of convergence, as is the case with other FL methods.  We appreciate the reviewer for raising this point, and we will add these discussions in our revision.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The paper claims that SFL is increasingly relevant for LLM pretraining and fine-tuning in the introduction, yet the experimental validation primarily relies on small-scale models like ResNet-18. This creates a mismatch between the stated motivation and empirical evidence. For instance, ResNet-18 can be easily deployed on edge devices in its entirety, rendering model splitting unnecessary and undermining the practicality of the proposed method for such scenarios. The lack of experiments with prevalent open-source LLMs (e.g., LLaMA 3, or Qwen) weakens the claim about SFL-SAGE’s applicability to modern foundation models. The authors should either substantiate their claims with experiments on larger, cutting-edge LLMs or refine their motivation to align with the evaluated scenarios.", "theoretical_claims": "I have carefully examined the correctness of the theoretical proofs and found no apparent errors. The authors have provided sufficient justifications and detailed explanations, ensuring the validity of their claims.", "experimental_designs_or_analyses": "The experimental designs and analyses exhibit both strengths and limitations in validity:  \nThe use of Dirichlet-sampled non-IID data** rigorously tests robustness to heterogeneity, and communication cost metrics align with real-world federated constraints.\nLimitations:  \n1. Experiments on ResNet-18 (edge-deployable without splitting) and GPT-2-medium (outdated for modern LLM scales) fail to validate claims about \"large-model\" efficacy.\n2. The auxiliary model architecture (a subset of the server model) is fixed without ablation studies. This leaves uncertainty about how auxiliary design (e.g., depth, width) affects performance, especially for larger models.  \n3. The alignment interval l = 10 is arbitrary; no sensitivity analysis explores how varying l impacts accuracy or communication trade-offs.  \n4. While FedAvg and SplitFed are included, newer hybrid FL-SL methods (e.g., AdaptSFL[1], CPSL[2]) are absent, raising questions about comparative novelty. \n\n[1] Lin, Zheng, et al. \"Adaptsfl: Adaptive split federated learning in resource-constrained edge networks.\" arXiv preprint arXiv:2403.13101 (2024).\n[2] Wu, Wen, et al. \"Split learning over wireless networks: Parallel design and resource management.\" IEEE Journal on Selected Areas in Communications 41.4 (2023): 1051-1066.", "relation_to_broader_scientific_literature": "The key contributions of FSL-SAGE are situated within the evolving landscape of federated and split learning, addressing critical gaps identified in prior work. FL, epitomized by FedAvg (McMahan et al., 2016), prioritized data parallelism and privacy but assumed clients could train full models—an impracticality for modern large-scale architectures. Split learning (SL) (Vepakomma et al., 2018) relaxed client memory constraints via model splitting but introduced sequential bottlenecks, as seen in SplitFed (Thapa et al., 2022). Recent efforts like CSE-FSL (Mu & Shen, 2023) reduced communication via local losses but lacked server feedback, risking accuracy degradation. FSL-SAGE bridges these paradigms by introducing auxiliary models to estimate server gradients locally, enabling parallelism while preserving server guidance. The theoretical O(1/√T) convergence rate aligns with FedAvg’s guarantees, demonstrating that split training need not sacrifice convergence speed despite added complexity.", "essential_references_not_discussed": "Many recent works have explored issues related to client-server communication and convergence in SFL; however, these aspects are not addressed in the manuscript.\n\n[1] Lin, Zheng, et al. \"Adaptsfl: Adaptive split federated learning in resource-constrained edge networks.\" arXiv preprint arXiv:2403.13101 (2024).\n[2] Wu, Wen, et al. \"Split learning over wireless networks: Parallel design and resource management.\" IEEE Journal on Selected Areas in Communications 41.4 (2023): 1051-1066.\n[3] Oh, Seungeun, et al. \"Locfedmix-sl: Localize, federate, and mix for improved scalability, convergence, and latency in split learning.\" Proceedings of the ACM Web Conference 2022. 2022.\n[4] Lin, Zheng, et al. \"Hierarchical split federated learning: Convergence analysis and system optimization.\" arXiv preprint arXiv:2412.07197 (2024).", "other_strengths_and_weaknesses": "Strengths:\n1. The use of an auxiliary model to approximate activation gradients instead of directly computing the loss locally is an innovative approach.\n2. The experiments on communication overhead are comprehensive and provide valuable insights.\n\nWeaknesses:\n1. The manuscript lacks sufficient experiments on truly large models that necessitate split training. Additionally, the chosen models and datasets are somewhat outdated.\n2. The baseline selection is relatively limited, lacking comparisons with closely related studies on SFL.\n3. The experiments are insufficient, particularly in terms of exploring the impact of hyperparameters and conducting ablation studies. These aspects should be addressed to strengthen the overall analysis.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. It would be beneficial to indicate the second-best results in the experimental figures, as this could provide additional insights for the result analysis. However, this aspect is not discussed in the manuscript.\n2. The manuscript would benefit from a discussion on how system data heterogeneity impacts the final convergence speed and convergence stability. This aspect is crucial for understanding the robustness of the proposed method under varying conditions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Reviewer bte9 gave a strong reject score and it seems that the rebuttal did not change his mind. However, other reviewers tend to accept this paper. So, I think I need help fron the SAC.
Decision: Accept (poster)