FSL-SAGE: Accelerating Federated Split Learning via Smashed Activation Gradient Estimation

We thank the reviewer for their comments. Please refer to the abbreviations in reviewer HNw2's rebuttal.

Comment 1: Experimental validation relies on small-scale models. Weakens claim that FSL-SAGE applies to large models.

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.

Comment 2: AM architecture is fixed without ablation study. How does AM size affect performance?

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.

Comment 3: The alignment interval $l = 10$ chosen without ablation study. Sensitivity analysis missing.

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$.

Comment 4: Newer hybrid methods like AdaptSFL, CPSL, not included in baselines. Important references missing.

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.

Comment 5: 1. Experiments use outdated models/datasets. 2. Baselines are limited. 3. Insufficient ablation studies.

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.

Comment 6: Second-best results in experimental figures.

Response: Thank you, we will indicate the second-best results in our figures.

Comment 7: How does system/data heterogeneity affect convergence speed?

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.

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.

Comment 1: Marginal contribution; server feedback already mentioned in [1]. Periodic alignment is related to LSL-PGG[1] which should be cited.

Thank you for pointing us to Ref. [1]. We would like to clarify that our core contributions are distinct from Refs. [1-2] as follows:

• 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.

• 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.

• 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.

We will cite [1] and clarify these differences in the next revision

[1] Bhatti el al. "Locally Supervised Learning with Periodic Global Guidance." arXiv preprint arXiv:2208.00821 (2022). [2] Mu el al. "Communication and storage efficient federated split learning." ICC 2023-IEEE International Conference on Communications. IEEE, 2023.

Comment 2: Single local epoch setting, small number of clients $m$, and manual splitting impractical.

1. 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.

2. 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.

3. 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.

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.

• Please see our response to Comment 6 by iZqi.
• In the next revision, we will also compare our method against baselines using smaller AMs as you suggested.

Comment 4: Assumes honest server; Include discussion of potential privacy concerns. Cite works analyzing privacy risks. Does sending labels to server compromize privacy?

• 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.
• Sharing labels is the same in all FSL methods. Hence, our method achieves the same privacy performance as in other FSL methods.

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?

Storing 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.

Comment 6: Key hypothesis is periodic alignment; analysis of auxiliary gradient error is essential.

Thanks 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.

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.

Comment 1: Sending alignment dataset to server poses potential privacy risk.

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.

Comment 2: The AM size is larger than the CSM; imposes compute burden on clients.

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.

Comment 3: Optimization of $x_a$ requires double derivative, causing compute burden on server.

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. In our experiments with GPT2m, we used an AM of size 92.4M and didn't face very high compute burdens.

Comment 4: In theoretical results, does FSL-SAGE enjoy linear speedup?

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.

[1] Han, Pengchao, et al. "Convergence analysis of split federated learning on heterogeneous data." arXiv preprint arXiv:2402.15166 (2024).

We thank the reviewer for their comments. Please refer to the abbreviations in reviewer HNw2's rebuttal.

Comment 1: In Theorem 4.3, physical meaning of $c$?

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.

Comment 2: Theorem 4.3 is agnostic to $l$; does convergence depend on $l$?

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.

Comment 3: [1] though most relevant, is not included in baselines.

Response: We did not include [1] in our baselines for the following reasons:

• [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.
• 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].
• Lastly, our method is not a generalization of [1]. The following are two important differences between our method and [1]:
  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.
  2. Han et. al [1] used multiple SSMs, one for each client, while we use only 1 SSM.

[1] Han, Dong-Jun et al. “Accelerating Federated Learning with Split Learning on Locally Generated Losses.” (2021). [2] Mu, Yujia, and Cong Shen. "Communication and storage efficient federated split learning." ICC 2023-IEEE International Conference on Communications. IEEE, 2023.

Comment 4: Arbitrary choice of AMs has non-trivial impact on performance. Cut strategy should also be investigated.

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.

Comment 5: Training latency as a metric.

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.

Comment 6: Empirically justify $\epsilon$ in learnability assumption.

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.

Comment 7: Include references for convergence of SplitFed.

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.

Comment 8: Some typos; inclusion of source code.

Response: We thank the reviewer for detecting these oversights. In the next revision we will:

• Proofread and fix typos.
• Add numbers to Fig. 1 to indicate the order of operations.

We 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).

Comment 9: Technical challenges faced in proofs and implementation.

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.