Thinking Forward: Memory-Efficient Federated Finetuning of Language Models
Authors: Kunjal Panchal
Nisarg Parikh
Sunav Choudhary
Lijun Zhang
Yuriy Brun
Hui Guan
Venue: NeurIPS 2024 poster (2024)
Keywords: Federated Learning
Large Language Models
Forward-mode Automatic Differentiation
Forward-mode AD
Memory-efficient Finetuning
Memory-efficiency
Data Heterogeneity
Primary area: optimization_for_deep_networks
TL;DR: Spry is a federated learning algorithm that enables finetuning LLMs using Forward-mode Auto Differentiation; to achieve low memory footprint, high accuracy, and fast convergence.

Abstract: Finetuning large language models (LLMs) in federated learning (FL) settings has become increasingly important as it allows resource-constrained devices to finetune a model using private data. However, finetuning LLMs using backpropagation requires excessive memory (especially from intermediate activations) for resource-constrained devices. While Forward-mode Auto-Differentiation (AD) can significantly reduce memory footprint from activations, we observe that directly applying it to LLM finetuning results in slow convergence and poor accuracy. In this paper, we introduce Spry, an FL algorithm that splits trainable weights of an LLM among participating clients, such that each client computes gradients using forward-mode AD that are closer estimations of the true gradients. Spry achieves a low memory footprint, high accuracy, and fast convergence. We formally prove that the global gradients in Spry are unbiased estimators of true global gradients for homogeneous data distributions across clients, while heterogeneity increases bias of the estimates. We also derive Spry's convergence rate, showing that the gradients decrease inversely proportional to the number of FL rounds, indicating the convergence up to the limits of heterogeneity. Empirically, Spry reduces the memory footprint during training by 1.4-7.1$\times$ in contrast to backpropagation, while reaching comparable accuracy, across a wide range of language tasks, models, and FL settings. 
Spry reduces the convergence time by 1.2-20.3$\times$ and achieves 5.2-13.5\% higher accuracy against state-of-the-art zero-order methods. When finetuning Llama2-7B with LoRA, compared to the peak memory consumption of 33.9GB of backpropagation, Spry only consumes 6.2GB of peak memory. For OPT13B, the reduction is from 76.5GB to 10.8GB. Spry makes feasible previously impossible FL deployments on commodity mobile and edge devices. Our source code is available for replication at https://github.com/Astuary/Spry.

Thinking Forward: Memory-Efficient Federated Finetuning of Language Models
Kunjal Panchal University of Massachusetts Amherst, MA 01003-9264 kpanchal @umass. edu
Nisarg Parikh University of Massachusetts Amherst, MA 01003-9264 nkpari inhamasss edu
Sunav Choudhary Adobe Research Bangalore, India 560103
schoudha@adobe. com
Lijun Zhang
Yuriy Brun University of Massachusetts Amherst, MA 01003-9264 brunêcs umass. edu
Hui Guan
University Massachusetts Amherst, MA 01003-9264 1i l unzhanggcs. umass edu
University Massachusetts Amherst, MA 01003-9264 hui muiguan@cs. umass edu
Abstract
Finetuning large language models (LLMs) in federated learning (FL) settings has become increasingly important it allows resource-constrained devices finetune a model using private data. However, finetuning LLMs using backpropa- gation requires excessive memory (especially from intermediate activations) for resource-constrained devices. While Forward-mode Auto-Differentiation (AD) can significantly reduce memory footprint from activations, we observe that directly applying it to LLM finetuning results in slow convergence and poor accuracy. In this paper, we introduce SPRY, an FL algorithm that splits trainable weights of an LLM among participating clients, such that each client computes gradients using Forward-moot AD that are closer estimations of the true gradients. SPRY achieves low memory footprint, high accuracy, and fast convergence. We formally prove that the global gradients in SPRY are unbiased estimators of true global gradients for homogeneous data distributions across clients, while heterogeneity increases bias the estimates. We also derive SPRY'S convergence rate, showing that the gradients decrease inversely proportional to the number of FL rounds, indicating the convergence up to the limits of heterogeneity. Empirically, SPRY reduces the memory footprint during training by 1.4-7.1x in contrast to backpropagation, while reaching comparable accuracy, across a wide range language tasks, models, and FI. settings. SPRY reduces the convergence time by x and achieves 5.2-13.5% higher accuracy against state-of-the-arr zero-order methods. When finetuning Llama2-7B with LoRA, compared to the peak memory consumption of 33.9GB of backpropagation, SPRY only consumes 6.2GB of peak memory. For OPT13B, the reduction is from 76.5GB 10.8GB. SPRY makes feasible previously impossible FL deployments on commodity mobile and edge devices. Our source code is available for replication https: // //gi thub. com/ Astn/ay/turrrry
Introduction
In cross-device federated learning (FL), thousands of edge devices (called clients) collaborate through an orchestrator (called server) to jointly train a machine learning (ML) model 2 In each round of FL, the server sends an ML model to participating clients, who then update the model weights for several epochs on their individual data and send the new weights back to the server. The server aggregates the weights update the model and initiates the next round of FL training. Due to the
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

inherent privacy preseeving nature of FL, it has been adopted in many privacy sensitvve domains, such as healthcare 33. IoT
[ and e-commerce [
In parallel, large language models (LLMS) have demonstrated impressive performance on natural language processing tasks 8. creating a surge of interest in finetuning L.Mss in FL settings p9 10] However, the problem is challenging in practice because of the memory requirements in finetuning LLMs. LLMs can have billions of weights, and finetuning them with backpropagation requires dozens of GBs of memory. These requirements easily overwhelm edge devices, particularly mobile phones with limited memory. The memory footprint of finetuning LLMs mainly comes from model weights and their gradients, optimizer states, and intermediate activations.
There are three categories of existing algorithms that reduce the memory footprint of finetuning LLMs: parameter-efficien finetuning (PEFT) 16) quantization [17], and zero-order gradient estimation methods [0] Although PEFT and quantization can reduce the memory consumption from parameters and optimizer states, the memory consumed by intermediate activations remains significant bottleneck because these methods still use backpropagation finetune LLMs. Back- propagation requires storing all intermediate activations during the forward pass to estimate gradients in the backward pass. For example, finetuning 4-bit quantized Llama2-7B model [2] with LoRA techniques 12] requires -33. of RAM, with 83.8% used for intermediate activations. Zero-order methods leverage finite difference [22 to estimate gradients and thus reduce the memory consumption from intermediate activations [19] 20 233. However, these methods suffer from slow convergence and poor model quality because the accumulation of truncation and round-off errors (24] a fundamental issue of finite difference, leads to noisy estimation of weight gradients.
Forward-mode into-Differentiation (AD) 24 2 has the potential to address the memory consump-
tion problems of backpropagation without introducing the round-off errors of finite difference. It estimates gradients by computing a Jacobian-vector product (jvp) based on random perturbations of weights during the forward pass, alleviating the need to store all intermediate activations, similar to zero-order methods. jvp represents how much changing the weights in the direction of random perturbation affects the outputs. However, simply replacing backpropagation with Forward AD in FL settings does not produce convergence speed and accuracy performance comparable to established federated optimizers based on backpropagation, such as FEDAVG [.. FEDYOGI [25], FEDSGD [. Forward gradients are computationally inefficient and inaccurate estimating true gradients when the number of trainable weights is large. We empirically observed that, for LLMS finetuned with the LoRA technique, Forward-mode AD suffers from 8.3-57.2% accuracy loss and is 1.4-3.9% slower to converge for models whose number of trainable weights exceed approximately 1.15M (see Appendix G)).
In this paper, we propose SPRY an FL algorithm for finetuning LLMs using Forward-mode AD while achieving low memory footprint, high accuracy, and fast convergence. SPRY tackles the shortcomings of Forward- per FL round. SPRY improves computation efficiency as each client only needs to perturb a small fraction of trainable weights to derive their gradients, reducing the number of computations in each forward pass. SPRY achieves higher accuracy and faster convergence because the smaller number of trainable weights for each participating client allows computing gradients that are closer estimations of the true gradients. In contrast to zero-order methods where one training iteration requires 20-100 forward passes, each with different perturbation 20 19 to estimate weight gradients well, SPRY allows computing weight gradients from only one forward pass per iteration for each client. Unlike split learning (26), SPRY does not need to transfer intermediate activations among clients. The union of the partial weights trained from each participating client in an FL round updates all the trainable weights of the language model. Since only a subset of weights is finetuned per client, SPRY also saves client-to-serve communication bandwidth.
AD by splitting the trainable weights among participating clients
We formally prove that the global gradients aggregated on the server side in SPRY are unbiased estimators of the true global gradients in case of homogeneous data distributions across clients, while the heterogeneity increases the bias of the estimations. We also derive the convergence rate of SPRY, showing that the norm of global gradients decreases linearly with the inverse the number of FL rounds. We further discuss how configurations in SPRY affect the convergence behavior of the algorithm and empirically validate the theoretical analysis.
Named after its light memory consumption and speed.

---

We empirically evaluate SPRY'S memory efficiency, accuracy, computation efficiency, and communi- cation efficiency through experiments on wide range of language tasks, models, and FL settings. SPRY achieves within 0.6-6.2% of the accuracy of the best-perfore rest-performing FL. backpropagation, with 1.4-7.1x less memory consumption for each client and comparable time convergence. SPRY also outperforms zero-order-basee baselines with 2..5%%5%% higher accuracy. an average of 1.5-28.6x faster per-round computation time, and 1.2-20.3x faster convergence. We also compare SPRY'S communication efficiency to that of FEDAVG (per-epoch communication) and FEDSGD (per-iteration communication). For communication frequency of per-epoch, SPRY reduces the number of model weights sent from a client to the server by M times, where M is the number of participating clients per round. For per-iteration communication frequency, each client of SPRY only needs to send back a scalar to the server, fixing the client-to- to-server total communication cost to M scalar values.
We make the following contributions:
SPRY, the first work that demonstrate the potential of Forward-mode AD for finetuning language models (with 18M to 13B parameters) in FL settings with low memory footprint, high accuracy, and fast convergence.
2. A federated optimization strategy that only requires single forward pass per batch on each client to finetune a language model.
3. A theoretical analysis of how SPRY'S global gradients estimate true gradients based on the heterogeneity of FL clients, and a proof that SPRY'S convergence is linearly dependent on the number of FL rounds when client's learning rate is inversely proportional the size of perturbations and client data heterogeneity.
4. An empirical evaluation shows that SPRY consumes 1.4-7.1x less memory than its backpropagation-basee counterparts, and converges 1 faster with 5.2-13.5% higher accuracy compared to zero-order counterparts.
2 Forward-mode Automatic Differentiation
This section presents the background on Forward-mode Auto- Auto-Differentiation (AD) necessary to follow the work. Related works on zero-order optimization methods, Forward-mode AD, and FL for LLMs are discussed in detail in Appendix A
Forward-mode AD computes gradients by measuring how changes in model weights, in the direction of random perturbation, affect the loss. Since these gradients are derived from forward pass, they referred to forward gradients (23] In contrast, backpropagation (also Reverse-mode AD) calculates a direction to adjust weights in, to decrease the loss. Formally, for each training iteration, given the trainable weights w, Forward-mode AD generates random perturbation whose size is the same as w. one forward pass, given training data D, Forward-mon -wodddddddddeeeooe AD computes the value of the objective function fww; D) and the Jacobian- product (jvp) as follows:
J V fo (w; D) a 8w1
af(::))
wa]""
(I)
Ows
This jvp is scalar for neural networks since the output of the objective function f is a scalar. Multiplying jvp with the perturbation v gives us the unbiased estimate of true gradients (23
F(w)) E. 1
[[[ (1=:))
(f(::: Dwa
Dw:
ED
Df(w:D) Oura
ED
/ [V)w;;))
(3)
The partial derivative @5(:: D) is computed by chain rule on intermediate activations in the forward pass (24). Unlike backpropagation where all the intermediate activations need to be stored during the forward pass, Forward-mode AD only stores the previous layer's activations in the forward pass for the chain rule derivatives. Hence, the memory overhead deriving gradients would be the size of the largest activation in the forward pass.
3 SPRY: Memory-Efficient Federated Finetuning of Language Models
While Forward-mode AD can decrease memory usage during model finetuning, merely substitut- ing backpropagation with Forward-mode AD in FL scenarios often results in poor accuracy and

---

3a. Compute JVP
Send frozen and assgged jesiiiits and
forvard AD
jwP
1. Spli layers across participant clients
Weight perturbation
HIHH
Send updated assigned weights iNN0eee
Client #2
3b. Update the weights
Language Model with Bockbone Trairable PEFT Weights
optim atop
Updated weight
Forward gradient
Legend
Vector
Client #4
Aggregate layer to update the model
Scalar
Matrix
Figure I: Overview of SPRY, federated learning framework to finetune language models with low memory footprint. The term "PEFT" stands for parameter- efficieent fine-tuning.
computational inefficiency. To address this challenge, SPRY recognizes that Forward-mode AD operates more efficiently and yields better gradient estimations when the trainable weight count minimized. Therefore, SPRY optimizes performance by distributing trainable weights among participating clients, assigning each client a responsibility to compute gradients for only subset of trainable weights. SPRY is compatible with PEFT methods such as IA3 [TS], ADAPTER based meth- ods [27]. BITFIT (28). and LORA [12), which mitigate the memory consumption from gradients and optimizer states. this work, we focus on LORA due its demonstrated superiority, as highlighted in Appendix G
Figure [ gives overview of SPRY It includes 5 main steps: (1) At the beginning of each FL round, the server assigns a few trainable layers to each of the participating clients of that round. (2) Each client is sent (i) the trainable weights of the layers assigned to it, (ii) frozen weights of the remaining layers if not previously received, and (iii) a scalar seed value. (3) On the client side, weight perturbations for the allocated layers are generated based on the received seed. These perturbations are utilized to update the assigned weights through computation of the forward gradients. only transmit the updated weights back the server. (5) The server aggregates all the trained layer weights and updates the language model for the subsequent round.
(4) Clients
Next, discuss the two key steps of SPRY in detail: Step (1), where the server assigns trainable weights to the participating clients and Step (3), where each client finetunes the assigned trainable weights using Forward-mode AD.
3.1 Assigning Trainable Layers to Clients at the Server-side
To enable closer gradient estimations through forward gradients, the server reduces the number of trainable weights per client by selecting layer and assigning to client in cyclic manner. With LORA, the server selects a LoRA layer, which consists of pair of weights (w and Wp matrices) for each client. When trainable layers # participating clients, each client will be assigned more than one layer. Otherwise, each layer will be assigned to more than one client. The server aggregates the trained weights from each client and updates the model using adaptive optimizers such as FEDY YOG. Adaptive optimizers are shown to be less prone to noisy updates compared to FEDAVG in the literature [25 29]. The server keeps a mapping of layer names client IDs, hence can gather updated layer weights from all clients and update the model.
FL often faces the data heterogeneity issue, where the data distribution of one client differs from another, leading to poor model accuracy. While the primary aim of SPRY does not directly tackle this issue, it seamlessly integrates with existing finetuning-base mitigate it. SPRY distributes trainable classifier layers to all participating clients, enabling each client to finetune these layers to personalize the jointly trained model.
personalization techniques [30] to
3.2 Finetuning Weights with Forward Gradients the Client-side
Clients update the assigned trainable weights with gradients estimated through Forward-mode AD. Specifically, each participating client will get copy of the trainable weights assigned to it and

---

scalar seed value from the server. Using the seed value, the client generates a random perturbation for each trainable weight, following a normal distribution with mean of 0 and a standard deviation of 1. Forward gradients are obtained during forward pass, as shown in Eq. 3 The subsequent steps depend on the communication frequency.
Per-Epoch Communication. Per-epoch communication means that each client transmits the updated trainable weights to the server after every one or more epochs. Locally, the trainable weights are updated using optimizers such as SGD and ADAM. Only the updated trainable weights are sent back to the server, which reduces communication costs. Each client transmits the weights of ] layers. This means that if there are more trainable layers than participating clients, each client sends back weights for multiple layers.
max #Trainable Layers
#participating clients
Per-Iteration Communication. Unlike FEDSGD, where gradients are sent back to the server and aggregated after each iteration, SPRY offers additional communication cost savings. this communi- cation mode, SPRY only requires sending the jvp scalar value back to the server after each iteration of fine-tuning. Since the server has the seed value, it can generate the same random perturbation used by each client. Using the received jvp values and the generated random perturbations, the server can then compute the gradients and update the model weights.
A detailed breakdown of communication and computation costs is given in Appendix F
4 Theoretical Analysis
This section theoretically analyzes convergence behaviors of SPRY. SPRY has unique aggregation rule for the trainable weights as each client trains a subset of weights. It is also the first work to utilize forward gradients train LLMs in FL settings where clients could have heterogeneous data distributions. Therefore, we theoretically analyze the effects of (a) data heterogeneity on gradient estimations of SPRY, and (b) configurations in SPRY, including the number of FL rounds, dimension of perturbations, the number of perturbations per iteration, data heterogeneity, and the number of participating clients, on the convergence of SPRY. The proofs are detailed in Appendix !|
Theorem 4.1 (Estimation of the Global Gradient). In SPRY, global forward gradients f of the trainable weights w E Rdd with the corresponding weight perturbations v E IRd, computed by M participating clients is estimated in terms of true global gradients f as,
EmeM, Ee am E V (uT $) "I.. (1. Ho.)] (1, Ee um B WED V Sun(u (w
[V)
1 M
1 +1. 3/) " +1. Wo))
(4)
where the expectation is under the randomness of sampled data and random perturbation C is total number of classes and Omm =
R0m. For a class ne is its sample count, 0g is its Dirichlet concentration parameter. For a client m, Omm is the sample count of the tt class and
The [2
[Dal
Dm is the size of the data of client m. The global data is D =EmEM? Dm M the set of clients training an arbitrury subset of weights, M MMa,, E [M/d].
Discussion. We focus on analyzing how data heterogeneity affects the estimation error between the global forward gradient and the global true gradient. Specifically, the estimation error of SPRY depends on the coefficient Olm,c clients share the same data distribution, the Dirichlet concentration parameter az. is for any class As the ratio of ne/ (total samples in class total samples globally) matches the ratio of Tmm,e/ (Dmll (total samples in a class to total samples for a client m), the bias term becomes 0 since am, 0, Vm, Hence, the global forward gradients are unbiased estimators of the true global gradients. In data heterogeneous settings, when data across clients becomes more heterogeneous, 0e and thus Orn,c -> ne/ increasing the estimation error. Hence the global forward gradients are biased estimators that depend on the distances of the data distributions across participating clients.
In data homogeneous settings, where all
ml
Theorem 4.2 (Convergence Analysis). Under the assumptions on L-smoothness (Asmp L.T). bounded global variance of between global true gradients and aggregated expected gradients of each client
5

---

(AsmpI.2 and bound on gradient magnitude G (Asmpl 1.3 learning rate ne.
and the following conditions on the local
= min (0
(
7
VBnGGL
VBaG
//12(1 82)
MK
(5)
B2G(3d - - mm Ece(cr onne
The global true gradients of SPRY satisfies the following bound,
5 f(
BBIl((w
min
E-
V
OSTSR
RR
2
MeL
vI BrGre ()dd + K
Σ mEMc
E a
(6)
2T
3
MK
where R is the total number of FI. rounds, w E Rdd are trainable weights, v. E TRd are random perturbation, K count of random perturbations per iteration, n is global learning rate, T is adaptability constant, and is client sampling rate. Rest of the symbols are defined in Theorem 4.1
Discussion. We focus on analyzing how different configurations in SPRY affect its convergence. (a) The number of FL rounds (R): The upper bounds of the norm of the global forward gradient in Eq. [e]eccas in proportion to the inverse of R, indicating the convergence of the algorithm up (b) Dimension of perturbations (d): ae x d shows that as the number of weights be perturbed increases, the learning rate must decrease. A lower learning rate can make convergence slower, or worse, as our empirical experiments in Appendix G will show, not converge at all. (c) The number of perturbations per iteration (K): We observe both in nominator and denominator, which indicates that increasing K brings little advantage in convergence speed. Results Appendix G confirm the above statement. (d) Data heterogeneity: "e x a shows that more homogeneous data distributions across clients allow higher learning rate, and hence faster error reduction. This observation corroborated by the comparison convergence speeds between homogeneous and heterogeneous clients in Appendix H (e) The number of clients training same subset of weights: ne M shows that more clients training the same subset of weights is beneficial for faster convergence, similar observation is shown in Appendix G
to the limits of data heterogeneity.
The theorem also sheds light the accuracy performance gap between Forward- -owaardddodddae.o AD and backpropagation in FL settings. The upper bounds on the global forward gradient norm includes a second term that increases with oz and variance oz but does not decrease with R. This results in gap between estimation errors of backpropagation- -aase methods like FEDAVG and SPRY.
5 Empirical Evaluation
We empirically evaluate SPRY on language tasks, 5 medium, and 3 large language models, un- der various FL settings. Our evaluation measures SPRY'S prediction performance, peak memory consumption, and time-to-convergence. We also ablate SPRY'S components to study the impact of communication frequency, the number of trainable parameters, the number of perturbations per iteration, the number participating clients, and the importance of splitting layers on performance.
Datasets, Tasks, and Models. Our evaluation uses 8 datasets: AG News [31 (4-class classifica- tion), SST2 (32 (2-class classification). Yelp (31] (2-class classification), Yahoo 31 (10-class classification), SNLI [3 (3-class classification). MNLI (34) (3-class classification). SQuADv2 (35]
(Closed-book question answering), and MultiRC [36] (2-class classification). We chose these datasets because they allow to generate heterogencous splits in FL settings using Dirichlet distri- bution B7]
The default dataset split is across 1,000 clients, except the smallest datasets SST2 and MultiRC, where there are 100 clients. SQuADv2 has 500 total clients. Each dataset has two versions: (i) Dirichlet = 1.0 (Homogeneous split), and (ii) Dirichlet a = 0.1 (Heterogeneous split).
Our evaluation uses the following language models: OPT13B 38. Llama2-7B (21] OPT6.7B 38
RoBERTa Large (355M)
39 BERT Large (336M)
40
BERT Base (110M) 40, Distil-
BERT Base (67M) 41, and Albert Large V2 (17.9M)
42
For the billion-sized models, we
use 4-bit quantization. For all the models, we use LoRA as the PEFT method. Appendix describes
the datasets and hyperparameters in more detail.

---

Table Generalized accuracy for SPRY and its backpropagation- and zero-order order-based counterparts on RoBERTa Large and LLMs. SQuADv2 uses Fl score. t shows that higher values are better.
The datasets are split with Dir a 0.1. o Llama2-7B. * OPT..7B. D OPT13B.
SPRY outperforms the best-pent best-performing 7ero-order-basel methods by 5.15-13.50% and approaches the performance of backpropagation baaed methods, with a difference of 0.60-6.16%
Backpropagation-kan based
Zero-order- based
First-order
Difference between
Methods f
Methods 1
Forward Mode AD
performances of SPRY and
best performinoo backpropagation method f
best-performing zero-order method
FEDAVG FEDYOG!
FWDLLM+ FEDMEZO BATTLE+
SPRY
AG News
93.07%
92.77%
76.94%
70.56%
57.69%
87.89%
-5.18%
10.95%
SST2
88.00%
92.14%
84.41%
72.17%
61.57%
91.54%
-0.60%
7.13%
SNLI
86.45%
79.31
74.3%%
69.57%
62.10%
82.66%
-3.79%
8.36%
MNLI
84.29%
84.98%
72.66%
66.66%
62.85%
80.32%
-4.66%
7.66%
Yahoo
67.37%
63.08%
56.06%
44.69%
37.81%
61.21%
-6.16%
5.15%
Yelp
90.48%
79.10%
71.83%
65.10%
55,99%
85.33%
-5.15%
13.50%
MultiRC
47.56%
72.53%
64,58%
N/A
58.12%
68.65%
-3,88%
4.07%
SQuADv2 *
19.06
19.91
13.46
13.09
11.09
16.75
-3.16
3.29
SQuADv2 C
11.88
11.30
7.85
8.07
6.92
8.84
-3.04
0.77
Comparison Counterparts and Metrics, We compare SPRY to (a) Backpropagation-bassee fed-
erated optimizers FEDAVG [ FEDYOGI 25 FEDSGD (Variant of FEDAVG with per-iteration communication) [1. (b) Zero-order federated methods FEDMEZO (federated version of MEZO [18). BAFFLE 200 FWDLLM [9]. all based on finite difference. MEZO uses prompt- based finetuning to improve the performance of finite differences. FWDLLM generates a random perturbation that has high cosine similarity to the global gradients of the previous rounds. BAFFLE generates ~100-500 perturbations per iteration. More details of these methods are in Appendix A The original imple- mentations of FWDLLM and BAFFLE had excessive memory usage in their implementations. We improve their codebase to be memory-efficient by perturbing only the trainable weights, similar to SPRY. We refer to our implementation as FWDLLM+ and BAFFLE+.
Evaluation of SPRY and its counterparts for classification tasks is on generalized accuracy Accg and personalized accuracy Acep, which are metries measured on server-side aggregated model and client-side locally updated model, respectively. Similarly, for question-answeriig measure Exact Matches and F1 Score. We also measure time to convergence and peak memory consumption during training. Our convergence criterion is the absence of change in the variance of performance metric, assessed at intervals of 50 rounds.
tasks, we
SPRY is implemented in Flower 433 library. Quantization is done using AutoGPTQ [44]. For the zero-order methods, we used their respective client-side implementations with the server simulation structure of Flower. We utilized two Nvidia 1080ti to conduct all experiments of sub-billion sized models and billion-sized models for SPRY and its zero-order methods. We used two RTX8000s and two A100s for Llama2-7B and OPT models on backpropagation-baat based methods respectively. Each experiment was run thrice with 0, 1, and 2 as seeds.
5.1 Accuracy Performance Comparison
Table I reports the accuracy performance of SPRY and its backpropagation- and zero-order- counterparts on heterogeneous datasets for million-sizer RoBERTa Large, and billion-sized Llama2- 7B, OPT6.7B, and OPT13B. Similar results on personalized performance is shown in Appendix HH Figure 5 Results on MultiRC for FEDMEZO are absent since the prompt-based finetuning variant of Llama2-7B was unavailable. Results on more model architectures and dataset combinations are available in Appendix G Details on the learning curves, homogeneous dataset splits, and variance across 3 runs are in Appendix H
Overall, SPRY achieves 5.15-13.50% higher generalized accuracy and 4.87-12.79% higher personal- ized accuracy over the best pest-performing zero-order- based methods across all datasets. FWDLLM+ the best- performing zero-order counterpart, attempts reduce the effect of numeric instability of finite differences by (a) Sampling K perturbations (default K = 10) per batch for each client and picking perturbation per batch that has the highest cosine similarity with the previous round's aggregated gradients and (b) Only picking trained weights from clients whose computed gradients have variance lower than set threshold. However, we posit that this strategy leads to some clients getting excluded due to low variance threshold or outlying clients getting included due to a high
7

---

variance threshold. Besides, picking new perturbations based on the previous round's aggregated gradients in the initial rounds can damage the learning trajectory. While BAFFLE+ samples more perturbation for each batch to make zero-order finite differences more tractable, the scale of language models demands perturbations on the scale of 500-1 per batch, which becomes computationally infeasible. FEDMEZO manages to outperform BAFFLE+ due to its prompt-based finetuning trick but still falls short due to only using perturbation per batch for finite differences on each client. In contrast, Forward-mode AD used in SPRY avoids the numerical instability from finite differences and improves accuracy by reducing trainable weights assigned to each client.
Compared to backpropagation as close as 0.60-6. 6.6 of generalized accuracy and 2.50-14, 14.14111 12% of personalized accuracy. The performance gap between backpropagation and horward-mode AD arises because in backpropagation, weight updates are more accurate as all gradients are computed directly using the error signal of the objective function. In contrast, Forward-mods AD relies on random perturbations, which is re relatively less accurate for gradient estimation. Nonetheless, the advantages of SPRY become evident when we see the peak memory consumption of Forward-more AD compared to backpropagation, which we will discuss next.
bassed based methods, FEDAVG and FEDYOGI, SPRY manages to come
5.2 Peak Memory Consumption Comparison
Parameters Activations
Gradients + Optimizer
States + Misc
Figure 2 shows peak memory consump- tion of backpropagation (used in FEDAVG, FEDYOGI), zero-order finite differences (used in FWDLLM+, BAFFLE+, FEDMEZO), and first- order forward mode AD (used in SPRY). The methods are profiled for a single client.
Memory Footprint (in GBx)
RoBERTa Large
Llama2-78
2.222a
34.003
5,480
5,200
Back- ampagattion
Zemo-Dnder Onder Differace
Tirsk Ordwr Mide AD
Back- propagation
Zero Orter Daffrence
Fire- Order Morde AD
Compared to backpropagation, Forward-mode AD reduces peak memory usage of RoBERTa Large by 27.90%, Llama2-7B by 81.73%, OPT6. 7 by 86.26% and OPT13B by 85.93%.
Meroory Footprint
OPT6.78
OPT13B
60
(In GBs) %
50
The sizes of trainable parameter (colored O) and gradient + optimizer state (colored O are consis- tent across the 3 modes of computing gradients for all 4 models. Hence the savings come from the reduced memory footprint related to activa- tions (colored O) in Forward-mode AD. Com- pared to backpropagation baseed methods, the memory cost related to activations is decreased by 22.12-49.25 in SPRY. Unlike storing all the intermediate activations in backpropagation, Forward-mode AD only has store the previous layer's activation in a forward pass. The activa- tion footprint of Forward-mode AD is equal to the size of the largest activation.
6,24GB
7//305
5,2308
15 Poc
Baco: prosagation
Zero- Order
Riss Orduu
Back propagation
Zero Order
Rly. Ordor
Diffenence
Mide AD
Deffernnce
Mode Ab
Figure 2: Peak memory consumption of SPRY'S Forward-mode AD versus backpropgation- and 7ero-order-based methods. RoBERTa Large, Llama2-7BB 7B, and OPT6.7B are profiled with a batch size of 8, and OPT13B with a batch size of
4. SPRY reduces total memory usage by 27. 77.90-86.20 26%% compared backpropagation- based methods. The 1 1.54-1.96 1.96x additional memory SPRY uses, compared to pero-order-ie based methods, is offset by the accuracy gains ($ 5.1)
Against zero-order methods, Forward-mode AD activations cost 1.96x, 1.95x, 1.83x, and 1.54x more for RoBERTa Large, Llama2-7B, OPT6.7B, and OPT13B respectively. The increasing cost comes from parallel evaluations of (a) the objective function on the original weights and (b) jvp computation on the perturbations in a single forward pass. However, as discussed in $ 5. the increased memory cost is offset by a boost of up 13.50% in accuracy performance. And as 5.3 will discuss, Forward-mode AD reaches convergence faster than zero-order methods since it takes fewer steps to compute closer gradient estimation.
5.3 Time to Convergence Comparison
Figure 3 shows wallclock time-to- co-ooorreeeeeeeeeeeeeeeeeeeeeeee for SPRY and its counterparts. We observe that SPRY 1.15-1155x 6. 16.2228, and 1.34-2.98x faster than the zero-order methods FWDLLM++ BAFFLE+. and FEDMEZO respectively. For a client in each round, SPRY achieves faster round computation time of 1.46x, 28.57 x, and 1.80x on average, against FWDLLM+, BAFFLE+, and FEDMEZO. Forward-mode AD achieves faster convergence and faster rer-round computation by providing a more accurate gradient estimation through a single perturbation per batch, leading to fewer steps needed to reach convergence. Since each client in SPRY only trains partial weights,
8

---

it gains a speedup of 1.14x over backpropagation-basee FEDAVG, FEDYOGI, and FEDSGD for methods, SPRY slows down for
RoBERTa Large. However, compared to the backpropagatton billion-sized LMs. We attribute this loss of speedup to the way jup is computed in Forward-mode AD. jvp is computed column- wise, while its counterpart vjp in backpropagation are computed row-wise. The column-wise computation incurs time overhead.
5,4 Ablation Studies
Fedhyg
FedMeZo
Fedfoggi
Baffle+
We summarize the ablation experiments on various com- ponents of SPRY. Further discussions are in Appendix [
Fed5.dd
Spry
Test Test Accuracy AcCa
0.8-
FwdLLM
0.6
SPRY can generalize to other language model architec- tures. Similar the observations from Table I we see a trend of SPRY outperforming the best performing zero- order method FWDLLM+ by 3. 3.15-10.25% 25% for generalized accuracy, demonstrating that SPRY can generalize to other language model architectures.
0.2
202
400
2000
1800
2000
2200
2420
Time (in minutes)
(a) AG News with RoBERTa Large
SPRY compatible with other PEFT methods. We integrate SPRY with different PEFT methods like IA3, BITFIT, and CLASSIFFER-ONN shows that LORA with SPRY performs the best, with accuracy improvements of 10.60-16.53%.
12
FINETUNING. Results
Test Score Fl,
Effects of the number of trainable weights. We change the number of trainable weights by controlling the rank and scale hyperparameters in LORA. Results show that SPRY achieves the highest accuracy with (r=1, a=1) setting, which has the smallest trainable weight count. The result is consistent with our theoretical analysis in $
203 2000
Time (in minutes)
(b) SQuADv2 with OPT13B
Figure 3: Time to convergence for SPRY and its counterparts, SPRY achieves faster convergence than zero-order methods due more accurate gradient estimations a single perturbation.
Effects of communication frequeney. Per-iteration com- munication SPRY has been shown to boost accuracy by 4.47% compared to per-epoch communication. This improvement brings the accuracy within 0.92% and 0.96% of FEDAVG and FEDSGD, respectively.
Effects of perturbation count per batch. We observe that increasing the number of perturbations per batch (K) for Forward-mode AD has little to no impact on the end prediction performance of SPRY, with K- 100 improving the generalized test accuracy by 1.1% over K 1. The benefits of increasing K are however seen in the convergence speed. Setting K = 10 achieves steady state (of accuracy ~86%) around the 200th round, while the setting with K takes 500 rounds.
Effects of participating client count. Increasing the client count increases the prediction performance of SPRY. For the SST2 dataset, with the total client count fixed to 100, the three settings C' 10, C²= 50, and C 100 produce accuracies of 85. 85.14%, 86.56% and 88.08%, respectively. We also an improvement in the convergence speed as the participating client count increases. To achieve an accuracy of ~85% C = 10, C 50, C 100 require 500, 450, and 150 rounds, respectively.
Importance of splitting weights. To understand the effects of splitting, conduct the following two experiments: (a) With PEDAYT we apply the strategy of splitting trainable layers across clients (see 3.1 backpropagation--aasee FEDAVG, and (b) With FEDFGD, omit the splitting strategy of SPRY for FL. with Forward-mode AD. We observe that FEDAVGSPLIT fails to achieve similar accuracy to FEDAVG with drop of 2.60-10.00% FEDEGD fails to converge the size of trainable weights increases, e.g., with RoBERTa Large with LORA. A, that has 1.15M trainable weights. This proves the necessity of splitting trainable weights for Forward-mode AD in SPRY.
5.5
Communication and Computation Costs
Tables Bn Appendix | shows the communication and computation overhead of SPRY against all its baselines. We summarize the main observations here:
SPRY has lower communication cost due to the splitting strategy and scalar jvp. Suppose Wy is the total trainable parameter count of model to be trained in federated setting. The

---

backpropagation-base based baselines; FEDAVG (per-epoch communication), FEDSGD (per-iteration communication), and FEDYOGI (per-epoch communication) transmit the entire set of trainable parameters to each participating client, and receives the same from each participating client. That results in "client to server" communication cost of W and "server to client" communication cost of We x M. The per-epoch versions of the zero-order baselines (FEDMEZO, BAFFLE, FWDLLM) follow a similar logic due to all the parameters needing to be transmitted to each client, with "client to server" communication costing W'g. and "server to client" communication taking the cost of Wy x M. However, the per-iteration versions of the zero-order baselines fare better, with "client to server" communication only requiring each client sending scalar finite difference (cost of 1), and "server to client" communication accruing (Wg + 1) x M cost, due to the server also needing to send scalar seed to each client now.
Meanwhile, SPRY only needs to send max L 1) layers (where I. is the total layer count of model, and M is the participating count of clients for a round), each layer of size wy. Hence, for per-epoch "client to server" communication accrues wy max [M 1). which is smaller than Wg. Similarly, per- epoch "server to client" communication costs Wy max (17 1). which is also smaller than the cost of Wig related to the baselines. For per-iteration communication, clients only need to send scalar jvp (cost of 1) to the server; this matches the communication cost of per-iteration zero-order methods. Server needs send total of w,M max L, M) (derivation given in Table 2 which a smaller cost than the costs of backpropagation and zero-order methods.
SPRY accrues lower computation cost due to lower trainable parameter count. SPRY'S client- side computation cost is traded off by a faster convergence to higher accuracy through better gradient approximations compared to finite difference-basee methods. And SPRY is the least computationally expensive on the server side due to needing to aggregate fewer parameters from the clients.
Let's assume that matrix multiplication costs for each layer, resulting in a forward pass cost of c. The cost of backpropagation is 2e because the computation of the current layer's weight gradient and the cost of computing the previous layer's activation gradient is another c. jvp computation in SPRY takes additional cost of c for each layer. Moreover, since jvp calculation happens through column-by-colume vector multiplications, the related overhead is quantified by
Hence backpropagation- based methods FEDAVG, FEDSGD, and FEDYOGI computationally costs 3c at client, and costs wy (M- - 1) at server (due to additions). Note that L amounts to all layers in the model here. FEDMEZO costs L L(2c- 3wy) at client through two forward passes, and generating perturbations three times. FWDLLM and BAFFLE costs K LL(2e + wy) due to K perturbations for all L. layers, with two forward passes and generation perturbations once. Against that, SPRY costs 2 max L 37. 1)(c + v) + W.L. for a smaller count of T., traded-off by the jvp computation cost of e.
On the server side, SPRY is the least computationally demanding. SPRY needs to aggregate a subset of layer weights from only the clients that were assigned to those layers. Computation cost on the server-side changes based on the communication frequency per-iteration communication incurs an additional overhead of wy L( M and multiplying those perturbations with aggregate of the jvp values received from the clients) for SPRY and its zero-order counterparts respectively.
7, + 1) and wy LMM + 1) (generation of perturbations at the server-side,
6 Conclusion
SPRY enables finetuning medium and large language models in cross-device FL.. It introduces a training strategy where trainable weights are split across federated clients, so each client only applies Forward-mode AD to fraction of the weights. This approach significantly reduces the memory footprint compared to backpropagation and achieves better gradient estimation, resulting in higher accuracy and faster convergence than zero-order methods. Experiments on various language tasks and models demonstrate SPRY'S effectiveness in reducing memory usage while maintaining accuracy comparable to backpropagation. We formally prove that the estimation bias of the global forward gradients in SPRY depends on data heterogeneity across clients. We also analyzed how the convergence rate of SPRY relates to the configurations of SPRY and FL settings including properties of weight perturbations, data heterogeneity, and the number of clients and FL rounds.
A Acknowledgments and Disclosure of F Funding
This material is based upon work supported by the National Science Foundation under grant no. 2CF-2210243, DMSS-220211L CNS-2224054,
and CNS-233851 and by Adobe.
10

---

References
[1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. communnication-fficienn Learning of Deep Networks from Decentralized Data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, Proceedings of Machine Learning Research. PMLR, 2017.
[2] Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mchdi Bennis, Arjun Nitin Bhagoji, Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. Oliveiraa, Hubert Eichner, Salim EI Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser. Zaïd Harchaoui, Chaoyang He, Lie He. Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi. Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konecný. Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mchryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song, Weikang Song, Scbastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Pranecth Vepakomma, Jianyu Wang. Li Xiong. Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, 2021.
[3] Tyler J Loftus, Matthew M Ruppert, Benjamin Shickel, Tezcan Ozrazgat-Baslanti, Jeremy A Balch, Philip A Efron, Jr. Gilbert Upehurch, Parisa Rashidi, Christopher Tignanelli, Jiang Bian, and Azra Bihorac. Federated learning for preserving data privacy in collaborative healthcare research. Digital Health, 2022.
[4] Li Li. Xi Yu, Xuliang Cai, Xin He, and Yanhong Liu. Contract theory based incentive mechanism for federated learning in health crowdsensing. IEEE Internet of Things Journal, 2022.
[5] Suresh Dara. Ambedkar Kanapala, A. Ramesh Babu, Swetha Dhamercherala. Ankit Vidyarthi, and Ruchi Agarwal. Scalable sederated-learning and internet-of-tling enabled architecture for chest computer tomography image classification. Computers and Electrical Engineering, 2022.
[6] Fabio Pinelli, Gabriele Tolomei, and Giovanni Trappolini. Flirt: retrieval. In Proceedings of the 46th International ACM SIGIR Conference on Research and Development in Information Retrieval, 2023.
Federated learning for information
[7] OpenAL, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, llge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shapiro. Christopher Bemer, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd. Anna-Luisa Brakman, Greg Brockman, Tim Brooks, Miles Brundage. Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann. Brittany Carey. Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decarcaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Elcti, Tyna Eloundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gibson, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Gontijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hallacy, Jesse Han, Jeff Harris, Yuchen He, Mike Hcaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey, Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga. Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang. Roger Jiang. Haozhun Jin. Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Fukasz Kaiser, Ali Kamali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, fukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosie, Gretchen Krueger, Vishal Kuo, Michael Lampe. Ikai Lan, Teddy Lee. Jan Leike, Jade Leung. Daniel Levy, Chak Ming Li, Rachel Lim. Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez. Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning, Todor Markov. Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mchta, Jacob Mcnick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reichiro Nakano, Rajeev Nayak, Arvind Neelakantan, Richard Ngo, Hyeonwoo Noh Long Ouyang, Cullen OeKeefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr H. Pong. Tolly Powell, Alethea Power, Boris Power, Elizabeth Prochl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross, Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish
11

---

Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Selsam. Kyla Sheppard, Toki Sherbakow, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler. Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky. Yang Song. Natalie Staudacher, Felipe Petroski Such, Natalic Summers, Ilya Sutskever, Jie Tang. Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng. Preston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vijayvergiya, Chelsea Voss, Carroll Wainwright, Justin Jay Wang. Alvin Wang. Ben Wang, Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng. Lilian Weng, Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong. Lauren Workman, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao. Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. Gpt-4 technical report. arXiv 2303. 08774, 2024.
[8] Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent EI Shafey, Yanping Huang, Kathy Mcier-Hellstem, Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder. Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hemandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette-Choee Crepy, Shachi Dave, Mostafa Dehghani, Sunipa Dev, Jacob Devlin, Mark Díaz, Nan Du, Ethan Dyer, Vlad Feinberg, Fangxiaoyu Feng. Vlad Fienber, Markus Freitag, Xavier Garcia, Sebastian Gehrmann, Lucas Gonzalez, Guy Gur-Ari, Ari, Steven Hand, Hadi Hashemi, Le Hou, Joshua Howland, Andrea Hu, Jeffrey Hui, Jeremy Hurwitz, Michael Isard, Abe Ittycheriah, Matthew Jagielski, Wenhao Jia, Kathleen Kenealy, Maxim Krikun, Sneha Kudugunta, Chang Lan, Katherine Lee, Benjamin Lee, Fric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontack Lim, Hanzhao Lin, Zhongtao I Liu, Frederick Liu, Marcello Maggioni, Aroma Mahendru. Joshua Maynez, Vedant Misra, Maysam Moussalem, Zachary Nado, John Nham, Eric Ni. Andrew Nystrom, Alicia Parrish, Marie Pellat, Martin Polacek, Alex Polozov, Reiner Pope, Siyuan Qiao, Emily Reif, Bryan Richter, Parker Riley, Alex Castro Ros, Aurko Roy, Brennan Saeta, Rajkumar Samuel, Renee Shelby. Ambrose Slone, Daniel Smilkov, David R. So, Daniel Sohn. Simon Tokumine. Dasha Valter, Vijay Vasudevan, Kiran Vodrahalli, Xuczhi Wang, Pidong Wang, Zirui Wang. Tao Wang. John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengeheng Yin, Jiahui Yu, Qiao Zhang, Steven Zheng, Ce Zheng, Weikang Zhou, Denny Zhou, Slav Petrov, and Yonghui Wu. Palm 2 technical report, arXiv 2305.10403. 2023.
Aakanksha Chowdhery. Clément
[9] Bill Yuchen Lin, Chaoyang He, Zihang Ze, Hulin Wang, Yufen Hua, Christophe Dupuy, Rahul Gupta, Mahdi Soltanolkotabi, Xiang Ren, and Salman Avestimehr. FedNLP: Benchmarking federated learning methods for natural language processing tasks, In Findings of the Association for Computational Linguistics: NAACL 2022. Association for Computational Linguistics. 2022.
[10] Yuanyishu Tian, Yao Wan, Lingjuan Lyu, Dezhong Yao, Hai Jin, and Lichao Sun. Fedhert: When federated learning meets pre-training. ACM Transactions on Intelligent Systems and Technology (TIST), 2022.
[11] Alexander Borzunov, Max Ryabinin, Artem Chumachenko, Dmnity Baranchuk, Tim Dettmers, Younes Belkada, Pavel Samygin, and Colin A Raffel. Distributed inference and fine-tuning of large language models over the internet. Advances in Neural Information Processing Systems, 2024.
[12] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen -hu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
[13] Donggi Cai, Yaozong Wu, Shangguang Wang, and Mengwei Xu. Fedadapter: Efficient federated learning for mobile nip. In Procediings of the ACM Turing Award Celebration Conference- China 2023. Association for Computing Machinery, 2023.
[14] Zeju Qiu, Weiyang Liu, Haiwen Feng. Yuxuan Xue. Yao Feng, Zhen Liu, Dan Zhang, Adrian Weller, and Bernhard Schölkopf. Controlling text-to-image diffusion by orthogonal finctuning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[15] Haokun Liu, Derek Tam, Mohammed Mugeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin A Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. In Advances in Neural Information Processing Systems, 2022.
[16] Xiang Lisa Li and Percy Liang. teninntunnnn:: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers). Association for Computational Linguistics, 2021.
12

---

[17] Benoit Jacob. Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer- arithmetic-only inference. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018.
[18] Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lec, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[19] Mengwei Xu, Dongqi Cai, Yaozong Wu, Xiang Li, and Shangeuang Wang. Fwdllm: Efficient fedllm using forward gradient. arXiv 13894, 2024.
[20] Haozhe Feng, Tianyu Pang, Chao Du Wei Chen, Shuicheng Yan, and Min Lin. Does federated learning really need backpropagation? arXiv 2301. 12195, 2023.
[21] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Ann Lachaux, Thibaut Lavril, Jenya Lee. Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan 7arov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama Open foundation and fine-tuned chat models. arXiv 2307.09288 2023.
[22] C. H. Richardson. An introduction to the calculus of finite differenees. by c.h. richardson pp. vi, 142. 28s. 1954. (van nostrand, new york; macmillan, london). The Mathematical Gazette, 39(330), 1955, doi: 10.2307/3608616.
[23] Atllm Güneş Baydin, Barak A. Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without backpropagation. arXiv 2202.08587, 2022.
[24] Attlum Günnes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey. The Journal of Machine Learning Research, 2017.
[25] Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. International Conference on Learning Representations, 2021.
[26] Chandra Thapa. Pathum Chamikara Mahawaga Arachchige. Seyit Camtepe, and Lichao Sun. Splitfed: When federated learning meets split learning. Proceedings of the AAAI Conference on Artificial Intelligence, 2022.
(27] Renrui Zhang, Jiaming Han, Chris Liu, Aojun Zhou, Pan Lu, Hongsheng Li, Peng Gao, and Yu Qiao. LLaMA adapter: Efficient fine- moeituning of large language models with zero-inie aro-iintialized attention. In The Twelfth International Conference on Learning Representations, 2024.
[28] Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel. BitFit: Simple parameter-effiiien fine-tuning for transformer- based masked language-modes In Proceedings the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers). Association for Computational Linguistics, 2022.
[29] Kunjal Panchal, Sunav Choudhary, Subrata Mitra, Koyel Mukherjee, Somdeb Sarkhel, Saayan Mitra, and Hui Guan. Flash: concept drift adaptation in federated learning. In International Conference on Machine Learning, pages 26931-26962. PMLR, 2023.
[30] Shanshan Wu, Tian Li, Zachary Charles, Yu Xiao, Ziyu Liu, Zheng Xu, and Virginia Smith. Motley: Benchmarking heterogeneity and personalization in federated learning. arXiv 2206. 2022.
[31] Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for
text classification. In Neural Information Processing Systems, 2015. Available at lhtpss http:///gggggggacoootooaaaaaaaaa sittcastssyeepppllleee attps://huggingfactt httpee/eetttttttt//aaaaaaa http:yyawoooooooeeooooooeeooooee
co/ htt.://sttttttttttt///aa/aaaoeee news wnwwwrs Accessed 15 May, 2024.
13

---

[32] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang. Christopher D. Manning, Andrew Ng. and Christo- pher Polts. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Prucessing. Association for Com- putational Linguistics, 2013. Available at atthhtps://mmie hhap::::///ggnggaaceee wt//aaasett///ttaaaaaiaaaaaaaa http://///ttaoooo///sssss////s// Accessed on 15 May, 2024.
[33] Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher D. Manning. A large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference on Empirical Methods Natural Language Processing. Association for Computational Linguisties, 2015. Available ace.
/nnli Accessed on 15 May, 2024.
[34] Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Longuage Technologies. Volume I (Long Papers). Association for Computational Linguisties, 2018. Available at https: http:///hgggnngaeeeeee ca/ cataatastss/ SetFit/mnl Accessed on 15 May, 2024.
[35] Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for SQuAD. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Pupers). Association for Computational Linguistics, 2018. Available at https://
huggingface. ct/ /tt////aaatttrrafarrrraaauaaaaaa v2
Accessed on 15 May. 2024.
[36] Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long Papers). Association for Computational I Linguistics, 2018. Available at
http:///ggggggaaecoaott/ttttttts so//aatasets//ttc//mllaae /tt//itircc////////mo/// Accessed on 15 May, 2024.
[37] Kunjal Panchal, Sunav Choudhary, Nisarg Parikh, Lijun Zhang, and Hui Guan. Flow: Per-instance personalized federated learning. In Thirry--event Conference on Neural Information Processing Systems, 2023.
[38] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models. arXiv 2205.01068, 2022.
[39] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. arxiv 1692, 2019.
[40] Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. arXiv 1810.04805, 2018.
[41] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv 1910.01 108, 2019.
[42] Zhenzhong Lan, Mingda Chen, Schastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. ALBERT: A lite BERT for self-supervised learning of language representations. arXiv 1999,11942, 2019.
[43] Daniel Beutel, Taner Topal, Akhil Mathur, Xinchi Qiu, Titouan Parcollet, and Nicholas D Lane. Flower: A friendly federated learning research framework. arXiv preprint 2007. 143901 2020.
[44] AutoGPTO, 2024. htlphattp::/// /tstta://gittubbl http:////gghhuubbcocoooaateeaeae
com/ com/AuttooPT//uutooooooooooooooo
[45] Richard L. Burden and J. Douglas. Faires. Numerical analysis/ / Richard L. Burden, J. Douglas Faires. Thomson Brooks/Cole, 8th ed. edition, 2005. ISBN 0534392008.
[46] Károly Jordán. Calculus of finite differences. American Mathematical Soc., 1965.
[47] Wenzhi Fang. Ziyi Yu, Yuning Jiang, Yuanming Shi, Colin N. Jones, and Yong Zhou. Communication- efficient stochastic zeroth-order optimization for federated learning. IEEE Transactions on Signal Prucexs- ing, 2022.
[48] Rui Ye, Wenhao Wang. Jingyi Chai, Dihan Li, Zexi Li, Yinda Xu, Yaxin Du Yanfeng Wang, and Siheng Chen. Openfedllm: Training large language models on decentralized private data via federated learning. arXiv 06954, 2024.
14

---

[49] Tao Fan, Yan Kang, Guoqiang Ma. Weijing Chen, Wenbin Wei, Lixin Fan, and Qiang Yang. Fate-Ilm: A industrial grade federated learning framework for large language models. arXiv 2310.10049, 2023.
[50] Fan Lai, Yinwei Dai, Sanjay S. Singapuram, Jiachen Liu, Xiangfeng Zhu, Harsha V. Madhyastha, and Mosharaf Chowdhury. Fedscale: Benchmarking model and system performance of federated learning scale. arXiv 2105.11367. 2022.
[51] Jae Ro, Theresa Breiner, Lara McConnaughey, Mingqing Chen, Ananda Suresh, Shankar Kumar, and Rajiv Mathews. Scaling language model size cross-device federated learning. Proceedings of First Workshop on Federated Learning for Natural Language Processing (FLANLP 2022). Association for Computational Linguistics, 2022.
[52] Shubham Malaviya, Manish Shukla, and Sachin Lodha. Reducing communication overhead in federated learning for pre-trained language models using parameter-cfficient finetuning. In Conference on Lifelong Learning Agents. PMLR, 2023.
[53] Zhuo Zhang, Yuanhang Yang. Yong Dai, Qifan Wang, Yue Yu. Lizhen Qu, and Zenglin Xu. FedPETuning: When federated learning meets the parameter-effeieet tuning methods of pre-trained language models. In Findings of the Association for Computational Linguisties: ACI. 2023. Association for Computational Linguistics, 2023.
[54] Seonghwan Park, Dahun Shin, Jinseok Chung, Namhoon Lee. Fedfwd: Federated learning without backpropagation. arXiv 2309.0/ 150, 2023.
[55] Geoffrey Hinton. The forward-forward algorithm: Some preliminary investigations. arXiv 2212.13345, 2022.
[56] Ning Ding. Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang. Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et Parameter-efficient fine-tuning of large-scale pre-trained language models. Nature Machine Intelligence, 2023.
[57] Zeyu Han, Chao Gao, Jinyang Liu, Jeff Zhang, and Sai Qian Zhang. Parameter-efficient fine-tume fine-tuning for large models: A comprehensive survey. arXiv 2403. 14608, 2024.
[58] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuning of quantized LLMs. In Thirty- Conference on Neural Information Processing Systems, 2023.
15

---

A
Related Work
SPRY uses first-order forward gradients for federated finetuning of language models. Hence, we review the literature on estimating gradients with low memory consumption and show how SPRY represents a significant advancement in finetuning language models in FL.
Zero-order Gradients. Gradients derived from finite difference 45 22 46) methods are called zero-order gradients since they don't involve Taylor expansion of the objective function MEZO 18] has shown that finite difference with one perturbation per batch does not reach conver- gence on own without additional tricks like prompt-based finetuning, which are highly specific to the tasks. For a more accurate gradient approximation, an average of zero-order gradients de- rived from multiple (~10 to 100) random perturbations on the same input batch required [24], leading to slow convergence. BAFFLE 20 and FEDZO 47 utilize zero-order gradients in feder- ated settings. To train vision models (of parameter count S 13M), BAFFLE requires (a) ~1000500 perturbations per batch for each client respectively, and (b) per-iteration communication among clients like FEDSGD [[] FEDZO also requires ~20 perturbations per batch for a vision model of size 25M. Besides, numerical errors associated with finite difference make MEZO, BAFFLE, and FEDZO suffer from sub-optimal predictions compared to mackpropagaion-baan using Forward-mode AD, computes gradients more accurately in single forward pass compared to averaged gradients obtained through finite difference methods. This higher accuracy is achieved without needing modifications model architectures or task structures, while also maintaining memory footprint similar to that of finite difference methods.
counterparts. SPRY,
First-order Forward Gradients. Gradients derived from Forward-mode Auto-Differentiation (AD) are considered first-order since it involves computing partial derivatives of the intermediate activations with respect to the input. We guide interested readers to the survey on different modes of automatic differentiation [24]. FGD [23] shows preliminary results on the speedup and comparable accuracy achieved by Forward-mode AD against backpropagation. The challenge that makes forward-mode AD less popular than backpropagation is that gradients derived from forward mode require more jvp column- column-by-collne by-column evaluations per input batch the number of trainable weights increases. Moreover, evaluation of FGD is limited to multi-layer perceptron of size d 11.M. Direct use of FGD finetune language models leads to slow or no convergence. SPRY splits the trainable layers of large language model across multiple clients in FL., letting each client finetune only a small subset of weights through forward gradients.
Training or Finetuning Language Models in Federated Learning. In recent years, several frameworks have been proposed to train or finetune LLMs in FL 4 [50]. The backpropagation-ban based
methods I0 [1] even with parameter efficient finetuning (PEFT) and quantization [13 52 53), have
large memory footprints due to the overhead related to activations, gradients, and gradient history storage for adaptive optimizers [18].
FEDFWD [54] applies FWDFWD (55) (which measures "goodness" of forward pass activations
to judge which perturbations are useful) in FL, but FWDFWD struggles as model size scales up. FWDLLM uses zero-order gradients to finetune language models. It samples ~10 perturbations per batch. For each batch, it picks perturbation that has the highest cosine similarity with the previous round's gradients. Sampling new perturbations based on aggregated gradients from previous rounds during the initial stages can disrupt the learning trajectory. SPRY requires I perturbation (without resampling) per batch to reach a higher prediction performance faster than FWDLLM.
16

---

B Datasets and Hyperparameters
Here we provide details of the datasets and their corresponding training hyperparameters used in this work.
Simulating Heterogeneity through Dirichlet Distribution. For each of the tasks, the class dis- tribution each client gets depends on the Dirichlet distribution, where a parameter a regulates the concentration of samples of a particular class for specific client. Dir 1.0 means all clients have homogeneous datasets where each class is equally likely to be on each client. With Dir a + O, the datasets of each client get more heterogeneous, where the sample distribution of each class is more likely to be concentrated on only a subset of clients.
Default Hyperparameters. Unless otherwise mentioned in dataset-specific paragraphs, the default hyperparameters for each method and for all datasets are stated here. For the backpropagation-base methods FEDAVG, FEDYOGI, and FEDSGD; we will fix the number of epochs to since the goal of this work is to inch closer to backpropagation-lik prediction performance while reducing the memory footprint. All the experiments have been run for 1 1500 FL rounds, except the experiments on OPT models, which are run for 600 FL rounds. Our observation from hyperparameter tunnig shows that the learning rate that gives the best performance is the same for all studied methods. BAFFLE and memory-effiieen improvement BAFFLE+ made by us, can perform better as the number of perturbations per batch increases, but due to the scale of experiments with 10-100 per round and up to 1500 rounds in the FL setting, we limit the total perturbations per batch BAFFLE+ to 20 perturbations per batch and fixed finite difference step size o 1e-4. FWDLLM+ samples 10 perturbations for each batch, finite difference step size of a le-2. FEDMEZO samples perturbation for each batch, with finite difference step size of a 1e-3. FEDMEZO also requires 3-5 epochs for each client. SPRY has 1 perturbation per batch for each client. For SPRY and its zero-order counterparts (BAFFLE+, FWDLLM+, and FEDMEZO), perturbations are sampled for normal distribution with 0 mean and variance. Default LoRA r and a are and I, respectively. All methods use ADAMW client-side optimizer. Besides FEDAVG, all methods use FEDY as server-side optimizer.
AG News. AG News dataset 31 has been derived from a corpus of 496.835 categorized news articles. A subset of the corpus is used that has 120,000 total training samples and 7,600 total testing samples spread equally across 4 classes. The news articles are classified into 4 classes: World, Sports, Business, and Sci/Tech. We split this data across 1000 clients. Each client gets equal number of samples for train and test datasets. This dataset is under Creative Commons CCZero(CCO) public domain dedication.
Learning rate for backpropagation--basd (FEDAVG, FEDYOGI, and FEDSGD), bero-order-based (FwDLLM, BAFFL,F, and FEDMEZO), and first-order-based SPRY is (1e-3, 5e-4, 1e-4, le-5). The batch size is set to The max sequence length 128. All methods use ADAMW as a client-side optimizer, while SPRY performs better with SGD. Variance threshold of FWDLLM+ is le+1.
SST2. Stanford Sentiment Treebank Binary 32 (or SST2) dataset is for binary sentiment classifi- cation task. The dataset has 1 1,855 sentences derived from a set of movic reviews. The corpus was parsed using the Stanford Parser into 215,154 discrete phrases annotated by 3 human judges. This dataset contains 67.349 training samples, 872 validation samples, and 1821 testing samples. This sentiment classification dataset has the following sentiments classes: Positive, and Negative. These samples are equally split between 100 clients, depending on the Dirichlet distribution. This dataset is under Creative Commons CCZero public domain dedication.
Learning rate for backpropagatton--base (FEDAvG, FEDYOGI, and FEDSGD), zero-order-based (FwDLLM, BAFFLE, and FEDMEZO), and first-order-base SPRY is (1e-3, 5e-4, 1e-4, le-5). The batch size is set to 8. The max sequence length is 64. Variance threshold of FWDLLM+ is 5e+0.
The Yelp reviews dataset 31 a binary classification dataset gathered during the 2015 Yelp Dataset Challenge. The full dataset has 1,569,264 total samples, and it defines classification tasks. We use the polarity classification task, which a binary classification problem. It considers 1-2 stars negative polarity and 3-4 stars positive polarity. This dataset has 280,000 training samples and 19,000
Yelp.
17

---

test samples in each polarity, We split this data into 1000 clients. This dataset is under the Apache License. Version 2.0.
Learning rate for backpropagatton--aased (FEDAvG, FEDYOGI, and FEDSGD), zero-order-hased (FwDLLM, BAFFLE, and FRDMEZO), and fist-order-based SPRY is (1e-3, 5e-4, le-4, 5e-5, 1e-5). The batch size is set to 8. The max sequence length is 128. Variance threshold of FWDLLM+ is 5e+1.
Yahoo. Yahoo! Answers Comprehensive Questions and Answers dataset 31 was gathered via the Yahoo! webscope program. The corpus itself contains 4,483,032 question/answer pairs, which were then formulated as a 10-class classification task. Each class in this dataset has 140,000 training and 5,000 testing samples.are The question/answer pairs are split into the following classes: 1. Society Culture, 2. Science & Mathematics, 3. Health, 4. Education & Reference, 5. Computers & Internet, Sports, 7. Business & Finance, 8. Entertainment & Music, 9. Family & Relationships, 10. Politics & Government. We split this dataset between 1000 clients, where each client gets equal amount of data samples, where the data distribution is set by changing the Dirichlet distribution a to range from most homogencous (a = 1) to least homogeneous (a = 0). This dataset is under the Apache License, Version 2.0.
Learning rate for backpropagation-basse (FEDAvG, FEDYOGI, and FEDSGD), zero-order-based (FwDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is {1e-3, 5e-4, 1e-4, le-5). The batch size is set to 8. The max sequence length is 128. Variance threshold of FWDLLM+ is 5e+1.
SNLI. The Stanford Natural Language Inference corpus 33 has 570,152 total sentence pairs. is natural language inference dataset, where the task is identifying if one sentence infers another. It has 550,152 training samples, 10,000 testing samples, and 10,000 evaluation samples. This dataset is split among 1,000 clients. The dataset has 3 classes: 1) The first sentence entails the second sentence, 2) The first sentence is neutral to the second sentence and 3) The first sentence contradicts the second sentence. This dataset is under CC BY-SA 4.0.
Learning rate for backpropagattion-bsse based (FEDAvG, FEDYOGI, and FEDSGD), cero-order-base (FwDLLM, BAFFLE, and FEDMEZO), and first-orderbased SPRY is (1e-3, 5e-4, lc-4, 5e-5, le-5). The batch size set to 8. The max sequence length is 80. Variance threshold of FWDLLM+ is le+2.
corpus contains 432.702 sen- tence pairs which were crowd-sourced and then annotated with textual entailment information. This dataset is also Natural Language Inference dataset with SNLI. This dataset has 392,702 training, 20,000 evaluation, and 20,000 testing samples. These samples are split among 1,000 clients. This dataset draws from multiple sources, most of which are under the Open American National Corpus (OANC) license. The rest are under either the CC BY 3.0 Unported licenses or the CC BY-SA 3.0 licenses.
MNLI.
The Mente Natural Language Inference (MNLI) 34
Learning rate for backpropagatton-bassed (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FwDLLM, BAFFLE, and FEDMEZO), and fist-order-based SPRY is (le-3, 5e-4, le-4, 5e-5, le-5). The batch size is set to 8. The max sequence length is 80. Variance threshold of FWDLLM+ is le+2.
SQuADv2. The Stanford Question Answering Dataset (SQuADv2) 35 consists of crowd-sourced questions about a set of Wikipedia articles. It is a reading comprehension dataset, where the answer to a question is a section (or a span) from the passage. It is also possible for the question to be unanswerable. The dataset contains 100,000 answerable and 50,000 unanswerable questions. This dataset was split into 500 clients. The heterogeneity is generated based on the topic labels (or "titles") associated with each question in the dataset. There are 35 titles available. This dataset is under the CC BY-SA 4.0 license.
Learning rate for backpropagattion-bseee bassd based (FEDAvG, FEDYOGI, and FEDSGD), cero-order-based (FwDLLM, BAFFLE, and FEDMEZO), and firstoorder-based SPRY is (le-3, 5e-4, le-4, 1e-5). The batch size is set to 8 for Forward-mode AD and zero-order differentiation-basee methods. Backpropagation-baan methods use a batch size of The max sequence length is 400. Variance threshold of FWDLLM+ is 5c+1.
18

---

MultiRC. Multi-sentence Reading Comprehension (MultRCC corpus is dataset that contains short paragraphs with questions, whose answers can be found the paragraph itself. We consider task where we classify if the input question given as input is correct or false. It has 6k nulti-sentence questions about 800+ paragraphs. Due to the scale of the dataset, we split into 100 clients. This dataset is under the MIT license.
Learning rate for backpropagatioon-besed (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is {le-3, 5e-4, 1e-4, le-5). The batch size is set 8. The max sequence length is 256. Variance threshold of FwDLLM+ is le+2.
C Limitations and Future Work
SPRY achieves remarkable drop in memory consumption due to Forward-mode AD not having to store activations during the forward pass. However, the current implementation of Forward-mode AD by PyTorch is still suboptimal in terms of computation time. The high computation time is attributed to the column-by- -ouummb computation of the intermediate results of jvp. Improving the computation of jvp such that it takes significantly less time (almost half) than backpropagation remains an open and interesting problem. Moreover, there's 's room for improvement reducing the memory consumption to that of zero-order methods. In zero-order methods, the weights perturbed and forward pass computed on the perturbed weights. However, with the current implementation of Forward-mode AD, perturbations create separate copy from the original weights, which accrues additional overhead. To further utilize the computation capacity clients in FL, revice-heterogeneity-eee
strategies splitting and mapping layers clients can be explored for SPRY, e.g., layer selection could happen on the client-side according to their data distributions.
D Broader Impact
Through SPRY, we aspire to bring impact in terms of data privacy and accessible finetuning of large language models (LLMs) on edge devices.
For small and medium organizations and individual users, the cost finetune these language models can be prohibitively expensive as the size of the trainable weights increases. With SPRY, we enable finetuning LLMSs feasible for a wider range of users. Moreover, the federated setting also provides benefits of data privacy. This can be ideal for use cases where LLMs can bring improved performance for plethora of personalized downstream tasks, but the finetuning data containing sensitive and confidential information is never shared with third party. With SPRY, such data remains on the user's s device, ensuring privacy while still allowing for effective finetuning of the LLM.
resource-oonstrained edge devices with low memory consumption, making
However, making LLM finetuning accessible comes with its own challenges like a spread of biases or misinformation. Without preprocessing and filtering client data their devices, the LLMS can be fed harmful and misleading information. Hence, is necessary to develop guardrails on what kind of data should be filtered in, to finetune LLMs with sourced compute resources.
broader range of organizations and individuals
19

---

E
SPRY Pseudocode
Algorithm
shows the workflow of SPRY.
Algorithm 1: SPRY
Input:
Total number rounds, r E [R]] Round index, Number of clients per round, m é [M]] Client index, .M: Set of available clients, s: Client sampling rate, Dm: Dataset of client m., IJc
Local learning rate, (r); Model weights 4h round, wht) Subset of assigned trainable weights for client m at r round, f: Objective function.
Output: (R+1)); Model the end the training
Main SPRY
Server loads pre-trained LM w with PEFT
for E [R] round
Sample M clients from M at rate of
C t list of trainable PEFT param names
crient_layers www.er_mappinggeeeee.m... +- MAAPLAYERSTOCUEETSSCCCC..
for m [M] in parallel do
CLEENTRRAIN client_ client__ayerrmappinngii
Build w (r) with wh! using seient_layes aayee__apping/nt Vm [M]
10
Use adaptive optimizer like FEDYOGI to build w (++1) based on the aggregated (()
11
end
13 end
Function MAPLAYYESSTOCLLENTSSS.. MM)
Function CLIENTTRAIN trainable layers)
client_idx f- 0; mapping () for name E do
Freeze whn parameters & trainable, crainable_layers
16
Generate perturbations + (00 Sbupe
17
rollover_idx: + client_idddd % M
18
juniiggrooloover__d_llllllplpli client_ idx +- client_idy
V trainable w
jvp - f F.FOWARDDDDDD um v; Dm) W.
29
end
W ml
(r) Me jvp)
21
return mapping
return all trainable /.
22 end
end
The function MAPLAVERSTOCLIENTS on Line 14 in Algorithm D shows how we have assigned only few trainable layers to each client in FL. to make Forward-mode AD more effective at generating better estimation of the gradients.
function CLIENTTRAIN on Line 23 Algorithm [ each client m gets copy of the entire language model (r) and a list trainable lonolle_lleeer layers of parameter names the client m has to train. client m freezes the parameters that are not included in its trainable layers. And for each of the parameter wh which need to be trained, SPRY generates corresponding random perturbation v using normal distribution A/ ((t) shape
Once client has obtained forward gradients of all the trainable parameters wh! wir) those parameters locally updated with optimizers like SGD ADAM. The updated trainable parameters w(r) sent back to the server. The server has a mapping of parameter names to client IDs, and hence builds w (r) by using i Vm E [M]. If there are multiple clients mapped to the same parameter, then we take a weighted average (similar to FEDAVG) of all the parameters to build w(r), SPRY uses adaptive optimizers like FEDY at server-side on effective gradients A = w"(r) - w(() to reduce the noise of forward gradients.
20

---

Communication and Computation Costs
F.1 Communication Costs
Table 2 illlstrates communication costs of SPRY and its backpropagation- and finite difference- based baselines. A discussion on communication modes of SPRY is also given in Section 3.2, "Per-Epoch Communication" and "Per-Iteration Communication".
Table 2: Communication cost of SPRY and all baselines. M is the count of participation clients. Total count of trainable parameters of a global model is wy w. (for simplicity, we assume that each layer has wy parameters).
Communication cost (in parameter count) from each client to server for each communication round
Communication cost
(in parameter count) from server all clients for each communication round
Method
Gradient computation
(Comm. frequency)
FEDAVG FEDYOGI (Per-epoch)
Backpropagation (First-order gradients)
W0
Wg x M
FEDSGD (Per-iteration)
W0
Wg x M
FEDM:ZO / /WDLLM/ BAFFLE (Per-epoch)
Finite differences (Zero-order gradients)
W0
Wg x M
(w, + 1) M
FEDMEZO/ FEDMEZO/FWDLLM BAFFLE (Per-iteration)
FWDLLM
(of finite
("11" for perturbation sced)
difference scalar)
We X max(L/M,
SPRY (Per-epoch)
max (L/M. 1) x M we max(L, M)
Forward-mode AD (First-order gradients)
(Assuming L%M = for cach of exposition)
(we X max (L/M 1) x M)
SPRY (Per-iteration)
1 (of jvp scalar)
+(1 x M)
we max M) M
Here we discuss the costs related to those communication modes:
Per-epoch Communication. SPRY'S client-to- server communication cost does not scale linearly with clients like in its backpropagation and finite-differenee counterparts, but instead decreases or stays constant for L as more clients are present. Server-to-client communication cost lower in L SPRY due to only sending one layer per client when M > L or 37 layers per client otherwise. This result follows from the below observation:
Backpropagation based and diitte-dfference--ee based methods have communication cost of Wg, where Wg represents the global model size. Each client in [M]] (set of participating clients) receives all trainable parameters from the server, requiring the server to send total of wy x M parameters each round.
SPRY'S communication cost per epoch is wi max( L Mf 1). where L is the layer count and my is the count of parameters for each layer. Each client sends subset of trainable parameters, incurring communication cost of WeL. parameters for L > M,, and we for L < M. When L < M, each client gets layer, hence the communication for each client is we.
Per-iteration Communication. SPRY accrues lower communication cost than the finite difference and backpropagation counterparts due to the layer splitting strategy, and the server's 's ability to compute gradients based on the jvp value. This is because:
The communication cost from client to server for forward-mode AD and finite differences is 1. This due an FL round that involves (1) server selecting a random seed, (2) server sending with trainable parameters to clients, (3) clients generating perturbations based on the seed, (4) deriving and
21

---

sending back scalar or finite difference scalar to the server, and then (5) server computing gradients by multiplying the derived perturbations with the seed.
The server to client communication is (Wy + 1) x M, where the "+1" is due the randomness seed.
F.2 Computation Costs
Table shows the computation costs of SPRY and its baselines, where the client-side cost is for cach iteration, and the server-side cost is for each round.
Briefly, SPRY'S client-side computation cost is traded off by a faster convergence to higher accuracy through better gradient approximations compared to finite difference- methods. Furthermore, SPRY is the least computationally expensive on the server side due to needing to aggregate fewer parameters from the clients.
Table 3 assumes that matrix multiplication costs c for each layer, resulting in a forward pass cost of c. The cost of backpropagation is 2c because the computation of the current layer's weight gradient is e, and the cost of computing the previous layer's activation gradient another c. jvp computation in SPRY takes additional cost of for each layer. Moreover, since jvp calculation happens through column-by-columi vector multiplications (Sec 3.1 of (24). the related overhead quantified by v.
Table 3: Computation cost of SPRY and all its baselines. The client-side cost is for each iteration, and the server-side cost for each round. cost of matrix multiplication for each layer. is the overhead related to column-by-colume vector multiplications jvp. wy is the size of each layer and hence size each layer's s perturbation too. K is the perturbation count per iteration. K for SPRY and FEDMEZO, and 20 for BAFTLE and FWDLLM.
the layer count, M is the participating client count, e the
Computation cost of each client for each iteration
Computation cost of the server for each round
Method (Comm. frequency)
Gradient computation
(Aggregating layer weights from M clients) (M 1) w,L
FEDAVG / FEDYOG! (Per-epoch)
Backpropagation
3Lc
(First-order gradients)
FEDSGD (Per-iteration)
3Le
(M 1) weL
FEDMEZO (Per-epoch)
L(2c+ +
(M - 1) x weL
Finite differences Zero-order gradients)
FWDLLM
BALLM//BAIII BAFFLE (Per-epoch)
KL
(M - 1) x weL
weT. + MT. + we 1)T
2M
FEDMFZO (Per-iteration)
T.(2c 3we)
(perturbation generation + gradient calculation
+ weight update)
FWDLLM/ / BAFFLE (Per-iteration)
KL + wi)
2Mw,L
SPRY (Per-epoch)
2 max( 4 $,11 x(c+ v) + weh
EMCiM (C.M) L)w max (,4 1))
Forward-mode AD (First-order
(usually, |M max (#, 1))
gradients)
SPRY (Por-iteration)
2 x max( $1 1) x(e+ v) + WeL
Emela (2|M/ max (g..
Client-side per-iteration computation cost. Backpropagation needs 3 matrix multiplication op- erations per layer. For zero-order methods, there are matrix multiplications (incurred due to two forward passes) per layer, and per perturbation within a training iteration; and additional overhead uy T. for perturbation generation. MFZO requires generation of perturbations thrice for the same seed (Algorithm 1 in MEZO (18)..
22

---

SPRY 's computation cost is 2 x max( M 1) x (c+ + wyL. Since SPRY allocates at most M layers each client, the computation cost only scales with max( & 1), against its counterparts scaling with L. However, ford AD computes jvp column-wise, while its counterpart vip in backpropagation is computed row-wise. This results in time overhead () if the number of trainable parameters exceeds the output size (1 as loss is scalar), which is the case for neural networks. Therefore, SPRY'S per-iteration computation cost higher compared to other approaches.
Note that the per-iteration computation cost of SPRY is not the whole picture. It takes fewer communication rounds to reach higher accuracy due to better gradient approximation of forward- mode AD than finite difference methods. This is why "Time to Convergence" (Section discusses fair comparison of SPRY'S 's runtime and prediction performance.
Server-side per-round computation cost. the server side, SPRY is the least computationally demanding. SPRY needs to aggregate a subset of layer weights from only the clients that were assigned to those layers, while its counterparts need to aggregate all layers from all clients.
Computation cost on the server-side changes based on the communication frequency per-iteration communication incurs an additional overhead of wy I. M + 1) and w.L((M 1) (generation of perturbations at the server-side, and multiplying those perturbations with aggregate of the jvp values received from the clients) for SPRY and its zero-order counterparts respectively.
23

---

Table 4: Generalized (Accg) and personalized (Accy) accuracies (the higher, the better) for SPRY and its backpropagation and zero-order based counterparts on various language model architectures. The datasets are split with Dir 0.1.
Zero-order based Method
First-order Forward-mode AD
Backpropgation based
FEDAVG
FEDYOGI
FWDLLM
SPRY
Acc,
Acc,
Accg
Acep
Acc,
Accp
Acc,
Accp
AG News BERT Base
93.00% 93.34% 93.31 % 93.88% 83.41% % 83.42% 86. 74%
92.42%
SST2 on DistiBERT Base
91.47% 95.28% 87.95% 92.97% 79.09% 80.94% 84.90%
87.12%
SNLJ on BERT Large
85.79% 89.36% 86.72% 90.33% 66.76% 64.84% 77.01% %
77.21% %
Yahoo on DistilBERT Base
69.13% 74.75% 63.84% 71.25% 54.29% 55.87% 61.17%
61.47%
Yelp on AlbertV2 Large
90.17% 92.78% 90.24% 94.00% 82.65%
83.25% 85.80%
86.00%
0.8
0.8
0.6-
Gener Text
Acturery Atre
0.8
Generalized Test
Accuracy Acts
0.6
Generaizzd
a=1 8, a=16
0.4
LoRA
0.4
0.4
IA3
Spry the Communication
FodSpd Peerttcaaton Communicat
Fedlag (Per Eandh Conmunication
32, a-32
16, a-16
Classiflar-orly
SpryPertar Per ktaasion Com municationn
0.0
200
400
669 ado 100012001400 Rounds
0.0
200
600
800
1400
200
10001200
Rounds
Rounds
(a) Effect of different PEFT methods on SPRY
(b) Effect of per-opoch and per-iteration communication
(c) Changing LoRA and a for SPRY
Figure 4: Ablation studies on PEFT methods, communication frequency, and LoRA hyperparameters.
G
Ablation Studies
Here we will dive into various components of SPRY and see their impact on the performance of SPRY.
SPRY can Generalize to Different Language Model Architectures. Shows generalized and personalized accuracies Acc, and Acc, for homogeneous data splits for language model architectures other than RoBERTa Large.
For zero-order gradients, we show the results of the best-performing method, FWDLLM+. We see similar trend of SPRY outperforming FWDLLM+ by 3.15-10.25% for generalized accuracy, and by 2.75-12.37% for personalized accuracy, exhibiting how SPRY is independent of model architectures. SPRY also comes as close as 4.44-9.71 % to the best-performing backpropagation- based method.
SPRY Supports Other PEFT Methods. Figure 4a shows the generalized accuracy of SPRY using three different PEFT methods: LORA, IA3, and BITHIT. We also experiment with finetuning only classifier layers, calling it CLASSIFIER- -ONLY. LoRA (with 0.3241 % of the total parameters of RoBERTa Large) outperforms IA3 (with 0.3449% of the total parameters) by 10.60% while BITFFr fails to converge for all datasets. Only training classifier layers is worse than finetuning LoRA weights by 16.53%. The observation on comparison of LORA with IA3 is consistent with the work benchmarking various PEFT methods (56). BITFIT has been observed fail LLMs (57). Furthermore, unlike 1 IA3, LORA A has been shown to be successful at finetuning quantized billion-sized models in QLORA 58). making a strong candidate for our work.
Effects of Communication Frequency. One way to reduce the noise introduced by the random perturbations in gradient computation is, to communicate the gradients back to the server every iteration instead of every few epochs. Figure 4b shows results of per-epoch and per-iteration communication variants of SPRY and FEDAVG.
24

---

0.8
0.2
0.3-
Greemn DeA Tee Aczuneyy Ade
Acculacy
Genernlted
Accu/acy
0.4
0.4
C
0.2
0.2
0.0
100
0.0
100
leang
Tacheppplk
Bounds
00020000000000000000000
Bounds
200 426
600
1200
220 3400
200 400
500
830
LOGE LOOGI400
(a) Changing K (Forward-mode AD
(b) Changing C (Per-round participating client count)
(c) RoBERTa Large (I.eft) and BERT Base (Right)- Splitting parameters across clients Backpropagation vs Forward-mode AD
perturbation count per iteration)
Figure 5: Ablation studies perturbation counts, participating client counts, and layer splitting strategy.
By communicating each iteration, we see boost of 4.47% in the accuracy of SPRY, only 0.92% and 0.96% away from the accuracy of FEDAVG and FEDSGD respectively. Furthermore, as shown in the convergence speed of SPRY also improves by communicating each iteration.
Figure 4b
As discussed in Section 3 to reduce the trade-off between performance gains and communication cost per iteration, each client can send only the jvp scalar to the server instead of transmitting all the assigned trainable weights 20). With the seed value of the randomness, the server can then generate the random perturbation vector which was used by all the clients to generate their respective jvp values. Then the random perturbation multiplied with the received j jvp values of all clients, to compute gradients.
Effects of the Number of Trainable Weights. Figure 4c displays the effect of changing trainable LORA parameter count on the prediction performance of SPRY. For DistiBBERT Base (total parameter count of 66M), LORA reduces the trainable parameter count to 0.61M (0.91%), O. 0.74M (1.11%) 0.89M (1.33%) and 1.18M (1.77%) with LoRA A hyperparameter settings of (r=1, o=1), (r=8, a=16), (r=16, l=16), and (r=32, a=32).
SPRY achieves the highest accuracy of 84.90% with (r=1, a=1) setting, which has the smallest trainable parameter count. The accuracy increases as the layer size decreases since fewer perturbed weights provide less noisy gradients.
Effects of the Number of Perturbations per Batch. The effect of increasing the number of perturbations per batch and hence the number of jvp evaluations for a batch is shown in Figure55a for the SST2 dataset. Here, gradients generated from each random perturbation and their corresponding jvp values are averaged to update the model. We observe that increasing K (perturbations per batch) for Forward- mode AD has little no impact on the end prediction performance of SPRY, with K = 100 improving the generalized test accuracy by 1.1% over K = 1. However, the benefits of increasing the perturbation count per batch are seen in the convergence speed. Setting K 10 achieves a steady state (of accuracy ~86% around th round, while the setting with K = 1 takes 500 rounds. The improvements in convergence speed are saturated for K > 10. This shows that more perturbations reduce the gradient estimation noise only to an extent.
Effects of the Number of Participating Clients per Round. Figure 5b shows how changing the per-round number of participating clients C' influences SPRY on the SST2 dataset. Increasing client count increases the prediction performance of SPRY. With the total client count fixed to 100, the three settings C = 10, C = 50, and C 100 produce accuracies of 85.14%, 86.56% and 88.0%%, respectively. Similar to the findings of Section G we also see an improvement in the convergence speed as the participating client count increases. To achieve an accuracy of 855% C'= = 10, = 50, C = 100 require 500, 450, and 150 rounds respectively. The performance gains and faster convergence are due to more clients training the weights of the same layers.
The Importance of Splitting Layers. To understand the effects of splitting, we compare the results of the following two experiments: (a) With FEDAVGSPLIT, we apply the strategy of splitting trainable layers across clients (Section 3.T) backpropagation-baal based FEDAVG, and (b) With FFDFGD, we omit the splitting strategy of SPRY.
25

---

Figure 5c shows the performance of FEDAVG and PEDAVGSPLIT against FEDFGD and SPRY for two LMs: RoBERTa Large (355M) and BERT Base (110M). We observe that FEDAVGSPLIT fails to achieve similar accuracy for both models with a drop of 2.60% and 10.00%. This is because in FEDAVGSPLIT, fewer clients are training each subset of weights. Moreover, we see similar accuracy with an absolute difference of 2. 2.70-3.61 % between FEDAVGSPI.IT and SPRY, since the trainable weight count per client is low. FEDYOGISPLIT follows the same observation of not achieving similar accuracy to FEDYOGI if the trainable weights split across clients. On the contrary, FEDFGD converges for the smaller model BERT base, albeit 150 rounds slower than SPRY, and with 2.87% accuracy drop. But as the size of trainable weights increases, e.g., for RoBERTa Large, FEDFGD fails to converge. This proves the necessity of splitting layers for Forward mode AD so that each client has fewer trainable weights to perturb.
26

---

H
Additional Results
H.1 Generalized Performance Curves
Generalized results on homogeneous and heterogeneous clients with Dir a 1.0 and a = 0.1 are shown in (a) Figures and [ for RoBERTa Large, Llama2-7B, OPT6.7B, OPT13B; and (b) Figure for BERT Large, BERT Base, DistilBert Base, Albert Large v2.
1.0
1.0
1.0
Gomeraltonn Tosr
Acturacy AtC.
Generalized
0.6
Gomeralized
FedMazo
0.
FedEgd
Spry
0.24
0.2
0.0
200
400
600
1000 Rounds
500
12001400
0.0
200
400
600
800 100012001400
0.0
200
400
600
800 1000
Rounds
Rounds
(a) AG News with RoBERTa Large
(b) SST2 with RoBERTa Large
(c) SNLI with RoBERTa Large
Genaratzed lest
0.8
Generatize tast
Accurncy AtCe
0.8
Genaratzed Test
0.8
0.6
0.6
0.6
0.4
0.4
0.4
120012001400
800
100012001400
(d) MNLI with RoBERTa Large
Yahoo with RoBERTa Large
(I) Yelp with RoBERTa Large
1.0
35
Generalized Test F1 Score F1,
30
2.00
Ganeralzeed
Exack Masthes
25
Generalized
20
15
10
0.2
0.0
0-
200
400
100012001400
100
200
300
400
500
600
100
200
300 Rounds
400
500
600
Rounds
Rounds
(g) MultiRC with Llama2-7B
(h) SQuADv2 with OPT6.7B
(i) SQuADv2 with OPT6.7B
20-
25.0-
Goneralzeed
15-
Cenirallzed
12.5
10.0
7.5-
0:
100
200
300 Rounds
400
Rounds
(i) SQuADv2 with OPT13B
(k) SQuADv2 with OPT13B
Figure 6: Generalized accuracy F1 score Exact matches for
homogeneous clients (Dirichlet 1.0) setting
27

---

1.0
1.0
Corneralieed Tast
Accurscy aCC.
Caneraliavd
Generalize
0.4
0.4
FwdLLM+
0.2
Fedhtgi
Spry
0.24
0.21
0.0
200
400
600
500
1000
0.0
200
400
600 800 100012001800
0.0 200
460
660
800 100012001400
Rounds
Rounds
Rounds
(a) AG News with BERT Base
(b) SST2 with DistilBERT Base
(c) SNLI with BERT Large
1.0
Generalized lest
0.8
Genaralzzed Test
0.8
0.6
0.6
100012001460
100012h014ón
Yahoo with DistilBERT Base
(e) Yelp with Albert Large v2
Figure 7: Generalized accuracy /FI score Exact matches for
homogeneous clients (Dirichlet 1.0) setting for variety of language models
28

---

1.0
1.0
1.0
0.8
Ceneralzed Test
Accufary Act.
0.0
0.0
Ceneralzzed
0.6
Genwrllzzed
0.4
Fedding
Baffle
Fectigd PM
Spry
0.0
200
400
B00
1000
12001400
0.0
200
400
800
820
1000
1200140
0.0
200
100012001000
Rounds
Rounds
Rounds
(a) AG News with RoBERTa Large
(b) SST2 with RoBERTa Large
(c) SNLI with RoBERTa Large
Caneraltzed lest
Ceneraltzed Best
Accureyy Arts
0.8
Generalized Test
0.6
0.6
0.6
0.4
0.4
0.4
0.2
200 400
800 300012001400
200
460
600
800
300012001400
200 400
800 100012001400
(d) MNLI with RoBERTa Large
Yahoo with RoBERTa Large
(f) Yelp with RoBERTa Large
1.0
20.9
Geveraliee "eet
15.3
Generalized Exact Matches
Cansraaiaed
12.5
Score
7.5
5.0
0.0
200
3000 12011000100
0.0
100
200
350
450
500
650
1.00
200
400 Rounds
300
500
600
Rounds
Rounds
(g) MultiRC with Llama2-7B
(h) SQuADv2 with OPT6.7B
(i) SQuADv2 with OPT6.7B
10-
Exort Marchug EM,
Gemeralized
Generalized
2.
100
200
300 Rounds
400
500
100
200
300 Rounds
400
500
600
(i) SQuADv2 with OPT13B
(k) SQuADv2 with OPT13B
Figure 8: Generalized accuracy FI score Exact matches for
heterogeneous clients (Dirichlet = 0.1) setting
29

---

H.2 Personalized Performance Curves
Personalized results on homogeneous and heterogeneous clients with Dir 1.0 and a 0.1 shown in (a) Figures and for BERT Large, BERT Base, DistilBert Base, Albert Large v2.
and 11 for RoBERTa Large, Llama2-7B, OPT6.7B, OPT13B; and (b) Figure I0
Table shows accuracy and FI scores of SPRY and its backpropagation and zero-order based counterparts.
1.0
1.0
penooaaized Test
0.8
Personalzzed Test
0.8
0.8
0.6
Accurscy
Rensonalzedd
0.6
0.4
Fectirg
RedMeedd
Fedtog
Boffle
0.2
Fectigf
Spry
FadLLM
0.0
1200140
Rounds
(a) AG News with RoBERTa Large
(b) SST2 with RoBERTa Large
(c) SNLI with RoBERTa Large
1.0
Persoraline lest
Penooallzed Best
Accufary ACS,)
Persorelized Test
0.6
0.4
0.4
0.4
0.2
0.2
0.2-
200
400
bod Bo0 1000 19001000
0.0
200
400
800
500
1000
12001400
200
400
bo0 see 3000 190010001000
Rounds
Rounds
Rounds
(d) MNLI with RoBERTa Large
(c) Yahoo with RoBERTa Large
(f) Yelp with RoBERTa Large
1.0
35
Pertonaizedd Test
0.0
Percomalized Test
0.6
F1 Score F1.
parsonalized
Matches
25
20
20
15.
FedAvg FedVogl
FwdLLM+
10
Baffic+
FedSgd
Spry
Rounds
(g)
MultiRC with Llama2-7B
(h) SQuADv2 with OPT6.7B
SQuADv2 with OPT6.7B
20
15.0
1.5
Pesonalized
10-
7.5
5-
160
200
300
658
(i) SQuADv2 OPT13B
(k) SQuADv2 with OPT13B
Figure 9: Personalized accuracy / FI score Exacc matches for
homogeneous clients (Dirichlet a = 1.0) setting
30

---

Penoratzed ledt
0.8
Persoraazed Test Accurccy ACEp
0.8
Resoraiizdd Text
0.0
0.8
0.8
0.8
0.4
0.4
0.4
FwILUM+
0.2
ledhogi
Spry
0.2
0.2
FedSgd
200
200
400
800
120012001160
200
100012001400
(a) AG News with BERT Base
SST2 DistilBERT Base
SNLI with BERT Large
1.0
1.0
AcC.
Personation
Personaliee
0.0
0.0
200
400
100012001400
200
400
3000 300012001400
Rounds
Rounds
(d) Yahoo with DistilBERT Base
(c) Yelp with Albert Large v2
Figure 10: Personalized accuracy / score Exact matches for
homogeneous clients (Dirichlet 1.0) setting for variety of language models
Table 5: Personalized accuracy (Acc,) for SPRY and its backpropagation- and zero-order-based counterparts on RoBERTa Large and LLMs. SQuADv2 uses Fl score. 1 shows that higher values are better. The datasets are split with Dir 0.1. Llama2-7B.
7PT..7B. U OPT13B.
SPRY significantly outperforms the best pest-performing zero-order- oraee--asedddddd methods.
Backpropagation based
Zero-order based
First-odder
Difference between
Methods
Methods
Forward Made
performances
best perfo-peffrrningg huckpropaattion method
best-performing
FEDAVO FEDYOG!
FWDLLM+ FEDMEZO BATTLE+
SPRY
order method
AG News
97.76%
97.71%
79.94%
72.69%
60.89%
89.91%
-7.85%
9.97%
SS12
95.84%
95.90%
85.51%
73.26%
64.55%
93.40%
-2.50%
7.89%
SNLI
97.41%
97.57%
74.53%
71.54%
68.55%
83.45%
11.11%
8.92%
MNLI
90.38%
90.03%
72.71%
67.53%
63.58%
80.63%
-9.75%
7.92%
Yahuo
77.93%
67.64%
Yelp
93.44%
97.81%
73.04%.
68.77%
57.78%
85.81%
-11.98%
12.79%
MultRC
50.28%
75.21%
65.91%
N/A
71.20%
-4.11
5.29%
SQuADv2
20.49
21.25
14.43
14.04
12.27
17.73
-3.52
3.30
SQuADv2
13.06
12.49
8.96
9.10
8.17
-3.18
0.78
31

---

1.0
1.0
0.8
penonalized Test
0.8
0.0
Accuracy Ace,
Perooaatzed
0.6
Personalzeed
0.6
0.4
Festing
FedMazd
Fodtog
Soffie
Fectigd
Spry
PM
0.0
Rounds
(a) AG News with RoBERTa Large
(b) SST2 with RoBERTa Large
(c) SNLI with RoBERTa Large
Personaiin lest
Percoralteed Test
Accuracy AC5,
Personalzzed Test
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
200
400
660 800 100012001400
0.0
200
400
600 B00 160012001600
0.0
200
460
660
800 100012001400
Rounds
Rounds
Rounds
(d) MNLI with RoBERTa Large
(c)
Yahoo with RoBERTa Large
(f)
Yclp with RoBERTa I Large
1.0
Pertonaliaed Test
acc-
0.8
Panconalioed Test
f1 Score Fls
20-
Matches EM;
20
0.6
Parsonalized
15
10-
Rounds
Rounds
MultiRC with Llama2-7B
(h) SQuADv2 with OPT6.7B
SQuADv2 with OPT6.7B
12
Personalzed bett
Parsonalteed Test
8-
01
100
200
300
400
500
600
100
200
300
400
500
600
Rounds
Rounds
( SQuADv2 with OPT13B
(k) SQuADv2 with OPT13B
Figure 11: Personalized accuracy FI score / Exact matches for
heterogencous clients (Dirichlet a = 0.1) setting
32

---

H.3 Experiment Variance
Tables 6 and 7 show the variance of running the same experiments thrice with the random seeds 0, 1, and 2.
Table 6: Experimental variance (+) for SPRY and its counterparts on RoBERTa Large.
Backpropgation-ben based
Zero-order-based
First-order
Methods
Methods
Forward-mode AD
FEDAVG
FEDYOGI
FWDLLM+ FEDMEZO
BAFFLE+
SPRY
Acc,
Accp
Accg
Accp
Accy
Accp
Acco
Acep
Accy
Acep
Accg
Accp
AG News
0.51% 0,44% 0.26% 0.21% 0.73%
0.71% 1.34%
1.27%
0.68%
0.51%
1.16%
1.08%
SST2
0.37% 0.47% 0,43% 0.37% 1.26%
1.14%
1.13%
1.06%
0.41% 0.45%
0.77%
0.95%
SNLI
0.45%
0.39% 0.17% 0.11% 0.82%
0.67%
0.74%
0.65%
0.84% 0.78% 0.51%
0.45%
MNLI
0.63%
0.58% 0.29% 0.24% 1.39%
1.25%
1.98%
1.85%
1.15%
1.01% 1.45%
1.32%
Yahoo
0.24%
0.33% 0.53%
0.47%
0.47%
0.44%
1.06%
0.93%
0.59%
0.48% 0.99%
0.85%
Yelp
0.22% 0.25% 0.36%
0.27%
0.54%
0.49%
0.82%
0.66%
0.83%
0.71%
0.76%
0.61%
Table 7: Experimental variance (+) for generalized (Acc, for MultiRC/ F1, for SQuADv2) and personalized (Acc, for MultiRC Flp for SQuADv2) accuracy or Fl score for SPRY and its counterparts. o = Llama2 7B. * OPT 6.7B. O = OPT 13B.
Backpropgation-ben based Methods
Zero-order brae-odder-based based
First-order
Methods
Forward-movd AD
FEDAVG
FEDYOGI
FWDLLM+ FEDMEZO BAFFLE+
SPRY
Acng
Accp
Accg
Accp
Accg
Accp
Accy Accp
Accy
Accp
Acco
Accp
MultiRC
0.58%
0.43%
0.34%
0.29%
0.89%
0.74%
N/A
N/A
1.34%
1.02%
0.97%
0.81%
SQuADv2 1.73
1.07
1.42
0.84
2.17
1.87
1.78
1.51
2.78
2.01
1.76
0.99
SQuADv2
0.86
0.43
0.61
0.45
1.16
0.97
1.03
0.87
1.34
1.14
0.97
0.81
33

---

I Proofs
I.1 Basics
Server Update. The server update of SPRY uses adaptive optimizer FEDYOGI, However, to simplify the proofs without losing generality, we the server update of FEDADAM (2 FEDADAM has the exact update rule as FEDYOGI but without a sign function in its calculation of the second moment of the gradients (See Algorithm 2 of AFO [25]
Hence, the server update SPRY
(r)
w
(1
y trainable weights w(() [d].
VAr)) +
A (r) is the square of accumulated gradients from all clients. A(r) is the second moment of A (r)), is a small positive real number, to prevent division by zero errors. Note that we are assuming flattened weights R without the loss of generality.
With SPRY, the aim to solve the following optimization problem:
M
min wERd
f(w)
1
Σ
Fm (Wm);
(8)
where m E (m-1)d M
+ 1, ad (um) = (rr.g)-D... Im (Wi (x, v))l is the objective function, M
Dm is the dataset, and fm the loss function of a client m [M].
Accumulated Gradients. With SPRY, each client trains a subset of weights Wa.. In a low participa- tion rate setting, each Wan is only trained by one of the participating clients from the set of available clients M. Although we make our analysis more generally applicable by showing multiple clients training the same Wms.
The true global gradients can be written as,
vf(w)=
Σ VF(Wm)
m
E
(m 1)d
+1,
md]
.M CM
(9)
M
M
M
mEM
where M is set of clients training the subset of the weights Wm- M is the size of M.
Client Update. The directional derivative of Forward-mode AD is denoted as fo (x, v)).
where Rd is the random perturbation of weights w and (x, y) are sampled from dataset D. For each client m, through Forward-mode AD, have V fm (Wrr, Umi (z, y)) (V fm. (Wmi (x, Um)) to estimate the true gradient F FF.(W..) (Wmm):
E Dm V (w Umi Dm))
1 DK
Σ
E Ba..
(w.y) V (wmi (z.,
(2,5))>DDn
1=1
(10)
E Dm
Fm (Wm, Umi Dm))
V (Wm)
Theorem 1 of FGD 23
(11)
Here, the expectation is under the randomness of sampled data D, and random perturbation K is the number perturbations per batch. SPRY uses K by default, but here we aim to make our analysis more general the impact of K various properties of SPRY.
I.2 Assumptions
We will be using the following assumptions to derive the first and second moment of the forward gradient and to calculate the rate of convergence of SPRY, Assumption I.1 (Smoothness). The gradient of function F,, is L-Lipschitz,
V (wi) V.m. (w>)l < L|]wi wall: Vm [M] and W1, we gd
34

---

Assumption 1.2 (Bounded Global Variance (Assumption 2 in AFO (2 The variance of the global objective function bounded by og as
M M
1
F [w]]]
V [V(m))]] 112 oog Vm E [M], E I2 and Vj [d].
where Vm is the true gradient of client m. As defined earlier, m E
[m-1)d M
+ 1,
md M
Assumption I.3 (Bounded Gradients (Assumption in AFO (25 The function fm (w; (x, y)) has G-bounded gradients such that for any client m E [M] weights w E IRd, and sampled data (3r, v) Dms have
Im (w; (z, w)l, < G, Vjeld
I.3 First and Second Moments Forward Gradients
We first prove that in SPRY, estimation of V by the accumulated forward mode gradients V across all clients of an arbitrary round r, depends on the heterogeneity across client datasets. Statements on homogeneous data have been proven for FGD ( and MEZO [88 in single-client centralized settings. Here we focus on the specific federated setting of SPRY, where gradients are accumulated differently than in traditional FEDAYG [ In SPRY, we have
[(m 1)d
md
Vf(w) E. vf(
Σ E
V fm (Wir, Umi Dmn)
AM M:m
+
1,
M
M
mEM
(12)
where w represents weights, are their corresponding perturbations, and Dm the dataset of client m.
We omit the round index of the model weights and their perturbations for this section since the same relationship will hold for any arbitrary round r.
Theorem L.4 (Estimation of the Global Gradient). In SPRY, global forward gradient V of the trainable weights w e TRd, with the corresponding weight perturbations v E IRd, computed by participating clients estimated in terms of true global gradient V as,
Ev, (V. sw, u; D)} = [[(w
"p. (a, (a, EmEM, E am BLN..DD fm 21 "p
1
EmeMs E am (wr a +1. "[ +1 (Ir, Mc))
+
M
where C is total number of classes and am,c count, Ore is its Dirichlet concentration parameter. For client m; nmm is the sample count of the
For a class The is its sample
th class, and Dm is the size of the data of client m. The global data is D = EmEm Dm- M is the set of clients training an arbitrary subset of weights, M -Mali Vi [M/d].
Proof. Suppose the global dataset D is defined as combination of all D,, Hence, D = UmE Di
In SPRY, the global forward gradient is defined
M Emem, Vim [1 "1,,)1 1 E Emema V (w) u +1, 24) 2 +1, +1 2 D)
E E F Vf(W, w; D)]
M
(13)
V V im(wwau-ne Eme
+1,d] +1,0) لبييي +1.0] D)
35

---

To combine the gradient estimates from the above equation for all clients m E .,, we first consider the bias bm between the gradient estimate under (a) Globally combined data D and (b) Data D,. of an arbitrary client m.
Measuring the dataset bias. Note that we consider samples (1r, y) of DDn to be sampled from D, since Dm C D.
Enga. fm (W'm Um; D) Ev Drm. fi (Wy:: Um; Dm)) + bm
(14)
Computing the bias term b requires information on dataset distribution. In FL settings, Dirichlet distribution to spread the data into heterogeneous splits is a popular way to simulate heterogeneity (37).
The biased dataset D. is sampled from the combined dataset D using the Dirichlet distribution. This allows us to utilize the properties of the distribution to derive the relationship between forward gradients on global data and on individual local data:
For classification tasks, let's say D has total classes: yi..... .Cc We have ai..... ag as concen- tration parameters of the Dirichlet distribution. The expected forward gradient for the combined dataset can be expressed as a weighted sum of the expected forward gradients for each class C:
E Expry
WiT, Unii D) mm
Re E Umm. (2. D
f (Wis Uai (xt, wel)
(15)
where ne is the sample count of class e. And for the biased dataset Dm sampled with Dirichlet concentration parameters a1...... ac,, the expected forward gradient can be expressed as,
Bu De V Fm (WiT, Um; Dim) Σ
(ome)) Vnn(um Umi (2, vcll]. (16)
[Dml
where nm, is the sample count of class e for client m.
Subtracting Equation T6 from Equation
15
bm EvmrD D
V fm WWm, Um: D) - Bum.D Dm f (Wm, Um: Dm))
(17)
Σ
V V (WTr. Umii (2r, Nc))
(18)
D
1Dml
For any 77 E
(m- 1)d
+ 1, md
we have Eugr. [ Fum Umm Dm)) = Fm((w) from Lemma
M
M
in FGD 23
Plugging in the above result and Equation 18 Equation 14
E Evm:D D
V (Wm, Umi D)
E Eup,
De
Vim (W', Umi Dm) + dm
c The
F (Wm)
E Eym.
Sm (Wir, Pam (z, v.)
(19)
1mml
For 7th row in Equation
13
Emm:D D
V (WT, Uzi D)
1
Σ
Ew
E
Vmn (Wrre Umn DI)
(20)
M
meMi
Σ
(Fm(Wum +
The
E Evm.
Vim(Wm (wi,, Umi (I,, mcl)
D
Dmml
mEM,
c=1
(21)
C
[)((m) + Σ
Tle
a [Dml
E Evm (2.
V Umi (2, yol)]
M
D
MEM,
(22)
36

---

Putting Equation 22 in Equation 13
L. (w, v; D)]
(23)
x/ Vf(w[1, [1 M
1
EmEM, Ee
( n2
E (L..))ED
V fm (w0 1. 1..1 (z. Be))
[D]
Pm
Vf(W [ (1+
M Eme.. E
Ro Daal
E ((D.o)EDD
V
Vmm (t +1, +1 #1 ve))
(D]
(24)
(t) 7
Eme Ec=1
Dr
Den 1m
E(
V fm
"I i "[I # (2, wel)
T
[2
V # ( +1.
Eme M2
Ee
The 27
V fm #1 " -1, #5 (2), wol))
D
M
(25)
Emehi EEsl EmeM, Ecsl
( nge
Pe 1mml
B B(I.DJDD
V (2 #1 "p, 1 (z. vcll
[D]
= fw)
nre
V fm (wr a +1 "[ 3 (cc, Mel)
|D]
M
(26)
Next, we formulate the norm of the forward gradient vff
Lemma I,5 (Norm of the Forward Gradient). Under Assumption 1.2, and at the participation rate of 8, M participating clients training weights w E Rd through random perturbations IRd in SPRY derives the accumulated forward gradient f(w, D, such that,
E = [[w))]]
1+
2(3d + 1I-1- -1
Σ
The
MK
D
D
mEM
220 s) + K 1))
Σ
The
(27)
MK
[D]
[Dml
MEM
where D is the combination of datasets of all M clients, K is the number of perturbations per batch, M is the count of clients training a particular weight subset, o²g is the upper bound of the global gradient variance. For total of C classes in D, we define ne as the sample count of the tt class, Oa Dirichlet concentration parameter for the th class, is the size of the global data. For a client m; Tm the sample count of the th class for client m, and Dm is the size of the data of client
Proof. The proof follows similar style of Lemma 2 in MEZO [88. The difference in our setting is that we have V. fwe u: D) which is an aggregate of V fm (w, v; D) derived from the federated clients, while MEZO has results under the setting of single client.
From Theorem we have,
E..D
Vww (w, ;; D
V + 1 1
EmeE, D
(
E ELWDEED
V (10 [1. ("T " "[. (a, vc)) Sm +1, #1 "[ ye))
T
M
(1))
Vs(w [ ++1 + Eme C
am.ccc
E EiEyy.DD
[2]
Dell
(28)
37

---

where
V I (Wrr, Umm (2r Mol)]
EE
[(V5 (Wiri (Ir (e) Vi.m) Ul
|D|K
(2.9e)~D SE[K]
using Equation T0 The right hand side expectation is also under the randomness of partial perturba- tions. Here, K is perturbation count, and (D) is the size of the combined dataset 'D = UmE[M] mm-
To compute the second moment for client m,,
V I'm (Wir, Uamis
(2,, yc))
f (Wm, Uaii (I, y
Σ
Σ
E
(V (I, u (V (umm (e,))e)) Vy f.
DDaK2
(D iE[K (D) jE[K]
(29)
For simplicity, let and b be arbitrary vectors representing E [m (ar, (e),))]]] and E Vm (Wmi (Z. T.))" respectively.
For the sum over all perturbations, we have two cases (all the expectations are under the randomness of partial v),
1. If # (Occurs K (K 1) times)
E [ ab" (since and By are independent of each other and b don't ton't depend on v))
U) Uym fjmmffmm = [ [ ab"
j.m]
i.
E [e).m²
p.m
(30)
=I ab" I ab?
(31)
2. If j. (Oceurs K times)
E Pimm abT Uyome [ua.m oo. abT Usmu
(32)
Er, [U]]] (a, b)
(33)
For all such v; with E [K].
E E. = 3d Sym (y
(34)
2d abT + d 1 l
(35)
Plugging in the results of the above two cases in Equation 29
E V
Fmm(ue (Wm, Vmm (xe, Uc))
V I (Wm, Uni (x, Uc))
1
Σ
[KK - 1)ab" + 2d + dK I- bTb]
(36)
22
(D.
(2.0.~D)
Σ
(2d + K - 1)B [V fm (Wmi (x, vc)) smn (wii (z,
D
(n.Je~Di) (T.J_~D)
+ d. R. Vmm(wmm (e, SmnTmmmmmmm (Wii (I, v.))
(37)
Here, the randomness is under the sampled data. For the sum samples of D, we have two cases,
1. If (I, No) A (I. V.). (Occurs [DI([P) 1) times)
E Vm fmmmmi (z, v V fm (wmi (I, T.) V.F. (Wm V Fmm(um))
(38)
38

---

2. If (r, ye) = (T, T.). (Occurs
D times)
B fm (Wm: (z, Be)) (umm (I, Vm (Wm)) VFm (Wm) ) (wm)) (39)
Combining both the cases, we get
(( V (wmm (x, yc)) V amm (Wis (z, (x, 1 Fm (10m) Fm (wm) )²m
T + (D) {(wm)
(40)
Plugging in Equation 40 in Equation 37
Bu (2.80) V Fm (Wm: Uni (I, Mc.) V fm (om(umm (Wi, Uais (I, Uc))
(2d + K L)I Fm (Wm) V Fm (Wix) T + [(wm)
D K
+ + [D] (PD)-I|VPnn (wm (F(wm)))
(41)
(2d + K -
1
Vm (wm) Vmm (wm)² ( +
E( E(wm)
K
[D]
(IVV (wm))) + D tr (I(w)))
+
(42)
K
3d + K
(awa)sDlll Vm (wmi (xr, yc))ll
(43)
K
Hence for client m, the expected norm of forward gradients under randomness of perturbations is,
Evwr (2.g.. V fmmmumm (Wm, umi (I, v=))11 3d + K
1
((,,,o))~ ~ {mmm (I,, v.))
(44)
K
3d K
V Fm (wm))1
(45)
K
For 7th row of f( D)
of Equation 28
V (wi se- +1, " +1, #1
D)|
E
Vi +1, #
Σ
am
E EI
V V (w
( معر 1. 1 " +1, # (x,, He
(M)² mEM, <E(C
[D]
[mm
(46)
The left-hand side expectation under the randomness of sampled data D and subset of random perturbations v.
Plugging in Equation 45 in the above equation and using 7
[i- 1)
1,
id AV
v f (w, 2g
D) || =E [V(
(3d + K -1 (M)2K
Σ IIV (w:))l
Σ
Te
e
i m
mEM.,
0E(C]
(47)
39

---

Since ff. AM C M are independent of each other, we can compute the norm of ff as follows,
E Vf(w, n; D)]1 B /
(3d+K--) K- -
Σ MEM
V Fm (Wm)
Σ cE|c]
T1c
mm
(M)2K
D
[Dml
(48)
E [ 3d K K- (M)K
Σ
V (Fm((u)) (Wm)
Vf + (
Σ
The |D]
(Dml
mEm
CE|C1
(49)
= / +
2(3d -)
Σ
V Fmmmm)) - ([((m))
Σ
nCe [
m0m
(M)2K
[Dm
mE.M
CE(C]
223d K K-1)
2
Σ McM
V (((m))
Σ
The
amc (Dml
(50)
(M)*
[D]
eE|C]
Using Assumption 1.2 and limited participation rate of s,
220g(1 8) s3d + K - 1)
2
E V [) = E| E/VS
E Σ
Σ
e00
MK
[D]
[Dml
mEMoc(c
2(3d + - 1)
2
Σ Σ
(51)
MK
[D]
ID
mEM CE(O)
Rearranging the terms, we get,
VV.plV/ll,, u;
B B(w)]]
1+ 2(3d + + -1
Σ
The D
nem.ooee [Dmml
MK
me.
ME|ocE(C
20
5)d + K -11)
Σ Σ
Toe
)
(52)
MK
|D]
[Dml
mEMocE(C
L.4 Convergence Rate of SPRY
The general template for the convergence analysis of SPRY is similar to FEDADAM. hence we will follow Theorem of AFO (25]. Our aim here to highlight the differences in treatment of our gradient estimator V fm Vm E [M), and the aggregate global gradient f as shown Equation Theorem I.6. Under the assumptions on L--smoothness (Asmp 1.1 bounded global variance g of accumulated gradients (Asmp]I.2 and bound on gradient magnitude G (Asmpl I.3) and the following conditions on the local learning rate Ne.
min {0
VB&G
/BEVI O26
MK
0
(53)
3rGl3d + K. -1) 1) EmEIM ExEloy un
SPRY satisfies the following bound,
E, (w) < S(w(0)) [En[((u(RR)]
min 0SPSR
nR
VI 3T-35mm
(3)(1-))33d s))dss(3dd))) + K -1)
2
Σ Σ
2 a
(54)
272
33
MK
cEM
40

---

where R is the total round count, w E R are the trainable weights, E R are the random perturbations, K is the count of random perturbations per batch. n the global learning rate, is adaptability hyperparameter, & is client sampling rate. The rest of the symbols are defined in Theorem 1.4
Proof. As shown in Equation 7 an update of the model weights w at server-side in SPRY looks like,
A
((++1) w(r)
(55)
V +
Using Assumption L.] and then the server-side update rule, we have
f(wlrt)) s f(w(r)) + V (V(w(()), ( 1
I.
[1(((+t))
11
(56)
=((w(() + n ( V
(r)
n L
(1(r)
(57)
VA() +
(VA!T +
iE|d
Taking expectation over randomness of round and simplifying the terms.
(r)
+ T E,,
(A
E, B.[f(w!rt)))) s f(w(r)) + ()(w(()) E.
VBA +T
2 EE|d]
[(VA +
R2
A (r)
A (7)
(l( By
(58)
VA(r) +T
3AAC-D +T
R2
Bounds for R2 are derived in the exactly same manner as "Bounding R2" in Theorem 2 of FEDADAM (25]
G
R2 v1 ByE
(59)
LV +
j=1
Bounding R has a different treatment due to the distinct aggregation strategy of SPRY:
A)
R1
V (() E,
(60)
BA
Since (r) piece-wise made of aggregations of forward gradients for several parts of the model weights, as shown in Equation 12] squared norm,
we first center each gradient piece for the computation of the
ne f(u(t) ( v (A (D1))
V.((w())))
Rı
(61)
Using ab S (a² + 12/
V )+
118
ne E, | / ,",
< R1
(62)
2 jEld
2
BA
Using the result of Lemma 1.5
V w (r) 12
20
s) + K -K1)
2
Ri <
17
+ N1
Σ
The
e
2
%EI4
(r-
+
2
MK
MEM CE(C)
[
D
1+
2(3d + K
Σ
Σ
7m
Er [vll
MK
D
Dmal
mEM
CE(C)
(63)
41

---

Putting R1 and Rg bounds in Equation 58
V (1c(()))
E, (f
f(w(1))
mp
Σ
2
3A 11- +
jeld
Me
20
s3d + K
Ne |D]
e.c 1Pm
2
2
MK
mEM cE|C
Me
1+
2(3d + K - 1)
Σ
Σ
The
Bem,Clte (Dml
Er
[f(
2
MK
|D]
MEM CE|C1
B2G
Σ
ET
(64)
2
iE|4]
+
T
dd]
A +
Summing over r= 0 to R and using telescoping sum, we get
[f(w(1)]
[x(f(w(())) s f(w(())
me
BA/r-11 +
r=0 jE|d
me R
20 (1 8) (3d + K - 1)
Σ
2
+
2
MK
D
1Dml
mEM CE(C]
14
(2(3d + K --1) 1)
Σ
T1e (D]
Bp
V }r
f
MK
1Dmll
mEM cE|CI
L
32G
R-1
+
(65)
2
72
KE(M
R4
Bounding R4 follows a similar derivation as "Bounding R₁ "
R
(A (
V (V)(()),,)))
R4 E
E
(66)
+ +2
A + +
r= veld
p=
KE(d]
R 1
v f((w())) 2,1)), D)
<
E
(67)
72
r=
Using Lemma 1.5 once again,
R4
2 2
1+
2(3d [2(dd+K +K KK
Σ
The
2
E
R-1
S
(1)
MK
[D]
1Dmml
mEA cE|O)
v/2
220 - R)d + K - 1)
Σ
Σ
(68)
?
MK
(
D m
MEM cE
42

---

Updating Equation 65 with the bounds of Rs,
R-1
V Vf(w(r)112
[x(f(w(()) 1 f(w(())
m
2 r=0 jEld 131 +T
me R
21 - 8) s3d + K 1
E mEMceic1
The
2
+
2
MR
[D]
[Dml
+
Me
1+
(2(3d + K --1)-
ΣΣ CE
a
R-1
E,
f(
2
MK
[D]
1Dml
ME.
+
L
M/ - B2G
1+
2(3d + K - -1)
Σ
The [
7mm
2
R
E,
1 V / (1))
72
MK
1Dml
mE
r=
112
1BG)
T
20
R(3d + K - 1)
Σ
nhe
+
72
MK
(
1Dml
mEM
(69)
Rearranging the terms,
R-1
(1 V
[(wl0)) En (z)(
I
BA +
17
JEId
21 s)(3d + K 11
Σ
The
7lm 1Dml
MK
[D]
mEM CE|C)
+
1+
2(3d + (11)
Σ
Σ
The [D]
Ex
(f(
MK
|Dml
MEM CE|C
+
nL +
2VI B&
Ne
1+
2(3d + K -1)
Σ
nhe [D]
c
R 1
Σ
Ex
V
7
72
MK
Deml
cE(C
r=0
v1 32G
N/8
Ro
s) + K -
Σ
T1e
+
+
72
MK
DD]
Dml
CE
(70)
Getting a lower bound for the left hand side term through using the fact (r--1) S nCG from Theorem 2 of AFO 225).
R-1 d E, (f
R-1 d E.
(f
R
2
min USrSit
E,
E,
324 111 +T
VBnKG++
VB +
(71)
43

---

R
min / (10 & /(wlo)) -ER(((u((1))
OSPSR
+
2
MeL 2T²
+
vI vT- BrGmm)
Roo(1 s)(3d K -K-1)
Σ
cm
2
73
MK
[D]
12ml
mEM
cE|c
1+
123d + K -
Σ Σ
1ce
E
R-1
MK
(D)
1Dml
mEM CE|C)
nT. 2/1 B2G
ne
1+
2(3dd + K
Σ
Σ
E
R
)1
72
MK
[D]
DDml
mEMocEO)
(72)
Considering the coefficients of E [[((r)))] terms, conditioning on following inequality,
R
MeL
2n 1 32G
+1)
(
2(3d + K-KK- 1)
E
Σ
The
2
V3mMKK+ +T
73
MK
D
1Dml
mE(M) EE(C)
(73)
We get the following condition the local learning rate,
Te min {o
(V
72
1 (V22G) VB&G)
3
144
.0 V32v1 - 32G2
VBenGL
MK
(74)
32G(3d K --1) 1) EmElM EeElcy
Bm 1m
[D)
for the following bound the gradient norm,
(( f(w(())
(xll(((())))
min |r 0SrsR
nR
+(2 2+
MeL
VI 5gGm
o s3 + - 1)
Σ
Σ
2-
33
MK
mEM CE|C] [D]
Dm
(75)
44

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims on (a) SPRY consuming significantly less memory than its backpropagation-been baaag counterparts, (b) SPRY achieving higher accuracy than its zero- order counterparts, and comparable accuracy to the backpropagation- based counterparts, (c) SPRY running faster in terms of wallclock time to convergence than its rero-order counterparts are proved Section 5.2 5.1 and 5.3 respectively. The theoretical claims are discussed and proved in Section and Appendix | The claims on direct application Forward- mode AD in an FL setup being inefficient and inaccurate are supported by
experiments shown in Appendix G Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings. It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? [Yes]
Answer:
Justification:
The limitations and potential paths for future work are discussed in Appendix
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
45

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification:
The assumptions and proofs are listed in detail in Appendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The details about implementation and experimental setup are given at the onset The pseudocode is in Appendix E and each part of SPRY is explained in
of Section 5
Section 3
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility.
46

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: The source code is currently available as supplementary material accompanying this work and an anonymized version is also available here
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
http://Seeuesbsssssooooooioooooo p
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurIPS code and data submission guidelines (https
CodeSuumissseeeeeeee ssioiooicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental setting is described in Section Full details about datasets, their splits, and hyperparameters are given in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Experimental variance related to the main results in the paper, across differently seeded runs is given in Appendix H.3
Guidelines:
The answer NA means that the paper does not include experiments.
47

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources are discussed at the onset Section Details on memory consumption are given in Section 5.2 Time to convergence is discussed in Section 5.3
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: Our work does not involve human subjects. All the experiments are done open-source datasets, which are widely used in the natural language community. Moreover, the goal of SPRY is enable crowd-sourcing compute resources to finetune L.LMs, without making private data accessible to a third party.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
48

---

Answer: [Yes]
Justification:
The positive and negative broader impact is discussed in Appendix D
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deg.oymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We have not created any new models or new data. Our main goal is to finetune the existing pre-trained language models. The datasets have used open-source, and free-to-use. The license information of datasets is given in Appendix B
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors
should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Licenses for all the datasets stated in Appendix Our work is within the usage bounds of each of the datasets. Dataset URLs are also provided under the references section.
Guidelines:
49

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paaassttssppersseaaes paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
The source code (also available as a supplemental material) of SPRY comes with detailed README md, explaining how to the presented experiments of our algo- rithm and its counterpart baselines.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing and research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
50

---

Answer:
[NA]
Justification:
Our work does not involve research with or on human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
51


Reviews:
[{"review_id": "G8nmBss9wN", "reviewer": "Reviewer_RCwU", "summary": "This paper introduces a forward-mode AD federated learning algorithm (SPRY). They use SPRY to finetune LLMs and demonstrate a low memory footprint compared to backpropagation-based federated learning algorithms. The authors also derive SPRY’s convergence rate and provide theory behind why SPRY’s global gradients are unbiased for homogeneous data distributions across clients. The authors empirically evaluate SPRY on 8 language tasks as well as perform ablation studies.", "questions": "* Are the authors aware of AsyncFGD (https://openreview.net/pdf?id=45RBLZBJid)? They also implement FGD in parallel by splitting weights across workers. They do not reach the model and experiment size of this paper, but there does appear to be some similarity in the approach. Would it be possible to identify these differences? \n* A general limitation of FGD as a learning method is the variance of the gradient estimator with increasing dimension of the parameter space. How does SPRY overcome this challenge? Is there a theory that updating individual layers on different clients reduces this variance?\n* For per-iteration communication, is there any advantage to having clients, since the server needs to generate all the tangent vectors and then update all the weights of the model. Since there is no matrix multiplication, perhaps there is no need for a GPU on the server-side? Also, generally what would the communication and compute cost be for the server?\n* What is the current state-of-the-art in federated learning?", "limitations": "Authors have done a good job at describing limitations and memory usage of their approach.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* The empirical evaluation on eight language tasks while testing multiple different language models is a key strength of the paper. \n* The paper is organized clearly with nice figures and a clear structure to the sections. The empirical evaluation section is particularly well structured. \n* The ablation studies further strengthen the results of the paper.\n* The section of peak memory consumption is clear and highlights the superior memory performance of both zero-order and first-order forward mode AD compared to backpropagation.\n* Splitting forward-mode across multiple layers seems a novel way to perform federated learning.", "weaknesses": "* The baseline FWDLLM is described as a zero-order-based approach but upon reading this paper, it seems like it uses a first-order forward mode AD update rule in the federated setting (https://arxiv.org/pdf/2308.13894). Why does the paper categorize this approach as a zero-order-based method? Also, Figure 10 in that paper shows a better performance of all the baselines on the different data sets. Why is there such a significant discrepancy? E.g. RoBERTa-large gets a performance of much greater than 80 % for FwdLLM in the original paper for AG News and is over 60 % for Yahoo.\n* Equation (4) is challenging to interpret. Why is the summation over classes? What do the subscripts in square brackets mean? Does the paper define data distribution only according to how class labels are distributed among the clients? This did not appear obvious in the paper.\n* The implications of theorem 4.2 are not clear. Is this bound related to the variance of the gradient estimator? Generally the variance is related to $E[x^2]$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# W1. Comparison to FwdLLM\n\nFwdLLM[1] shows the equation of finite difference (which involves only function evaluation and no gradient computation) in their paper’s Eq1. And their experiment scripts[2] refer to the user of finite differences as well. Hence, we categorized FwdLLM as a zero-order method.\n\n[1]FwdLLM: Efficient FedLLM using Forward Gradient (ATC 24)\n\n[2]FwdLLM Code: https://tinyurl.com/2cwuhjfu\n\nThe results differ from our Table 1 due to variations in participating and total client numbers in each FL round. FwdLLM[1] originally used 100 clients per round, but we used 10 clients due to limited compute resources. Previous studies[3, 4] have also used 10 clients per round. This results in an accuracy gap of approximately 3.06-8.17% for AGNews, Yelp, and Yahoo datasets.\n\n[3]Ditto: Fair and Robust Federated Learning Through Personalization (ICML 21)\n\n[4]Federated Optimization in Heterogeneous Networks (MLSys 20)\n\n# W2. Interpretation of Eq4\nIn Eq4 we quantify the gradient estimation bias caused by heterogeneous data across clients using Dirichlet distribution, a popular method in FL to simulate heterogeneity in classification. It divides data using $\\alpha_c$, higher values indicating more samples of class $c$.\n\nSummation over classes\n\nThe bias of gradient estimation is defined by difference of gradient expectations over global data randomness. This difference can be computed per class and summed over all classes to get the total difference, as shown in Eqs. 17, 18 of Thm H.4 proof, Adx H. That’s why we sum over classes. \n\nSubscripts in square brackets\n\nSubscripts under the model params $w$ and perturbations $v$ show the parameters a client $m$ uses to perturb and fine-tune, with weight perturbations $v_{[a, b]}$ being used to fine-tune a weight subset $w_{[a, b]}$. \n\nDefining data distributions\n\nWhile our homogeneous data experiments are not limited to classification, it’s common in FL works to define heterogeneity by a Dirichlet distribution on classification tasks[5, 6]. Dir distribution simulates real-world data distributions since a client usually has data distributed disproportionally across classes.\n\nHence, to express the bias of gradient estimations for theoretical analysis, we used Dir distribution. More details are given in proof of Theorem H.4.\n\n[5]Federated Learning Based on Dynamic Regularization (ICLR 21)\n\n[6]Personalized Federated Learning under Mixture of Distributions (ICML 23)\n\n# W3. Implication of Theorem 4.2\n\nThm 4.2 derives Spry's error bounds based on: a) total comm. rounds, b) perturbation size for forward-mode AD, c) perturbation count per iteration, d) data heterogeneity, & e) number of clients training a subset of weights. The standard form of error bounds in case of non-convex objective $f$ is L2 norm of $\\nabla f$[7,8].\n\nWe show that increasing training rounds or number of clients training the same parameters decreases convergence error, while decreasing training parameter count or decreasing data heterogeneity also reduces convergence error.\n\n[7]SCAFFOLD: Stochastic Controlled Averaging for Federated Learning (ICML 20)\n[8]Adaptive Federated Optimization (ICLR 21)\n\n# Q1. Comparison to AsyncFGD\nAsyncFGD[9] targets an efficient resource utilization for gradient descent with gradients derived by forward-mode AD (also called “forward gradient descent”) by parallelizing jvp computations across various iterations on a single device. We point out the major differences as follows: \n\nSpry aims to improve memory efficiency in fine-tuning large models through forward-mode AD, while AsyncFGD aims to enhance resource utilization by parallelizing jvp computations across multiple iterations.\n\nLimited model sizes for AsyncFGD\n\nAsyncFGD is designed for a single client and has limited experiments for models with 13M parameters. Applying AsyncFGD to larger models like RoBERTa Large (355M params) for a single client would fail, similar to the failures of FGD in preliminary experiments (Sec 5.4 and Adx F on “The Importance of Splitting Layers”). \n\nAsyncFGD doesn't address high-dimensional parameter space issues, while Spry does by splitting the space across multiple clients. AsyncFGD's resource utilization can be applied to each client in Spry, making them orthogonal.\n\n[9]Accelerated On-Device Forward Neural Network Training with Module-Wise Descending Asynchronism (NeurIPS 2023)\n\n# Q2. Variance of the gradient estimator\nIt is true that variance of estimated gradients increases with param dimension. This led us to the development of Spry, which uses participating clients to train a small set of parameters, reducing the impact of gradient estimation variance faster in FL rounds.\n\nThm 4.2 shows that increasing the number of clients training the same layers and reducing the trainable parameter count decreases convergence error due to the reduction in the impact of global gradient's bounded variance by smaller perturbation size and larger number of clients training the same layer.\n\n# Q3a. Per-iteration communication advantages\nSpry's per-iteration communication variant requires clients to derive jvp values from local data, while the server generates tangent vectors using random tensors from a normal distribution. Once clients provide jvp values, the server multiplies the average scalar with the tangent vector. The reviewer points out that there's no need for a GPU on the server side, which is not an assumption for Spry. \n\n# Q3b. Server communication and compute cost\nPlease refer to the global “Author Rebuttal”.\n\n# Q4. Current state-of-the-art in FL\nAdx A provides an overview of recent methods for fine-tuning large language models in federated settings, as it gains popularity in this context. Due to space limits, we refer to recent surveys[10,11] which encompass progress made on solving FL challenges.\n\n[10] Recent advances on federated learning: A systematic survey (Neurocomp 24)\n[11] Federated Learning for Generalization, Robustness, Fairness: A Survey and Benchmark (TPAMI 24)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "34yDln0zYn", "reviewer": "Reviewer_KHem", "summary": "This manuscript is focused on the memory-efficient federated finetuning of LLMs. The author first uses Forward-mode Auto-Differentiation to reduce memory. Then, the author observes that merely substituting backpropagation with Forward-mode AD in FL scenarios often results in poor accuracy and computational inefficiency. To address this challenge, the author recognizes that Forward-mode AD operates more efficiently and yields better gradient estimations when the trainable weight count is minimized. Therefore, SPRY only assigns each client a responsibility to compute gradients for only a subset of trainable weights. The experimental results show that the memory overhead can be reduced significantly. The topic is of interest and the presented numerical results seem, indeed, promising. However, there are still some questions/comments/suggestions for the current version of the paper, please refer to my comments under Questions.", "questions": "* What is the computational overhead situation in each step (e.g., the training time for each gradient step)? A comparison of computational overheads is suggested in the experimental section.\n* Is each client's trained layer the same from start to finish? Or are the layers reassigned each communication round?\n* It seems that the layers for each client are fixed at the beginning. If we consider the dynamic distribution of layers, it might alleviate the issues caused by heterogeneous devices.\n* Communication overhead is an important metric in federal systems, and it is suggested that the authors give the communication data size in each communication round.", "limitations": "The authors give a limitation analysis in the checklist section.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper performs a number of experimental verifications as well as theoretical proofs.\n* The experimental results show that the method is enabled to significantly reduce GPU memory consumption while ensuring a competitive performance.\n* Additionally, the appendix contains valuable results.", "weaknesses": "* This manuscript does not analyze the computational load of Forward-mode.\n* There are still some issues with the experimental setup, such as the need to consider the impact of the number of clients on performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **W1. Computational load of Forward-mode AD and**\n### **Q1. Computational overhead**\n\nCompute cost of each client and the server is given in “2. Computation costs” under the global “Author Rebuttal”. We also state time per iteration cost here followed by the result analysis:\n\nThe computational overhead (in terms of seconds) is measured on Nvidia 1080ti and RTX8000 for RoBERTa Large and Llama2 respectively. For a fixed batch size of 8 is as follows,\n| Method | Time per iteration (in seconds) for RoBERTa Large with LoRA on AGNews | Time per iteration (in seconds) for Quantized Llama 2 7B with LoRA on MultiRC |\n|---|---|---|\n| Backpropagation [First-order] (used in FedAvg, FedYogi, FedSGD) | 0.1207 | 0.0167 |\n| Zero-order Finite Difference [Zero-order]  (used in FedMeZO) | 0.0843 | 0.0115 |\n| Zero-order Finite Difference [Zero-order]  (used in FwdLLM+) | 0.7593 | 0.2383 |\n| Zero-order Finite Difference [Zero-order]  (used in Baffle+) | 2.683 | 1.3175 |\n| Forward-mode AD for all weights [First-order] (used in FGD) | 0.3215 | 0.1237 |\n| Forward-mode AD for 1/3rd of the weights [First-order] (used in Spry) | 0.1036 | 0.0209 |\n\nTo summarize, while backpropagation is faster than forward-mode AD due to the column-wise gradient computation of Pytorch, forward-mode AD proves to be more feasible for edge devices due to its superiority in memory footprint reduction. And finite differences are more unstable, resulting in the methods based on the finite difference having to evaluate 20 to 100 perturbations per iteration, while achieving sub-optimal performance compared to Spry due to numerically unstable gradients. The explanation is as follows:\n\n&nbsp;\n\n*Compared to backpropagation-based methods and FedMeZO*, forward-mode AD is slower than backpropagation due to how the gradients are currently computed (discussed in Section 5.3 of our work, and further elaborated in Section 3.1 of [1]): Forward-mode AD computes \\texttt{jvp} column-wise, while its counterpart \\texttt{vjp} in backpropagation is computed row-wise. The column-wise computation incurs time overhead if the number of trainable parameters >> output size (loss is scalar, hence the output size is 1) – which is the case for neural networks.\n\nAlthough Spry’s forward-mode AD is slower than backpropagation and FedMeZO’s finite difference, its chief advantages lie in its much-reduced peak memory consumption (shown in Section 5.2) with comparable time to convergence (shown in Section 5.3). Spry has consumed 27.90% to 86.26% lower peak memory compared to backpropagation because it does not need to store activations of the entire forward pass. And Spry is 1.34-2.98x faster than FedMeZO to achieve convergence due to better gradient approximations.\n\n&nbsp;\n\n*Compared to zero-order finite difference methods FwdLLM and Baffle*, forward-mode AD is faster per iteration. FwdLLM uses a variance control mechanism to pick the perturbations which have smaller variance with the previous round’s aggregated gradient. This adds an overhead in FwdLLM for each iteration to sample and pick the appropriate perturbations. \n\nThe zero-order finite difference gradients derived from Baffle, and FwdLLM are numerically unstable due to the accumulation of truncation and round-off errors, resulting in having to take average of the gradients derived across 20 to 100 finite difference evaluations per iteration. \n\nCompared to 20 to 100 evaluations of finite differences per iteration, forward-mode AD can get a better approximation of the true gradient in 1 forward pass per iteration given that it has to perturb and train a fraction of the trainable weights of a large neural network.\n\n[1] Automatic Differentiation in Machine Learning: a Survey (Baydin et al, JMLR 2018)\n___\n### **W2. Impact of the number of clients on performance**\nResults on scaling up participating client count are discussed in Section 5.4, under “Effects of participating client count” (more details available in Appendix F “Effects of the Number of Participating Clients per Round”). \n\nThe experiment involves raising the participating client count from 10 (participation ratio of 0.1) to 100 (participation ratio of 1.0), for a total of 100 clients on the SST2 dataset. We observe that as the participation ratio increases from 0.1 to 1.0, accuracy increases by 2.94%. Furthermore, to achieve an accuracy of ~85%, Spry needs only 150 FL rounds for participation ratio 1.0, compared to 500 FL rounds with participation ratio of 0.1.\n\nThis shows that as more clients train the same model parameters through Forward-mode AD, the gradient estimates are closer approximations of the true gradient, leading to a higher accuracy and faster convergence of Spry, which is corroborated in Theorem 4.1.\n\nIn fact, as the number of participating clients increase, not only there would be more clients to train each layer (improving accuracy as discussed in the initial paragraph), but each client would have to train fewer layers (reducing computation and communication cost for all the clients).\n___\n### **Q2 and Q3. Layer assignment to the participating clients**\nEach client’s trained layer is not the same from start to finish. At the beginning of a communication round, trainable layers are randomly assigned to the participating clients of that particular round (Line 6 of Algo 1, Appendix E).\n\nIn our cross-device FL setting, we do not make any assumptions on which clients will be participating for a given round, and for our main experiments (in Table 1 of Spry), the participation ratio (#participating clients/#total clients) is 0.01 to 0.1. Therefore, it’s unlikely that a client will get randomly picked for subsequent rounds, and it’s also unlikely that the same client will be randomly assigned to the same layers for multiple rounds.\n\nHence, as the reviewer has suggested, we are indeed doing layer assignment dynamically.\n___\n### **Q4. Communication overhead**\nPlease see our reply under “1. Communication overhead” in the global “Author Rebuttal”.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "4i8ssIrx5E", "reviewer": "Reviewer_sp7C", "summary": "This paper introduces SPRY, a federated learning (FL) algorithm designed to finetune large language models (LLMs) on resource-constrained devices by addressing the excessive memory requirements of traditional backpropagation methods. SPRY tackles the challenge of high memory usage from intermediate activations by utilizing Forward-mode Auto-Differentiation (AD) and splitting trainable weights among participating clients, allowing each client to compute accurate gradients with reduced memory. Theoretical analysis shows SPRY's global gradients are unbiased for homogeneous data and provides a convergence rate dependent on FL rounds and data heterogeneity. Empirical results demonstrate SPRY's efficiency, reducing memory usage by 1.4-7.1× and achieving faster convergence and higher accuracy compared to existing methods. This makes feasible the finetuning of LLMs on mobile and edge devices, significantly impacting FL deployments.", "questions": "* **Q.1.** Will the SPRY algorithm work in a large scale setting of 1k+ GPUs?", "limitations": "n/a", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* **S.1.** The proposed SPRY algorithm tackles a difficult task and is backed up with both empirical and theoretical results.\n* **S.2.** The paper is well written and the illustrations are helpful.\n* **S.3.** The empirical experiments are conducted on multiple datasets, models, and hardware configuration, while showing that SPRY outperforms previous works.\n* **S.4.** The paper provides an anonymous code repository of SPRY.", "weaknesses": "* **W.1.** While SPRY shows promising results with a relatively small number of clients, the scalability to a larger number of clients, which is typical in FL scenarios, is not thoroughly investigated. The communication and computational overheads associated with increasing the number of participating clients, especially in terms of synchronization and gradient aggregation, are not discussed. This raises concerns about the practicality and efficiency of SPRY in large-scale deployments.\n* **W.2.** The empirical evaluation, although thorough in certain aspects, is limited in terms of the variety of datasets and models used. The evaluation focuses primarily on specific language tasks and a narrow range of LLMs. A broader evaluation including more diverse datasets and model architectures would strengthen the generalizability of the findings. Additionally, the impact of SPRY on tasks beyond natural language processing, such as computer vision or other modalities, is not explored.\n* **W.3.** While SPRY does outperform existing zero-order methods, it converges slower (time-wise) compared to traditional gradient based methods such as FedAvg.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **W1. Scalability to a larger number of clients**\nPlease refer to our answer to reviewer KHem under “W2. Impact of the number of clients on performance”.\n\n&nbsp;\n___\n### **W1. Communication and computational overheads**\nFor communication and computational overheads, please see “1. Communication overhead” and “2. Computation costs” in the global “Author Rebuttal”. \n\nTo summarize, \n\n(a) As the participating client count increases, both per-epoch and per-iteration modes of Spry have a lower communication cost are communicatively lower-cost than classic backprop-based FL algorithm FedAvg and other finite-differences based baselines. In FedAvg, the communication cost of the server to clients scales linearly with the number of participating clients for the entire global model. In contrastOn the bright side, Spry’s communication cost only scales with max(layer count, number of participating clients) for a subset of the global model.\n\n(b) Spry is more computationally efficient than FedAvg at the server-side, since Spry aggregates only a subset of model weights updated by each client, instead of aggregating all model weights from all the clients.     \n\n\n&nbsp;\n\nWe next discuss in detail the server-side computation costs related gradient aggregation and synchronization at the server-side in the following two parts: \n\n### • *Synchronization*\nA server in Spry needs to assign layers randomly to participating clients for each round, and it needs to keep that mapping of layers-to-client-ids for aggregation. As shown in Algorithm 1, Appendix E; this mapping is not a bottleneck for the server since it’s a simple loop to map layer ids to client ids (See lines 14 to 22 in Algorithm 1). \n\nMoreover, for per-iteration communication frequency, the server incurs an additional cost of $w_\\ell \\max(L/M, 1)$, where $w_\\ell$ is the size of a layer (we have assumed the same layer size for each layer for ease of exposition). As shown in Table 2 of the PDF attached to the global “Author Rebuttal”, this cost is also a factor for per-iteration finite differences methods. For Spry, the additional cost comes from having to generate the random perturbations at the server-side and multiplying it with \\texttt{jvp} scalar values received from the client.\n\n&nbsp;\n\n### • *Gradient Aggregation*\nThe only major difference between Spry and backpropagation-based FedAvg at server-side is that with Spry, instead of aggregating all layer weights from all clients, the server only needs to aggregate layer weights from the clients assigned to that layer. Hence Spry has lower computation cost than FedAvg by not having to aggregate all layer weights from all clients.\n\n&nbsp;\n___\n### **W2. Generalizability of Spry**\nSpry focuses on fine-tuning language models of sizes between 17.9M and 13B mainly because of their popularity and practicality on edge devices. This goal and the scope of the experimentations align well with recent efforts published at both NeurIPS and other top venues  [1, 2, 3].  \n\nNonetheless, we value the reviewer’s feedback and aim to do more evaluation on vision and a wider variety of language tasks in the future.\n\n[1] MeZO: Fine-Tuning Language Models with Just Forward Passes (Malladi et al., NeurIPS 2024)\n\n[2] Distributed inference and fine-tuning of large language models over the internet (Borzunov et al., NeurIPS 2023)\n\n[3] Pockengine: Sparse and efficient fine-tuning in a pocket (Zhu et al., MICRO 2023)\n\n&nbsp;\n___\n### **W3. Slow convergence of Spry compared to FedAvg**\nWe have acknowledged the slower convergence of Spry compared to FedAvg in general. However, we would like to emphasize that Spry’s much reduced peak memory consumption makes it feasible for language model finetuning on devices, which is infeasible with FedAvg. \n\nWe must note that a faster gradient calculation method would have no utility for fine-tuning or training large models on edge devices, if its peak memory footprint is larger than available memory of that device. This is the trade-off between runtime and peak memory consumption of Spry, which sacrifices time to convergence by 2.65x to gain 27.90% to 86.26% memory reduction, while achieving 0.6% to 6.2% of the accuracy of the best-performing FL backpropagation.\n\nAs a side note, we also reiterate that Spry does gain a speedup of 1.14x against FedAvg for a medium-sized language model, RoBERTa Large. This speedup is attributed to the fact that clients in Spry only need to train a subset of layers, unlike all FedAvg clients training all the layers.  \n\n&nbsp;\n___\n### **Q1. Large scale setting of 1k+ GPUs**\nFor now it’s computationally slower for us to simulate federated training with 1k participating clients. It takes ~4h to simulate and execute 1k clients per round, running 500 such rounds would take ~83 days. However, given our experiments on 1k total clients (See Table 1 of Spry: rows of AG News, SNLI, MNLI, Yahoo, Yelp) and ablation studies (“Effects of participating client count” in Section 5.4 and Appendix F) on increasing the participating client ratio, we anticipate Spry achieving better performance as the participating client count scales up. Our expectations are based on the observations that more clients training the same model weights decreases the forward gradient’s approximation error, leading to a better prediction performance. \n\nHence, the layer-to-client mapping is not a bottleneck at the server for Spry; and gradient aggregation and synchronization of Spry can be faster than methods which let all the clients compute all gradients. We observe that Spry can be more scalable than backpropagation and finite-difference-based methods, while being 1.46x to 28.57x faster than finite-difference methods and takes 27.90% to 86.26% less memory than backpropagation-based methods, making it a feasible and scalable way to train larger models on edge devices.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We thank the reviewers for their suggestion on adding information on communication and computation costs, and we will update the manuscript with detailed explanations of the following:
____
### **1. Communication overhead**

Table 1 of the PDF attached to this response illustrates communication costs of Spry and its backpropagation and finite difference based baselines. A discussion on communication modes of Spry is also given in Section 3.2, “Per-Epoch Communication” and “Per-Iteration Communication”. Here we discuss the costs related to those communication modes:

&nbsp;
### • *Per-epoch Communication*

Spry's client-to-server communication cost does not scale linearly with clients like backpropagation and finite-difference counterparts, but instead decreases or stays constant for $L$ as more clients are present. Server-to-client communication cost is lower in Spry due to only sending one layer per client when $M > L$ or $\frac{L}{M}$ layers per client otherwise. This result follows from the below observation:

Backpropagation-based and finite-difference-based methods have a communication cost of $w_g$, where $w_g$ represents the global model size. Each client in $[M]$ (set of participating clients) receives all trainable parameters from the server, requiring the server to send a total of $w_g \times M$ parameters each round.

Spry's communication cost per epoch is $w_\ell \max(\frac{L}{M}, 1)$, where $L$ is the layer count and $w_\ell$ is the count of parameters for each layer. Each client sends a subset of trainable parameters, incurring a communication cost of $\frac{w_\ell L}{M}$ parameters for $L > M$, and $w_\ell$ for $L \leq M$. When $L \leq M$, each client gets 1 layer, hence the communication for each client is $w_\ell$.

&nbsp;
### • *Per-iteration Communication*

Spry accrues lower communication cost than the finite difference and backpropagation counterparts due to the layer splitting strategy, and the server’s ability to compute gradients based on the jvp value. This is because:

The communication cost from client to server for forward-mode AD and finite differences is 1. This is due to an FL round that involves 1. server selecting a random seed, 2. server sending it with trainable parameters to clients, 3. clients generating perturbations based on the seed, 4. deriving and sending back a scalar or finite difference scalar to the server, and then 5. server computing gradients by multiplying the derived perturbations with the seed.

The server to client communication is $(w_g + 1) \times M$, where the “+1” is due the randomness seed.

&nbsp;
____
### **2. Computation costs**

Table 2 of the PDF attached this response shows the computation costs of Spry and its baselines, where the client-side cost is for each iteration, and the server-side cost is for each round. 

Briefly, Spry’s client-side computation cost is traded off by a faster convergence to higher accuracy through better gradient approximations compared to finite difference-based methods. Furthermore, Spry is the least computationally expensive on the server side due to needing to aggregate fewer parameters from the clients.

Table 2 assumes that matrix multiplication costs $c$ for each layer, resulting in a forward pass cost of $c$. The cost of backpropagation is $2c$ because the computation of the current layer’s weight gradient is $c$, and the cost of computing the previous layer’s activation gradient is another $c$. jvp computation in Spry takes additional cost of $c$ for each layer. Moreover, since jvp calculation happens through column-by-column vector multiplications (Section 3.1 of [1]), the related overhead is quantified by $v$. 

&nbsp;
### • *Client-side per-iteration computation cost*

Backpropagation needs 3 matrix multiplication operations per layer. For zero-order methods. There are 2 matrix multiplications (incurred due to two forward passes) per layer, and per perturbation within a training iteration; and additional overhead $w_\ell KL$ for perturbation generation. MeZO requires generation of perturbations thrice for the same seed (Algorithm 1 in MeZO [2]). 

Spry’s computation cost is $2 \times \max(\frac{L}{M} , 1) \times (c + v) + w_\ell L$. Since Spry allocates at most $\frac{L}{M}$ layers to each client, the computation cost only scales with $\max(\frac{L}{M}, 1)$, against its counterparts scaling with $L$. However, forward-mode AD computes jvp column-wise, while its counterpart vjp in backpropagation is computed row-wise. This results in time overhead ($v$) if the number of trainable parameters exceeds the output size (1 as loss is scalar), which is the case for neural networks. Therefore, Spry's per-iteration computation cost is higher compared to other approaches.

Note that the per-iteration computation cost of Spry is not the whole picture. It takes fewer communication rounds to reach higher accuracy due to better gradient approximation of forward-mode AD than finite difference methods. This is why "Time to Convergence" (Section 5.3) discusses a fair comparison of Spry's runtime and prediction performance.

&nbsp;
### • *Server-side per-round computation cost*

On the server side, Spry is the least computationally demanding. Spry needs to aggregate a subset of layer weights from only the clients that were assigned to those layers, while its counterparts need to aggregate all layers from all clients.

Computation cost on the server-side changes based on the communication frequency
Per-iteration communication incurs an additional overhead of $w_\ell L (\frac{M}{L} + 1)$ and $w_\ell L (M+1)$ (generation of perturbations at the server-side, and multiplying those perturbations with aggregate of the jvp values received from the clients) for Spry and its zero-order counterparts resp.
___
[1] Automatic Differentiation in Machine Learning: a Survey (Baydin et al, JMLR 2018)

[2] MeZO: Fine-Tuning Language Models with Just Forward Passes (Malladi et al., NeurIPS 2024)

Meta-review: Fine-tuning large language models in federated learning settings has become increasingly important as it allows resource-constrained devices to fine-tune a model using private data. However, fine-tuning LLMs using backpropagation requires excessive memory (especially from intermediate activations) for resource-constrained devices. While Forward-mode Auto-Differentiation can significantly reduce the memory footprint from activations, the authors observe that directly applying it to LLM fine-tuning results in slower convergence and lower accuracy. They introduce SPRY, which reduce the memory footprint, speed up convergence and combine it with LORA, which leads to more scalable FL training algorithms.

Generally, the reviewers were impressed by the numerical results, and found interesting the convergence analysis. Furthermore, it studies an important topic. I thus recommend accepting this paper.
Decision: Accept (poster)