Exploring Selective Layer Fine-Tuning in Federated Learning

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?

• 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

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.

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.

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.