Paper deep dive
ZorBA: Zeroth-order Federated Fine-tuning of LLMs with Heterogeneous Block Activation
Chuiyang Meng, Ming Tang, Vincent W. S. Wong
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 88%
Last extracted: 7/21/2026, 12:29:43 AM
Summary
The paper introduces ZorBA, a zeroth-order optimization-based federated fine-tuning framework for Large Language Models (LLMs) that addresses high VRAM usage and communication overhead. ZorBA employs heterogeneous block activation, where the central server allocates different subsets of transformer blocks to clients, and utilizes shared random seeds to estimate gradients via forward passes only, eliminating gradient storage and reducing communication costs.
Entities (15)
Relation Signals (13)
ZorBA → uses → Zeroth-order Optimization
confidence 95% · ZorBA leverages zeroth-order optimization to eliminate the storage of gradients at the clients by forward passes.
ZorBA → uses → Heterogeneous Block Activation
confidence 95% · ZorBA includes a heterogeneous block activation mechanism in which the central server allocates different subsets of transformer blocks to clients
ZorBA → reduces → VRAM
confidence 92% · ZorBA reduces the VRAM usage by using zeroth-order optimization.
ZorBA → reduces → Communication Overhead
confidence 90% · ZorBA utilizes shared random seeds and the finite differences of gradients in order to reduce the communication overhead.
ZorBA → solvedby → epsilon-constraint lexicographic algorithm
confidence 88% · We propose an epsilon-constraint lexicographic algorithm to solve this problem.
ZorBA → evaluatedon → SST-2
confidence 85% · We conduct experiments on AG-News, SST-2, and SNLI datasets
ZorBA → evaluatedon → SNLI
confidence 85% · We conduct experiments on AG-News, SST-2, and SNLI datasets
ZorBA → evaluatedon → AG News
confidence 85% · We conduct experiments on AG-News, SST-2, and SNLI datasets
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Federated fine-tuning of large language models (LLMs) enables collaborative tuning across distributed clients. However, due to the large size of LLMs, local updates in federated learning (FL) may incur substantial video random-access memory (VRAM) usage. Moreover, frequent model exchange may lead to significant communication overhead. To tackle these challenges, in this paper we propose ZorBA, a zeroth-order optimization-based federated fine-tuning framework with heterogeneous block activation. ZorBA leverages zeroth-order optimization to eliminate the storage of gradients at the clients by forward passes. ZorBA includes a heterogeneous block activation mechanism in which the central server allocates different subsets of transformer blocks to clients in order to accelerate the convergence rate and reduce the VRAM usage. Furthermore, ZorBA utilizes shared random seeds and the finite differences of gradients in order to reduce the communication overhead. We conduct theoretical analysis to characterize the effect of block activation decisions on the convergence rate and VRAM usage. To jointly enhance the convergence rate and reduce the VRAM usage, we formulate an optimization problem to optimize the block activation decisions. We propose an $\epsilon$-constraint lexicographic algorithm to solve this problem. Experimental results show that ZorBA outperforms three federated fine-tuning baselines in VRAM usage by up to 62.41% and incurs a low communication overhead.
Tags
Links
- Source: https://arxiv.org/abs/2603.04436v1
- Canonical: https://arxiv.org/abs/2603.04436v1
Trouble viewing inline? Open PDF directly →
Full Text
100,659 characters extracted from source content.
Expand or collapse full text
ZorBA: Zeroth-order Federated Fine-tuning of LLMs with Heterogeneous Block Activation Chuiyang Meng1, Ming Tang2, and Vincent W.S. Wong1 Abstract Federated fine-tuning of large language models (LLMs) enables collaborative tuning across distributed clients. However, due to the large size of LLMs, local updates in federated learning (FL) may incur substantial video random-access memory (VRAM) usage. Moreover, frequent model exchange may lead to significant communication overhead. To tackle these challenges, in this paper we propose ZorBA, a zeroth-order optimization-based federated fine-tuning framework with heterogeneous block activation. ZorBA leverages zeroth-order optimization to eliminate the storage of gradients at the clients by forward passes. ZorBA includes a heterogeneous block activation mechanism in which the central server allocates different subsets of transformer blocks to clients in order to accelerate the convergence rate and reduce the VRAM usage. Furthermore, ZorBA utilizes shared random seeds and the finite differences of gradients in order to reduce the communication overhead. We conduct theoretical analysis to characterize the effect of block activation decisions on the convergence rate and VRAM usage. To jointly enhance the convergence rate and reduce the VRAM usage, we formulate an optimization problem to optimize the block activation decisions. We propose an ϵε-constraint lexicographic algorithm to solve this problem. Experimental results show that ZorBA outperforms three federated fine-tuning baselines in VRAM usage by up to 62.41% and incurs a low communication overhead. I Introduction Large language models (LLMs) [zhao2023survey, achiam2023gpt, grattafiori2024llama, team2024gemma, 11044734] have shown remarkable performance across numerous natural language processing tasks. Fine-tuning techniques further enhance LLMs’ ability to handle specific downstream applications [ding2023parameter, hu2022lora, liu2022few]. Meanwhile, federated learning (FL) [mcmahan2017communication, li2020federated, karimireddy2020scaffold, wang2020tackling] has emerged as a promising paradigm for collaborative training of machine learning models. By allowing multiple resource-constrained clients to locally train the models and aggregate their parameters without sharing raw data, FL facilitates the deployment of large-scale machine learning applications across distributed environments. By combining these two approaches, federated fine-tuning empowers LLMs to adapt to diverse and decentralized datasets across clients [qin2023federated, 10447454, wang2024flora, bai2024federated, 10855336, 11044514, panchal2024thinking]. However, most conventional FL approaches use backpropagation (BP) to compute the first-order gradients for model updates. This leads to two critical challenges for directly applying FL to fine-tune LLMs. First, LLMs typically contain a number of transformer blocks, resulting in model sizes ranging from hundreds of millions to billions of parameters. Storing the corresponding gradients during fine-tuning requires substantial video random-access memory (VRAM) on graphics processing units (GPUs). This requirement may exceed the VRAM capacities of resource-constrained clients, thus hindering the deployment of FL for LLM fine-tuning. Second, conventional FL cannot be used when the first-order gradients are unavailable, such as in models with non-differentiable operators or in black-box systems. Figure 1: VRAM usage during zeroth-order optimization with OPT-125M. The forward-pass activations per block include tensors for hidden tensors, projections for Q, K, and V in transformer blocks, and feed-forward networks (FFNs). To address the aforementioned challenges, zeroth-order optimization has emerged as an effective alternative to the conventional first-order FL approaches in LLM fine-tuning [shu2023zeroth, malladi2023fine, chen2023deepzero, zhang2024revisiting, ma2025revisiting, chen2024enhancing]. In particular, zeroth-order optimization replaces BP-based first-order optimization approaches with a forward-pass-only approach (i.e., a BP-free approach). Zeroth-order optimization estimates the true gradients by using finite differences of loss function values generated by random perturbation vectors. Several recent studies have successfully integrated zeroth-order optimization into federated fine-tuning [9917343, chen2023fine, qiu2023zeroth, ling2024convergence, li2024achieving, shu2024ferret]. Nevertheless, the aforementioned works applied zeroth-order optimization to all blocks, leading to three main limitations. First, as shown in [malladi2023fine], zeroth-order optimization approaches exhibit slower convergence rates when compared with first-order approaches in high-dimensional model parameter space. This is because the perturbation vectors introduce variance on the estimated gradients, which grows with the model dimension. Second, while zeroth-order approaches eliminate the need to store the gradients at the clients, they still require storing forward-pass activations. As shown in Fig. 1, these activations constitute a significant portion of VRAM usage and increase linearly with the number of activated blocks. Third, due to frequent client-server communications in FL, the high dimension of the estimated gradients may result in considerable communication overhead. Based on the aforementioned discussions, we focus on addressing the following question: Is there a zeroth-order approach enabling clients to activate fewer blocks while reducing the overhead during client-server communications? In response, we aim to develop a federated fine-tuning framework integrating zeroth-order optimization with a selective block activation strategy across distributed clients. Achieving such a goal is challenging due to the following unexplored questions: (i) How can this framework optimally allocate subsets of blocks to clients in order to jointly improve the convergence rate while reducing the VRAM usage across clients? (i) How can we improve the convergence rate of this federated fine-tuning framework? (i) Can we propose an efficient and scalable algorithm to determine the block activation decisions under multiple clients and blocks? In this work, we make the following contributions as answers to the aforementioned questions: • We propose ZorBA, a zeroth-order federated fine-tuning framework with heterogeneous block activation for LLMs. ZorBA reduces the VRAM usage by using zeroth-order optimization. It synchronizes the perturbation vectors across the clients and central server via shared random seeds, which prevents gradient leakage and reduces the communication overhead. ZorBA includes a heterogeneous block activation mechanism, enabling the clients to activate different subsets of blocks based on the convergence rate and VRAM constraints. • We analyze the convergence rate of ZorBA in the nonconvex setting. Through a toy example, we reveal key insights into optimizing block activation decisions. We theoretically quantify how these decisions impact the convergence rate and VRAM usage across clients. • We formulate a multi-objective optimization problem to jointly improve the convergence rate and reduce VRAM usage across clients. To solve this problem, we propose an ϵε-constraint lexicographic algorithm by decoupling the problem into two subproblems. First, we derive a closed-form expression for the maximal least activated block across all clients. Then, we design a greedy algorithm under VRAM constraints. It activates additional blocks to minimize the number of clients whose least popularity remains at this maximum. We obtain the Pareto front and select the block activation matrix on the front in ZorBA, balancing the convergence rate and total VRAM usage. • We conduct experiments on AG-News, SST-2, and SNLI datasets by using OPT-125M and OPT-1.3B models, respectively. We compare ZorBA with FedIT [10447454], FedZO [9917343], and DeComFL [li2024achieving]. Results show that ZorBA converges faster than FedZO and DeComFL. Moreover, ZorBA reduces the communication overhead significantly and reduces the total VRAM usage by up to 62.41%. The rest of this paper is organized as follows. Our proposed ZorBA framework is introduced in Section I. The convergence rate of ZorBA and theoretical insights are presented in Section I. The problem formulation and the proposed algorithms are shown in Section IV. Experimental results are presented in Section V. Conclusions are drawn in Section VI. Notations: We use italic upper case letters, boldface upper case letters, and boldface lower case letters to denote sets, matrices, and vectors, respectively. ℝM×NR^M× N denotes the set of real-valued matrices. d×d∈ℝd×dI_d× d ^d× d denotes the identity matrix. ∅ denotes the empty set. Mathematical operators |⋅||·|, [⋅]E[·], (⋅)⊺(·) , [⋅;⋅][·;·], d(⋅)d(·) ⟨⋅⟩ · , ⌊⋅⌋ · , , ∥⋅∥\|·\|, and ∥⋅∥op\|·\|_op denote the cardinality of a set, expectation, transpose, column-wise concatenation, derivative, inner product, floor function, 2-norm, and operator norm of a matrix, respectively. ∼ denotes “distributed as” and (⋅)U(·) denotes the uniform distribution. I ZorBA In this section, we present ZorBA, a federated fine-tuning framework. ZorBA (i) enables heterogeneous block activation in which each client updates different subsets of transformer blocks, (i) incorporates shared random seeds for independent perturbation vector generation between the clients and central server in order to reduce the communication overhead, and (i) uses zeroth-order optimization to estimate the gradients with only forward passes in order to reduce the VRAM usage. We consider a central server and N clients. Let =1,2,…,TT=\1,2,…,T\ and =1,2,…,NN=\1,2,…,N\ denote the set of T fine-tuning rounds and the set of N clients, respectively. Let ℳ=1,2,…,MM=\1,2,…,M\ denote the set of M transformer blocks of the pretrained model. Note that each transformer block is a matrix of model parameters. For illustration simplicity, we flatten all M blocks of the pretrained model and concatenate them into one vector as =[1;2;…;M]∈ℝdw=[w_1;w_2;…;w_M] ^d, where d is the model dimension. Let mw_m denote the model parameters of the m-th block, m∈ℳm . We consider that each client n∈n has a local dataset nD_n. Let n∼n ξ_n _n denote a mini-batch of training samples of nD_n. Let D denote the dataset of all clients. The central server broadcasts ¯1 w^1 to all clients at the beginning of the first fine-tuning round. The ZorBA framework is shown in Fig. 2. Figure 2: Illustration of our proposed ZorBA framework. Heterogeneous Block Activation: In practical systems, clients often have limited VRAM when fine-tuning large models. In addition, exchanging the full model update between the clients and central server may result in gradient leakage and privacy concerns. Therefore, in ZorBA, a heterogeneous block activation mechanism is introduced. In particular, based on the VRAM capacity of each client, the central server activates a subset of blocks for each client to update. Let n∈0,1Ma_n∈\0,1\^M denote the block activation decision vector for client n. we have am,n=1a_m,n=1 if the m-th block is activated for client n; otherwise, am,n=0a_m,n=0. Let =[1,2,⋯,N]∈0,1M×NA=[a_1,a_2,·s,a_N]∈\0,1\^M× N denote the block activation allocation decision matrix. Client n utilizes its decision vector na_n to determine the activated blocks of the local model. At the beginning of the t-th training round, all clients and the central server have the same model, which is obtained at the end of the previous fine-tuning round. It can be expressed in a “block form” as ¯t=[¯1t;¯2t;…;¯Mt]∈ℝd w^t=[ w_1^t; w_2^t;…; w_M^t] ^d. We denote the parameters of the m-th block of client n after block activation as ^m,nt=am,n¯mt,m∈ℳ,n∈ w_m,n^t=a_m,n w_m^t,m ,n . The activated model parameters of client n in the t-th fine-tuning round are written as ^nt=[^1,nt;^2,nt;…;^M,nt]=[a1,n¯1t;a2,n¯2t;…;aM,n¯Mt]∈ℝd w_n^t=\>[ w_1,n^t; w_2,n^t;…; w_M,n^t]=\>[a_1,n w_1^t;a_2,n w_2^t;…;a_M,n w_M^t] ^d. Shared Random Seeds: Due to the large model size, exchanging parameters between clients and the central server incurs significant communication overhead. Hence, ZorBA introduces a shared random seed mechanism that replaces model exchange with seed exchange. At the beginning of fine-tuning, the central server initializes a set of P random seeds =s1,s2,…,sPS=\s_1,s_2,…,s_P\ and shares it with all clients. These random seeds are used to initialize random number generators which sample perturbation vectors =1,2,…,PV=\v_1,v_2,…,v_P\ from a standard Gaussian distribution with zero mean and an identity matrix as the covariance matrix [panchal2024thinking]. Then, these perturbation vectors are normalized onto the unit sphere. Each perturbation vector in set V will be used for gradient estimation in zeroth-order estimation. To match the block structure of ^nt w_n^t, each perturbation vector is written as p=[p,1;p,2;…;p,M]∈ℝdv_p=[v_p,1;v_p,2;…;v_p,M] ^d, where p,mv_p,m is the perturbation vector for the m-th block initialized by random seed sps_p. Since all clients and the central server share the same set of random seeds, they can independently generate the same perturbation vectors based on the exchanged random seeds. Fine-tuning Process: Let Fn(^nt;n)F_n( w_n^t; ξ_n) denote the local loss function of client n∈n on n ξ_n with model ^nt w_n^t. In the t-th fine-tuning round, the central server randomly selects a subset of Q seeds from set S and sends them to all clients. Each client then determines the corresponding set of Q perturbation vectors, denoted by t⊆Q^t , for zeroth-order optimization. Client n∈n determines Fn(^nt;n)F_n( w_n^t; ξ_n), and then estimates the gradient of ^nt w_n^t via zeroth-order optimization by Q trials, each corresponding to a perturbation vector. Note that client n estimates only the gradients of model parameters in the activated blocks. For the frozen blocks, the corresponding gradient estimation is skipped. In each trial, client n estimates the gradient by performing two forward passes that evaluate the local loss at the current point and at a perturbed point. This is done by adding a sampled perturbation vector q∈tv_q ^t, scaled by a small smoothing constant μ, to the activated model parameters ^nt w_n^t. The constant μ controls the finite difference of the perturbation applied for gradient estimation. Then, client n calculates two losses as Fn(^nt+μq;n)F_n( w_n^t+ _q; ξ_n) and Fn(^nt;n)F_n( w_n^t; ξ_n), respectively. The finite difference of the estimated gradient of client n in the direction of the perturbation vector qv_q in the t-th fine-tuning round is given by ρq,nt=Fn(^nt+μq;n)−Fn(^nt;n)μ _v_q,n^t= F_n( w_n^t+ _q; ξ_n)-F_n( w_n^t; ξ_n)μ. By averaging all Q trials, the estimated gradient of the m-th block of client n in the t-th fine-tuning round can be determined by ∇~^m,ntFn(^nt;n)= ∇_ w_m,n^tF_n( w_n^t; ξ_n)= 1Q∑q∈tρq,ntq,m. \> 1Q _v_q ^t _v_q,n^tv_q,m. (1) After finishing local updates, the clients send the model update to the central server for aggregation. Instead of directly transmitting the full estimated gradient, client n transmits the finite differences of the estimated gradients corresponding to the sampled perturbation vectors. Since both the model ¯t w^t and the sampled perturbation vector subset tQ^t are identical between the clients and central server, transmitting only the finite difference of the estimated gradient is sufficient for the server to determine the estimated gradient. Upon receiving the finite differences of the estimated gradients from all clients, the central server updates the global model based on the estimated gradients over the activated blocks. In particular, the m-th block is updated as follows: ¯mt+1= w_m^t+1= ∑n∈am,n∑n′∈am,n′(¯mt−ηt∇~^m,ntFn(^nt;n)) \>Σ _n a_m,nΣ _n a_m,n ( w_m^t-η^t ∇_ w_m,n^tF_n( w_n^t; ξ_n) ) = = ¯mt−ηt(1Q∑n∈am,n∑n′∈am,n′∑q∈tρq,ntq,m) \> w_m^t-η^t ( 1QΣ _n a_m,nΣ _n a_m,n _v_q ^t _v_q,n^tv_q,m ) = = ¯mt−ηt∑q∈tρ¯qtq,m,m∈ℳ, \> w_m^t-η^t _v_q ^t ρ_v_q^tv_q,m, m , (2) where ηtη^t denotes the learning rate in the t-th fine-tuning round. ρ¯qt=1Q∑n∈am,n∑n′∈am,n′ρq,nt ρ_v_q^t= 1Q _n a_m,n _n a_m,n _v_q,n^t denotes the average finite difference of estimated gradients in the direction of qv_q. Let ¯t=ρ¯qt:q∈t P^t=\ ρ_v_q^t:v_q ^t\ denote the set of average finite differences for all perturbation vectors in the t-th fine-tuning round. The aggregated model is updated by ¯t+1=[¯1t+1;¯2t+1;…;¯Mt+1] w^t+1=[ w_1^t+1; w_2^t+1;…; w_M^t+1]. Then, the central server broadcasts the set of average finite differences of estimated gradients ¯t P^t to all clients. By using the shared random seeds, the server no longer broadcasts high-dimensional parameters and reduces the communication overhead significantly. Since both the model ¯t w^t and the sampled perturbation vector subset tQ^t are shared, upon receiving ¯t P^t, each client can obtain the model for the (t+1)(t+1)-th fine-tuning round as ¯t+1 w^t+1 by using eqn. (I). VRAM Usage Model: In zeroth-order optimization, the VRAM usage consists of the storage for both model parameters and forward-pass activations of all blocks. We characterize the total VRAM usage as the sum of these two components. We define ψmax,n _max,n as the VRAM capacity for client n∈n . The VRAM usage for model parameters is identical across clients, as they share the same pre-trained model architecture. Thus, for any client n∈n and fine-tuning round t∈t , the VRAM usage for model parameters is denoted as ψmd _md. Let B, L, H, and K denote the batch size, input size, hidden size, and number of attention heads, respectively. The VRAM usage for forward-pass activations of each block consists of three parts: (i) hidden state activations: ψhd=BLH _hd=BLH; (i) activations of Q, K, and V modules in the transformer block: ψqkv=3KBLH _qkv=3KBLH; (i) activations of feed-forward network (FFN): ψffn=αBLH _ffn=α BLH, where α is the FFN expansion ratio for the model. The VRAM usage for each block is given by ψblock=ψhd+ψqkv+ψffn=(α+3K+1)BLH _block= _hd+ _qkv+ _ffn=(α+3K+1)BLH. ψblock _block models the VRAM usage to store the intermediate states of the blocks at client n. We adopt an additive form since zeroth-order training evaluates multiple perturbations per round and reuses these stored states across evaluations. For client n∈n in each fine-tuning round, the VRAM usage of the total activations is the VRAM required to store the activations of all activated blocks, which is given by ψact,n(n)=(α+3K+1)BLH∑m∈ℳam,n _act,n(a_n)=(α+3K+1)BLH _m a_m,n. Thus, the total VRAM usage of client n∈n is ψtot,n(n)=ψact,n(n)+ψmd _tot,n(a_n)= _act,n(a_n)+ _md. I Theoretical Analysis of ZorBA In this section, we analyze how the block activation matrix affects the convergence rate of our proposed ZorBA. Without loss of generality, we conduct analysis under nonconvex loss functions. We denote the expected loss of client n with model ntw_n^t as fn(nt)=n∼nFn(nt;n)f_n(w_n^t)=E_ ξ_n _nF_n(w_n^t; ξ_n). We denote the global loss function as f(t)f(w^t). We denote the gradient of f(t)f(w^t) as ∇f(t)=n∈∇fn(t)∇ f(w^t)=E_n ∇ f_n(w^t). We first present the following assumptions which are widely used in the literature (e.g., [li2019convergence, wang2020tackling, 10971879]). Assumption 1. The loss function of each client n∈n is continuously differentiable and L-smooth. That is, for arbitrary two vectors ntw_n^t and ~nt w_n^t, we have fn(nt)≤fn(~nt)+⟨∇fn(~nt),nt−~nt⟩+L2‖nt−~nt‖2.f_n(w_n^t)≤ f_n( w_n^t)+ ∇ f_n( w_n^t),w_n^t- w_n^t + L2\|w_n^t- w_n^t\|^2. Assumption 2. The variance of the local stochastic gradient of each client n∈n is upper-bounded, i.e., n∼n[‖∇Fn(nt;n)−∇fn(nt)‖2]≤σ2E_ ξ_n _n [\|∇ F_n(w_n^t; ξ_n)-∇ f_n(w_n^t)\|^2 ]≤σ^2. Assumption 3. The dissimilarity between the local gradient of each client n∈n and the gradient of the global loss function is upper-bounded, i.e., ‖∇f(t)−∇fn(t)‖2≤σG2\|∇ f(w^t)-∇ f_n(w^t)\|^2≤ _G^2. Moreover, we introduce two lemmas to facilitate our convergence analysis. Lemma 1. Let n∼n[∇~Fn(t;n)]=∇~fn(t),n∈E_ ξ_n _n[ ∇F_n(w^t; ξ_n)]= ∇f_n(w^t),n . The dissimilarity between the zeroth-order and first-order gradient estimation of each client n∈n is upper-bounded, i.e., ‖∇~fn(t)−∇fn(t)‖≤12μL(d+3)32\| ∇f_n(w^t)-∇ f_n(w^t)\|≤ 12μ L(d+3) 32. This result follows from the analysis in [nesterov2017random]. Lemma 2. When μ is small enough (i.e., μ→0μ→ 0), the square norm of the zeroth-order gradient estimation of each client n∈n satisfies n∼n[‖∇~Fn(nt;n)‖2]=d+Q−1Qn∼n[‖∇Fn(nt;n)‖2]E_ ξ_n _n [\| ∇F_n(w_n^t; ξ_n)\|^2 ]= d+Q-1QE_ ξ_n _n [\|∇ F_n(w_n^t; ξ_n)\|^2 ]. This results from [malladi2023fine]. We denote Λ()=∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 (A)= _n _m \ ( a_m,n _n a_m,n )^2 \. Let ⋆w denote the optimal model. We present the convergence bound of ZorBA. Theorem 1. (Standard convergence bound of ZorBA) Under Assumptions 1−-3 and Lemmas 1−-2, if the learning rate ηtη^t is small such that ηt=η≤12LdN2η^t=η≤ 12LdN^2, the standard convergence bound of our proposed ZorBA satisfies: 1T∑t=1T‖∇f(¯t)‖2≤ 1T _t=1^T \|∇ f( w^t) \|^2≤ 1Ω(Λ())T∑t=1T(f(¯1)−f(⋆)) \> 1 ( (A))TΣ _t=1^T (f( w^1)-f(w ) ) +Θ1Ω(Λ())Λ(), + _1 ( (A)) (A), (3) where Ω(Λ())=(η2−3Lη2(d+Q−1)N2QΛ()) ( (A))= ( η2- 3Lη^2(d+Q-1)N2Q (A) ) and Θ1=14μ2L2η(d+3)3N+(ηN+3Lη2(d+Q−1)N2Q)σG2+3Lη2(d+Q−1)N2Qσ2 _1= 14μ^2L^2η(d+3)^3N+ (η N+ 3Lη^2(d+Q-1)N2Q ) _G^2+ 3Lη^2(d+Q-1)N2Qσ^2. Sketch of proof. The proof starts by showing the update of one fine-tuning round. By using Assumption 1, we expand the expected loss of the average model into two terms: a first-order descent term T1T_1 and a second-order variance term T2T_2. T1T_1 captures the inner product between the true gradient and the model update made with zeroth-order optimization. In particular, T1T_1 is upper-bounded by using the bounded stochastic-gradient variance (Assumption 2), the gradient-dissimilarity assumption (Assumption 3), and the zeroth-order estimation dissimilarity in Lemma 1, giving a bias that depends on the block-activation decisions A. T2T_2 is bounded by using Lemma 2 and Assumptions 2−-3. Lemma 2 draws a connection between the zeroth-order gradients to first-order ones. Assumptions 2−-3 characterize the gradient variance and data heterogeneity, respectively. By summing over all fine-tuning rounds, rearranging the inequality, and choosing the proper learning rate, the convergence bound is obtained. ∎ According to Theorem 1, the convergence bound of ZorBA consists of two terms: an optimality gap that diminishes as the number rounds T increases, and a non-diminishing bias term (i.e., Θ1Ω(Λ())Λ() _1 ( (A)) (A)). The bias term quantifies the impact of zeroth-order gradient estimation error, data heterogeneity, and local gradient variance on the convergence. However, to guarantee convergence of ZorBA, Theorem 1 requires a learning rate η=(1dN2)η=O( 1dN^2). In practical systems, the model dimension d can exceed billions, and the number of clients N may be large. Thus, the required learning rate may be impractically small. To overcome this challenge, we introduce a condition and adopt a supporting lemma from [malladi2023fine], allowing us to derive a convergence bound independent of d. Condition 1. Let G(t)=maxn∈maxn∼n‖∇Fn(nt;n)‖G(w^t)= _n _ ξ_n _n\|∇ F_n(w_n^t; ξ_n)\|, there exists a Hessian matrix (t)⪯Ld×dH(w^t) LI_d× d such that for arbitrary ∈ℝdw ^d with ‖−t‖≤ηtdG(t)\|w-w^t\|≤η^tdG(w^t), we have ∇2f()⪯(t)∇^2f(w) (w^t). The effective rank of (t)H(w^t) satisfies ((t))=Tr((t))‖(t)‖op≤κerank(H(w^t))= (H(w^t))\|H(w^t)\|_op≤κ. This is from [malladi2023fine]. Condition 1 shows that the Hessian of local loss can be approximated by a matrix such as its curvature is dominated by only κ effective directions. Hence, the convergence bound depends on κ instead of d. Lemma 3. Let F(;)F(w; ξ) denote the global stochastic loss such that f()=∼[F(;)]f(w)=E_ ξ [F(w; ξ)]. We define ()=[∇F(;)(∇F(;))⊺]−∇f()(∇f())⊺ (w)=E [∇ F(w; ξ) (∇ F(w; ξ) ) ]-∇ f(w) (∇ f(w) ) . The outer product of the global model difference between two consecutive rounds satisfies [(¯t+1−¯t)(¯t+1−¯t)⊺] [( w^t+1- w^t)( w^t+1- w^t) ] = = (1+d−2d+2)(ηt)2N(∇f(¯t)(∇f(¯t))⊺+1N(¯t)) \> (1+ d-2d+2 )(η^t)^2N (∇ f( w^t) (∇ f( w^t) ) + 1N ( w^t) ) +d+2(ηt)2N(‖∇f(¯t)‖2+1NTr((¯t)))d×d. + dd+2(η^t)^2N ( \|∇ f( w^t) \|^2+ 1N ( ( w^t) ) )I_d× d. (4) This result follows from [malladi2023fine]. Then, we present the dimension-free convergence bound of our proposed ZorBA. Theorem 2. (Dimension-free convergence bound of ZorBA) Under Assumptions 1−-3, Lemma 1, Condition 1, and Lemma 3, if the learning rate ηtη^t satisfies ηt=η≤d+22LdN2(κ+2)η^t=η≤ d+22LdN^2(κ+2), the dimension-free convergence bound of our proposed ZorBA is: 1T∑t=1T‖∇f(¯t)‖2≤ 1T _t=1^T \|∇ f ( w^t ) \|^2≤ 1Φ(Λ())T(f(¯1)−f(¯⋆)) \> 1 ( (A))T (f( w^1)-f( w ) ) +Θ2Φ(Λ())Λ()⏟bias term, + _2 ( (A)) (A)_bias term, (5) where Φ(Λ())=(η2−dη2NL(κ+2)2(d+2)Λ()) ( (A))= ( η2- dη^2NL(κ+2)2(d+2) (A) ) and Θ2=ηN(σG2+14μ2L2(d+3)3)+dη2L(κ+2)σ22(d+2) _2=η N ( _G^2+ 14μ^2L^2(d+3)^3 )+ dη^2L(κ+2)σ^22(d+2). Sketch of proof. To remove the dependency on the model dimension d, we expand the expected loss of the average model to second order and invoke the local κ-effective-rank condition on the Hessian (Condition 1). Lemma 3 shows the outer product of model differences, which further connects the Hessian term to the gradient norm and a trace term without d. The key insight is bounding this Hessian term by leveraging the gradient covariance matrix. The resulting convergence bound has a similar structure to Theorem 1, but with a new bias term which incorporates κ and eliminates the learning rate’s dependence on d. ∎ Theorem 2 introduces a non-diminishing bias term which reflects the impact of zeroth-order gradient error, data heterogeneity, and local gradient variance while ensuring the learning rate is independent of model dimension. To improve convergence, it is essential to design the block activation matrix A to minimize this bias. Although the involved constants (e.g., L, κ, σG _G, σ) are generally intractable, Lemma 4 shows that accelerating the convergence is equivalent to minimizing Λ() (A), regardless of these constants. Figure 3: An example of different block activation decisions on three clients. We consider that each client’s model has three blocks. Each column denotes the blocks for each client. The blue square denotes the activated block. The pink square denotes the frozen block. The value of the block in the m-th row and n-th column corresponds to the aggregation weight, i.e., am,n∑n′∈am,n′ a_m,n _n a_m,n . Lemma 4. (Monotonicity of bias term) The optimal block activation matrix A which minimizes the bias term Θ2Φ(Λ())Λ() _2 ( (A)) (A) is equivalent to the minimizer of Λ() (A). Proof. The first-order derivative of the bias term Θ2Φ(Λ())Λ() _2 ( (A)) (A) with respect to (w.r.t.) Λ() (A) can be derived as dΘ2Φ(Λ())Λ()dΛ()=Θ2Φ(Λ())+dη2NL(κ+2)2(d+2)Θ2Λ()(Φ(Λ()))2. d _2 ( (A)) (A)d (A)= _2 ( (A))+ dη^2NL(κ+2)2(d+2) _2 (A) ( ( (A)) )^2. (6) In particular, Θ2 _2 and Λ() (A) are both positive. When η≤d+22LdN2(κ+2)η≤ d+22LdN^2(κ+2), Φ(Λ()) ( (A)) is non-negative. The bias term is monotonically increasing w.r.t. Λ() (A) and is minimized at the lower bound of Λ() (A). ∎ To further demonstrate the impact of the block activation matrix A on the convergence rate, we present an example in Fig. 3. Based on Lemma 4, we use Λ() (A) to characterize the bias term and the convergence rate. We present several important observations as follows. Observation 1. At a large scale, increasing the number of activated blocks generally reduces Λ , thereby leading to faster convergence. The optimal convergence rate is achieved when all clients activate all the blocks, i.e., full-block fine-tuning. Observation 2. At a small scale, activating more blocks does not always improve the convergence. Observation 3. In some cases, activating more blocks can degrade the convergence. For example, case 5 activates one more block than case 11, but case 11 has a lower value of Λ . Observation 4. The convergence rate can vary even when the total number of activated blocks remains the same. For example, both cases 8 and 9 activate the same total number of blocks. Although case 9 has a more imbalanced block allocation, it achieves a lower value of Λ . Based on the above observations, we provide insights into minimizing Λ() (A). We define cm(m,:)=∑n∈am,nc_m(a_m,:)= _n a_m,n as the popularity of block m. It is the number of clients that activate block m. For each client n, we define its least popularity as the minimum popularity among all blocks activated by client n. The least popularity of client n satisfies c¯n()=minm:am,n=1cm(m,:)≥1 c_n(A)= _m:a_m,n=1c_m(a_m,:)≥ 1. Let ¯()=[c¯1,c¯2,…,c¯N]∈ℝN c(A)=[ c_1, c_2,…, c_N] ^N denote the vector of the least popularity of all clients. In particular, Λ() (A) satisfies Λ()=∑n∈1c¯n2() (A)= _n 1 c_n^2(A). To analyze the optimization of the convergence rate, we start by showing Schur-convexity and majorization. Lemma 5. (Schur-convexity and majorization of Λ() (A)) We define ()=sort↓(¯())l(A)= *sort ( c(A)) as a vector of the least popularities of all clients in a descending order. If there exists another block activation matrix ′A such that ()l(A) is majorized by (′)l(A ), we have Λ()<Λ(′) (A)< (A ). Proof. We start by recalling majorization and Schur-convexity. We define two descending vectors ↓,↓∈ℝNx ,y ^N, where xn↓x_n and yn↓y_n denote the n-th largest value of ↓x and ↓y , respectively. If for any k=1,2,…,Nk=1,2,…,N, ↓x and ↓y satisfy ∑n=1kxi↓≤∑n=1kyi↓ _n=1^kx_i ≤ _n=1^ky_i and ∑n∈xi↓=∑n∈yi↓ _n x_i = _n y_i , we say ↓x is majorized by ↓y , i.e., ↓≺↓x . Then, we introduce Schur-convexity. We say a function f:ℝN→ℝf:R^N is Schur-convex if ↓≺↓x , f(↓)≤f(↓)f(x )≤ f(y ). Due to the property of Schur-convexity, if f is a convex function, then ∑n∈f(xn↓) _n f(x_n ) is Schur-convex. In our settings, 1c¯n2() 1 c_n^2(A) is convex and monotonically decreasing w.r.t. c¯n() c_n(A) for any n∈n , Λ()=∑n∈1c¯n2() (A)= _n 1 c_n^2(A) is Schur-convex. Then, based on the definition of majorization, we have Λ()<Λ(′) (A)< (A ) if ()≺(′)l(A) (A ). ∎ Lemma 5 provides key insight into optimizing the convergence rate, which is summarized in the following theorem. Theorem 3. (Dominance of the least popular blocks) The value of Λ() (A) depends on ()l(A). Minimizing Λ() (A) is equivalent to maximizing the least popularity of all clients (i.e., maxminn∈c¯n() _A _n \ c_n(A)\) and then minimizing the number of clients which obtain this minimum. Proof. Since Λ() (A) is Schur-convex in the sorted vector of least popularities ()l(A), Lemma 5 shows that minimizing Λ() (A) is achieved by flattening ()l(A) as much as possible. Flattening occurs first by maximizing the minimal value across c¯n(),n∈ c_n(A),n (i.e., maximizing the least popularity across all clients). Then, with that minimum fixed, we reduce the number of clients which obtain this minimum. ∎ In addition, we present the following proposition. Proposition 1. (Λ -optimal multiplicity) Λ() (A) depends only on the multiset of ¯() c(A). Moreover, Λ() (A) is permutation invariant to ¯() c(A). That is, for any permutation π, if two block activation matrices A and ′A satisfy ¯(′)=π(¯()) c(A )=π( c(A)), they achieve a similar convergence rate. Proof. The value of Λ() (A) depends only on the multiset of least popularities c¯1(),c¯2(),…,c¯N()\ c_1(A), c_2(A),…, c_N(A)\. The ordering of its elements plays no role because Λ()=∑n∈1c¯n2() (A)= _n 1 c_n^2(A) is a symmetric sum. Therefore, for any permutation π, replacing ¯() c(A) with π(¯())π( c(A)) leaves the value of the sum unchanged. ∎ The above analysis uncovers a counterintuitive insight, which is summarized in the following remark. Remark 1. While the total number of activated blocks affects the convergence rate, the factor that controls the convergence is how the least popularities are distributed across clients. In general, increasing the total number of activated blocks across all clients enhances the convergence rate at the cost of higher VRAM usage. This reveals a fundamental trade-off between convergence performance and VRAM usage: activating more blocks accelerates optimization but demands more VRAM resources on each client. IV Problem Formulation and ϵε-constraint Lexicographic Algorithm IV-A Problem Formulation Based on the aforementioned theoretical analysis, we aim to jointly minimize Λ() (A) to improve the convergence rate of ZorBA and reduce each client’s VRAM usage. We formulate the problem as follows: 1: _1: minimize *minimize_A Λ(),ψtot,1(1),ψtot,2(2), (A), _tot,1(a_1), _tot,2(a_2), ⋯,ψtot,N(N) ·s, _tot,N(a_N) (7a) subjectto \>to ∑m∈ℳam,n≥1,n∈, _m a_m,n≥ 1, n , (7b) ∑n∈am,n≥1,m∈ℳ, _n a_m,n≥ 1, m , (7c) am,n∈0,1,m∈ℳ,n∈, a_m,n∈\0,1\, m ,n , (7d) where constraint (7b) ensures that each client activates and updates at least one block for fine-tuning. Constraint (7c) ensures that each block is activated and updated by at least one client. Constraint (7d) ensures that the block activation decisions are binary variables. The challenges of solving 1P_1 are twofold. First, 1P_1 is an NP-hard integer programming problem. Second, 1P_1 contains multiple objectives which depend on the block activation decisions. To address these challenges, we propose an ϵε-constraint lexicographic algorithm to achieve a close-to-optimal solution in the following subsection. IV-B ϵε-constraint Lexicographic Algorithm We adopt the ϵε-constraint method [haimes1971bicriterion] to transform the original multi-objective problem into a single-objective optimization problem, which allows us to balance the convergence performance with VRAM constraints. In particular, for each client n, we denote the VRAM usage reduction as ϵn=τψmax,n _n=τ _max,n, where τ denotes the desired reduction ratio in VRAM usage. To explore a diverse set of trade-offs, we sample τ independently from a uniform distribution as τ∼[0,1]τ [0,1]. Let ϵ=[ϵ1,ϵ2,…,ϵN]∈ℝN ε=[ _1, _2,…, _N] ^N denote the resulting VRAM reduction vector for all clients. We generate E such vectors and collect them in the set ℰE, which is used to evaluate the candidate solutions under different VRAM reduction scenarios across clients. We solve 1P_1 for each reduction vector ϵ∈ℰ ε . Under a given ϵ ε, problem 1P_1 is reformulated as 2: _2: minimize *minimize_A Λ() (A) (8a) subjectto \>to ψtot,n(n)=ψmax,n−ϵn,n∈, _tot,n(a_n)= _max,n- _n, n , (8b) constraints(7b)−(7d). \>( c12)\!-\!( c14). Problem 2P_2 is difficult to solve since matrix A is a binary matrix with dimension M×NM× N. For each ϵ ε, solving problem 2P_2 incurs an exponential computation complexity of up to (2MN)O(2^MN), which is not scalable as M and N grow. To reduce the complexity, we propose a lexicographic optimization algorithm. Based on Theorem 3, we equivalently transform 2P_2 into two subproblems, namely IP_I and IP_I. Problem IP_I aims to maximize the least popularity of all clients. Let γ⋆γ denote the optimal objective value to IP_I. Problem IP_I is the least popularity adjustment problem. The objective is to minimize the number of clients whose least popularity remains γ⋆γ after additional block activation. Least popularity maximization problem IP_I: IP_I is formulated as follows: I: _I: maximize *maximize_A minn∈c¯n() _n \ c_n(A)\ (9a) subjectto \>to constraints(7b)−(7d),(8b). \>( c12)\!-\!( c14),( c21). The optimal objective value of IP_I is given by Theorem 4. Theorem 4. (Optimal least popularity) Let ℬB denote any non-empty subset of blocks, i.e., ℬ⊆ℳB and ℬ≠∅B≠ . γ⋆γ satisfies γ⋆=minℬ⌊∑n∈minψmax,n−ψmd−ϵn(α+3K+1)BLH,|ℬ||ℬ|⌋. γ = _B \ _n \ _max,n- _md- _n(α+3K+1)BLH, \ \. (10) Proof. For client n∈n , given its VRAM capacity, the maximum block activation budget is given by λn⋆=ψmax,n−ψmd−ϵn(α+3K+1)BLH _n = _max,n- _md- _n(α+3K+1)BLH. We consider that each block in the subset ℬ⊆ℳB is to be activated by at least γ clients. Then, the total number of activated blocks in this subset ℬB is |ℬ|γ |γ. However, client n∈n can only activate at most minλn⋆,|ℬ| \ _n , |\ blocks. The total number of activated blocks of all clients for subset ℬB is ∑n∈minλn⋆,|ℬ| _n \ _n , |\. To make γ feasible, we must have |ℬ|γ≤∑n∈minλn⋆,|ℬ| |γ≤ _n \ _n , |\. Note that this inequality must hold for all possible subsets ℬB. Hence, the optimal γ is γ⋆= γ = minℬγ:|ℬ|γ≤∑n∈minλn⋆,|ℬ|,ℬ⊆ℳ,ℬ≠∅ \> _B \γ: |γ≤ _n \ _n , |\,B ,B≠ \ = = minℬ⌊∑n∈minψmax,n−ψmd−ϵn(α+3K+1)BLH,|ℬ||ℬ|⌋. \> _B \ _n \ _max,n- _md- _n(α+3K+1)BLH, \ \. (11) ∎ After γ⋆γ has been determined, we apply Dinic’s algorithm [dinitz2006dinitz] to construct the optimal block activation matrix for problem IP_I. We denote this matrix as ~∈0,1M×N A∈\0,1\^M× N, which serves as the initial block activation matrix. Based on ~ A, we obtain the initial popularity of each block and initial least popularity of each client as c~m,m∈ℳ c_m,m and c¯~n,n∈ c_n,n , respectively. The initial popularity vector and initial least popularity vector are denoted as ~∈ℝM c ^M and ¯~∈ℝN c ^N, respectively. Least popularity adjustment problem IP_I: Note that ~ A may not be optimal to problem 2P_2. Hence, based on Theorem 3, we formulate problem IP_I to activate as many additional blocks as possible on top of ~ A such that the number of clients with the least popularity γ⋆γ is minimized. Let xm,n∈0,1x_m,n∈\0,1\ denote the additional block activation decision. We set xm,n=1x_m,n=1 to additionally activate the m-th block of client n when a~m,n=0 a_m,n=0. Otherwise, xm,n=0x_m,n=0. That is, xm,na~m,n=0,m∈ℳ,n∈. x_m,n a_m,n=0, m ,n . (12) Let ∈0,1M×NX∈\0,1\^M× N denote the additional block activation decision matrix. Note that the total number of blocks that each client activates must satisfy the VRAM constraint. That is, ∑m∈ℳ(xm,n+a~m,n)≤λn⋆,n∈. _m (x_m,n+ a_m,n)≤ _n , n . (13) Then, given the decision matrix X, we obtain the popularity of each block and the least popularity of each client as cm c_m =c~m+∑n∈xm,n,m∈ℳ, = c_m+ _n x_m,n, m , (14) c¯n c_n =minm:a~m,n+xm,n=1cm,n∈. = _m: a_m,n+x_m,n=1c_m, n . (15) In addition, we introduce an auxiliary indicator variable yn∈0,1y_n∈\0,1\ to characterize whether client n’s least popularity is still γ⋆γ after additional block activation. In particular, if c¯n=γ⋆ c_n=γ , then yn=1y_n=1. Otherwise, yn=0y_n=0. If client n activates the m-th block after additional block activation and its least popularity is larger than γ⋆γ , then the popularity of the m-th block must be larger than γ⋆γ . That is, for xm,n+a~m,n=1x_m,n+ a_m,n=1 and yn=0y_n=0, cmc_m must satisfy cm≥γ⋆+1,m∈ℳ,n∈c_m≥γ +1,m ,n . We use the big-M method [griva2008linear] to linearize this condition as cm≥(γ⋆+1)−β1(1−(xm,n+a~m,n))−β2yn, c_m≥(γ +1)- _1 (1-(x_m,n+ a_m,n) )- _2y_n, m∈ℳ,n∈, m ,n , (16) where β1 _1 and β2 _2 are two positive constants. When β1=β2=γ⋆+1 _1= _2=γ +1, the aforementioned condition always holds. Thus, inequality (IV-B) is equivalent to cm≥ c_m\ ≥ (γ⋆+1)−(γ⋆+1)(1−(xm,n+a~m,n)) \>(γ +1)-(γ +1) (1-(x_m,n+ a_m,n) ) −(γ⋆+1)yn,m∈ℳ,n∈. -(γ +1)y_n, m ,n . (17) In summary, problem IP_I can be formulated as follows: I: _I: minimize,ynn∈ *minimize_X,\y_n\_n ∑n∈yn _n y_n (18a) subjectto \>to constraints(12)−(15),(IV-B), \>( eq:x_constraint)\!-\!( eq:c_n_constraint),\>( eq:big-M-v2), xm,n,yn∈0,1,m∈ℳ,n∈. x_m,n,y_n∈\0,1\, m ,n . (18b) To solve problem IP_I, we propose a greedy update algorithm. We start by introducing several terms. Let ℒ=m:c~m=γ⋆,m∈ℳL=\m: c_m=γ ,m \ denote the set of blocks with its initial popularity as γ⋆γ . Let ℱ=n:c¯~n=γ⋆,n∈F=\n: c_n=γ ,n \ denote the set of clients whose initial least popularity is γ⋆γ . Let ∈ℝNr ^N denote the vector of remaining VRAM budgets, in which each element rn=⌊λn⋆⌋−∑m∈ℳa~m,nr_n= _n - _m a_m,n represents the number of additional blocks client n can activate after the initial block activation. For client n∈n , we denote n=m∈ℒ:a~m,n=1W_n=\m : a_m,n=1\ as the “bottleneck” set. It includes those blocks which are activated while satisfying the initial popularity. The greedy update algorithm can be summarized as follows: We define Gain_ls[m]Gain\_ls[m] as the number of clients whose current least popularity is γ⋆γ and would increase to γ⋆+1γ +1 if their m-th block were activated. Among all blocks with popularity γ⋆γ (i.e., m∈ℒm ), we denote m⋆m as the “most valuable” block that maximally reduces the number of clients with least popularity γ⋆γ when activated. We then activate the m⋆m -th block by assigning it to any client n with remaining VRAM budget to activate more blocks (i.e., rn>0r_n>0) and has not yet activated m⋆m -th block (i.e., a~m⋆,n=0 a_m ,n=0). Thus, the popularity of the m⋆m -th block increases from γ⋆γ to γ⋆+1γ +1, i.e., cm⋆←γ⋆+1c_m ←γ +1. Next, for every client n∈n whose “bottleneck” set nW_n includes block m⋆m , we remove m⋆m from nW_n. If a client’s bottleneck set becomes empty, it is removed from the set of remaining clients ℱF. This procedure is repeated iteratively until the number of clients with the least popularity γ⋆γ is minimized. We summarize this algorithm in Algorithm 1. To solve problem 1P_1, we iterate over each ϵ∈ℰ ε . For each ϵ ε, we solve the corresponding subproblems IP_I and IP_I, and obtain the optimal block activation matrix ϵ⋆A_ ε . Based on ϵ⋆A_ ε , we can determine the corresponding value of Λ(ϵ⋆) (A_ ε ) and ψtot,n(ϵ⋆),n∈ _tot,n(A_ ε ),n . We denote the total usage of all clients as ∑n∈ψtot,n(ϵ⋆) _n _tot,n(A_ ε ). By going over all ϵ∈ℰ ε , we can form a Pareto front which characterizes the trade-off between Λ and total VRAM usage of all clients. We can choose the optimal block activation matrix ⋆A on the Pareto front to balance the VRAM usage and the convergence rate for ZorBA. The algorithm to solve problem 1P_1 is shown in Algorithm 2. Algorithm 1 Greedy update algorithm for IP_I 1: Input: Initial block activation matrix ~ A; initial popularity vector ~ c; initial least popularity vector ¯~ c; optimal least popularity γ⋆γ ; n,n∈W_n,n ; ℒL; ℱF; r; λn⋆,n∈ _n ,n ; Gain_lsGain\_ls = []. 2: While ℒ≠∅L≠ , ℱ≠∅F≠ , and ∑n∈λn⋆>0 _n _n >0 do 3: Create a candidate block index set :=m∈ℒ:∃n∈,suchthata~m,n=0andrn>0C:=\m :∃ n ,\>such\>that\> a_m,n=0\>and\>r_n>0\. 4: For m∈m do 5: Gain_lsGain\_ls[m] := |n∈ℱ:m∈n||\n :\>m _n\|. 6: End for 7: m⋆:=argmaxm∈Gain_ls[m]m := *arg\,max_m \Gain\_ls[m]\. 8: Randomly select client n such that rn>0r_n>0 and a~m⋆,n=0 a_m ,n=0. 9: xm⋆,n:=1x_m ,n:=1, rn:=rn−1r_n:=r_n-1. 10: Remove m⋆m from ℒL. 11: For n∈n and a~m⋆,n=1 a_m ,n=1 do 12: Remove m⋆m from nW_n. Remove n from ℱF if n=∅W_n= . 13: End for 14: End while 15: :=~+A:= A+X. 16: Output: Updated block activation matrix A. By using Dinic’s algorithm, solving problem IP_I incurs a computation complexity of (MNM+Nlog2N)O(MN M+N _2N). The greedy update algorithm in the second stage incurs a computation complexity of (minM,NMN)O( \M,N\MN). By going through ℰE, our proposed ϵε-constraint lexicographic algorithm incurs a total computation complexity of (EMN(M+Nlog2N+minM,N))O(EMN( M+N _2N+ \M,N\)), which is significantly lower than that of directly solving problem 1P_1. The overall workflow of our proposed ZorBA is shown in Algorithm 3. Algorithm 2 ϵε-constraint lexicographic algorithm for 1P_1 1: Input: VRAM capacities ψmax,n _ ,n, n∈n ; initial constraint parameter vector ℰE. 2: For ϵ∈ℰ ε do 3: Solve problem IP_I to determine γ⋆γ by using eqn. (10). 4: Construct the initial block activation matrix ~ A. 5: Obtain the updated block activation matrix ϵ⋆A_ ε using Alg. 1. 6: Determine Λ(ϵ⋆) (A_ ε ) and ψtot,n(ϵ⋆),n∈ _tot,n(A_ ε ),n . 7: Record tuple (ϵ ε, ϵ⋆A_ ε , Λ(ϵ⋆) (A_ ε ), ∑n∈ψtot,n(ϵ⋆) _n _tot,n(A_ ε )). 8: End for 9: Select ⋆A on the Pareto front ∑n∈ψtot,n(ϵ⋆),Λ(ϵ⋆)ϵ∈ℰ\ _n _tot,n(A_ ε ), (A_ ε )\_ ε . 10: Output: Optimal block activation matrix ⋆A . Algorithm 3 ZorBA Algorithm 1: Input: Initial model ¯1 w^1; local fine-tuning datasets nD_n, n∈n ; initial random seeds S; learning rate η. 2: The central server solves problem 1P_1 to determine the block activation matrix ⋆A by using Algorithm 2. 3: For client n∈n in parallel do 4: Activate local model parameters. 5: End for 6: For t∈t do 7: The central server selects a subset tQ^t from S. 8: For client n∈n in parallel do 9: Estimate the gradients ∇~Fn(^nt;n) ∇F_n( w_n^t; ξ_n) based on eqn. (1) and record the finite differences of the estimated gradients. 10: Transmit the finite differences of the estimated gradients to the central server. 11: End for 12: The central server performs aggregation based on eqn. (I), obtains and broadcasts ¯t P^t back to the clients. 13: For client n∈n in parallel do 14: Determine ¯t+1 w^t+1 based on ¯t P^t. 15: End for 16: End for 17: Output: fine-tuned model ¯T+1 w^T+1. V Performance Evaluation V-A Simulation Setup We conduct federated LLM fine-tuning over N=50N=50 clients and a central server. We use OPT-125M and OPT-1.3B [zhang2022opt] as the local models for experiments. In particular, OPT-125M and OPT-1.3B have 12 and 24 transformer blocks, respectively. We conduct experiments on text classification datasets, including AG-News [AG-news], SST-2 [socher2013recursive], and SNLI [bowman2015large]. We use the Dirichlet distribution Dir(1.0)Dir(1.0) to create non-independent and identically distributed data partitioning across clients’ local datasets. To characterize the VRAM capacity, we denote the maximum VRAM of all clients as the sum of the number of model parameters and activations of all blocks. We set the minimum VRAM of all clients to the sum of the number of model parameters and activations of one block. The VRAM capacity of each client follows a uniform distribution between the minimum and maximum of VRAM. We use the number of transmitted parameters to characterize the communication overhead. In addition, we set α=4α=4, η=5×10−5η=5× 10^-5, μ=10−4μ=10^-4, B=8B=8, E=1000E=1000, P=4096P=4096, Q=10Q=10, and T=500T=500. We compare the performance of ZorBA with the following baseline schemes: • FedIT [10447454]: Clients perform first-order BP with all blocks. • FedZO [9917343]: Clients sample one perturbation vector to perform zeroth-order optimization with all blocks during each fine-tuning round. Then, clients exchange the estimated gradients with the central server. • DeComFL [li2024achieving]: Clients use a shared perturbation vector to perform zeroth-order optimization with all blocks during each fine-tuning round. Then, clients exchange the finite differences of estimated gradients with the server. V-B Experiments V-B1 Comparison of the Convergence To show the convergence rate, we compare the number of training rounds each approach takes to achieve a target accuracy. In particular, for AG-News, SST-2, and SNLI, we set the target accuracies as 0.8, 0.8, and 0.72, respectively. Since the central server determines the block activation matrix in prior, we present the total VRAM usage of all clients per fine-tuning round, which remains consistent during fine-tuning. Furthermore, we use the number of parameters transmitted between clients and the central server to characterize the total communication overhead. Results in Table I show that our proposed ZorBA converges faster than DeComFL by up to 23.76%. ZorBA also converges faster than FedZO in most cases. These results highlight the benefit of optimizing the heterogeneous block activation matrix to accelerate the convergence. V-B2 Comparison of the Total VRAM Usage and Communication Overhead We compare the total VRAM usage and communication overhead in Table I. Our proposed ZorBA achieves a much lower total VRAM usage, which is up to 62.41%, 54.75%, and 54.75% lower than FedIT, FedZO, and DeComFL, respectively. Note that FedZO and DeComFL incur the same total VRAM usage since they both activate all blocks. In addition, ZorBA incurs negligible communication overhead when compared with FedIT and FedZO. ZorBA also achieves a comparable total communication overhead to DeComFL. Table I: Comparison of the number of fine-tuning rounds required to achieve the target average testing accuracy. Model Datasets ZorBA FedIT FedZO DeComFL OPT-125M AG-News 138 49 155 181 SST-2 167 70 161 205 SNLI 231 112 255 246 OPT-1.3B AG-News 61 23 79 68 SST-2 99 48 104 104 SNLI 157 82 186 173 Table I: Comparison of the total VRAM usage of all clients per round (in GB) and total communication overhead (i.e., Commun.) of all clients to achieve the target accuracy. Model Datasets Metrics ZorBA FedIT FedZO DeComFL OPT-125M AG-News VRAM 47.31 118.67 95.37 95.37 Commun. 1.24×1051.24× 10^5 6.13×10116.13× 10^11 1.94×10121.94× 10^12 1.36×1051.36× 10^5 SST-2 VRAM 59.32 118.67 95.37 95.37 Commun. 1.67×1051.67× 10^5 8.75×10118.75× 10^11 2.01×10122.01× 10^12 1.54×1051.54× 10^5 SNLI VRAM 131.39 190.74 167.44 167.44 Commun. 2.43×1052.43× 10^5 1.40×10121.40× 10^12 3.19×10123.19× 10^12 1.85×1051.85× 10^5 OPT-1.3B AG-News VRAM 538.04 1431.17 1189.03 1189.03 Commun. 7.93×1047.93× 10^4 2.99×10122.99× 10^12 1.03×10131.03× 10^13 9.18×1049.18× 10^4 SST-2 VRAM 794.49 1431.17 1189.03 1189.03 Commun. 1.34×1051.34× 10^5 6.24×10126.24× 10^12 1.35×10131.35× 10^13 1.40×1051.40× 10^5 SNLI VRAM 1583.55 2378.05 2315.90 2315.90 Commun. 2.27×1052.27× 10^5 1.07×10131.07× 10^13 2.42×10132.42× 10^13 2.34×1052.34× 10^5 (a) (a) Pareto front of Λ versus the total VRAM usage ratio. (b) Comparison between SC1, SC2, and SC3 for the number of rounds to achieve the target average testing accuracy and the total VRAM usage. V-B3 Trade-off between Λ and Total VRAM Usage To show the trade-off between the value of Λ and the total VRAM usage of all clients, we conduct experiments by using OPT-125M and SST-2 dataset. For better illustration, we use the fraction of the total VRAM usage of all clients (i.e., ∑n∈ψtot,n∑n∈ψmax,n _n _tot,n _n _ ,n). The Pareto front is shown in Fig. 5(a) (a). Results show that by activating more blocks across clients, Λ generally exhibits a descending trend, which substantiates our Observation 1 in Section I. Moreover, the value of Λ has a trend of initially slow decline, followed by a sharp drop in the middle, and gradually converging to the minimum. Based on this trend, we can empirically determine the block activation matrix such that the total VRAM usage is effectively reduced without degrading the convergence rate of our proposed ZorBA. V-B4 Case Study To further demonstrate the relationship between Λ and the convergence rate of our proposed ZorBA, we conduct three case studies. In particular, we select three scenarios from Fig. 5(a) (a), namely SC1, SC2, and SC3, respectively. The values of the data tuples of SC1, SC2, and SC3 are (0.71, 0.50), (0.86, 0.50), and (0.86, 0.95), respectively. The comparison of the number of fine-tuning rounds to achieve the target average testing accuracy of 0.8 and the total VRAM usage of all clients is shown in Fig. 5(a) (b). Results show that SC1 requires less total VRAM than SC2, indicating fewer activated blocks across all clients. Nevertheless, SC1 achieves a similar (a bit faster) convergence rate because both SC1 and SC2 yield comparable values of Λ . This supports Observation 2 and Proposition 1 in Section I. In contrast, SC3 uses more VRAM than SC1 and matches SC2 in the total number of activated blocks. However, it needs more fine-tuning rounds to reach the target accuracy. This slower convergence highlights the role of Λ , which validates Observations 3 and 4. VI Conclusion In this work, we proposed ZorBA, a zeroth-order optimization-based federated fine-tuning framework. ZorBA leverages zeroth-order optimization and a heterogeneous block activation mechanism to reduce the VRAM usage. It introduces shared random seeds to reduce the communication overhead. We theoretically analyzed the convergence rate of our proposed ZorBA and investigated the impact of block activation decisions on the convergence rate and VRAM usage. We proposed an ϵε-constraint lexicographic algorithm to optimize the block activation decisions. Experimental results show that ZorBA converges faster than zeroth-order optimization baselines, while reducing VRAM usage by up to 62.41% and incurring negligible communication overhead. References Proof of Theorem 1 Proof. According to the model update rule, the global model at the beginning of (t+1)(t+1)-th fine-tuning round can be expressed as ¯t+1=[∑n∈a1,n∑n′∈a1,n′(¯1,nt−ηt∇~Fn(¯1,nt;ξn)),…,∑n∈aM,n∑n′∈aM,n′(¯M,nt−ηt∇~Fn(¯M,nt;ξn))]. w^t+1= [Σ _n a_1,nΣ _n a_1,n ( w_1,n^t-η^t ∇F_n ( w_1,n^t; _n ) ),…,Σ _n a_M,nΣ _n a_M,n ( w_M,n^t-η^t ∇F_n ( w_M,n^t; _n ) ) ]. (19) Based on Assumption 1, we expand [f(¯t+1)]E[f( w^t+1)] as [f(¯t+1)]≤[f(¯t+1)]+[⟨∇f(¯t),¯t+1−¯t⟩]⏟T1+L2[‖¯t+1−¯t‖2]⏟T2. [f( w^t+1) ] [f( w^t+1) ]+ E [ ∇ f( w^t), w^t+1- w^t ]_T_1+ L2 E [ \| w^t+1- w^t \|^2 ]_T_2. (20) We analyze T1T_1 first. In particular, it satisfies T1 T_1 = = [∇f(¯t)⊺][¯t+1−¯t] \>E [∇ f( w^t) ]E [ w^t+1- w^t ] = = [∇f(¯t)⊺][−ηt[∑n∈a1,n∑n′∈a1,n′∇~Fn(¯1,nt;ξn),…,∑n∈aM,n∑n′∈aM,n′∇~Fn(¯M,nt;ξn)]] \>E [∇ f( w^t) ]E [-η^t [Σ _n a_1,nΣ _n a_1,n ∇F_n ( w_1,n^t; _n ),…,Σ _n a_M,nΣ _n a_M,n ∇F_n ( w_M,n^t; _n ) ] ] =(a) (a)= −ηt[∇f(¯t)⊺][[∑n∈a1,n∑n′∈a1,n′∇~fn(¯1,nt),…,∑n∈aM,n∑n′∈aM,n′∇~fn(¯M,nt)]] \>-η^tE [∇ f( w^t) ]E [ [Σ _n a_1,nΣ _n a_1,n ∇f_n ( w_1,n^t ),…,Σ _n a_M,nΣ _n a_M,n ∇f_n ( w_M,n^t ) ] ] =(b) (b)= −ηt2[‖∇f(¯t)‖2]−ηt2[‖[∑n∈a1,n∑n′∈a1,n′∇~fn(¯1,nt),…,∑n∈aM,n∑n′∈aM,n′∇~fn(¯M,nt)]‖2] \>- η^t2E [ \|∇ f( w^t) \|^2 ]- η^t2E [ \| [Σ _n a_1,nΣ _n a_1,n ∇f_n ( w_1,n^t ),…,Σ _n a_M,nΣ _n a_M,n ∇f_n ( w_M,n^t ) ] \|^2 ] +ηt2[‖[∇f(¯1t)−∑n∈a1,n∑n′∈a1,n′∇~fn(¯1,nt),…,∇f(¯Mt)−∑n∈aM,n∑n′∈aM,n′∇~fn(¯M,nt)]‖2]⏟T3, + η^t2 E [ \| [∇ f( w_1^t)-Σ _n a_1,nΣ _n a_1,n ∇f_n ( w_1,n^t ),…,∇ f( w_M^t)-Σ _n a_M,nΣ _n a_M,n ∇f_n ( w_M,n^t ) ] \|^2 ]_T_3, (21) where equality (a) results from Assumption 2. Equality (b) is obtained by using the fact that ⟨i,j⟩=12(‖i‖2+‖j‖2−‖i−j‖2) _i,w_j = 12(\|w_i\|^2+\|w_j\|^2-\|w_i-w_j\|^2). Then, we bound T3T_3. In particular, it satisfies T3 T_3 = = [∥[∑n∈a1,n∑n′∈a1,n′(∇f(¯1t)±∇fn(¯1t)−∇~fn(¯1,nt)),…, \>E [ \| [Σ _n a_1,nΣ _n a_1,n (∇ f( w_1^t)±∇ f_n( w_1^t)- ∇f_n( w_1,n^t) ),…, ∑n∈aM,n∑n′∈aM,n′(∇f(¯Mt)±∇fn(¯Mt)−∇~fn(¯M,nt))]∥2] 37.00002ptΣ _n a_M,nΣ _n a_M,n (∇ f( w_M^t)±∇ f_n( w_M^t)- ∇f_n( w_M,n^t) ) ] \|^2 ] ≤(a) (a)≤ 2[‖[∑n∈a1,n∑n′∈a1,n′(∇f(¯1t)−∇fn(¯1t)),…,∑n∈aM,n∑n′∈aM,n′(∇f(¯Mt)−∇fn(¯Mt))]‖2] \>2E [ \| [Σ _n a_1,nΣ _n a_1,n (∇ f( w_1^t)-∇ f_n( w_1^t) ),…,Σ _n a_M,nΣ _n a_M,n (∇ f( w_M^t)-∇ f_n( w_M^t) ) ] \|^2 ] +2[‖[∑n∈a1,n∑n′∈a1,n′(∇fn(¯1t)−∇~fn(¯1,nt)),…,∑n∈aM,n∑n′∈aM,n′(∇fn(¯Mt)−∇~fn(¯M,nt))]‖2] +2E [ \| [Σ _n a_1,nΣ _n a_1,n (∇ f_n( w_1^t)- ∇f_n ( w_1,n^t ) ),…,Σ _n a_M,nΣ _n a_M,n (∇ f_n( w_M^t)- ∇f_n ( w_M,n^t ) ) ] \|^2 ] ≤(b) (b)≤ 2N∑n∈[‖[a1,n∑n′∈a1,n′(∇f(¯1t)−∇fn(¯1t)),…,aM,n∑n′∈aM,n′(∇f(¯Mt)−∇fn(¯Mt))]‖2] \>2NΣ _n E [ \| [ a_1,nΣ _n a_1,n (∇ f( w_1^t)-∇ f_n( w_1^t) ),…, a_M,nΣ _n a_M,n (∇ f( w_M^t)-∇ f_n( w_M^t) ) ] \|^2 ] +2N∑n∈[‖[a1,n∑n′∈a1,n′(∇fn(¯1t)−∇~fn(¯1,nt)),…,aM,n∑n′∈aM,n′(∇fn(¯Mt)−∇~fn(¯M,nt))]‖2] +2NΣ _n E [ \| [ a_1,nΣ _n a_1,n (∇ f_n( w_1^t)- ∇f_n ( w_1,n^t ) ),…, a_M,nΣ _n a_M,n (∇ f_n( w_M^t)- ∇f_n ( w_M,n^t ) ) ] \|^2 ] ≤(c) (c)≤ 2N∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2[‖(∇f(¯t)−∇fn(¯t))‖2] \>2NΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \E [ \| (∇ f( w^t)-∇ f_n( w^t) ) \|^2 ] +2N∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2[‖(∇fn(¯t)−∇~fn(¯t))‖2] +2NΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \E [ \| (∇ f_n( w^t)- ∇f_n( w^t) ) \|^2 ] ≤(d) (d)≤ 2N(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2, \>2N ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \, (22) where inequalities (a) and (b) result from Jensen’s inequality. Inequality (d) is obtained from Assumption 3 and Lemma 1. Therefore, by combining inequality (VI) and inequality (VI), we have T1≤ T_1≤ −ηt2[‖∇f(¯t)‖2]−ηt2[‖[∑n∈a1,n∑n′∈a1,n′∇~fn(¯1,nt),…,∑n∈aM,n∑n′∈aM,n′∇~fn(¯M,nt)]‖2] \>- η^t2E [ \|∇ f( w^t) \|^2 ]- η^t2E [ \| [Σ _n a_1,nΣ _n a_1,n ∇f_n ( w_1,n^t ),…,Σ _n a_M,nΣ _n a_M,n ∇f_n ( w_M,n^t ) ] \|^2 ] +ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 +η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ≤(a) (a)≤ −ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2, \>- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \, (23) where inequality (a) is because [‖[∑n∈a1,n∑n′∈a1,n′∇~fn(¯1,nt),…,∑n∈aM,n∑n′∈aM,n′∇~fn(¯M,nt)]‖2]≥0E [ \| [Σ _n a_1,nΣ _n a_1,n ∇f_n ( w_1,n^t ),…,Σ _n a_M,nΣ _n a_M,n ∇f_n ( w_M,n^t ) ] \|^2 ]≥ 0. Then, we bound T2T_2. In particular, it satisfies T2 T_2 = = [‖−ηt[∑n∈a1,n∑n′∈a1,n′∇~Fn(¯1,nt;ξn),…,∑n∈aM,n∑n′∈aM,n′∇~Fn(¯M,nt;ξn)]‖2] \>E [ \|-η^t [Σ _n a_1,nΣ _n a_1,n ∇F_n ( w_1,n^t; _n ),…,Σ _n a_M,nΣ _n a_M,n ∇F_n ( w_M,n^t; _n ) ] \|^2 ] ≤(a) (a)≤ (ηt)2N∑n∈[‖[a1,n∑n′∈a1,n′∇~Fn(¯1,nt;ξn),…,aM,n∑n′∈aM,n′∇~Fn(¯M,nt;ξn)]‖2] \> (η^t )^2NΣ _n E [ \| [ a_1,nΣ _n a_1,n ∇F_n ( w_1,n^t; _n ),…, a_M,nΣ _n a_M,n ∇F_n ( w_M,n^t; _n ) ] \|^2 ] ≤(b) (b)≤ (ηt)2N∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2[‖∇~Fn(¯t;ξn)‖2] \> (η^t )^2NΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \E [ \| ∇F_n ( w^t; _n ) \|^2 ] ≤(c) (c)≤ (ηt)2dN∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2[‖∇Fn(¯t;ξn)‖2] \> (η^t )^2dNΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \E [ \|∇ F_n ( w^t; _n ) \|^2 ] = = (ηt)2dN∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2[‖∇Fn(¯t;ξn)±∇fn(¯t)±∇f(¯t)‖2] \> (η^t )^2dNΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \E [ \|∇ F_n ( w^t; _n )±∇ f_n ( w^t )±∇ f ( w^t ) \|^2 ] ≤(d) (d)≤ 3(ηt)2dN∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2([∥∇Fn(¯t;ξn)−∇fn(¯t)∥2] \>3 (η^t )^2dNΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ (E [ \|∇ F_n ( w^t; _n )-∇ f_n ( w^t ) \|^2 ] [∥∇fn(¯t)−∇f(¯t)∥2]+[∥∇f(¯t)∥2]) [ \|∇ f_n ( w^t )-∇ f ( w^t ) \|^2 ]+E [ \|∇ f ( w^t ) \|^2 ] ) ≤(e) (e)≤ 3(ηt)2dN∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2([‖∇f(¯t)‖2]+σ2+σG2), \>3 (η^t )^2dNΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ (E [ \|∇ f ( w^t ) \|^2 ]+σ^2+ _G^2 ), (24) where inequalities (a) and (d) result from Jensen’s inequality. Inequality (b) is obtained by using Lemma 2. Inequality (c) is derived by using Lemma 3. Inequality (e) results from Assumptions 2 and 3. Then, we combine inequalities (20), (VI), and (VI) as follows: [f(¯t+1)] [f( w^t+1) ] ≤ ≤ [f(¯t)]−ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +L2(3(ηt)2dN∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2([‖∇f(¯t)‖2]+σ2+σG2)). + L2 (3 (η^t )^2dNΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ (E [ \|∇ f ( w^t ) \|^2 ]+σ^2+ _G^2 ) ). (25) By rearranging inequality (VI), we have (ηt2−3L(ηt)2dN2∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2)[‖∇f(¯t)‖2] ( η^t2- 3L (η^t )^2dN2Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ )E [ \|∇ f ( w^t ) \|^2 ] ≤ ≤ [f(¯t)]−[f(¯t+1)]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]-E [f( w^t+1) ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +L2(3(ηt)2dN(σ2+σG2)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2). + L2 (3 (η^t )^2dN (σ^2+ _G^2 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ). (26) We define Ω()=(ηt2−3L(ηt)2dN2∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2) (A)= ( η^t2- 3L (η^t )^2dN2Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ). To guarantee Ω()≥0 (A)≥ 0, the learning rate in the arbitrary t-th fine-tuning round must be small enough such that ηt=η≤min13LdN∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2=13LdN2. η^t=η≤ _A \ 13LdNΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ \= 13LdN^2. (27) Hence, we have [‖∇f(¯t)‖2]≤1Ω()([f(¯t)]−[f(¯t+1)])+Θ1Ω()∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2, [ \|∇ f ( w^t ) \|^2 ]≤ 1 (A) (E [f( w^t) ]-E [f( w^t+1) ] )+ _1 (A)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \, (28) where Θ1=14μ2L2η(d+3)3N+(ηN+3Lη2dN2)σG2+3Lη2dN2σ2 _1= 14μ^2L^2η(d+3)^3N+ (η N+ 3Lη^2dN2 ) _G^2+ 3Lη^2dN2σ^2. To determine the convergence rate, we sum up both sides of inequality (28) for all T fine-tuning rounds and multiply both sides by 1T 1T. We have 1T∑t=1T[‖∇f(¯t)‖2]≤ 1TΣ _t=1^TE [ \|∇ f ( w^t ) \|^2 ]≤ 1Ω()T∑t=1T(f(¯t)−f(¯t+1))+Θ1Ω()∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \> 1 (A)TΣ _t=1^T (f( w^t)-f( w^t+1) )+ _1 (A)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ≤(a) (a)≤ 1Ω()T(f(¯1)−f(⋆))+Θ1Ω()∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2, \> 1 (A)T (f( w^1)-f(w ) )+ _1 (A)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \, (29) where inequality (a) is obtained by using f(¯T+1)≥f(⋆)f( w^T+1)≥ f(w ). This completes the proof of Theorem 1. ∎ Proof of Theorem 2 Proof. Based on second-order Taylor expansion, we expand [f(¯t+1)]E[f( w^t+1)] as [f(¯t+1)] [f( w^t+1) ] ≈ ≈ [f(¯t)]+[⟨∇f(¯t),¯t+1−¯t⟩]+12[(¯t+1−¯t)⊺∇2f(¯t)(¯t+1−¯t)] \>E [f( w^t) ]+E [ ∇ f( w^t), w^t+1- w^t ]+ 12E [( w^t+1- w^t) ∇^2f( w^t)( w^t+1- w^t) ] ≤(a) (a)≤ [f(¯t)]−ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +12[(¯t+1−¯t)⊺∇2f(¯t)(¯t+1−¯t)] + 12E [( w^t+1- w^t) ∇^2f( w^t)( w^t+1- w^t) ] ≤(b) (b)≤ [f(¯t)]−ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +12[⟨(¯t),(¯t+1−¯t)(¯t+1−¯t)⊺⟩]⏟T1, + 12 E [ ( w^t),( w^t+1- w^t)( w^t+1- w^t) ]_T_1, = = [f(¯t)]−ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +12[(¯t)][(¯t+1−¯t)(¯t+1−¯t)⊺]⏟T2, + 12E [H( w^t) ] E [( w^t+1- w^t)( w^t+1- w^t) ]_T_2, (30) where inequality (a) results from inequality (VI) in the proof of Theorem 1. Inequality (b) is obtained by using Assumption 4. Then, to bound T1T_1, we analyze T2T_2. We define ()=[∇F(;)(∇F(;))⊺]−∇f()(∇f())⊺ (w)=E [∇ F(w; ξ) (∇ F(w; ξ) ) ]-∇ f(w) (∇ f(w) ) . In particular, T2T_2 satisfies T2 T_2 ≤(a) (a)≤ (1+d−2d+2)(ηt)2N∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2(∇f(¯t)(∇f(¯t))⊺+1N(¯t)) \> (1+ d-2d+2 )(η^t)^2NΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ (∇ f( w^t) (∇ f( w^t) ) + 1N ( w^t) ) +d+2(ηt)2N∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2(‖∇f(¯t)‖2+1NTr((¯t)))d×d, + dd+2(η^t)^2NΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ( \|∇ f( w^t) \|^2+ 1N ( ( w^t) ) )I_d× d, (31) where inequality (a) results from Jensen’s inequality and Lemma 3. By combining inequalities (VI) and (VI), we have [f(¯t+1)] [f( w^t+1) ] ≤ ≤ [f(¯t)]−ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +12(1+d−2d+2)(ηt)2N∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2(∇f(¯t)⊺(¯t)∇f(¯t)+⟨1N(¯t),(¯t)⟩) + 12 (1+ d-2d+2 )(η^t)^2NΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ (∇ f( w^t) H( w^t)∇ f( w^t)+ 1N ( w^t),H( w^t) ) +d2(d+2)(ηt)2N∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2(‖∇f(¯t)‖2+1NTr((¯t)))Tr((¯t)) + d2(d+2)(η^t)^2NΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ( \|∇ f( w^t) \|^2+ 1N ( ( w^t) ) ) (H( w^t) ) ≤(a) (a)≤ [f(¯t)]−ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +12(1+d−2d+2)(ηt)2NL∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2(‖∇f(¯t)‖2+1NTr((¯t))) + 12 (1+ d-2d+2 )(η^t)^2NLΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ( \|∇ f( w^t) \|^2+ 1N ( ( w^t) ) ) +d2(d+2)(ηt)2NLκ∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2(‖∇f(¯t)‖2+1NTr((¯t))) + d2(d+2)(η^t)^2NLκΣ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ( \|∇ f( w^t) \|^2+ 1N ( ( w^t) ) ) ≤(b) (b)≤ [f(¯t)]−ηt2[‖∇f(¯t)‖2]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]- η^t2E [ \|∇ f( w^t) \|^2 ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +d(ηt)2NL(κ+2)2(d+2)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2(‖∇f(¯t)‖2+σ2N), + d(η^t)^2NL(κ+2)2(d+2)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ( \|∇ f( w^t) \|^2+ σ^2N ), (32) where inequality (a) results from Assumption 4. Inequality (b) is obtained by using the fact that Tr(())=[‖∇F(;)−∇f()‖2] ( (w) )=E [\|∇ F(w; ξ)-∇ f(w)\|^2 ] and Assumption 2. By rearranging inequality (VI), we have (ηt2−d(ηt)2NL(κ+2)2(d+2)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2)‖∇f(¯t)‖2 ( η^t2- d(η^t)^2NL(κ+2)2(d+2)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ) \|∇ f( w^t) \|^2 ≤ ≤ [f(¯t)]−[f(¯t+1)]+ηtN(σG2+14μ2L2(d+3)3)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2 \>E [f( w^t) ]-E [f( w^t+1) ]+η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ +d(ηt)2L(κ+2)σ22(d+2)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2. + d(η^t)^2L(κ+2)σ^22(d+2)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \. (33) We define Φ()=(ηt2−d(ηt)2NL(κ+2)2(d+2)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2) (A)= ( η^t2- d(η^t)^2NL(κ+2)2(d+2)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ ). To guarantee Φ()≥0 (A)≥ 0, the learning rate in the arbitrary t-th fine-tuning round must be small enough such that ηt=η≤mind+2LdN(κ+2)∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2=d+2LdN2(κ+2). η^t=η≤ _A \ d+2LdN(κ+2)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \ \= d+2LdN^2(κ+2). (34) Hence, we have [‖∇f(¯t)‖2]≤1Φ()([f(¯t)]−[f(¯t+1)])+Θ2Φ()∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2, [ \|∇ f ( w^t ) \|^2 ]≤\> 1 (A) (E [f( w^t) ]-E [f( w^t+1) ] )+ _2 (A)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \, (35) where Θ2=ηtN(σG2+14μ2L2(d+3)3)+d(ηt)2L(κ+2)σ22(d+2) _2=η^tN ( _G^2+ 14μ^2L^2(d+3)^3 )+ d(η^t)^2L(κ+2)σ^22(d+2). To determine the convergence rate, we sum up both sides of inequality (35) for all T fine-tuning rounds and multiply both sides by 1T 1T. We have 1T∑t=1T[‖∇f(¯t)‖2]≤(a)1Φ()T(f(¯1)−f(¯⋆))+Θ2Φ()∑n∈maxm∈ℳ(am,n∑n′∈am,n′)2, 1T _t=1^TE [ \|∇ f ( w^t ) \|^2 ] (a)≤\> 1 (A)T (f( w^1)-f( w ) )+ _2 (A)Σ _n _m \ ( a_m,nΣ _n a_m,n )^2 \, (36) where inequality (a) is obtained by using f(¯T+1)≥f(⋆)f( w^T+1)≥ f(w ). This completes the proof of Theorem 1. ∎