Exploring Selective Layer Fine-Tuning in Federated Learning

## Abstract

Abstract Federated learning (FL) has emerged as a promising paradigm for fine-tuning foundation models using distributed data in a privacy-preserving manner. Under limited computational resources, clients often find it more practical to fine-tune a selected subset of layers, rather than the entire model, based on their task-specific data. In this study, we provide a thorough theoretical exploration of selective layer fine-tuning in FL, emphasizing a flexible approach that allows the clients to adjust their selected layers according to their local data and resources. We theoretically demonstrate that the layer selection strategy has a significant impact on model convergence in two critical aspects: the importance of selected layers and the heterogeneous choices across clients. Drawing from these insights, we further propose a strategic layer selection method that utilizes local gradients and regulates layer selections across clients. The extensive experiments on both image and text datasets demonstrate the effectiveness of the proposed strategy compared with several baselines, highlighting its advances in identifying critical layers that adapt to the client heterogeneity and training dynamics in FL.

## 1 Introduction

Foundation models (Bommasani et al., ), including BERT (Devlin et al., ), GPT (Radford et al., ; Brown et al., ), CLIP (Radford et al., ; Dosovitskiy et al., ), LLaMA (Touvron et al., ), and so on (Ramesh et al., ; Chowdhery et al., ), have attracted considerable attention due to their exceptional ability in handling complex tasks (Eloundou et al., ).
When it comes to practical deployments of these models in specialized fields, fine-tuning with domain-specific data becomes critical.
Nevertheless, the distributed nature of data across various users and organizations presents a challenge for centralized storage and training, as it may lead to severe privacy concerns and incur additional transmission costs.
Such issues have positioned federated learning (FL) (McMahan et al., ) as a promising paradigm to fine-tune foundation models, aligning model enhancement with privacy preservation (Chen et al., ; Kuang et al., ).

Generally, FL aims to learn a global model through a collaborative process where clients perform local training and upload the parameter updates to a central server for aggregation.
Given that clients have limited resources (Bonawitz et al., ; Imteaj et al., ), such as computational power, communication bandwidth, and available memory, it becomes impractical for them to fine-tune the entire foundation model.
Two kinds of solutions have recently emerged to tackle this challenge. The first solution employs *parameter-efficient fine-tuning* techniques (Houlsby et al., ; Gao et al., ; Hu et al., ; Li & Liang, ), which introduces additional modules integrated into foundation models and updates these modules with domain-specific data while keeping the parameters of the foundation model frozen. The second one is *selective model fine-tuning* (Lee et al., ; Xu et al., ; Zhang et al., ; Shen et al., ), which only selects an impactful subset of parameters for optimization to streamline the fine-tuning process under resource constraints.

This study focuses on selective model fine-tuning as it is particularly well-suited to address the inherent heterogeneity in FL, i.e., the data heterogeneity and device heterogeneity (Yang et al., ; Chai et al., ; Li et al., ).
Specifically, clients involved in FL have non-independent and identically distributed (non-IID) data distribution and different system resources, leading to the need to customize fine-tuning strategies to such discrepancies.
For example, clients with limited computation resources may opt to update only a fraction of the model, while those with sufficient resources and high-quality data prefer fine-tuning a large portion of the model to enhance performance.
Selective model fine-tuning enables clients to adjust the chosen part of the model to be updated based on their own capabilities, providing a flexible and advanced solution to mitigating sub-optimal issues induced by the heterogeneity in FL.

The exploration of selective model fine-tuning within the context of FL, is still in its early stages.
Previous studies (Shen et al., ; Xu et al., ; Lee et al., ; Dun et al., ) have concentrated on designing static strategies for subnetwork selection to improve model fine-tuning performance, without adequately considering heterogeneity among clients.
To fulfill this gap, in this paper, we provide a comprehensive theoretical analysis on selective model fine-tuning in FL, focusing on a general scenario where clients are allowed to choose different layers for local training and vary their choices across different training epochs, called selective layer fine-tuning.
Specifically, we formulate the optimization objective of selective layer fine-tuning in FL, and provide insights on effectively determining critical layers to achieve model convergence. Building on these insights, we further propose a novel layer selection strategy that leverages local gradients and the regulation of unified selections.

Our main contributions are summarized as follows:

- •
We study a practical FL setup where clients choose to fine-tune some layers of the model, with arbitrary layer selection that may vary among clients and across different training epochs.
We theoretically analyze such a training scheme and investigate the impact of layer selection.
The analytical results show that the selected layers affect the convergence performance with two critical aspects, namely the importance of layers and heterogeneous choices across clients.
- •
Building on the theoretical analysis, we formulate the optimization problem of selective layer fine-tuning considering the limited and diverse resource budgets of clients.
Inspired by the solution to this optimization problem, we propose an effective strategy for selecting layers for fine-tuning that are well-suited for the local data and available resources at clients.
- •
We conduct experiments to compare the proposed layer selection strategy with baseline methods on both image and text datasets.
The experimental results demonstrate the superior performance of the proposed strategy in achieving better model performance, highlighting that the proposed strategy can find critical layers for fine-tuning while considering the client heterogeneity and training dynamics in FL.

## 2 Related Works

Various approaches have been proposed to properly select a subset of parameters for fine-tuning foundation models within centralized training, including optimizing a non-structured mask matrix (Lee et al., ; Xu et al., ; Zhang et al., ; Shen et al., ; Zaken et al., ; Zhang et al., ; Kovaleva et al., ; Lee et al., ) and adopting layer-wise selection strategies (Kovaleva et al., ; Lee et al., ; ; Kaplun et al., ). For example, Lee et al. () suggest updating the parameters in a stochastic manner based on the Bernoulli distribution, while Kovaleva et al. (); Lee et al. () showcase that fine-tuning the top few layers achieves competitive model performance in downstream tasks.
Moreover, Lee et al. () propose to select layers according to their gradient statistics.

Recent studies have extended the selective fine-tuning techniques to FL scenarios (Nguyen et al., ; Chen et al., ; Hilmkil et al., ; Zhang et al., ). Specifically, researchers (Lee et al., ; Dun et al., ) investigate layer-wise network decomposition to achieve selective model fine-tuning on clients.
However, these works fail to offer methodologies for adaptive and dynamic layer selection that take into account the heterogeneous characteristics of clients.
In addition, personalized FL algorithms (Pillutla et al., ; Chen et al., ) propose to train different subnetworks on clients towards better local models.
Different from previous studies, we focus on providing an in-depth understanding of selective layer fine-tuning in FL, considering the heterogeneity from the perspective of client resources and local data distributions.

## 3 Problem Formulation

##### Federated learning

We consider an FL system with a central server and $N$ clients (denoted by the set $\mathcal{N}=\{1,\dots,N\}$), where each client has a private dataset $\mathcal{D}_{i}$ consisting of $d_{i}=|\mathcal{D}_{i}|$ data instances.
The server owns a pretrained foundation model $\theta\in\mathbb{R}^{P}$, containing $P$ trainable parameters and $L$ layers with the index set $\mathcal{L}=\{1,2,\dots,L\}$.
The server aims to fine-tune this foundation model based on the clients’ datasets $\mathcal{D}=\{\mathcal{D}_{i}\}_{i\in\mathcal{N}}$ but does not directly access these datasets. The learning goal is formally given as:

$$ $\vspace{-5pt}\min_{\theta\in\mathbb{R}^{P}}f(\theta)=\sum_{i=1}^{N}\alpha_{i}f _{i}(\theta),$ (1) $$

where $\alpha_{i}=\frac{d_{i}}{\sum_{j=1}^{N}d_{j}}$ denotes the relative sample size, and $f_{i}(\theta)=\frac{1}{d_{i}}\sum_{\xi\in\mathcal{D}_{i}}F_{i}(\theta;\xi)$ denotes the local training objective of client $i$.
Here we use $F_{i}(\theta;\xi)$ to define the (possibly non-convex) loss function computed by the model $\theta$ on data sample $\xi$.
The training process of FL is divided into $T$ training epochs.
In each epoch $t\in[T]$, the server chooses a subset of clients $\mathcal{S}^{t}$, and sends the up-to-date global model $\theta^{t}$ to these clients for local training.

Figure: Figure 1: An overview of selective layer fine-tuning in FL. The colored layers are selected for fine-tuning.
Refer to caption: x1.png

##### Selective layer fine-tuning in FL

An overview of selective layer fine-tuning in FL is illustrated in Figure 1.
Due to resource limitations, clients tend to update some of the layers in the local training process rather than the entire global model.
Formally, we define a masking vector $\mathbf{m}_{i}^{t}\in\{0,1\}^{L}$ for each client $i$. The $l$-th element $\mathbf{m}_{i}^{t}(l)$ equals $1$ if the $l$-th layer is selected to be updated in the $t$-th training epoch, and $\mathbf{m}_{i}^{t}(l)=0$ otherwise.
Accordingly, the selected layer set of client $i$ is denoted by $\mathcal{L}_{i}^{t}\triangleq\{l\in\mathcal{L}|\mathbf{m}_{i}^{t}(l)=1\}$, and the set for all selected layers in the $t$-th training epoch is denoted by $\mathcal{L}_{t}=\bigcup_{i\in\mathcal{S}^{t}}\mathcal{L}_{i}^{t}$.
The choice of selected layer sets has a substantial effect on training performance, which will be discussed in detail later.

After determining the selected layer set $\mathcal{L}_{i}^{t}$, clients initialize the local model according to the global model sent by the server, i.e., $\theta_{i}^{t,0}=\theta^{t}$, and train the local model for $\tau$ local steps using the mini-batch SGD algorithm (McMahan et al., ; Wang et al., ; Karimireddy et al., ).
For local step $k\in[\tau]$, client $i$ samples a batch of data instances $\xi_{i}^{t,k}$, and calculates the gradients for the selected layers, which is given as(^1^11$\nabla_{l}F(\theta)$ represents the gradient of a function $F(\theta)$ w.r.t. the parameters of the $l$-th layer in model $\theta$.):

$$ $\sum_{l\in\mathcal{L}_{i}^{t}}g_{i,l}(\theta_{i}^{t,k};\xi_{i}^{t,k})=\sum_{l \in\mathcal{L}_{i}^{t}}\nabla_{l}F_{i}(\theta_{i}^{t,k};\xi_{i}^{t,k}).$ (2) $$

Notably, the local gradient calculation pertains solely to the layers within the subset $\mathcal{L}_{i}^{t}$.
Afterward, the local model is updated with the learning rate $\eta$:

$$ $\theta_{i}^{t,k}=\theta_{i}^{t,k-1}-\eta\sum_{l\in\mathcal{L}_{i}^{t}}g_{i,l}( \theta_{i}^{t,k-1};\xi_{i}^{t,k-1}),\forall k\in\{1,2,\dots,\tau\}.$ (3) $$

The accumulated model update in local training is summarized as:

$$ $\Delta_{i}^{t}=\frac{1}{\eta}(\theta_{i}^{t,0}-\theta_{i}^{t,\tau})=\sum_{k=0} ^{\tau-1}\sum_{l\in\mathcal{L}_{i}^{t}}g_{i,l}(\theta_{i}^{t,k};\xi_{i}^{t,k}).$ (4) $$

After local training, clients upload their model updates $\Delta_{i}^{t},i\in\mathcal{S}^{t}$ to the server.
The server performs federated aggregation among these model updates and optimizes the global model accordingly:

$$ $\Delta^{t}=\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{S}^{t}}w_{i,l}^{t}\sum _{k=0}^{\tau-1}g_{i,l}(\theta_{i}^{t,k};\xi_{i}^{t,k}),\vspace{-4pt}$ (5) $$

and

$$ $\theta^{t+1}=\theta^{t}-\eta\Delta^{t}.$ (6) $$

Inspired by previous studies (McMahan et al., ; Li et al., ), the aggregation weights in selective layer fine-tuning are defined based on the data ratio and the masking vectors, which are formally given as:

$$ $w_{i,l}^{t}=\left\{\begin{array}[]{ll}\frac{d_{i}}{\sum_{\{j\in\mathcal{S}^{t} |\mathbf{m}_{j}^{t}(l)=1\}}d_{j}},&\text{if }\mathbf{m}_{i}^{t}(l)=1,\\ 0,&\text{otherwise.}\end{array}\right.$ (7) $$

The details of the training process are summarized in Algorithm 1.

Figure: Algorithm 1 Selective Layer Fine-tuning in FL

## 4 Which Layers Should Be Selected For Fine-Tuning?

The aforementioned training process provides substantial flexibility in selective layer fine-tuning, namely, clients are allowed to select different layers for local training and adjust their choices in different training epochs.
Such flexibility enables clients to tailor their local training to their data and resources, providing feasible solutions for handling the heterogeneity in FL.

However, without a well-designed strategy for layer selection, the optimization of the global model in FL could be severely hindered, potentially leading to a suboptimal solution or even failure in convergence. As a result, researchers have proposed several useful strategies for layer selection in recent years, including:

- •
All clients select the same layer set for fine-tuning (Pillutla et al., ; Lee et al., ; Zhang et al., ; ; Lee et al., ), i.e., $\mathcal{L}_{i}^{t}=\mathcal{L}_{j}^{t},\forall i\neq j$;
- •
Clients fix their selections across different training epochs (Arivazhagan et al., ; Chen et al., ), i.e., $\mathcal{L}_{i}^{t_{1}}=\mathcal{L}_{i}^{t_{2}},\forall t_{1},t_{2}\in[T]$.

These strategies for layer selection are proposed based on the insights drawn from experts’ experience, serving as special instantiations of the selected layer sets $\mathcal{L}_{i}^{t}$.
It is worth noting that these experience-driven strategies might not consistently yield optimal results in various FL applications, particularly considering the client heterogeneity.
This leads to an essential question: How to effectively determine the task-specified layer selection strategy among a large search space of possible options?
In the rest of this section, we provide a theoretical analysis to answer this question.

### 4.1 Theoretical Analysis

Following previous theoretical analysis in FL (Wang et al., ; Karimireddy et al., ; Li et al., ), we begin with some necessary assumptions.

###### Assumption 4.1 .

($\gamma$-Smoothness)
There exists a constant $\gamma>0$ such that for any $\theta,\theta^{\prime}\in\mathbb{R}^{P}$, $\|\nabla f_{i}(\theta)-\nabla f_{i}(\theta^{\prime})\|_{2}\leq\gamma\|\theta-
\theta^{\prime}\|_{2},\forall i\in\mathcal{N}$.

For analyzing the effect of each layer on the model convergence, we give several assumptions for the gradient with respect to each layer $l$.

###### Assumption 4.2 .

(Unbiased and variance-bounded stochastic gradient)
The stochastic gradient $g_{i,l}(\theta^{t};\xi_{i}^{t})$ on a randomly sampled batch of data $\xi_{i}^{t}$ is an unbiased estimate of the full-batch gradient, i.e., $\mathbb{E}[g_{i,l}(\theta^{t};\xi_{i}^{t})]=\nabla_{l}f_{i}(\theta^{t})$.
Besides, there exist constants $\sigma_{l}>0,\forall l\in\mathcal{L}$ such that $\|g_{i,l}(\theta^{t};\xi_{i}^{t})-\nabla_{l}f_{i}(\theta^{t})\|^{2}\leq\sigma_
{l}^{2},\forall i\in\mathcal{N}$ and $\sum_{l\in\mathcal{L}_{t}}\sigma_{l}^{2}\leq\sigma^{2}$.

The non-IID data owned by clients causes diverse gradients.
In the following assumption, we state the diversity of each layer’s gradient.

###### Assumption 4.3 .

(Gradient diversity)
There exist constants $\kappa_{l}>0,\forall l\in\mathcal{L}$ such that $\left\|\nabla_{l}f(\theta^{t})-\nabla_{l}f_{i}(\theta^{t})\right\|^{2}\leq
\kappa_{l}^{2},\forall i\in\mathcal{N}$.

Here we first consider a case where $\tau\!=\!1$ to simplify the analysis without affecting the insights on layer selection. The detailed analysis for the generalized case, i.e., $\tau\!>\!1$, is provided in Section A.3.

Compared with the theoretical analysis for the standard FL settings (Wang et al., ; Li et al., ), there exist three additional challenges in selective layer fine-tuning.
Firstly, since each client only updates some layers during the local training process, the aggregated gradient is no longer an unbiased estimate of the local gradient $\nabla f_{i}(\theta^{t})$, i.e.,

$$ $\mathbb{E}[\Delta_{i}^{t}]=\sum_{l\in\mathcal{L}_{i}^{t}}\nabla_{l}f_{i}( \theta^{t})\neq\nabla f_{i}(\theta^{t}),$ (8) $$

where the inequality holds unless all layers are selected for fine-tuning, i.e., $\mathcal{L}_{i}^{t}=\mathcal{L}$.
Secondly, since a certain layer may not be selected by all the clients, the aggregated gradient of this layer is not equivalent to the gradient computed based on the global loss function ($\sum_{l\in\mathcal{L}_{t}}\nabla_{l}f(\theta^{t})$), which is given as:

$$ $\displaystyle\mathbb{E}[\Delta^{t}]=\sum_{l\in\mathcal{L}_{t}}\sum_{i\in \mathcal{S}^{t}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta^{t})\neq\sum_{l\in\mathcal{L }_{t}}\nabla_{l}f(\theta^{t}),$ (9) $$

where the inequality holds unless all clients select the same subset of layers.
Last but not least, the aforementioned gaps vary across different training epochs, making it rather complicated in the theoretical analysis.

To link the aggregated and desired gradients, we define a surrogate objective function representing the underlying loss function optimized by the clients, which is given as:

$$ $h_{l}^{t}(\theta)\triangleq\sum_{i\in\mathcal{S}^{t}}w_{i,l}^{t}f_{i}(\theta).$ (10) $$

In essence, the layer-wise gradient of this objective function represents the update of the aggregated global update $\Delta^{t}$. This relationship is elaborated in the following lemma.

###### Lemma 4.4 .

With 4.2, we have:

$$ $\mathbb{E}[\Delta^{t}]=\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t }),$ (11) $$

where the expectation is with respect to mini-batch data sampling.

###### Proof.

We rewrite both sides of Equation 11 by using the definitions in Equations 5 and 10, and apply 4.2 to obtain the result.
∎

As aforementioned, the underlying loss function $h_{l}^{t}(\theta)$ deviates from the desired global loss function $f(\theta)$, which hinders the optimization of the global model and may lead to suboptimal model performance.
Such deviation can be quantified by the difference between the underlying update $\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})$ and the global gradient $\nabla f(\theta^{t})$, i.e.,

$$ $\mathcal{E}_{t}\triangleq\left\|\nabla f(\theta^{t})-\sum_{l\in\mathcal{L}_{t} }\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}.$ (12) $$

The term $\mathcal{E}_{t}$ can be further decomposed using the Jensen’s inequality into two parts:

$$ $\displaystyle\mathcal{E}_{t}\leq 2\left\|\nabla f(\theta^{t})-\sum_{l\in \mathcal{L}_{t}}\nabla_{l}f(\theta^{t})\right\|_{2}^{2}+2\left\|\sum_{l\in \mathcal{L}_{t}}\nabla_{l}f(\theta^{t})-\nabla_{l}h_{l}^{t}(\theta^{t})\right \|_{2}^{2}.$ (13) $$

###### Remark 4.5 .

These two terms in the right-hand side (RHS) of (13) can be interpreted as follows:
(i) The first term is the difference between the gradient w.r.t. all layers and the gradient w.r.t. the selected layers.
The value of this term becomes smaller when the selected layers have *larger gradients*;
(ii) The second term represents the mismatch between the desired gradient computed by all clients (i.e., $\nabla_{l}f(\theta^{t})=\sum_{i\in\mathcal{N}}\alpha_{i}\nabla_{l}f_{i}(\theta
^{t})$) and the underlying update computed by partial clients (i.e., $\nabla_{l}h_{l}^{t}(\theta^{t})=\sum_{i\in\mathcal{S}_{t}}w_{i,l}^{t}\nabla_{l
}f_{i}(\theta^{t})$), resulting from *different layer choices* among clients.
If some layer is selected by all clients, its corresponding term in this term can be diminished.

For a better understanding, the following lemma shows an upper bound for the value of $\mathcal{E}_{t}$.

###### Lemma 4.6 .

With 4.3, we have:

$$ $\vspace{-5pt}\mathcal{E}_{t}\leq 2\underbrace{\Bigg{[}\Bigg{\|}\sum_{l\notin \mathcal{L}_{t}}\nabla f(\theta^{t})\Bigg{\|}^{2}\Bigg{]}}_{\mathcal{E}_{t,1}} +2\underbrace{\sum_{l\in\mathcal{L}_{t}}\chi_{\mathbf{w}_{t,l}\|\mathbf{\alpha }}\kappa_{l}^{2}}_{\mathcal{E}_{t,2}},$ (14) $$

where $\chi_{\mathbf{w}_{t,l}\|\mathbf{\alpha}}\triangleq\sum_{i\in\mathcal{N}}\frac{
(w_{i,l}^{t}-\alpha_{i})^{2}}{\alpha_{i}}$.

The proofs can be found in Section A.1.

Next we aim to analyze the impact of layer selection on the convergence of the global model.
Following previous studies (Bottou et al., ; Wang et al., ), we consider an algorithm to have achieved convergence if it converges to a stationary point of the global loss function, namely, if its expected squared gradient norm $\min_{t\in[T]}\mathbb{E}\left[\left\|\nabla f(\theta^{t})\right\|_{2}^{2}\right]$ is zero.
The following theorem and corollary show the convergence of the proposed selective layer fine-tuning framework for FL.

###### Theorem 4.7 .

Define a constant $C\triangleq 1-4\eta L>0$.
With Assumptions 4.1-4.3, we have:

$$ $\displaystyle\min_{t\in[T]}\mathbb{E}\left[\left\|\nabla f(\theta^{t})\right\| _{2}^{2}\right]\leq\frac{2}{\eta CT}\left[f(\theta^{0})\!-\!f(\theta^{*}) \right]+\frac{2\gamma\eta}{C}\sigma^{2}+\frac{1}{T}\!\sum_{t=1}^{T}\left(\frac {1}{\gamma\eta C}\!+\!2\right)(\mathcal{E}_{t,1}\!+\!\mathcal{E}_{t,2}),$ (15) $$

where $\theta^{*}=\arg\min_{\theta\in\mathbb{R}^{P}}f(\theta)$ is the best model with the minimal loss.

The proofs can be found in Section A.2.

###### Corollary 4.8 .

With the commonly selected learning rate $\eta=\mathcal{O}\left(\frac{1}{\sqrt{T}}\right)$, the RHS of (15) except the last term becomes zero as $T\rightarrow\infty$.
Therefore, FL with selective layer fine-tuning may only oscillate around a stationary point of the global loss function with a non-zero error floor $\mathcal{O}(\mathcal{E}_{t,1}+\mathcal{E}_{t,2})$.

According to Theorem 4.7 and Corollary 4.8, the training performance of the global model is degraded by the increase of $\mathcal{E}_{t,1}+\mathcal{E}_{t,2}$, in the following aspects:

- •
The term $\mathcal{E}_{t,1}$ indicates that it might lead to a suboptimal global model if layers with large gradient norms were not selected for fine-tuning.
- •
The term $\mathcal{E}_{t,2}$ shows that the consistent selections among clients promote the convergence of the global model. Specifically, if the $l$-th layer has large gradient diversity $\kappa_{l}$, implying significant objective bias among clients, reducing the weight divergence $\chi_{\mathbf{w}_{t,l}\|\mathbf{\alpha}}$ helps to alleviate this term.

These findings highlight that the layer selection strategy that minimizes both terms can achieve better convergence of the global model.
However, minimizing these two terms simultaneously may lead to contradictory selection decisions.
Moreover, the optimal solution for minimizing the sum of $\mathcal{E}_{t,1}+\mathcal{E}_{t,2}$ is inaccessible, since the ground-truth values are intractable in practice.
To resolve these challenges, in the next subsection, we propose a strategy to adaptively select layers for clients and promote the learning performance of the global model.

###### Assumption 4.1 .

###### Assumption 4.2 .

###### Assumption 4.3 .

###### Lemma 4.4 .

###### Proof.

###### Remark 4.5 .

###### Lemma 4.6 .

###### Theorem 4.7 .

###### Corollary 4.8 .

### 4.2 Layer Selection Strategy

Based on the above analysis, we need to determine the selected layer sets $\{\mathcal{L}_{t}^{i}\}$ that minimize the values of $\mathcal{E}_{t,1}$ and $\mathcal{E}_{t,2}$.
As both terms are hard to compute directly, we first design an approach to estimate their values.

To minimize $\mathcal{E}_{t,1}$, we prefer selecting layers with larger gradients, which can be achieved by maximizing the value of $\sum_{l\in\mathcal{L}_{t}}\|\nabla_{l}f(\theta^{t})\|_{2}^{2}$. Since the norm of the global gradient $\nabla_{l}f(\theta^{t})$ is unknown, we estimate it by using the sum of stochastic local gradients, expressed as $\sum_{i\in\mathcal{S}^{t}}\|g_{i,l}(\theta^{t};\xi_{i}^{t})\|_{2}^{2}$.
Meanwhile, forcing the same layer selection among clients can reduce the value of $\chi_{\mathbf{w}_{t,l}\|\mathbf{\alpha}}$ and thus alleviate the term $\mathcal{E}_{t,2}$.
For this purpose, we introduce the regularization term $\sum_{j\neq i}\|\mathbf{m}_{i}^{t}-\mathbf{m}_{j}^{t}\|_{1}$ into the optimization objective.
Therefore, the selection of layers is determined by solving the following optimization problem:

$$ $\text{(P1)}\quad\begin{aligned} \max_{\{\mathbf{m}_{i}^{t}\}}&\sum_{i\in \mathcal{S}^{t}}\sum_{l\in\mathcal{L}_{i}^{t}}\|g_{i,l}(\theta^{t};\xi_{i}^{t} )\|_{2}^{2}-\frac{\lambda}{2}\sum_{i\in\mathcal{S}^{t}}\sum_{j\neq i}\|\mathbf {m}_{i}^{t}-\mathbf{m}_{j}^{t}\|_{1}^{2},\\ \text{ s.t. }&\mathcal{R}(\mathbf{m}_{i}^{t})\leq R_{i}^{t},\quad\forall i\in \mathcal{S}^{t}.\end{aligned}$ $$

Here $\lambda\geq 0$ is a weighting constant. Specifically, a large $\lambda$ forces consistent selection across clients, while $\lambda=0$ allows for independent choices among clients.
The constraint in Problem (P1) ensures that the total cost of the selected layers meets the clients’ local resource budgets $R_{i}^{t}$, and the cost function $\mathcal{R}(\cdot)$ is typically a linear function of $\mathbf{m}_{i}^{t}$.

Solving Problem (P1) further gives an effective layer selection strategy for clients.
At the beginning of a training epoch, each participating client $i\in\mathcal{S}^{t}$ evaluates the current global model $\theta^{t}$ on a batch of local data and obtains the layer-wise gradient $g_{i,l}(\theta^{t};\xi_{i}^{t}),\forall l\in\mathcal{L}$. Subsequently, clients upload the norms of these gradients $\|g_{i,l}(\theta^{t};\xi_{i}^{t})\|_{2},\forall l\in\mathcal{L}$, which are $L$-dimensional vectors, to the server. With these values, the server can optimize the selected layer sets for clients by solving Problem (P1).

In general, the proposed layer selection strategy leads to client-specific layer sets, determined based on the estimated gradient norms.
Meanwhile, a hyper-parameter $\lambda$ is used to regulate the extent to which clients are encouraged to select the same layer.
In the next section, we empirically demonstrate the benefits of the proposed strategy in effectively identifying critical layers and achieving better model performance than existing methods.

### 4.3 Discussions on Computational and Communication costs

In this section, we provide discussions on the computational cost and communication costs, considering a case where each client selects $R$ layers to fine-tune a model with $L$ layers.

Computational cost
Since all methods require the same forward operations, we focus on comparing the computational cost of backward operations in different methods.
For simplicity, we assume each layer requires $b$ FLOPs of backward operations.
The average computational cost of the proposed layer selection method is calculated as:

$$ $\text{Cost}_{\text{sel}}=\underbrace{b(L-1)}_{\text{Select}}+\underbrace{bR \tau}_{\text{Fine-tune}}=b(R\tau+L-1),$ (16) $$

where $\tau$ represents the local training steps.
For comparison, fully fine-tuning a model requires the computational cost of:

$$ $\text{Cost}_{\text{full}}=bL\tau=\frac{L\tau}{R\tau+L-1}\text{Cost}_{\text{sel }}.$ (17) $$

Thus, the proposed method takes a much lower computational cost than the full model fine-tuning, and the cost reduction is proportional to the number of layers and local training steps.

Meanwhile, the layer selection step in the proposed method introduces slightly additional costs of $\frac{L-R}{\tau L}\text{Cost}_{\text{full}}$, which can be further reduced by evaluating the model on a smaller volume of data or making the selection decision at a lower frequency.

Communication time
The communication time is determined by the transmitted bits of different methods.
The proposed method only transmits the selected layers, which is much less than uploading the full model.
For example, assuming that different layers have the same number of parameters, the communication cost of the proposed method is $\frac{R}{L}$ of full model fine-tuning.

To summarize, the computational cost and communication cost (time) of the proposed method are much lower than the full model fine-tuning. In the following section, the empirical observations also verify the benefits of the proposed selective layer fine-tuning method.

## 5 Experiments

### 5.1 Settings

Datasets & Models
We conduct a series of experiments on several widely-used image classification datasets, including CIFAR-10 (Krizhevsky & Hinton, ) and DomainNet (Peng et al., ), text classification dataset, i.e., XGLUE-NC (Liang et al., ), and five benchmark question-answering (QA) datasets, including SCIQ (Welbl et al., ), OpenbookQA (Mihaylov et al., ), PIQA (Bisk et al., ), ARC-Easy and ARC-Challenge (Bhakthavatsalam et al., ) datasets. More details of the adopted datasets can be found in in Section B.1.

As for the splitting of datasets, inspired by previous studies (Zhu et al., ; Kim et al., ), we consider two commonly observed data heterogeneity among clients: (i) Label skew: adopting Dirichlet distribution to allocate data samples of the CIFAR-10 dataset; (ii) Feature skew: adopting naturally domain shift on the DomainNet, XGLUE-NC, and QA datasets.
Specifically, each QA dataset is equally divided into two subsets, with each client possessing one subset of samples from one of the five datasets.
We adopt the CLIP model (Radford et al., ) for image classification tasks and the multi-lingual XLM-Roberta-Base (Conneau et al., ) model on the XGLUE-NC dataset.
In addition, we train a LLaMA-2-7B (Touvron et al., ) model on the QA dataset.

Server & Clients
In the experiments, we set up an FL system with a central server and $N=100$ clients.
In each training epoch, the server randomly selects a subset of $20$ clients, and broadcasts the up-to-date model to these clients for local training.
Besides, there are $N=10$ clients in the QA task and five clients are randomly chosen for training in each epoch.
The resource budgets of clients are limited, which are quantified as the maximum number of layers they can fine-tune in the local training. For example, we use $R_{i}=1$ to indicate that the resource of client $i$ cannot afford fine-tuning more than 1 layer.
The resource budgets can be identical or heterogeneous among clients.

Baselines   We compare the proposed layer selection strategy with several competitive baselines, including: (i) Top (Kovaleva et al., ; Lee et al., ): Clients only fine-tune the top few layers (near the output) based on their task-specific data;
(ii) Bottom (Lee et al., ): Clients only fine-tune the bottom few layers (near the input) based on their task-specific data, which can be beneficial for the tasks with input-shift;
(iii) Both (Xiao et al., ): Clients fine-tune an equal proportion of both the top and bottom layers, which shows the effectiveness for large language models;
(iv) SNR (Mahsereci et al., ): Clients fine-tune the layers with higher signal-to-noise ratio (SNR) values, defined as the ratio of the mean of gradient elements to their variance;
(v) RGN (Cheng et al., ; Lee et al., ): Clients fine-tune the layers with higher relative gradient norm (RGN) values, defined as the ratio of gradient norm to the parameter norm;
(vi) Moreover, we consider Full model fine-tuning, i.e., training the entire model, as the performance benchmark.
More implementation details can be found in Section B.2.

### 5.2 Comparisons

We conduct experiments with the identical resource scenario and the heterogeneous resource scenario.

**Table 1: Test accuracy (%) on both image and text datasets, where each client selects $R$ layers for fine-tuning. The best results are highlighted in bold.**
|  | CIFAR-10 | DomainNet | XGLUE-NC | QA |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | $R=1$ | $R=2$ | $R=1$ | $R=2$ | $R=1$ | $R=2$ | $R=1$ | $R=2$ |
| Full | 95.43 | 90.27 | 82.11 | 65.98 |  |  |  |  |
| Top | 93.09 | 93.61 | 87.86 | 88.32 | 69.86 | 77.05 | 63.90 | 64.44 |
| Bottom | 27.38 | 32.81 | 13.80 | 18.63 | 40.43 | 40.60 | 64.18 | 64.60 |
| Both | - | 94.96 | - | 85.48 | - | 74.65 | - | 64.41 |
| SNR | 94.47 | 90.49 | 86.38 | 87.67 | 69.11 | 79.92 | 63.80 | 64.58 |
| RGN | 92.69 | 89.48 | 88.80 | 87.19 | 74.06 | 79.48 | 63.73 | 64.70 |
| Ours | 95.47 | 96.05 | 89.37 | 89.64 | 74.95 | 80.39 | 64.71 | 65.03 |

Identical resource scenario
We first consider clients with identical computational resources, i.e., clients select the same number of layers ($R_{i}=R,\forall i\in\mathcal{N}$) for fine-tuning. The experimental results are shown in Table 1.
From the model performance (accuracy) on the CIFAR-10 and DomainNet datasets, we observe that the proposed strategy demonstrates notable superiority over partial layer fine-tuning baselines.
Specifically, fine-tuning only one layer of the CLIP model achieves comparable performance with tuning the entire model, since the CLIP model is sufficiently powerful to extract useful features and thus requires less training on task-specific data.
This also reveals that selective layer fine-tuning well meets the performance requirement within the resources of clients.

On text datasets, including XGLUE-NC and QA, the proposed layer selection strategy and RGN demonstrate similar performance, both surpassing other baseline methods (especially Top and Both) by noticeable margins.
One potential explanation for this phenomenon could be that they result in similar layer selections, indicating that updating layers with higher relative gradient norms is more beneficial than other strategies, which is consistent with previous study (Lee et al., ).
Moving a forward step, the proposed method adopts a flexible and dynamic layer selection strategy instead of fixed strategies, which leads to competitive performance.

**Table 2: Test accuracy (%) on both image and text datasets, where clients have different resources ($R_{i}\in[1,4]$). The best results are highlighted in bold.**
|  | CIFAR-10 | DomainNet | XGLUE-NC | QA |
| --- | --- | --- | --- | --- |
| Full | 95.43 | 90.27 | 82.11 | 65.98 |
| Top | 91.22 | 89.29 | 78.17 | 64.10 |
| Bottom | 27.38 | 23.10 | 50.92 | 64.51 |
| Both | 89.91 | 86.27 | 73.01 | 64.64 |
| SNR | 75.72 | 87.34 | 78.24 | 64.51 |
| RGN | 93.83 | 88.19 | 79.36 | 64.56 |
| Ours | 95.57 | 89.39 | 80.18 | 65.80 |

Heterogeneous resource scenario
Further, we conduct experiments with heterogeneous clients, i.e., clients have different local resources and thus tend to select different numbers of layers for fine-tuning.
Such a heterogeneous resource scenario is more practical (Yang et al., ; Chai et al., ) and brings additional challenges for selective layer fine-tuning.
Inspired by previous studies (Wang et al., ; Nguyen et al., ), the number of layers to be fine-tuned, denoted as $R_{i}$ for client $i$, is sampled from a truncated half Normal distribution within $[1,4]$.

The experimental results are shown in Table 2, from which we observe that the proposed strategy consistently shows superiority over all the baseline methods on all the datasets.
Compared with baselines, the proposed strategy allows clients to flexibly determine the proper number of layers to be tuned and effectively find the most important layers.
This advantage arises from enabling clients with sufficient resources to prioritize the selection of more critical layers instead of being restricted to layers in fixed positions.
Overall, these experimental results demonstrate the advantage of the proposed strategy when handling heterogeneity in real-world FL applications.

### 5.3 Further Discussions

Visualization of selected layers

Figure: Figure 2: Visualization of selected layers ($R_{i}\in[1,4]$).
Refer to caption: x2.png

For a better understanding on the proposed layer selection strategy, we visualize the selected layers on different datasets in Figure 2.
When fine-tuning the CLIP model on the CIFAR-10 dataset, it can be observed that the focus is primarily on updating a few top layers, indicating that the low-level features (related to middle and bottom layers) are transferable from pre-trained data to downstream tasks.
In comparison, the DomainNet dataset, characterized by a significant domain shift, necessitates extensive tuning of the middle layers in the CLIP model.
Furthermore, on the XGLUE-NC dataset, we observe a clear progression of selected layers for fine-tuning, with a shift from the top layers progressively down to the bottom layers. Such a pattern is markedly different from the trend observed in the image datasets. One possible reason lies in the intrinsic differences between the modalities of text and image data.
These results highlight the necessity for adaptive layer selection and adjustment strategies in FL to accommodate varying dataset properties and domain shifts.

Comparisons regarding computational and communication costs
We compare the computational cost (in TFLOPs) and communication cost (in transmitted MBits) of the proposed method with the full model fine-tuning when adopting the CLIP model on the CIFAR-10.
For the proposed method, we consider fine-tuning only one layer, as it is sufficient to achieve comparable accuracy with full model fine-tuning according to Table 1.
The results in Table 3 evidence a substantial decrease in both computational and communication requirements when utilizing the proposed method.
Besides, the layer selection strategy takes as low as 3.4% of the computational cost.
Therefore, the aforementioned empirical observations demonstrate that the proposed method significantly reduces both the computational and communication costs.

**Table 3: Comparisons of computational and communication costs when fine-tuning the CLIP model on the CIFAR-10 ($R=1$). The numbers in brackets represent the cost of the layer selection strategy.**
|  | Computational cost (TFLOPs) | Ratio | Transmission (MBits) | Ratio |
| --- | --- | --- | --- | --- |
| Full Model Fine-tuning | 8.47 | 100% | 2,811 | 100% |
| Proposed Method | 2.24 (1.51) | 26% (17%) | 234 | 8.33% |
| Proposed (Sel. Period=2) | 1.46 (0.75) | 17% (9.5%) | 234 | 8.33% |
| Proposed (Sel. Batch=1) | 0.99 (0.30) | 12% (3.4%) | 234 | 8.33% |

## 6 Conclusions

In this paper, we study a practical FL setting for fine-tuning foundation models, where clients are allowed to optimize a subset of layers using their task-specific data.
We carefully consider the impact of both data heterogeneity and device heterogeneity across clients, providing a comprehensive theoretical analysis of the optimization objective of selective layer fine-tuning and global model convergence.
The theoretical analysis offers insights into how the selected layers influence global model training and highlights the role of layer importance and client heterogeneity.
We further propose a novel strategy for layer selection that considers the local data and available resources at clients.
The experimental results demonstrate that the proposed strategy outperforms baseline strategies in improving the global model training performance and even matches full model fine-tuning performance in some scenarios, showing the potential for more efficient and tailored real-world FL applications of the proposed layer selection strategy.

## Appendix A Convergence Analysis: Full Proofs

### A.1 Useful Lemmas

#### A.1.1 One-Round Loss Decay

###### Lemma A.1 .

With 4.1, we have:

$$ $\mathbb{E}[f(\theta^{t+1})]-\mathbb{E}[f(\theta^{t})]\leq\frac{1}{2\gamma} \mathcal{E}_{t}+\underbrace{\mathbb{E}\left\langle\sum_{l\in\mathcal{L}_{t}} \nabla_{l}h_{l}^{t}(\theta^{t}),\theta^{t+1}-\theta^{t}\right\rangle}_{ \mathcal{T}_{1}}+\gamma\underbrace{\mathbb{E}\left[\left\|\theta^{t+1}-\theta^ {t}\right\|^{2}\right]}_{\mathcal{T}_{2}}.$ (18) $$

###### Proof.

We begin with analyzing the loss decay by using $\gamma$-smoothness in 4.1 as follows:

$$ $\displaystyle\mathbb{E}[f(\theta^{t+1})]-\mathbb{E}[f(\theta^{t})]$ (19) $\displaystyle\leq$ $\displaystyle\mathbb{E}\langle\nabla f(\theta^{t}),\theta^{t+1}-\theta^{t} \rangle+\frac{\gamma}{2}\mathbb{E}[\left\|\theta^{t+1}-\theta^{t}\right\|^{2}]$ (20) $\displaystyle=$ $\displaystyle\mathbb{E}\left\langle\nabla f(\theta^{t})-\sum_{l\in\mathcal{L}_ {t}}\nabla_{l}h_{l}^{t}(\theta^{t})+\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^ {t}(\theta^{t}),\theta^{t+1}-\theta^{t}\right\rangle+\frac{\gamma}{2}\mathbb{E }[\left\|\theta^{t+1}-\theta^{t}\right\|^{2}]$ (21) $\displaystyle=$ $\displaystyle\underbrace{\mathbb{E}\left\langle\nabla f(\theta^{t})-\sum_{l\in \mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t}),\theta^{t+1}-\theta^{t}\right \rangle}_{\mathcal{T}_{0}}+\underbrace{\mathbb{E}\left\langle\sum_{l\in \mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t}),\theta^{t+1}-\theta^{t}\right \rangle}_{\mathcal{T}_{1}}+\frac{\gamma}{2}\underbrace{\mathbb{E}\left[\left\| \theta^{t+1}-\theta^{t}\right\|^{2}\right]}_{\mathcal{T}_{2}}.$ (22) $$

By Young’s inequality, we upper bound the term $\mathcal{T}_{0}$ as:

$$ $\displaystyle\mathcal{T}_{0}=$ $\displaystyle\mathbb{E}\left\langle\nabla f(\theta^{t})-\sum_{l\in\mathcal{L}_ {t}}\nabla_{l}h_{l}^{t}(\theta^{t}),\theta^{t+1}-\theta^{t}\right\rangle$ (23) $\displaystyle\leq$ $\displaystyle\frac{1}{2\gamma}\underbrace{\mathbb{E}\left[\left\|\nabla f( \theta^{t})-\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^ {2}\right]}_{\mathcal{E}_{t}}+\frac{\gamma}{2}\mathbb{E}\left[\left\|\theta^{t +1}-\theta^{t}\right\|^{2}\right]$ (24) $\displaystyle=$ $\displaystyle\frac{1}{2\gamma}\mathcal{E}_{t}+\frac{\gamma}{2}\mathcal{T}_{2}.$ (25) $$

Plugging (25) back into (22) gives the result in (18).

∎

###### Lemma A.1 .

###### Proof.

#### A.1.2 Analyzing ℰ t subscript ℰ 𝑡 \mathcal{E}_{t} caligraphic_E start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT : Proof of Lemma 4.6

In this subsection, we prove the result in Lemma 4.6.

We begin with decomposing the term $\mathcal{E}_{t}$ using the Jensen’s inequality as:

$$ $\mathcal{E}_{t}\leq 2\underbrace{\|\nabla f(\theta^{t})-\sum_{l\in\mathcal{L}_ {t}}\nabla_{l}f(\theta^{t})\|_{2}^{2}}_{\tilde{\mathcal{E}}_{t,1}}+2 \underbrace{\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}f(\theta^{t})-\nabla_{l}h_{l }^{t}(\theta^{t})\|_{2}^{2}}_{\tilde{\mathcal{E}}_{t,2}}.$ (26) $$

For the first term $\tilde{\mathcal{E}}_{1,t}$, we directly obtain:

$$ $\tilde{\mathcal{E}}_{1,t}=\mathbb{E}\left[\left\|\nabla f(\theta^{t})-\sum_{l \in\mathcal{L}_{t}}\nabla_{l}f(\theta^{t})\right\|^{2}\right]=\mathbb{E}\left[ \left\|\sum_{l\notin\mathcal{L}_{t}}\nabla_{l}f(\theta^{t})\right\|^{2}\right].$ (27) $$

Afterwards, we derive the value of $\tilde{\mathcal{E}}_{2,t}$ as follows:

$$ $\displaystyle\tilde{\mathcal{E}}_{2,t}=$ $\displaystyle\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}f( \theta^{t})-\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^ {2}\right]$ (28) $\displaystyle=$ $\displaystyle\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\sum_{i\in \mathcal{N}}\alpha_{i}\nabla_{l}f_{i}(\theta^{t})-\sum_{l\in\mathcal{L}_{t}} \sum_{i\in\mathcal{S}_{t}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta^{t})\right\|^{2}\right]$ (29) $\displaystyle=$ $\displaystyle\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\sum_{i\in \mathcal{N}}\alpha_{i}\nabla_{l}f_{i}(\theta^{t})-\sum_{l\in\mathcal{L}_{t}} \sum_{i\in\mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta^{t})\right\|^{2}\right]$ (30) $\displaystyle=$ $\displaystyle\sum_{l\in\mathcal{L}_{t}}\mathbb{E}\left[\left\|\sum_{i\in \mathcal{N}}\frac{w_{i,l}^{t}-\alpha_{i}}{\sqrt{\alpha_{i}}}\sqrt{\alpha_{i}} \left(\nabla_{l}f_{i}(\theta^{t})-\nabla_{l}f(\theta^{t})\right)\right\|^{2}\right]$ (31) $\displaystyle\leq$ $\displaystyle\sum_{l\in\mathcal{L}_{t}}\left[\sum_{i\in\mathcal{N}}\frac{(w_{i ,l}^{t}-\alpha_{i})^{2}}{\alpha_{i}}\right]\left[\sum_{i\in\mathcal{N}}\alpha_ {i}\mathbb{E}\left[\left\|\nabla_{l}f_{i}(\theta^{t})-\nabla_{l}f(\theta^{t}) \right\|^{2}\right]\right]$ (32) $\displaystyle\leq$ $\displaystyle\sum_{l\in\mathcal{L}_{t}}\chi_{\mathbf{w}_{t,l}\|\mathbf{\alpha} }\kappa_{l}^{2}.$ (33) $$

where (32) follows the Cauchy–Schwarz inequality and (33) applies 4.3.

By substituting the RHS of (27) and (33) into (26), we complete the proof.
∎

### A.2 Convergence Analysis for Single-Step Case

In this subsection, we consider $\tau=1$ and prove Theorem 4.7.

We derive the value of $\mathcal{T}_{1}$ as follows:

$$ $\mathcal{T}_{1}=\mathbb{E}\left\langle\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l }^{t}(\theta^{t}),-\eta\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t })\right\rangle=-\eta\mathbb{E}\left[\left\|\sum_{l\in L_{t}}\nabla_{l}h_{l}^{ t}(\theta^{t})\right\|^{2}\right].$ (34) $$

Afterwards, we give an upper bound for the term $\mathcal{T}_{2}$ as follows:

$$ $\displaystyle\mathcal{T}_{2}=$ $\displaystyle\left[\left\|\eta\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{S}^ {t}}w_{i,l}^{t}g_{i,l}(\theta^{t};\xi_{i}^{t})\right\|^{2}\right]$ (35) $\displaystyle\leq$ $\displaystyle\eta^{2}\mathbb{E}\left[\left\|\sum_{l\in L_{t}}\nabla_{l}h_{l}^{ t}(\theta^{t})\right\|^{2}\right]+\eta^{2}\sigma^{2},$ (36) $$

where (36) follows 4.2.

Using the result in Lemma A.1, we have:

$$ $\displaystyle\mathbb{E}[f(\theta^{t+1})]-\mathbb{E}[f(\theta^{t})]$ $\displaystyle\leq$ $\displaystyle\frac{1}{2\gamma}\mathcal{E}_{t}-\eta\mathbb{E}\left[\left\|\sum_ {l\in L_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+\gamma\left[ \eta^{2}\mathbb{E}\left[\left\|\sum_{l\in L_{t}}\nabla_{l}h_{l}^{t}(\theta^{t} )\right\|^{2}\right]+\eta^{2}\sigma^{2}\right]$ (37) $\displaystyle=$ $\displaystyle\frac{1}{2\gamma}\mathcal{E}_{t}-\eta(1-\gamma\eta)\mathbb{E} \left[\left\|\sum_{l\in L_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2} \right]+\gamma\eta^{2}\sigma^{2}.$ (38) $$

We define a constant $C\triangleq 1-\gamma\eta>0$ and arrange the terms in (38) as follows:

$$ $\displaystyle\mathbb{E}\left[\left\|\sum_{l\in L_{t}}\nabla_{l}h_{l}^{t}( \theta^{t})\right\|^{2}\right]\leq\frac{1}{\eta C}\left[\mathbb{E}[f(\theta^{t })]-\mathbb{E}[f(\theta^{t+1})]\right]+\frac{1}{2\gamma\eta C}\mathcal{E}_{t}+ \frac{\gamma\eta}{C}\sigma^{2}.$ (39) $$

By Jensen’s inequality, we have:

$$ $\displaystyle\mathbb{E}\left[\left\|\nabla f(\theta^{t})\right\|^{2}\right]$ (40) $\displaystyle=$ $\displaystyle\mathbb{E}\left[\left\|\nabla f(\theta^{t})-\sum_{l\in L_{t}} \nabla_{l}h_{l}^{t}(\theta^{t})+\sum_{l\in L_{t}}\nabla_{l}h_{l}^{t}(\theta^{t })\right\|^{2}\right]$ (41) $\displaystyle\leq$ $\displaystyle 2\mathbb{E}\left[\left\|\nabla f(\theta^{t})-\sum_{l\in L_{t}} \nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+2\mathbb{E}\left[\left\| \sum_{l\in L_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]$ (42) $\displaystyle=$ $\displaystyle 2\mathcal{E}_{t}+2\mathbb{E}\left[\left\|\sum_{l\in L_{t}}\nabla _{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right].$ (43) $$

Combining (39) and (43) gives:

$$ $\displaystyle\mathbb{E}\left[\left\|\nabla f(\theta^{t})\right\|^{2}\right] \leq\frac{2}{\eta C}\left[\mathbb{E}[f(\theta^{t})]-\mathbb{E}[f(\theta^{t+1}) ]\right]+\left(\frac{1}{\gamma\eta C}+2\right)\mathcal{E}_{t}+\frac{2\gamma \eta}{C}\sigma^{2}.$ (44) $$

We sum up both sides of (44) over $t=0,1,\dots,T-1$ and divide them by $T$ to obtain the following result:

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\left[\left\|\nabla f(\theta^{ t})\right\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{2}{\eta CT}\left[\mathbb{E}[f(\theta^{0})]-\mathbb{E}[f( \theta^{T})]\right]+\frac{1}{T}\sum_{t=1}^{T}\left(\frac{1}{\gamma\eta C}+2 \right)\mathcal{E}_{t}+\frac{2\gamma\eta}{C}\sigma^{2}$ (45) $\displaystyle\leq$ $\displaystyle\frac{2}{\eta CT}\left[f(\theta^{0})-f(\theta^{*})\right]+\frac{1 }{T}\sum_{t=1}^{T}\left(\frac{1}{\gamma\eta C}+2\right)\mathcal{E}_{t}+\frac{2 \gamma\eta}{C}\sigma^{2}$ (46) $\displaystyle\leq$ $\displaystyle\frac{2}{\eta CT}\left[f(\theta^{0})-f(\theta^{*})\right]+\frac{2 \gamma\eta}{C}\sigma^{2}+\frac{1}{T}\sum_{t=1}^{T}\left(\frac{1}{\gamma\eta C} +2\right)(\mathcal{E}_{t,1}+\mathcal{E}_{t,2}).$ (47) $$

∎

### A.3 Convergence Analysis for Multi-Step Case

Consider the general case where $\tau>1$. We characterize the convergence in the following theorem and note that the impact of $\mathcal{E}_{t,1}+\mathcal{E}_{t,2}$ is similar to that in Theorem 4.7.

###### Theorem A.2 .

Let $C^{\prime}\triangleq 1-4\eta\tau-8\eta^{2}\gamma^{2}\tau(\tau-1)-32\eta^{3}
\gamma^{2}\tau^{2}(\tau-1)>0$ and $A_{\tau}\triangleq\eta+2\eta^{2}\gamma^{2}\tau(\tau-1)+8\eta^{3}\gamma^{2}\tau
^{2}(\tau-1)$. With Assumptions 4.1-4.3, we have:

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\left[\left\|\nabla f(\theta^{ t})\right\|^{2}\right]\leq$ $\displaystyle\frac{2}{\eta\tau C^{\prime}T}\left[f(\theta^{0})-f(\theta^{*}) \right]+\frac{4A_{\tau}}{C^{\prime}}\sigma^{2}+\frac{1}{T}\sum_{t=1}^{T}\left( \frac{1}{\eta\tau\gamma C^{\prime}}+2\right)(\mathcal{E}_{t,1}+\mathcal{E}_{t, 2}).$ (48) $$

###### Proof.

In Lemma A.1, the term $\mathcal{T}_{1}$ is related to client drift caused by multiple local SGD steps, which can be upper bounded as follows:

$$ $\displaystyle\mathcal{T}_{1}$ $\displaystyle=$ $\displaystyle-\eta\sum_{k=0}^{\tau-1}\mathbb{E}\left\langle\sum_{l\in\mathcal{ L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t}),\sum_{l\in\mathcal{L}_{t}}\sum_{i\in \mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta_{i}^{t,k})\right\rangle$ (49) $\displaystyle=$ $\displaystyle-\eta\sum_{k=0}^{\tau-1}\mathbb{E}\left\langle\sum_{l\in\mathcal{ L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t}),\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{ l}^{t}(\theta^{t})\right\rangle$ $\displaystyle+\eta\sum_{k=0}^{\tau-1}\mathbb{E}\left\langle\sum_{l\in\mathcal{ L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t}),\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{ l}^{t}(\theta^{t})-\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t} \nabla_{l}f_{i}(\theta_{i}^{t,k})\right\rangle$ (50) $\displaystyle\leq$ $\displaystyle-\frac{\eta}{2}\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left\|\sum_{l \in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+\frac{ \eta}{2}\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}} \nabla_{l}h_{l}^{t}(\theta^{t})-\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N }}w_{i,l}^{t}\nabla_{l}f_{i}(\theta_{i}^{t,k})\right\|^{2}\right]$ (51) $\displaystyle=$ $\displaystyle-\frac{\eta\tau}{2}\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{ t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+\frac{\eta}{2}\sum_{k=0} ^{\tau-1}\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N }}w_{i,l}^{t}\nabla_{l}f_{i}(\theta^{t})-\sum_{l\in\mathcal{L}_{t}}\sum_{i\in \mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta_{i}^{t,k})\right\|^{2}\right]$ (52) $\displaystyle\leq$ $\displaystyle-\frac{\eta\tau}{2}\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{ t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+\frac{\eta\gamma^{2}}{2} \underbrace{\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t }}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\left(\theta^{t}-\theta_{i}^{t,k}\right) \right\|^{2}\right]}_{\mathcal{T}_{4}},$ (53) $$

where (51) uses the inequality $\left\langle\mathbf{a},\mathbf{b}\right\rangle\leq\frac{\|\mathbf{a}\|^{2}}{2}
+\frac{\|\mathbf{b}\|^{2}}{2}$, and (53) follows 4.1.

Then we analyze the term $\mathcal{T}_{2}$ as follows:

$$ $\displaystyle\mathcal{T}_{2}$ $\displaystyle\leq$ $\displaystyle\eta^{2}\mathbb{E}\left[\left\|\sum_{k=0}^{\tau-1}\sum_{l\in \mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta_{i}^{t ,k})\right\|^{2}\right]+\eta^{2}\tau\sigma^{2}$ (54) $\displaystyle\leq$ $\displaystyle\eta^{2}\tau\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left\|\sum_{l\in \mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta_{i}^{t ,k})-\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i }(\theta^{t})+\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t} \nabla_{l}f_{i}(\theta^{t})\right\|^{2}\right]$ $\displaystyle+\eta^{2}\tau\sigma^{2}$ (55) $\displaystyle\leq$ $\displaystyle 2\eta^{2}\tau\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left\|\sum_{l \in\mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta_{i} ^{t,k})-\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\nabla_{l}f _{i}(\theta^{t})\right\|^{2}\right]$ $\displaystyle+2\eta^{2}\tau\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left\|\sum_{l \in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+\eta^{2} \tau\sigma^{2}$ (56) $\displaystyle\leq$ $\displaystyle 2\eta^{2}\gamma^{2}\tau\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left \|\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\left(\theta_{i}^ {t,k}-\theta^{t}\right)\right\|^{2}\right]+2\eta^{2}\tau^{2}\mathbb{E}\left[ \left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2} \right]+\eta^{2}\tau\sigma^{2}$ (57) $\displaystyle=$ $\displaystyle 2\eta^{2}\gamma^{2}\tau\mathcal{T}_{4}+2\eta^{2}\tau^{2}\mathbb{ E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right \|^{2}\right]+\eta^{2}\tau\sigma^{2},$ (58) $$

where (54) follows 4.2, (55)-(56) apply the Jensen’s inequality, and (57) follows 4.1.

Following Lemma 22 in (Pillutla et al., ), $\mathcal{T}_{4}$ can be upper bounded as:

$$ $\displaystyle\mathcal{T}_{4}\leq$ $\displaystyle\sum_{k=0}^{\tau-1}\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{ t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\theta^{t}-\theta_{i}^{t,k}\right\|^{2}\right]$ (59) $\displaystyle\leq$ $\displaystyle 8\eta^{2}\tau^{2}(\tau-1)\mathbb{E}\left[\left\|\sum_{l\in \mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t}\nabla_{l}f_{i}(\theta^{t}) \right\|^{2}\right]+\sum_{l\in\mathcal{L}_{t}}\sum_{i\in\mathcal{N}}w_{i,l}^{t }4\eta^{2}\tau^{2}(\tau-1)\sigma_{l}^{2}$ (60) $\displaystyle=$ $\displaystyle 8\eta^{2}\tau^{2}(\tau-1)\mathbb{E}\left[\left\|\sum_{l\in \mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+4\eta^{2} \tau^{2}(\tau-1)\sigma^{2}.$ (61) $$

Plugging (54),(58) and (61) back into (18), we have the following result:

$$ $\displaystyle\mathbb{E}[f(\theta^{t+1})]-\mathbb{E}[f(\theta^{t})]$ (62) $\displaystyle\leq$ $\displaystyle\frac{1}{2\gamma}\mathcal{E}_{t}-\frac{\eta\tau}{2}\mathbb{E} \left[\left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\| ^{2}\right]+\frac{\eta\gamma^{2}}{2}\mathcal{T}_{4}+2\eta^{2}\gamma^{2}\tau \mathcal{T}_{4}+2\eta^{2}\tau^{2}\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_ {t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+\eta^{2}\tau\sigma^{2}$ (63) $\displaystyle=$ $\displaystyle\frac{1}{2\gamma}\mathcal{E}_{t}-\frac{\eta\tau}{2}\left(1-4\eta \tau\right)\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t }(\theta^{t})\right\|^{2}\right]+\eta^{2}\tau\sigma^{2}+\left(\frac{\eta\gamma ^{2}}{2}+2\eta^{2}\gamma^{2}\tau\right)\mathcal{T}_{4}$ (64) $\displaystyle=$ $\displaystyle\frac{1}{2\gamma}\mathcal{E}_{t}-\frac{\eta\tau}{2}\left(1-4\eta \tau\right)\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t }(\theta^{t})\right\|^{2}\right]+\eta^{2}\tau\sigma^{2}$ $\displaystyle+\left(\frac{\eta\gamma^{2}}{2}+2\eta^{2}\gamma^{2}\tau\right) \left\{8\eta^{2}\tau^{2}(\tau-1)\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{ t}}\nabla_{l}h_{l}^{t}(\theta^{t})\right\|^{2}\right]+4\eta^{2}\tau^{2}(\tau-1 )\sigma^{2}\right\}$ (65) $\displaystyle=$ $\displaystyle\frac{1}{2\gamma}\mathcal{E}_{t}-\frac{\eta\tau}{2}\left[1-4\eta \tau-8\eta^{2}\gamma^{2}\tau(\tau-1)-32\eta^{3}\gamma^{2}\tau^{2}(\tau-1) \right]\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}( \theta^{t})\right\|^{2}\right]$ $\displaystyle+\left(\eta^{2}\tau+2\eta^{3}\gamma^{2}\tau^{2}(\tau-1)+8\eta^{4} \gamma^{2}\tau^{3}(\tau-1)\right)\sigma^{2}.$ (66) $$

Let $C^{\prime}\triangleq 1-4\eta\tau-8\eta^{2}\gamma^{2}\tau(\tau-1)-32\eta^{3}
\gamma^{2}\tau^{2}(\tau-1)>0$ and $A_{\tau}\triangleq\eta+2\eta^{2}\gamma^{2}\tau(\tau-1)+8\eta^{3}\gamma^{2}\tau
^{2}(\tau-1)$. We have:

$$ $\mathbb{E}\left[\left\|\sum_{l\in\mathcal{L}_{t}}\nabla_{l}h_{l}^{t}(\theta^{t })\right\|^{2}\right]\leq\frac{2}{\eta\tau C^{\prime}}\left[\mathbb{E}[f( \theta^{t})]-\mathbb{E}[f(\theta^{t+1})]\right]+\frac{1}{\eta\tau\gamma C^{ \prime}}\mathcal{E}_{t}+\frac{2}{C^{\prime}}A_{\tau}\sigma^{2}.$ (67) $$

Using the result in (43), we have:

$$ $\mathbb{E}\left[\left\|\nabla f(\theta^{t})\right\|^{2}\right]\leq\frac{4}{ \eta\tau C^{\prime}}\left[\mathbb{E}[f(\theta^{t})]-\mathbb{E}[f(\theta^{t+1}) ]\right]+\left(\frac{1}{\eta\tau\gamma C^{\prime}}+2\right)\mathcal{E}_{t}+ \frac{4A_{\tau}}{C^{\prime}}\sigma^{2}.$ (68) $$

We sum up both sides of (68) over $t=0,1,\dots,T-1$ and divide them by $T$ to obtain the following result:

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\left[\left\|\nabla f(\theta^{ t})\right\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{2}{\eta\tau C^{\prime}T}\left[\mathbb{E}[f(\theta^{0})]- \mathbb{E}[f(\theta^{T})]\right]+\frac{1}{T}\sum_{t=1}^{T}\left(\frac{1}{\eta \tau\gamma C^{\prime}}+2\right)\mathcal{E}_{t}+\frac{4A_{\tau}}{C^{\prime}} \sigma^{2}$ (69) $\displaystyle\leq$ $\displaystyle\frac{2}{\eta\tau C^{\prime}T}\left[f(\theta^{0})-f(\theta^{*}) \right]+\frac{1}{T}\sum_{t=1}^{T}\left(\frac{1}{\eta\tau\gamma C^{\prime}}+2 \right)\mathcal{E}_{t}+\frac{4A_{\tau}}{C^{\prime}}\sigma^{2}$ (70) $\displaystyle\leq$ $\displaystyle\frac{2}{\eta\tau C^{\prime}T}\left[f(\theta^{0})-f(\theta^{*}) \right]+\frac{4A_{\tau}}{C^{\prime}}\sigma^{2}+\frac{1}{T}\sum_{t=1}^{T}\left( \frac{1}{\eta\tau\gamma C^{\prime}}+2\right)(\mathcal{E}_{t,1}+\mathcal{E}_{t, 2}).$ (71) $$

∎

###### Theorem A.2 .

###### Proof.

## Appendix B Experimental Details

We implement all methods with PyTorch and run experiments on Nvidia V100 GPUs.
For fair comparisons, we adopt the same training epochs and hyper-parameters for all methods.

### B.1 Training Tasks

To evaluate the methods in various scenarios with different non-IID patterns, we consider two image classification tasks and two text classification tasks.

**Table 4: Summary of datasets.**
| Dataset | Data Type | Non-IID Type | Partition |
| --- | --- | --- | --- |
| CIFAR-10 | Image | Label skew | $Dir(0.1)$ |
| DomainNet | Image | Feature skew | Domain |
| XGLUE-NC | Text | Feature skew | Domain |
| QA | Text | Feature skew | Domain |

The image classification tasks include:

- •
CIFAR-10 (Krizhevsky & Hinton, ): In this training task, we consider the label-skewed case where $P(y_{i})$ is different among clients. Following previous works (Li et al., ), we adopt Dirichlet allocation with concentration parameter $\alpha=0.1$ among clients.
- •
DomainNet (Peng et al., ): DomainNet contains six domains of data samples, including clipart, real, sketch, infograph, painting, and quickdraw. In this training task, we consider the varying feature case where $P(x_{i}|y_{i})$ is different among clients. Following (Li et al., ), each client is allocated random samples from only one domain.

On both tasks, we fine-tune a CLIP Vision Transformer (CLIP) model (Radford et al., ).

Besides, we fine-tune an XLM-Roberta-Base model on the following text dataset:

- •
XGLUE-NC (Liang et al., ): This is a news classification task consisting of 10 classes. The news texts comprise five languages (English, Spanish, French, German, and Russian). We allocate one random language to each client, which naturally introduces domain shift among clients.

In addition, we evaluate a LLaMA-2-7B  model on the QA datasets.

- •
QA: The QA datasets consist of four commonly used question-answering datasets, i.e., SCIQ (Welbl et al., ), OpenbookQA (Mihaylov et al., ), PIQA (Bisk et al., ), ARC-Easy and ARC-Challenge (Bhakthavatsalam et al., ). The datasets are transformed into classification tasks where the model determines the correct answer for each question and corresponding choices.
We allocate the samples from one random dataset to each client, indicating the domain shift among clients.

### B.2 Implementation Details

The CLIP model is pre-trained on the DataComp dataset and is adapted from [https://github.com/openai/CLIP](https://github.com/openai/CLIP);
the XLM-Roberta-Base model is adapted from [https://huggingface.co/xlm-roberta-base](https://huggingface.co/xlm-roberta-base); the LLaMA-2-7B model is adapted from [https://huggingface.co/meta-llama/LLaMA-7b-chat-hf](https://huggingface.co/meta-llama/LLaMA-7b-chat-hf).
In all training tasks, we freeze the embedding layers of the model and fix the classifier as commonly selected layers (Lee et al., ).
The values of adopted hyperparameters are summarized in Table 5.
For the proposed method, we tune the value of $\lambda$ from {1, 5, 10, 100, 500, 1000}.

**Table 5: Implementation details.**
| Dataset | CIFAR-10 | DomainNet | XGLUE-NC | QA |
| --- | --- | --- | --- | --- |
| Model | CLIP | CLIP | XLM-Roberta-Base | LLaMA-2-7B |
| Batch size | 64 | 64 | 32 | 16 |
| Learning rate | 0.01 | 0.01 | 0.01 | 2e-5 |
| Local steps^∗ | 5 | 1 | -1 | -1 |
| Total epochs | 50 | 30 | 20 | 2 |
| $\lambda$ | 1000 | 10 | 10 | 5 |
| ^∗The local steps -1 means clients iterate all training samples (i.e., one single local training epoch). |  |  |  |  |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
questions: 1. How problem 1 (line 344) is optimized? It is an integer programming problem, which can require exponential complexity to get the globally optimal solution. 
2. Problem 1 uses $$\sum_{i}\sum_{l} \| g_{i, l} \|_2^2$$ to approximate $$\mathcal{E}_{t, 1}$$. It looks weird to me since $$\sum_{l} \| \sum_{i} g_{i, l} \|_2^2$$ seems to be a closer approximation of the true $$\sum_l \| \nabla_l f \|_2^2$$. What is the beneficial of using the approximation in the paper? 
3. See Weakness 1, what is the exact algorithm used? 
4. Communication cost: since problem 1 requires pairwise comparison of client masks, it needs to be solved on server. It seems like in each communication round, each client needs to (1) compute the local gradient for all layers (2) upload all layers' gradients to the server, waiting for it to compute the mask m, (3) download mask m and perform selective layer updates, and (4) upload the gradients for partial layers. If it is true, the algorithm actually increase the communication cost. Please correct me if I am wrong. 
5. For the CLIP model, which part of the model is updated? Only the vision encoder, or both vision and text encoder? Please provide these details in the Appendix if possible.
soundness: 2
strengths: 1. The theoretical part is very well-presented. The notations, definitions, and assumptions of this paper is very rigorous, and each lemma or theorem has good intuition. 
2. The proposed method has stronger performance than baselines with fixed selection and baselines based on heuristics.
summary: This paper consider the selective layer update in federated learning. It first theoretically analyzes the impact of layer selection on the model convergence. Then, motivated by the form of the error term, it proposes a layer selection algorithm that minimizes a reformation of the error term as a surrogate loss. The proposed method has comparable model performance to full update, and saves computation and communication cost.
weaknesses: 1. The exact training process is not very clear. In Algorithm 1, it seems each client only compute the gradient once, and the model is updated on the server, which is similar to FedSGD. However, In the analysis of computational cost, it seems like each client will first compute a step of full layer update, do the layer selection, and then do $$\tau$$ steps of selective layer update. 
2. The computational cost seems over-simplified. In (16), the author simplify the cost of back-propagation as a simple sum of cost for each layer. However, this is generally not true. For example, for a ResNet network, finetuning the last linear layer is very efficient since one only needs to run back-propagation on the final layer. However, if one chooses to fine-tune the first convolutional layer, they need to back-propagate through the whole network to the very first layer, during which computing the gradient to intermediate feature maps. Therefore, just finetuning the last layer can save a lot of computation, while just finetuning the first layer may not save many computation. Also, the result in Table 3 may largely result from this, since in Figure 2(a), only the last few layers are selected after 10 epochs. Does the algorithm still save computation on other settings which may select the first layers (e.g., on XGLUE-NC)? 
3. Lack of hyperparameter sensitivity w.r.t. $$\lambda$$, which seems to be related to gradient norms. 
4. [Minor] Typos: Lemma 4.6 $$\nabla$$ in $$\mathcal{E}_{t,1}$$ should be $$\nabla_l$$​. 
5. [Minor] Corollary 4.8 may not be very rigorous. Typically the learning rate is either a constant or a function of $$t$$, not a function of $$T$$. Also when $$\eta = \frac{1}{\sqrt{T}}$$, the third term seems limit to infinity. I believe it is better to just set learning rate to constant, and claim the third term dominate the RHS.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See the weaknesses.
soundness: 2
strengths: 1. The paper provides a robust theoretical exploration of selective layer fine-tuning, enhancing understanding of how layer selection impacts convergence in FL.

2. By allowing clients to customize their fine-tuning strategies based on local resources and data, the proposed approach effectively addresses the inherent heterogeneity present in FL scenarios.
summary: This study investigates the concept of selective layer fine-tuning within the federated learning (FL) paradigm, where clients fine-tune only specific layers of foundation models based on their local data and resource constraints. By allowing clients to choose which layers to update, the paper emphasizes a flexible approach that adapts to the diverse capabilities and data distributions of clients. Theoretical analysis shows that the selection of layers significantly influences model convergence, focusing on the importance of chosen layers and the heterogeneity in client selections. An effective layer selection strategy leveraging local gradients is proposed, demonstrating enhanced performance in extensive experiments across image and text datasets compared to baseline methods.
weaknesses: 1. The proposed layer selection strategy may require sophisticated implementations and tuning, which could complicate its deployment in practice, particularly for clients with limited expertise or resources.

2. The effectiveness of selective layer fine-tuning in larger, more complex models or across a significantly larger number of clients remains to be fully explored, raising questions about its scalability.

3. While the paper addresses client heterogeneity, it does not fully consider the impact of changing data distributions over time, which could affect the effectiveness of the fine-tuning strategy as client environments evolve.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Layers with larger gradient norms are assumed to be more relevant to client tasks. However, could selecting these layers for global aggregation and fine-tuning exacerbate the discrepancy between global and local parameters?

2. Due to non-iid distributions, different clients should ideally select different layers. However, this work enforces the same layer selection across clients using regularization terms. Please explain the rationale behind optimization objective P1.

3. According to Section 3, only the selected layers are updated during global fine-tuning, while other layers remain unchanged. Could updating the unselected layers based on local gradients provide additional performance gains?

4. Could Figure 2 be supplemented with visualizations of the selected layers for the RGN and SNR baseline?

5. Regarding the proposed method’s reduction in computational costs in table 3, is this saving limited to the parameter updates on the client side? In the process of calculating local gradients, are gradients for all layers still required?
soundness: 2
strengths: 1. The idea of selectively fine-tuning specific layers in federated learning is innovative. 
2. The paper sets up various Non-IID scenarios on four benchmark datasets.
3. The paper is well written and easy to follow.
summary: This work proposes a layer selection strategy based on client gradient norms, where layers with larger gradient norms are fine-tuned rather than adjusting the entire model. This selective strategy considers practical factors such as data heterogeneity and client resources, achieving improved model performance on image and text datasets.
weaknesses: 1. The paper focuses on layer selection for fine-tuning but lacks a detailed discussion of experimental results in Tables 1 and 2. For instance, in the setup of selecting the same number of layers in Table 1, the differences in selected layers between the proposed method and RGN or SNR methods should be elaborated.

2. The comparative experiments lack stronger baselines, such as adding the same regularization terms to SNR and RGN as in the proposed method.

3. The paper lacks robust conclusions regarding layer selection. For instance, a dynamic layer selection strategy where bottom layers have larger gradient norms in early training stages and top layers have larger norms in later stages could be considered. Additionally, there is an absence of discussion on whether forcing all clients to select the same layers or allowing them to select different layers could be more effective.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: - Based on your theoretical findings, what is the specific reason on focusing on a layer-wise selection as opposed to element-wise selection?
- How does the pretrained foundation model's zero shot ability or local training performance compare to the federated methods?
- How does $\lambda$ affect the performance numerically? Since $\varepsilon_{t,1}$ and $\varepsilon_{t,2}$ can sometimes conflict with each other, is the performance bad when $\lambda$ is either too large or too small?
soundness: 3
strengths: - The problem this paper focuses on is important. Leveraging pre-trained foundation models in federated learning often does require careful training strategies due to clients' resource limitations.
- Theoretical analysis provides insights towards the proposed layer selection strategy, indicating the effect of large gradient norms and consistent selection among clients.
- Experiments include different modalities on various datasets.
summary: This paper examines selective layer fine-tuning in federated learning (FL) to adapt foundation models efficiently across clients with varying resources. It presents a theoretical analysis, providing insights towards the proposed gradient-based layer selection strategy over heterogeneous clients. This strategy enables clients to choose layers based on data and resource constraints. Experiments validate its efficiency, demonstrating performance on par with full model fine-tuning in resource-limited and heterogeneous FL settings.
weaknesses: - While selective layer fine-tuning and gradient magnitude-based selection are useful, these approaches are both well studied in the field (e.g. [1,2]). This paper lacks a distinct methodological innovation that would set it apart from merely combining the two existing techniques.

- The theoretical insights presented do not appear specific to the layer-selection strategy, making them applicable to more granular selection approaches, such as element-wise weight selection [2]. Without a clear motivation for focusing on layer-based selection, the paper does not sufficiently justify the connection between the theoretical findings and the layer-wise selection method.

- The experimental gains reported, typically less than 1%, may not justify the complexity of implementing this approach in real-world applications. Additionally, the lack of baseline results from purely local training (without FL) limits the evaluation of the added value of the proposed method over simple local fine-tuning.

- There is no ablation study on $\lambda$. Since $\lambda$ plays a key role in controlling the degree of uniformity in layer selection across clients, an ablation study would have provided critical insights into its impact on model performance. Exploring how $\lambda$ values affect the trade-off between personalization and global model alignment could help validate the robustness of the approach, making the results more interpretable and the method more adaptable to diverse FL settings.

- I do not think $\sum_{l\in \mathcal{L}_i^t}$ is accurate. E.g. in Eq.(2) how do you add the gradient with respect to different layers together?

[1] Surgical fine-tuning improves adaptation to distribution shifts

[2] FedSelect: Personalized Federated Learning with Customized Selection of
Parameters for Fine-Tuning

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1) The authors mention two types of solutions for finetuning LLMs in a resource-constrained setting: (a) parameter-efficient finetuning (PEFT) and (b) Selective model finetuning. Since they are both orthogonal, I wonder if they can be combined for an even better performance. Also, what is the rationale behind not comparing the proposed method against PEFT methods like LoRA? 

2.a) Under "Introduction," right above the list of "main contributions," there's a sentence stating "a novel layer selection strategy that leverages local gradients and the regulation of unified selections." Can the authors please clarify how their strategies are novel in terms of leveraging local gradients? Don't the SNR and RGN baselines also leverage the same local gradients? 

2.b) Also, what does "unified selections" mean in this context? Is it the same client picked by the same client every round? 

3) Layer selection should depend on the heterogeneity among selected clients for the current round. If clients are selecting which layers they want to finetune, is there any way for the clients to know which layers are getting selected by other clients?

4) Referring to Line 200-201, what is the rationale behind "clients being able to select which layers to finetune" relating to "handling heterogeneity in FL"? Can't/won't the clients just select the "easy to achieve high accuracy on" layers for local finetuning (which might be detrimental to minimizing the global loss)? In other words, clients (in contrast to server) selecting which layers to finetune should lead to clients being selfish and only caring about personalized accuracy, which is opposite to how global balance can be achieved in the face of heterogeneity.

5) May I know what does a "Layer" in the context of language models and the CLIP model consists of? Is it one linear layer (e.g., one of the query, key, or value linear layers) or a module containing multiple linear or other layers?

6) Theorem 4.7 has good findings (in Lines 317-323) related to the terms $\varepsilon_{t, 1}$ and $\varepsilon_{t, 2}$; I am wondering if there's an empirical way to observe these two phenomena:

(i) The selected layers have large gradient norms based on the gradients derived by the current round's clients.

(ii) Inconsistent selection of clients not leading to a global convergence.

7) What is the reasoning behind defining a resource budget in terms of how many layers a client can finetune? What if the budget is 1 layer, but that layer has large dimensions which cannot be fit into a device like a resource-constrained mobile phone? Who decides that budget and how?

8.a) Does Equation 16 have cost of $b$ in "Select" due to clients having to compute the gradient norm? 

8.b) Line 383 states, "The proposed method only needs to transmit the selected layers." However, don't the clients also need to transmit the gradient norms to the server?
soundness: 2
strengths: 1. The use case is relevant and timely; and the motivation makes sense. 

2. The work considers data and device heterogeneity under the FL setting, making it more practical.

3. The manuscript gives new insights on layer selection strategies related to picking layers with large gradient norms and consistency in layer selection by a client.

4. The writing is mostly clear (clarifying questions on a few confusing points have been asked in the questions section).
summary: This work proposes a layer-selection-based strategy to finetune large models in a federated setting. "Which layers to finetune" is selected based on the resource budget and data distribution of each client. Empirical evaluation shows the improvements in model accuracy, and theoretical analysis shows how heterogeneity based on data and hence derived gradients plays a role in global model convergence.
weaknesses: 1) Lack of strong baselines. The work is compared against only finetuning top or bottom layers and against methods that select layers based on signal-to-noise ratio (SNR) and relative gradient norm (RGN); however, there are methods like FedBiot [1], which uses LoRA adapters to finetune LLMs in FL. (Although FedBiot was quite recently published, a similar baseline that includes parameter-efficient finetuning should be an essential method to compare against.)

[1] Fedbiot: LLM local fine-tuning in federated learning without full model (ACM SIGKDD 2024)

2) Which party (server or clients) selected which layers should be trained by which client is unclear from algorithm 1 and its corresponding discussion text. Line 42 and 199 say "clients tend to update some of the layers" and "clients are allowed to select different layers", but Line 166 in Algorithm 1 indicates that the server is sending the selected layers to each client, and Line 355 states "the server can optimize the selected layer sets." What would be the difference in performance and variables to consider if it's (a) the server or (b) the clients picking which layers to finetune? (I have follow-up questions in Q3 and Q4 under the "Questions" as well.)

3) Section 4.3 can be massively improved by drawing comparisons of the computational and communication costs of the proposed method with all its baselines instead of just picking the most intensive baseline of full model finetuning.

4.a) The insights into why the proposed method finetunes the CLIP model so well compared to other baselines do not make sense. The line 460 "since the CLIP model is sufficiently powerful to extract useful features..." is the property of the model, not of the proposed method. And if that's the reasoning, at least one of the Top/Bottom baselines should also have great performance, but that is not the case. According to Figure 2, Top baseline should have worked way better.

4.b) Lines 465-467 state that RGN and the proposed layer selection strategy result in similar layer selections, so why does RGN not have higher resultant accuracy? What are the differences between RGN and your proposed method?

Meta-review: **Summary:** The paper investigates the impact of selective layer fine-tuning in federated learning settings. It provides a theoretical analysis of the relationship between selective layer updates and model convergence and proposes a gradient-based layer selection strategy to balance computational costs and heterogeneity across clients. Experiments conducted on both image and text datasets demonstrate the proposed approach's potential to save computation and improve adaptability compared to heuristic baselines.

**Decision:** This paper tackles an important problem in federated learning by exploring selective layer fine-tuning as a means to balance computation and heterogeneity. However, the proposed method lacks significant methodological innovation, as it builds on existing techniques without introducing distinct new ideas. Additionally, the theoretical contributions are generic and not directly tied to the proposed approach, weakening the overall impact. The experimental validation is limited to small-scale datasets and does not include comparisons to recent state-of-the-art methods. Besides, some reviewers also raise concerns about the further analysis of the results (e.g., sensitivity analysis of $\lambda$). These limitations outweigh the contributions, leading to the decision to reject. During the reviewer-AC discussion period, the reviewers unanimously agreed with this decision.

Decision: Reject