Title: Thinking Forward: Memory-Efficient Federated Finetuning of Language Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Federated Learning

Large Language Models

Forward-mode Automatic Differentiation

Forward-mode AD

Memory-efficient Finetuning

Memory-efficiency

Data Heterogeneity
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

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

1 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 [1, 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 to update the model and initiates the next round of FL training. Due to the

Sunav Choudhary

Hui Guan

inherent privacy-preserving nature of FL, it has been adopted in many privacy-sensitive domains, such as healthcare [3], IoT [4, 5], and e-commerce [6].

In parallel, large language models (LLMs) have demonstrated impressive performance on natural language processing tasks [7, 8], creating a surge of interest in finetuning LLMs in FL settings [9, 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-efficient finetuning (PEFT) [11-16], quantization [17], and zero-order gradient estimation methods [18-20]. Although PEFT and quantization can reduce the memory consumption from parameters and optimizer states, the memory consumed by intermediate activations remains a significant bottleneck because these methods still use backpropagation to finetune LLMs. Backpropagation requires storing all intermediate activations during the forward pass to estimate gradients in the backward pass. For example, finetuning a 4-bit quantized Llama2-7B model [21] with LoRA techniques [12] requires ∼ 33.9GB 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, 23]. 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 Auto-Differentiation (AD) [24, 23] has the potential to address the memory consumption 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 a random perturbation affects the outputs. However, simply replacing backpropagation with Forward-mode AD in FL settings does not produce convergence speed and accuracy performance comparable to established federated optimizers based on backpropagation, such as FEDAVG [1], FEDYOGI [25], FEDSGD [1]. Forward gradients are computationally inefficient and inaccurate in 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 1 , 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-mode AD by splitting the trainable weights among participating clients 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 a 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-server communication bandwidth.

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

We empirically evaluate SPRY's memory efficiency, accuracy, computation efficiency, and communication efficiency through experiments on a wide range of language tasks, models, and FL settings. SPRY achieves within 0.6-6.2% of the accuracy of the best-performing FL backpropagation, with 1.4-7.1 × less memory consumption for each client and comparable time to convergence. SPRY also outperforms zero-order-based baselines with 5.2-13.5% higher accuracy, an average of 1.5-28.6 × faster per-round computation time, and 1.2-20.3 × 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-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.

A federated optimization strategy that only requires a single forward pass per batch on each client to finetune a language model.

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 a client's learning rate is inversely proportional to the size of perturbations and client data heterogeneity.

An empirical evaluation shows that SPRY consumes 1.4-7.1 × less memory than its backpropagation-based counterparts, and converges 1.2-20.3 × faster with 5.2-13.5% higher accuracy compared to its zero-order counterparts.

2 Forward-mode Automatic Differentiation

This section presents the background on Forward-mode 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 a random perturbation, affect the loss. Since these gradients are derived from a forward pass, they are referred to as 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 a random perturbation v whose size is the same as w . In one forward pass, given training data D , Forward-mode AD computes the value of the objective function f ( w ; D ) and the Jacobian-vector product ( jvp ) as follows:

This jvp is a 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],

The partial derivative ∂f ( w ; D ) /∂w 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 of 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 substituting backpropagation with Forward-mode AD in FL scenarios often results in poor accuracy and

Client #1

Client #2

Client #3

Client #4

Server

A Language Model with

Trainable PEFT Weights

Weight

perturbation

4.

Send updated

assigned

weights

3a.

Compute JVP

5.

Aggregate trained layer

to update the model

,

1.

Split layers

across participant

clients

Forward

gradient

Updated

weight

3b.

Update the weights

(  ,  )

= optim_step

·

(     )

= forward_AD

2.

Send frozen and

assigned weights,

and a seed

Scalar

Vector

or

Matrix

Legend

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 is minimized. Therefore, SPRY optimizes performance by distributing trainable weights among participating clients, assigning each client a responsibility to compute gradients for only a subset of trainable weights. SPRY is compatible with PEFT methods such as IA3 [15], ADAPTER-based methods [27], BITFIT [28], and LORA [12], which mitigate the memory consumption from gradients and optimizer states. In this work, we focus on LORA due to its demonstrated superiority, as highlighted in Appendix G.

Figure 1 gives an 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. (4) Clients only transmit the updated weights back to the server. (5) The server aggregates all the trained layer weights and updates the language model for the subsequent round.

Next, we 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 a layer and assigning it to a client in a cyclic manner. With LORA, the server selects a LoRA layer, which consists of a pair of weights ( w A and w B 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 FEDYOGI. 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 to client IDs, hence it 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-based personalization techniques [30] to mitigate it. SPRY distributes trainable classifier layers to all participating clients, enabling each client to finetune these layers to personalize the jointly trained model.

3.2 Finetuning Weights with Forward Gradients on the Client-side

Clients update the assigned trainable weights with gradients estimated through Forward-mode AD. Specifically, each participating client will get a copy of the trainable weights assigned to it and a

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 a 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 max { #Trainable Layers #participating clients , 1 } layers. This means that if there are more trainable layers than participating clients, each client sends back weights for multiple layers.

Per-Iteration Communication. Unlike FEDSGD, where gradients are sent back to the server and aggregated after each iteration, SPRY offers additional communication cost savings. In this communication 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 a 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 to 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 I.

Theorem 4.1 (Estimation of the Global Gradient) . In SPRY , global forward gradients ∇ ˆ f of the trainable weights w ∈ R d , with the corresponding weight perturbations v ∈ R d , computed by M participating clients is estimated in terms of true global gradients ∇ f as,

where the expectation is under the randomness of sampled data and random perturbation v . C is total number of classes and α m,c = ( n c |D| -n m,c α c |D m | ) . For a class c , n c is its sample count, α c is its Dirichlet concentration parameter. For a client m , n m,c is the sample count of the c th class and D m is the size of the data of client m . The global data is D = ∑ m ∈M D m . ˜ M is the set of clients training an arbitrary subset of weights, ˜ M = | ˜ M i | , ∀ i ∈ [ 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 α m,c = ( n c |D| -n m,c α c |D m | ) . In data homogeneous settings, where all clients share the same data distribution, the Dirichlet concentration parameter α c , is 1 for any class c . As the ratio of n c / |D| (total samples in a class to total samples globally) matches the ratio of n m,c / |D m | (total samples in a class to total samples for a client m ), the bias term becomes 0 since α m,c = 0 , ∀ m,c . Hence, the global forward gradients are unbiased estimators of the true global gradients. In data heterogeneous settings, when data across clients becomes more heterogeneous, α c → 0 and thus α m,c → n c / |D| , increasing the estimation error. Hence the global forward gradients are biased estimators that depend on the distances of the data distributions across participating clients.

Theorem 4.2 (Convergence Analysis) . Under the assumptions on L -smoothness (Asmp I.1), bounded global variance σ 2 g between global true gradients and aggregated expected gradients of each client

(Asmp I.2), and bound on gradient magnitude G (Asmp I.3); and the following conditions on the local learning rate η ℓ ,

The global true gradients of SPRY satisfies the following bound,

where R is the total number of FL rounds, w ∈ R d are trainable weights, v ∈ R d are random perturbation, K is count of random perturbations per iteration, η is global learning rate, τ is adaptability constant, and s 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. 6 decrease in proportion to the inverse of R , indicating the convergence of the algorithm up to the limits of data heterogeneity. (b) Dimension of perturbations ( d ): η ℓ ∝ 1 d shows that as the number of weights to 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 K both in nominator and denominator, which indicates that increasing K brings little advantage in convergence speed. Results in Appendix G confirm the above statement. (d) Data heterogeneity: η ℓ ∝ 1 α 2 m,c shows that more homogeneous data distributions across clients allow higher learning rate, and hence faster error reduction. This observation is corroborated by the comparison of convergence speeds between homogeneous and heterogeneous clients in Appendix H. (e) The number of clients training same subset of weights: η ℓ ∝ ˜ M shows that more clients training the same subset of weights is beneficial for faster convergence, similar observation is shown in Appendix G.

The theorem also sheds light on the accuracy performance gap between Forward-mode AD and backpropagation in FL settings. The upper bounds on the global forward gradient norm includes a second term that increases with α 2 m,c and variance σ 2 g , but does not decrease with R . This results in a gap between estimation errors of backpropagation-based methods like FEDAVG and SPRY.

5 Empirical Evaluation

We empirically evaluate SPRY on 8 language tasks, 5 medium, and 3 large language models, under 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 of participating clients, and the importance of splitting layers on performance.

Datasets, Tasks, and Models. Our evaluation uses 8 datasets: AG News [31] (4-class classification), SST2 [32] (2-class classification), Yelp [31] (2-class classification), Yahoo [31] (10-class classification), SNLI [33] (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 us to generate heterogeneous splits in FL settings using Dirichlet distribution [37]. 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 α = 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], DistilBERT 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 B describes the datasets and hyperparameters in more detail.

The datasets are split with Dir α = 0 . 1 . ⋄ = Llama2-7B. ⋆ = OPT6.7B. □ = OPT13B. SPRY outperforms the best-performing zero-order-based methods by 5.15-13.50% and approaches the performance of backpropagation-based methods, with a difference of 0.60-6.16%.

Comparison Counterparts and Metrics. We compare SPRY to (a) Backpropagation-based federated optimizers FEDAVG [1], FEDYOGI [25], FEDSGD (Variant of FEDAVG with per-iteration communication) [1], (b) Zero-order federated methods FEDMEZO (federated version of MEZO [18]), BAFFLE [20], FWDLLM [19], all based on finite difference. MEZO uses prompt-based finetuning to improve the performance of finite differences. FWDLLM generates a random perturbation that has a 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 implementations 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 Acc g and personalized accuracy Acc p , which are metrics measured on server-side aggregated model and client-side locally updated model, respectively. Similarly, for question-answering tasks, we 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 a performance metric, assessed at intervals of 50 rounds.

SPRY is implemented in Flower [43] 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-based methods respectively. Each experiment was run thrice with 0, 1, and 2 as seeds.

5.1 Accuracy Performance Comparison

Table 1 reports the accuracy performance of SPRY and its backpropagation- and zero-order-based counterparts on heterogeneous datasets for million-sized RoBERTa Large, and billion-sized Llama27B, OPT6.7B, and OPT13B. Similar results on personalized performance is shown in Appendix H, 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 personalized accuracy over the best-performing zero-order-based methods across all datasets. FWDLLM+, the best-performing zero-order counterpart, attempts to reduce the effect of numeric instability of finite differences by (a) Sampling K perturbations (default K = 10 ) per batch for each client and picking 1 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 a set threshold. However, we posit that this strategy leads to some clients getting excluded due to a low variance threshold or outlying clients getting included due to a high

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-1000 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 1 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-based methods, FEDAVG and FEDYOGI, SPRY manages to come as close as 0.60-6.16% of generalized accuracy and 2.50-14.12% of personalized accuracy. The performance gap between backpropagation and Forward-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-mode AD relies on random perturbations, which is 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.

5.2 Peak Memory Consumption Comparison

Figure 2 shows peak memory consumption of backpropagation (used in FEDAVG, FEDYOGI), zero-order finite differences (used in FWDLLM+, BAFFLE+, FEDMEZO), and firstorder forward mode AD (used in SPRY). The methods are profiled for a single client.

Compared to backpropagation, Forward-mode AD reduces peak memory usage of RoBERTa Large by 27.90%, Llama2-7B by 81.73%, OPT6.7B by 86.26% and OPT13B by 85.93%. The sizes of trainable parameter (colored ) and gradient + optimizer state (colored ) are consistent across the 3 modes of computing gradients for all 4 models. Hence the savings come from the reduced memory footprint related to activations (colored ) in Forward-mode AD. Compared to backpropagation-based methods, the memory cost related to activations is decreased by 12.12-49.25 × in SPRY. Unlike storing all the intermediate activations in backpropagation, Forward-mode AD only has to store the previous layer's activation in a forward pass. The activation footprint of Forward-mode AD is equal to the size of the largest activation.

Back-

propagation

Zero-Order

Finite

Difference

First-Order

Forward

Mode AD

0

1

2

Memory Footprint

(in GBs)

2.23GB

1.56GB

RoBERTa-Large

Back-

propagation

Zero-Order

Finite

Difference

First-Order

Forward

Mode AD

0

20

40

34.0GB

5.48GB

6.21GB

Llama2-7B

Back-

propagation

Zero-Order

Finite

Difference

First-Order

Forward

Mode AD

0

20

40

60

Memory Footprint

(in GBs)

56.44GB

6.04GB

7.75GB

OPT6.7B

Back-

propagation

Zero-Order

Finite

Difference

First-Order

Forward

Mode AD

0

25

50

75

76.45GB

9.21GB

10.75GB

OPT13B

Parameters

Activations

Gradients + Optimizer

Against zero-order methods, Forward-mode AD activations cost 1.96 × , 1.95 × , 1.83 × , and 1.54 × 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.1, the increased memory cost is offset by a boost of up to 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 a closer gradient estimation.

5.3 Time to Convergence Comparison

Figure 3 shows wallclock time-to-convergence for SPRY and its counterparts. We observe that SPRY is 1.15-1.59 × , 6.16-20.28 × , and 1.34-2.98 × faster than the zero-order methods FWDLLM+, BAFFLE+, and FEDMEZO respectively. For a client in each round, SPRY achieves a faster per-round computation time of 1.46 × , 28.57 × , and 1.80 × on average, against FWDLLM+, BAFFLE+, and FEDMEZO. Forward-mode AD achieves faster convergence and faster per-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,

it gains a speedup of 1.14 × over backpropagation-based FEDAVG, FEDYOGI, and FEDSGD for RoBERTa Large. However, compared to the backpropagation-based methods, SPRY slows down for billion-sized LMs. We attribute this loss of speedup to the way jvp 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

We summarize the ablation experiments on various components of SPRY. Further discussions are in Appendix G.

SPRY can generalize to other language model architectures. Similar to the observations from Table 1, we see a trend of SPRY outperforming the best performing zeroorder method FWDLLM+ by 3.15-10.25% for generalized accuracy, demonstrating that SPRY can generalize to other language model architectures.

SPRY is compatible with other PEFT methods. We integrate SPRY with different PEFT methods like IA3, BITFIT, and CLASSIFIER-ONLY FINETUNING. Results shows that LORA with SPRY performs the best, with accuracy improvements of 10.60-16.53%.

Effects of the number of trainable weights. We change the number of trainable weights by controlling the rank r and scale α hyperparameters in LORA. Results show that SPRY achieves the highest accuracy with ( r =1, α =1) setting, which has the smallest trainable weight count. The result is consistent with our theoretical analysis in §4.

Effects of communication frequency. Per-iteration communication in 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.

0

200

400

0.0

0.2

0.4

0.6

0.8

1600

1800

2000

2200

2400

Time (in minutes)

Test Accuracy

Accg

FedAvg

FedYogi

FedSgd

FwdLLM+

FedMeZO

Baffle+

Spry

(a) AG News with RoBERTa Large

0

50

100

150

200

0

2

4

6

8

10

12

2700

2750

2800

Time (in minutes)

Test F1 Score

F

1

g

(b) SQuADv2 with OPT13B

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 a steady state (of accuracy ∼ 86%) around the 200 th round, while the setting with K = 1 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 . 14% , 86 . 56% , and 88 . 08% , respectively. We also see 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, we conduct the following two experiments: (a) With FEDAVGSPLIT, we apply the strategy of splitting trainable layers across clients (see § 3.1) to backpropagation-based FEDAVG, and (b) With FEDFGD, we omit the splitting strategy of SPRY for FL with Forward-mode AD. We observe that FEDAVGSPLIT fails to achieve similar accuracy to FEDAVG with a drop of 2.60-10.00%. FEDFGD fails to converge as the size of trainable weights increases, e.g., with RoBERTa Large with LORA, 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 2 and 3 in Appendix F 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 w g is the total trainable parameter count of a model to be trained in federated setting. The

backpropagation-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 g and 'server to client' communication cost of w g × 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 w g × M . However, the per-iteration versions of the zero-order baselines fare better, with 'client to server' communication only requiring each client sending a scalar finite difference (cost of 1), and 'server to client' communication accruing ( w g +1) × M cost, due to the server also needing to send a scalar seed to each client now.

Meanwhile, SPRY only needs to send max ( L M , 1 ) layers (where L is the total layer count of a model, and M is the participating count of clients for a round), each layer of size w ℓ . Hence, for per-epoch 'client to server' communication accrues w ℓ max ( L M , 1 ) , which is smaller than w g . Similarly, perepoch 'server to client' communication costs w ℓ M max ( L M , 1 ) , which is also smaller than the cost of w g M related to the baselines. For per-iteration communication, clients only need to send a scalar jvp (cost of 1) to the server; this matches the communication cost of per-iteration zero-order methods. Server needs to send a total of w ℓ M max( L, M ) (derivation given in Table 2), which is 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 clientside computation cost is traded off by a faster convergence to higher accuracy through better gradient approximations compared to finite difference-based 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 c for each layer, resulting in a forward pass cost of c . The cost of backpropagation is 2 c 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, the related overhead is quantified by v .

Hence backpropagation-based methods FEDAVG, FEDSGD, and FEDYOGI computationally costs 3 Lc at client, and costs w g ( M -1) at server (due to additions). Note that L amounts to all layers in the model here. FEDMEZO costs L (2 c +3 w ℓ ) at client through two forward passes, and generating perturbations three times. FWDLLM and BAFFLE costs KL (2 c + w ℓ ) due to K perturbations for all L layers, with two forward passes and generation perturbations once. Against that, SPRY costs 2 max( L M , 1)( c + v ) + w ℓ L for a smaller count of L , traded-off by the jvp computation cost of v .

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 w ℓ L ( M L +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.

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

Acknowledgments and Disclosure of Funding

This material is based upon work supported by the National Science Foundation under grant no. CCF-2210243, DMS-2220211, CNS-2224054, CNS-2312396, and CNS-2338512, and by Adobe.

References

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-Efficient 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.

Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D'Oliveira, Hubert Eichner, Salim El 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, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Praneeth 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.

Tyler J Loftus, Matthew M Ruppert, Benjamin Shickel, Tezcan Ozrazgat-Baslanti, Jeremy A Balch, Philip A Efron, Jr. Gilbert R Upchurch, Parisa Rashidi, Christopher Tignanelli, Jiang Bian, and Azra Bihorac. Federated learning for preserving data privacy in collaborative healthcare research. Digital Health , 2022.

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.

Suresh Dara, Ambedkar Kanapala, A. Ramesh Babu, Swetha Dhamercherala, Ankit Vidyarthi, and Ruchi Agarwal. Scalable federated-learning and internet-of-things enabled architecture for chest computer tomography image classification. Computers and Electrical Engineering , 2022.

Fabio Pinelli, Gabriele Tolomei, and Giovanni Trappolini. Flirt: Federated learning for information retrieval. In Proceedings of the 46th International ACM SIGIR Conference on Research and Development in Information Retrieval , 2023.

OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge 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 Berner, 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 Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, 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 Heaton, 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, Łukasz 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, Łukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, 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 Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reiichiro Nakano, Rajeev Nayak, Arvind Neelakantan, Richard Ngo, Hyeonwoo Noh, Long Ouyang, Cullen O'Keefe, 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 Proehl, 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

Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Selsam, Kyla Sheppard, Toki Sherbakov, 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, Natalie 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.

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 El Shafey, Yanping Huang, Kathy Meier-Hellstern, Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hernandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette-Choo, Aakanksha Chowdhery, Clément 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, 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, Eric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontaek Lim, Hanzhao Lin, Zhongtao 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, Xuezhi Wang, Pidong Wang, Zirui Wang, Tao Wang, John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengcheng 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.

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.

Yuanyishu Tian, Yao Wan, Lingjuan Lyu, Dezhong Yao, Hai Jin, and Lichao Sun. Fedbert: When federated learning meets pre-training. ACM Transactions on Intelligent Systems and Technology (TIST) , 2022.

Alexander Borzunov, Max Ryabinin, Artem Chumachenko, Dmitry 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.

Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations , 2022.

Dongqi Cai, Yaozong Wu, Shangguang Wang, and Mengwei Xu. Fedadapter: Efficient federated learning for mobile nlp. In Proceedings of the ACM Turing Award Celebration Conference - China 2023 . Association for Computing Machinery, 2023.

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 finetuning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Haokun Liu, Derek Tam, Mohammed Muqeeth, 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.

Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) . Association for Computational Linguistics, 2021.

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 integerarithmetic-only inference. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2018.

Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Mengwei Xu, Dongqi Cai, Yaozong Wu, Xiang Li, and Shangguang Wang. Fwdllm: Efficient fedllm using forward gradient. arXiv 2308.13894 , 2024.

Haozhe Feng, Tianyu Pang, Chao Du, Wei Chen, Shuicheng Yan, and Min Lin. Does federated learning really need backpropagation? arXiv 2301.12195 , 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, 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-Anne 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 Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. arXiv 2307.09288 , 2023.

C. H. Richardson. An introduction to the calculus of finite differences. 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.

Atılım Güne¸ s Baydin, Barak A. Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without backpropagation. arXiv 2202.08587 , 2022.

Atılım Günes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. The Journal of Machine Learning Research , 2017.

Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koneˇ cný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations , 2021.

Chandra Thapa, Pathum Chamikara Mahawaga Arachchige, Seyit Camtepe, and Lichao Sun. Splitfed: When federated learning meets split learning. In Proceedings of the AAAI Conference on Artificial Intelligence , 2022.

Renrui Zhang, Jiaming Han, Chris Liu, Aojun Zhou, Pan Lu, Hongsheng Li, Peng Gao, and Yu Qiao. LLaMA-adapter: Efficient fine-tuning of large language models with zero-initialized attention. In The Twelfth International Conference on Learning Representations , 2024.

Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel. BitFit: Simple parameter-efficient fine-tuning for transformer-based masked language-models. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers) . Association for Computational Linguistics, 2022.

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.

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.09262 , 2022.

Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In Neural Information Processing Systems , 2015. Available at https: //huggingface.co/datasets/ag_news , https://huggingface.co/datasets/yelp_polarity , https://huggingface.co/datasets/yahoo_answers_topics , Accessed on 15 May, 2024.

Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing . Association for Computational Linguistics, 2013. Available at https://huggingface.co/datasets/stanfordnlp/sst2 , Accessed on 15 May, 2024.

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 in Natural Language Processing . Association for Computational Linguistics, 2015. Available at https://huggingface.co/datasets/stanfordnlp/snli , Accessed on 15 May, 2024.

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 Language Technologies, Volume 1 (Long Papers) . Association for Computational Linguistics, 2018. Available at https://huggingface.co/datasets/ SetFit/mnli , Accessed on 15 May, 2024.

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 Papers) . Association for Computational Linguistics, 2018. Available at https:// huggingface.co/datasets/rajpurkar/squad_v2 , Accessed on 15 May, 2024.

Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: A 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 1 (Long Papers) . Association for Computational Linguistics, 2018. Available at https://huggingface.co/datasets/mtc/multirc , Accessed on 15 May, 2024.

Kunjal Panchal, Sunav Choudhary, Nisarg Parikh, Lijun Zhang, and Hui Guan. Flow: Per-instance personalized federated learning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

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.

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 1907.11692 , 2019.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. arXiv 1810.04805 , 2018.

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv 1910.01108 , 2019.

Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. ALBERT: A lite BERT for self-supervised learning of language representations. arXiv 1909.11942 , 2019.

Daniel J Beutel, Taner Topal, Akhil Mathur, Xinchi Qiu, Titouan Parcollet, and Nicholas D Lane. Flower: A friendly federated learning research framework. arXiv preprint 2007.14390 , 2020.

AutoGPTQ, 2024. URL https://github.com/AutoGPTQ/AutoGPTQ .

Richard L. Burden and J. Douglas. Faires. Numerical analysis / Richard L. Burden, J. Douglas Faires. Thomson Brooks/Cole, 8th ed. edition, 2005. ISBN 0534392008.

Károly Jordán. Calculus of finite differences . American Mathematical Soc., 1965.

Wenzhi Fang, Ziyi Yu, Yuning Jiang, Yuanming Shi, Colin N. Jones, and Yong Zhou. Communicationefficient stochastic zeroth-order optimization for federated learning. IEEE Transactions on Signal Processing , 2022.

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 2402.06954 , 2024.

Tao Fan, Yan Kang, Guoqiang Ma, Weijing Chen, Wenbin Wei, Lixin Fan, and Qiang Yang. Fate-llm: A industrial grade federated learning framework for large language models. arXiv 2310.10049 , 2023.

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 at scale. arXiv 2105.11367 , 2022.

Jae Ro, Theresa Breiner, Lara McConnaughey, Mingqing Chen, Ananda Suresh, Shankar Kumar, and Rajiv Mathews. Scaling language model size in cross-device federated learning. In Proceedings of the First Workshop on Federated Learning for Natural Language Processing (FL4NLP 2022) . Association for Computational Linguistics, 2022.

Shubham Malaviya, Manish Shukla, and Sachin Lodha. Reducing communication overhead in federated learning for pre-trained language models using parameter-efficient finetuning. In Conference on Lifelong Learning Agents . PMLR, 2023.

Zhuo Zhang, Yuanhang Yang, Yong Dai, Qifan Wang, Yue Yu, Lizhen Qu, and Zenglin Xu. FedPETuning: When federated learning meets the parameter-efficient tuning methods of pre-trained language models. In Findings of the Association for Computational Linguistics: ACL 2023 . Association for Computational Linguistics, 2023.

Seonghwan Park, Dahun Shin, Jinseok Chung, and Namhoon Lee. Fedfwd: Federated learning without backpropagation. arXiv 2309.01150 , 2023.

Geoffrey Hinton. The forward-forward algorithm: Some preliminary investigations. arXiv 2212.13345 , 2022.

Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et al. Parameter-efficient fine-tuning of large-scale pre-trained language models. Nature Machine Intelligence , 2023.

Zeyu Han, Chao Gao, Jinyang Liu, Jeff Zhang, and Sai Qian Zhang. Parameter-efficient fine-tuning for large models: A comprehensive survey. arXiv 2403.14608 , 2024.

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuning of quantized LLMs. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

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 f . MEZO [18] has shown that finite difference with one perturbation per batch does not reach convergence on its 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 derived from multiple ( ∼ 10 to 100) random perturbations on the same input batch is required [24], leading to slow convergence. BAFFLE [20] and FEDZO [47] utilize zero-order gradients in federated settings. To train vision models (of parameter count ≤ 13M), BAFFLE requires (a) ∼ 100-500 perturbations per batch for each client respectively, and (b) per-iteration communication among clients like FEDSGD [1]. 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 backpropagation-based counterparts. SPRY, using Forward-mode AD, computes gradients more accurately in a single forward pass compared to averaged gradients obtained through finite difference methods. This higher accuracy is achieved without needing modifications to model architectures or task structures, while also maintaining a memory footprint similar to that of finite difference methods.

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-by-column evaluations per input batch as the number of trainable weights increases. Moreover, evaluation of FGD is limited to a multi-layer perceptron of size d ≈ 1.8M. Direct use of FGD to finetune language models leads to slow or no convergence. SPRY splits the trainable layers of a 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 [48-50]. The backpropagation-based methods [10, 51], 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 1 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 1 perturbation (without resampling) per batch to reach a higher prediction performance faster than FWDLLM.

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 distribution each client gets depends on the Dirichlet distribution, where a parameter α regulates the concentration of samples of a particular class for a specific client. Dir α = 1 . 0 means all clients have homogeneous datasets where each class is equally likely to be on each client. With Dir α → 0 , 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-based methods FEDAVG, FEDYOGI, and FEDSGD; we will fix the number of epochs to 1 since the goal of this work is to inch closer to backpropagation-like prediction performance while reducing the memory footprint. All the experiments have been run for 1500 FL rounds , except the experiments on OPT models, which are run for 600 FL rounds. Our observation from hyperparameter-tuning shows that the learning rate that gives the best performance is the same for all studied methods. BAFFLE and its memory-efficient 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 of BAFFLE+ to 20 perturbations per batch and fixed finite difference step size σ = 1 e-4. FWDLLM+ samples 10 perturbations for each batch, finite difference step size of σ = 1 e-2. FEDMEZO samples 1 perturbation for each batch, with finite difference step size of σ = 1 e-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 a normal distribution with 0 mean and 1 variance. Default LORA r and α are 1 and 1, respectively. All methods use ADAMW as client-side optimizer. Besides FEDAVG, all methods use FEDYOGI 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 an equal number of samples for train and test datasets. This dataset is under Creative Commons CCZero(CC0) public domain dedication.

Learning rate for backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-4 , 1e-5 } . The batch size is set to 8 . The max sequence length is 128. All methods use ADAMW as a client-side optimizer, while SPRY performs better with SGD. Variance threshold of FWDLLM+ is 1e+1.

SST2. Stanford Sentiment Treebank Binary[32] (or SST2) dataset is for a binary sentiment classification task. The dataset has 11,855 sentences derived from a set of movie 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 as 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 backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-4 , 1e-5 } . The batch size is set to 8 . The max sequence length is 64. Variance threshold of FWDLLM+ is 5e+0.

Yelp. The Yelp reviews dataset[31] is a binary classification dataset gathered during the 2015 Yelp Dataset Challenge. The full dataset has 1,569,264 total samples, and it defines 2 classification tasks. We use the polarity classification task, which is 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

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 backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-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, 6. Sports, 7. Business & Finance, 8. Entertainment & Music, 9. Family & Relationships, 10. Politics &Government. We split this dataset between 1000 clients, where each client gets an equal amount of data samples, where the data distribution is set by changing the Dirichlet distribution α to range from most homogeneous ( α = 1 ) to least homogeneous ( α = 0 ). This dataset is under the Apache License, Version 2.0.

Learning rate for backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-4 , 1e-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. It is a 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 backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-4, 5e-5 , 1e-5 } . The batch size is set to 8 . The max sequence length is 80. Variance threshold of FWDLLM+ is 1e+2.

MNLI. The Multi-Genre Natural Language Inference (MNLI) [34] corpus contains 432,702 sentence pairs which were crowd-sourced and then annotated with textual entailment information. This dataset is also a Natural Language Inference dataset as 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.

Learning rate for backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-4, 5e-5 , 1e-5 } . The batch size is set to 8 . The max sequence length is 80. Variance threshold of FWDLLM+ is 1e+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 backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-4, 1e-5 } . The batch size is set to 8 for Forward-mode AD and zero-order differentiation-based methods. Backpropagation-based methods use a batch size of 4 . The max sequence length is 400. Variance threshold of FWDLLM+ is 5e+1.

MultiRC. Multi-sentence Reading Comprehension (MultiRC) corpus is a dataset that contains short paragraphs with questions, whose answers can be found in the paragraph itself. We consider a task where we classify if the input question given as input is correct or false. It has 6k multi-sentence questions about 800+ paragraphs. Due to the scale of the dataset, we split it into 100 clients. This dataset is under the MIT license.

Learning rate for backpropagation-based (FEDAVG, FEDYOGI, and FEDSGD), zero-order-based (FWDLLM, BAFFLE, and FEDMEZO), and first-order-based SPRY is { 1e-3, 5e-4, 1e-4 , 1e-5 } . The batch size is set to 8 . The max sequence length is 256. Variance threshold of FWDLLM+ is 1e+2.

C Limitations and Future Work

SPRY achieves a 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-column 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 room for improvement in reducing the memory consumption to that of zero-order methods. In zero-order methods, the weights are perturbed and a forward pass is computed on the perturbed weights. However, with the current implementation of Forward-mode AD, perturbations create a separate copy from the original weights, which accrues additional overhead. To further utilize the computation capacity of clients in FL, device-heterogeneity-aware strategies on splitting and mapping layers to 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 to finetune these language models can be prohibitively expensive as the size of the trainable weights increases. With SPRY, we enable finetuning LLMs on resource-constrained edge devices with low memory consumption, making it 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 a plethora of personalized downstream tasks, but the finetuning data containing sensitive and confidential information is never shared with a third party. With SPRY, such data remains on the user's device, ensuring privacy while still allowing for effective finetuning of the LLM.

However, making LLM finetuning accessible to a broader range of organizations and individuals comes with its own challenges like a spread of biases or misinformation. Without preprocessing and filtering client data on their devices, the LLMs can be fed harmful and misleading information. Hence, it is necessary to develop guardrails on what kind of data should be filtered in, to finetune LLMs with crowd-sourced compute resources.

E SPRY Pseudocode

Algorithm 1 shows the workflow of SPRY.

The function MAPLAYERSTOCLIENTS on Line 14 in Algorithm 1 shows how we have assigned only a few trainable layers to each client in FL, to make Forward-mode AD more effective at generating better estimation of the gradients. In function CLIENTTRAIN on Line 23 of Algorithm 1; each client m gets a copy of the entire language model w ( r ) m and a list trainable_layers of parameter names the client m has to train. A client m freezes the parameters that are not included in its trainable_layers . And for each of the parameter w ( r ) m which need to be trained, SPRY generates a corresponding random perturbation v using a normal distribution N (0 , I w ( r ) m . shape ) .

Once a client m has obtained forward gradients of all the trainable parameters w ( r ) m , those parameters are locally updated with optimizers like SGD or ADAM. The updated trainable parameters w ( r ) m are sent back to the server. The server has a mapping of parameter names to client IDs, and hence it builds w ′ ( r ) by using w ( r ) m ∀ m ∈ [ 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 FEDYOGI at server-side on effective gradients ∆ = w ′ ( r ) -w ( r ) to reduce the noise of forward gradients.

F Communication and Computation Costs

F.1 Communication Costs

Table 2 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:

Per-epoch Communication. SPRY's client-to-server communication cost does not scale linearly with clients like in its 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 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 × M parameters each round.

SPRY's communication cost per epoch is w ℓ max( L M , 1) , where L is the layer count and w ℓ is the count of parameters for each layer. Each client sends a subset of trainable parameters, incurring a communication cost of w ℓ L M parameters for L > M , and w ℓ for L ≤ M . When L ≤ M , each client gets 1 layer, hence the communication for each client is w ℓ .

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) × M , where the '+1' is due the randomness seed.

F.2 Computation Costs

Table 3 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 3 assumes that matrix multiplication costs c for each layer, resulting in a forward pass cost of c . The cost of backpropagation is 2 c 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 (Sec 3.1 of [24]), the related overhead is quantified by v .

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 ℓ KL for perturbation generation. MEZO requires generation of perturbations thrice for the same seed (Algorithm 1 in MEZO [18]).

SPRY 's computation cost is 2 × max( L M , 1) × ( c + v ) + w ℓ L . Since SPRY allocates at most L M layers to each client, the computation cost only scales with max( 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 ( ) 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 forwardmode 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.

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 ℓ L ( M L + 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.

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

Generalized Test

Accg

LoRA

Bitfit

IA3

Classifier-only

0

200

400

600

800

1000

1200

1400

Rounds

0.0

0.2

0.4

0.6

0.8

Generalized Test

Accg

FedAvg (Per-Epoch Communication)

FedSgd (Per-Iteration Communication)

Spry (Per-Epoch Communication)

SpryPerIter (Per-Iteration Communication)

0

200

400

600

800 10001200

Rounds

0.0

0.2

0.4

0.6

0.8

Generalized Test

Accg

r

= 1,  =1

r

= 8,  =16

r

= 16,  =16

r

= 32,  =32

(a) Effect of different PEFT methods on SPRY

(b) Effect of per-epoch and per-iteration communication

(c) Changing LORA

SPRY

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. Table 4 shows generalized and personalized accuracies Acc g and Acc p 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 a 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 BITFIT. 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 BITFIT 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 to fail on LLMs [57]. Furthermore, unlike IA3, LORA has been shown to be successful at finetuning quantized billion-sized models in QLORA [58], making it 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.

r

and

α

for

0

100 200 300 400 500 600 700

Rounds

0.0

0.2

0.4

0.6

0.8

Generalized Test

Accg

K = 1

K = 10

K = 20

K = 50

K = 100

(a) Changing

K

(Forward-mode AD

perturbation count per

0

100200300400500600700800

Rounds

0.0

0.2

0.4

0.6

0.8

Generalized Test

Accg

C = 10

C = 50

C = 100

(b) Changing

C

(Per-round participating

client count)

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

Generalized Test

Accg

FedAvg

FedAvgSplit

FedFgd

Spry

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

Generalized Test

Accg

FedAvg

FedAvgSplit

FedFgd

Spry

(c) RoBERTa Large (Left) and BERT Base (Right) - Splitting parameters across clients on Backpropagation vs Forward-mode AD

iteration)

By communicating each iteration, we see a 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 Figure 4b, the convergence speed of SPRY also improves by communicating each iteration.

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 v , which was used by all the clients to generate their respective jvp values. Then the random perturbation is multiplied with the received 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 DistilBERT Base (total parameter count of 66M), LORA reduces the trainable parameter count to 0.61M (0.91%), 0.74M (1.11%), 0.89M (1.33%), and 1.18M (1.77%) with LORA hyperparameter settings of ( r =1, α =1), ( r =8, α =16), ( r =16, α =16), and ( r =32, α =32).

SPRY achieves the highest accuracy of 84.90% with ( r =1, α =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 Figure 5a 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 to 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 200 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 . 08% , 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 ∼ 85%; C = 10 , C = 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.1) to backpropagation-based FEDAVG, and (b) With FEDFGD, we omit the splitting strategy of SPRY.

Figure 5c shows the performance of FEDAVG and FEDAVGSPLIT 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 a similar accuracy with an absolute difference of 2.70-3.61% between FEDAVGSPLIT 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 are 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.

H Additional Results

H.1 Generalized Performance Curves

Generalized results on homogeneous and heterogeneous clients with Dir α = 1 . 0 and α = 0 . 1 are shown in (a) Figures 6 and 8 for RoBERTa Large, Llama2-7B, OPT6.7B, OPT13B; and (b) Figure 7 for BERT Large, BERT Base, DistilBert Base, Albert Large v2.

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

FedAvg

FedYogi

FedSgd

FwdLLM+

FedMeZO

Baffle+

Spry

(a) AG News with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(b) SST2 with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(c) SNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(d) MNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(e) Yahoo with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(f) Yelp with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(g) MultiRC with Llama2-7B

0

100

200

300

400

500

600

Rounds

0

10

20

30

40

Generalized Test

F

1

g

(h) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

5

10

15

20

25

30

35

Generalized Test

Exact Matches

EMg

(i) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

5

10

15

20

Generalized Test

F

1

g

0

100

200

300

400

500

600

Rounds

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

Generalized Test

Exact Matches

EMg

(j) SQuADv2 with OPT13B

(k) SQuADv2 with OPT13B

0

200

400

600

800

1000

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

FedAvg

FedYogi

FedSgd

FwdLLM+

Spry

(a) AG News with BERT Base

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(b) SST2 with DistilBERT Base

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(c) SNLI with BERT Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(d) Yahoo with DistilBERT Base

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

FedAvg

FedYogi

FedSgd

FwdLLM+

FedMeZO

Baffle+

Spry

(a) AG News with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(b) SST2 with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(c) SNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(d) MNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(e) Yahoo with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(f) Yelp with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Generalized Test

Accg

(g) MultiRC with Llama2-7B

0

100

200

300

400

500

600

Rounds

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

Generalized Test

F

1

g

(h) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

5

10

15

20

Generalized Test

Exact Matches

EMg

(i) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

2

4

6

8

10

12

Generalized Test

F

1

g

(j) SQuADv2 with OPT13B

0

100

200

300

400

500

600

Rounds

0

2

4

6

8

Generalized Test

Exact Matches

EMg

(k) SQuADv2 with OPT13B

H.2 Personalized Performance Curves

Personalized results on homogeneous and heterogeneous clients with Dir α = 1 . 0 and α = 0 . 1 are shown in (a) Figures 9 and 11 for RoBERTa Large, Llama2-7B, OPT6.7B, OPT13B; and (b) Figure 10 for BERT Large, BERT Base, DistilBert Base, Albert Large v2.

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

FedAvg

FedYogi

FedSgd

FwdLLM+

FedMeZO

Baffle+

Spry

(a) AG News with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(b) SST2 with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(c) SNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(d) MNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(e) Yahoo with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(f) Yelp with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(g) MultiRC with Llama2-7B

0

100

200

300

400

500

600

Rounds

0

10

20

30

40

Personalized Test

F

1

p

(h) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

5

10

15

20

25

30

35

Personalized Test

Exact Matches

EMp

FedAvg

FedYogi

FedSgd

FwdLLM+

Baffle+

Spry

(i) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

5

10

15

20

Personalized Test

F

1

p

(j) SQuADv2 with OPT13B

0

100

200

300

400

500

600

Rounds

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

Personalized Test

Exact Matches

EMp

(k) SQuADv2 with OPT13B

0

200

400

600

800

1000

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

FedAvg

FedYogi

FedSgd

FwdLLM+

Spry

(a) AG News with BERT Base

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(b) SST2 with DistilBERT Base

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(c) SNLI with BERT Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(d) Yahoo with DistilBERT Base

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

FedAvg

FedYogi

FedSgd

FwdLLM+

FedMeZO

Baffle+

Spry

(a) AG News with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(b) SST2 with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(c) SNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(d) MNLI with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(e) Yahoo with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(f) Yelp with RoBERTa Large

0

200

400

600

800 100012001400

Rounds

0.0

0.2

0.4

0.6

0.8

1.0

Personalized Test

Accp

(g) MultiRC with Llama2-7B

0

100

200

300

400

500

600

Rounds

0

5

10

15

20

Personalized Test

F

1

p

(h) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

5

10

15

20

Personalized Test

Exact Matches

EMp

(i) SQuADv2 with OPT6.7B

0

100

200

300

400

500

600

Rounds

0

2

4

6

8

10

12

14

Personalized Test

F

1

p

(j) SQuADv2 with OPT13B

0

100

200

300

400

500

600

Rounds

0

2

4

6

8

10

Personalized Test

Exact Matches

EMp

(k) SQuADv2 with OPT13B

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.

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 use the server update of FEDADAM [25]. 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 of SPRY is

∆ ( r ) is the square of accumulated gradients from all clients. Λ ( r ) is the second moment of ∆ ( r ) . τ is a small positive real number, to prevent division by zero errors. Note that we are assuming flattened weights w ∈ R d without the loss of generality.

With SPRY, the aim is to solve the following optimization problem:

where m ∈ [ ( m -1) d M +1 , md M ] , F m ( w m ) = E ( x,y ) ∼D m [ f m ( w m , ( x, y ))] is the objective function, D m is the dataset, and f m is the loss function of a client m ∈ [ M ] .

Accumulated Gradients. With SPRY, each client trains a subset of weights w m . In a low participation rate setting, each w m 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 w m .

The true global gradients can be written as,

where ˜ M is a set of clients training the subset of the weights w m . ˜ M is the size of ˜ M .

Client Update. The directional derivative of Forward-mode AD is denoted as ∇ ˆ f v ( w ; ( x, y )) , where v ∈ R d is the random perturbation of weights w and ( x, y ) are sampled from a dataset D . For each client m , through Forward-mode AD, we have ∇ ˆ f m ( w m , v m ; ( x, y )) = ( ∇ ˆ f m,v ( w m ; ( x, y )) · v m ) to estimate the true gradient ∇ F m ( w m ) :

Here, the expectation is under the randomness of sampled data D , and random perturbation v . K is the number of perturbations per batch. SPRY uses K = 1 by default, but here we aim to make our analysis more general to see the impact of K on 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 m is L -Lipschitz,

Assumption I.2 (Bounded Global Variance (Assumption 2 in AFO [25])) . The variance of the global objective function f is bounded by σ g as

where ∇ F m is the true gradient of client m . As defined earlier, m ∈ [ ( m -1) d M +1 , md M ] .

Assumption I.3 (Bounded Gradients (Assumption 3 in AFO [25])) . The function f m ( w ; ( x, y )) has G -bounded gradients such that for any client m ∈ [ M ] , weights w ∈ R d , and sampled data ( x, y ) ∼ D m ; we have

I.3 First and Second Moments of Forward Gradients

We first prove that in SPRY, estimation of ∇ f by the accumulated forward mode gradients ∇ ˆ f across all clients of an arbitrary round r , depends on the heterogeneity across client datasets. Statements on homogeneous data have been proven for FGD [23] and MEZO [18] in single-client or centralized settings. Here we focus on the specific federated setting of SPRY, where gradients are accumulated differently than in traditional FEDAVG [1]. In SPRY, we have

where w represents weights, v are their corresponding perturbations, and D m is the dataset of client m .

We omit the round index of the model weights w and their perturbations v for this section since the same relationship will hold for any arbitrary round r .

Theorem I.4 (Estimation of the Global Gradient) . In SPRY , global forward gradient ∇ ˆ f of the trainable weights w ∈ R d , with the corresponding weight perturbations v ∈ R d , computed by M participating clients is estimated in terms of true global gradient ∇ f as,

where C is total number of classes and α m,c = ( n c |D| -n m,c α c |D m | ) . For a class c , n c is its sample count, α c is its Dirichlet concentration parameter. For a client m ; n m,c is the sample count of the c th class, and D m is the size of the data of client m . The global data is D = ∑ m ∈M D m . ˜ M is the set of clients training an arbitrary subset of weights, ˜ M = | ˜ M i | ; ∀ i ∈ [ M/d ] .

Proof. Suppose the global dataset D is defined as a combination of all D m . Hence, D = ∪ m ∈ [ M ] D m . In SPRY, the global forward gradient is defined as

To combine the gradient estimates from the above equation for all clients m ∈ M , we first consider the bias b m between the gradient estimate under (a) Globally combined data D and (b) Data D m of an arbitrary client m .

Measuring the dataset bias. Note that we consider samples ( x, y ) of D m to be sampled from D , since D m ⊂ D .

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 m 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 C total classes: y 1 , . . . , y C . We have α 1 , . . . , α C as concentration parameters of the Dirichlet distribution. The expected forward gradient for the combined dataset D can be expressed as a weighted sum of the expected forward gradients for each class C :

where n c is the sample count of class c . And for the biased dataset D m sampled with Dirichlet concentration parameters α 1 , . . . , α C , the expected forward gradient can be expressed as,

where n m,c is the sample count of class c for client m .

Subtracting Equation 16 from Equation 15,

For any m ∈ [ ( m -1) d M +1 , md M ] , we have E v m , D m [ ∇ ˆ f m ( w m , v m ; D m ) ] = F m ( w ) from Lemma 1 in FGD [23].

Plugging in the above result and Equation 18 in Equation 14,

For i th row in Equation 13,

Putting Equation 22 in Equation 13,

Next, we formulate the norm of the forward gradient ∇ ˆ f .

Lemma I.5 (Norm of the Forward Gradient) . Under Assumption I.2, and at the participation rate of s , M participating clients training weights w ∈ R d through random perturbations v ∈ R d in SPRY derives the accumulated forward gradient ∇ ˆ f ( w,v ; D ) such that,

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, σ 2 g is the upper bound of the global gradient variance. For a total of C classes in D , we define n c as the sample count of the c th class, α c is Dirichlet concentration parameter for the c th class, D is the size of the global data. For a client m ; n m,c is the sample count of the c th class for client m , and D m is the size of the data of client m .

Proof. The proof follows a similar style of Lemma 2 in MEZO [18]. The difference in our setting is that we have ∇ ˆ f ( w,v ; D ) which is an aggregate of ∇ ˆ f m ( w,v ; D ) derived from the federated clients, while MEZO has results under the setting of a single client.

From Theorem I.4 we have,

where

using Equation 10. The right hand side expectation is also under the randomness of partial perturbations. Here, K is perturbation count, and |D| is the size of the combined dataset D = ∪ m ∈ [ M ] D m .

To compute the second moment for a client m ,

For simplicity, let a and b be two arbitrary vectors representing E [ ∇ f m ( w m ; ( x, y c ))] and E [ ∇ f m ( w m ; ( x, y c )) T ] , respectively.

For the sum over all perturbations, we have two cases (all the expectations are under the randomness of partial v ),

̸

If i = j . (Occurs K times)

For all such v i with i ∈ [ K ] ,

Plugging in the results of the above two cases in Equation 29,

Here, the randomness is under the sampled data.

For the sum over samples of D , we have two cases,

̸

Combining both the cases, we get

Plugging in Equation 40 in Equation 37,

Hence for client m , the expected norm of forward gradients under randomness of perturbations v is,

For i th row of E [ ∇ ˆ f ( w,v ; D ) ] of Equation 28,

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 i = [ ( i -1) d M +1 , id M ] ,

Since ∇ ˆ f i , ∀ ˜ M⊂M are independent of each other, we can compute the norm of ∇ ˆ f as follows,

Using Assumption I.2 and limited participation rate of s ,

Rearranging the terms, we get,

I.4 Convergence Rate of SPRY

The general template for the convergence analysis of SPRY is similar to FEDADAM, hence we will follow Theorem 2 of AFO [25]. Our aim here is to highlight the differences in treatment of our gradient estimator ∇ ˆ f m ∀ m ∈ [ M ] , and the aggregate global gradient ∇ f as shown in Equation 9.

Theorem I.6. Under the assumptions on L -smoothness (Asmp I.1), bounded global variance σ 2 g of accumulated gradients (Asmp I.2), and bound on gradient magnitude G (Asmp I.3) and the following conditions on the local learning rate η ℓ ,

SPRY satisfies the following bound,

where R is the total round count, w ∈ R d are the trainable weights, v ∈ R d are the random perturbations, K is the count of random perturbations per batch, η is the global learning rate, τ is adaptability hyperparameter, s is client sampling rate. The rest of the symbols are defined in Theorem I.4.

Proof. As shown in Equation 7, an update of the model weights w at server-side in SPRY looks like,

Using Assumption I.1 and then the server-side update rule, we have

Taking expectation over randomness of round r and simplifying the terms,

Bounds for R 2 are derived in the exactly same manner as 'Bounding R 2 ' in Theorem 2 of FEDADAM [25],

Bounding R 1 has a different treatment due to the distinct aggregation strategy of SPRY:

Since ∆ ( r ) is piece-wise made of aggregations of forward gradients for several parts of the model weights, as shown in Equation 12, we first center each gradient piece for the computation of the squared norm,

Using ab ≤ ( a 2 + b 2 ) / 2 ,

Using the result of Lemma I.5,

Putting R 1 and R 2 bounds in Equation 58,

Summing over r = 0 to R -1 and using telescoping sum, we get

Bounding R 4 follows a similar derivation as 'Bounding R 1 ',

Updating Equation 65 with the bounds of R 4 ,

Rearranging the terms,

Getting a lower bound for the left hand side term through using the fact √ Λ ( r -1) ≤ η ℓ KG from Theorem 2 of AFO [25],

Considering the coefficients of E ||∇ f ( w ( r ) ) || 2 terms, conditioning on the following inequality,

We get the following condition on the local learning rate,

for the following bound on the gradient norm,

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-based counterparts, (b) SPRY achieving higher accuracy than its zeroorder counterparts, and comparable accuracy to the backpropagation-based counterparts, (c) SPRY running faster in terms of wallclock time to convergence than its zero-order counterparts are proved in Section 5.2, 5.1, and 5.3 respectively. The theoretical claims are discussed and proved in Section 4 and Appendix I. The claims on direct application of 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 in 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 or 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 can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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? Answer: [Yes]

Justification: The limitations and potential paths for future work are discussed in Appendix C. 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 in 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The assumptions and proofs are listed in detail in Appendix I.

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

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 experimental results of the paper to the extent that it 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 of Section 5. The pseudocode is in Appendix E, and each part of SPRY is explained in 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 be 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to 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 to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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 it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to 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 instructions to faithfully reproduce the main experimental results, as 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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 new proposed method and baselines. If only a 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 as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experimental setting is described in Section 5. 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 of 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, or 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 3 differently seeded runs is given in Appendix H.3.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).

It should be clear whether the error bar is the standard deviation or 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 a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Compute resources are discussed at the onset of Section 5. 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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our work does not involve human subjects. All the experiments are done on open-source datasets, which are widely used in the natural language community. Moreover, the goal of SPRY is to enable crowd-sourcing compute resources to finetune LLMs, without making private data accessible to a third party.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The positive and negative broader impact is discussed in Appendix D. Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 as 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a 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 we have used are 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 a 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 the 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 a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Licenses for all the datasets are stated in Appendix B. Our work is within the usage bounds of each of the datasets. Dataset URLs are also provided under the references section.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: The source code (also available as a supplemental material) of SPRY comes with a detailed README.md , explaining how to run the presented experiments of our algorithm and its counterpart baselines.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 screenshots, 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 not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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?

Answer: [NA]

Justification: Our work does not involve research with or on human subjects.

Guidelines:

The answer NA means that the paper does not 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 NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* **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.\n\nWEAKNESSES:\n* **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.\n\nQUESTIONS:\n* **Q.1.** Will the SPRY algorithm work in a large scale setting of 1k+ GPUs?\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThe authors give a limitation analysis in the checklist section.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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]$.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nAuthors have done a good job at describing limitations and memory usage of their approach.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}