A Closer Look at Personalized Fine-Tuning in Heterogeneous Federated Learning

Response to Q3

We appreciate the suggestion to provide a concise summary of the theoretical analysis. We have included this summary at the beginning of the theoretical analysis section to offer a clear road map of the theory section.

Here is the overview: To compare the global performance of LP-FT and FT, we make assumptions about the clients’ data-generating function (Assumption 4.1) and the model structure (Assumption 4.2). Both assumptions involve a two-layer network to distinguish between the linear head and the feature extractor layer. We then define the global loss function as MSE loss. Based on these assumptions, we compare the global performance of LP-FT and FT under concept shift (Theorem 4.4) and combined concept-and-covariate shift (Theorem 4.5). For this comparison, we use one step of gradient descent (a single pass through the entire dataset) to evaluate the global loss of both methods.

Response to Q4

Thank you for the suggestion. We have revised Figure 2 to provide a detailed visualization of the individual global and local accuracy trends across various distribution shift scenarios. The updated figure includes:

• (a) Global and local accuracy under different learning rates for full-parameter fine-tuning.
• (b) Accuracy trends under varying sparsity rates for sparse fine-tuning.
• (c) Accuracy behavior with different regularization strengths in proximal fine-tuning.

In all subplots, global accuracy is represented by solid lines and local accuracy by dashed lines. The results show that while global accuracy remains stable, local accuracy consistently decreases across all hyperparameter settings, highlighting overfitting in the local models.

We believe this revision addresses the concern and provides a clearer understanding of how personalization overfitting manifests. Thank you for the valuable feedback.

Response to Q1 & Q2

Thank you for the question. The superior performance of LP-FT (Linear-Probing-then-Fine-Tuning)—which involves fine-tuning the linear head followed by fine-tuning the full model (see Section 1 Line 83, and Figure 1(b); further details are provided in Appendix B.4 of the revised version)—can be explained both empirically and theoretically. Below, we first provide a high-level explanation and then point to the sections where these reasons are explored in depth.

1. High-Level Intuition When applying LP-FT, the features (last-layer representations before the linear head) are fixed during the linear head fine-tuning (LP stage), while only the linear head is updated. In the subsequent full fine-tuning stage, since the linear head is already adapted to local performance, the features experience less distortion compared to directly applying FT without LP. This reduced feature distortion leads to better performance for LP-FT (stated in Section 1 Line 82-90). The pre-trained model already performs well on global data, so minimizing feature distortion ensures better global performance after fine-tuning. This concept is illustrated graphically in Fig. 3.

2. Empirical Evidence In Sec. 3.5, we present experiments to explore the connection between federated feature distortion and the performance of LP-FT and FT. As shown in Fig. 4(a), FT results in significant feature distortion, causing a substantial decline in global performance. In contrast, LP-FT maintains higher global performance with significantly reduced feature distortion.

3. Theoretical Insight In Sec. 4, we provide a theoretical analysis comparing the global performance of LP-FT and FT under two distribution shift settings: concept shift and combined concept-and-covariate shifts. Using a two-layer neural network as the data-generating function and model structure, we simulate LP-FT and FT, where layer 1 represents the features and layer 2 represents the linear head. In Lemma 4.3, we derive the changes in features under FT inspired by feature distortion. In Theorems 4.4 and 4.5, we show that LP-FT achieves better global performance than FT under concept shift and combined concept-and-covariate shifts, respectively. These results demonstrate how reduced feature distortion in LP-FT contributes to its superior performance.

Summary LP-FT outperforms FT in global performance by significantly reducing federated feature distortion, leading to improved overall results.

[Continued from above]

Response to W1

Empirical Contributions (See Section 1 Contribution Evaluation):,

1. New Observations: Detailed Analysis on Overfitting in PFT: As shown in Figure 2, we analyze commonly used fine-tuning strategies and demonstrate their tendency to overfit, even with carefully tuned hyperparameters, across three different types of distribution shifts. These findings highlight a critical gap in prior FL research, where such tendencies have not been systematically explored.

2. Comprehensive Benchmarking PFT: Our experiments span seven diverse datasets (CIFAR-10-C, CIFAR-10, Digit5, DomainNet, Chexpert, CelebA, CIFAR-10), addressing four types of distribution shifts (input-level shift, feature shift, spurious correlation, and label shift, as detailed in Appendix Table 6). We evaluate five Personalized Fine-Tuning (PFT) methods using five robust metrics, providing a thorough and comprehensive analysis.

3. Establishing New Ablation Studies to Support Theoretical Results: In Figure 4, we conduct detailed analysis on feature distortion relation to overfitting and in Table 2, we conduct large ranges of label flipping to support the theoretical analysis validating our theoretical insights in low heterogeneity regimes and further demonstrating LP-FT’s superior performance in scenarios with higher heterogeneity.

Reference

• [1] Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution. In ICLR. OpenReview.net, 2022.

We sincerely thank the reviewer for their valuable comments and questions. Below, we provide our responses to the concerns and inquiries raised:

Response to W1

We would like to clarify that the novelty and contribution of our work extend beyond simply adapting an existing method (LP-FT [1]) to a new scenario (PFT). Our key contributions lie in the methodological and theoretical novelty, as well as empirical findings, benchmarking, and analysis. These efforts provide deeper insights and fine-grained evaluations of diverse distribution shifts, demonstrating the effectiveness of LP-FT in the new PFT setting.

Methodological Novelties (See Section 1 Contribution Insight)

1. Simple yet Non-Trivial Method: Novel Focus on Final-Stage LP-FT in FL

While our method may appear straightforward, its novelty lies in strategically shifting the insertion phase of LP-FT to the final local training stage in FL. Unlike most existing PFL methods that require modifications to the entire training process, including both General FL (GFL) training and local personalization, our approach focuses on the critical yet underexplored final local training stage. This makes our method a versatile solution, independent of the GFL training process. This shift not only introduces a fresh perspective on FL personalization but also demonstrates significant theoretical and empirical insights.

We also respectfully wish reviewers to agree that Simplicity does not diminish the novelty of our method; instead, it enhances its practical impact and potential for widespread applicability in real-world FL scenarios. A simple yet effective method, like ours, is a strength that enables broader adoption while providing substantial improvements over state-of-the-art techniques.

2. Pioneering LP-FT to FL:

To the best of our knowledge, this work is the first to systematically explore and leverage the potential of the simple yet effective LP-FT method within the FL framework. This represents a significant contribution, as it opens up a new line of research exploring the application and adaptation of LP-FT in addressing the unique challenges of FL, such as data heterogeneity and distribution shifts.

Theoretical Novelties (See Section 1 Theory Contribution),

1. Our theoretical setting and assumption differ significantly from that of [1]. While [1] assumes a centralized setting, we focus on personalized fine-tuning (PFT) in FL, a fundamentally different scenario. In our case, the pre-trained model is trained on data from all clients, reflecting a global client distribution. This contrasts with the centralized setting in [1], where the pre-trained model is trained without any notion of a global client distribution.

2. Our theoretical analysis diverges significantly from that of [1]. In [1], the assumption is that a single ground-truth data-generating function underlies all data points. This assumption makes it challenging to study LP-FT in different distribution shift settings in PFT for FL (e.g., concept shift). In contrast, we assume that each client has a distinct ground-truth function, enabling the exploration of personalized fine-tuning in FL under varying distribution shifts. Specifically, the way we define the data-generating function (Assumption 4.1), with unique linear heads before pre-training $V_i^* = \begin{bmatrix} {V_{com}^*}^T & \lambda e_i^T \end{bmatrix}^T$ for each client and the way we define data heterogeneity \Bigl( x_i = e_i + \epsilon n, \text{ with } \( n \sim \mathcal{N}(0, I) \Bigr) , allow us to study both concept shift and covariate shift using gradient descent to compare FT and LP-FT in terms of global performance. This approach provides new insights into the challenges of PFT in FL. In other words, the main challenge we address in our theoretical analysis lies in incorporating assumptions that align the theoretical framework with our specific setup. Our analysis connects the superior global performance of the LP-FT method to the concept of feature distortion. This novel perspective offers valuable insights into the effectiveness of LP-FT in FL through the lens of feature distortion.

3. In-depth Integrating Theoretical Analysis and Empirical Studies: We have conducted additional empirical studies to substantiate the theoretical analysis on the heterogeneity level $\lambda$. These include a comprehensive investigation with diverse ablation studies across varying levels of label-flipping ratios (Section 5). This significantly extends the scope and granularity of distribution shifts explored in the original LP-FT paper [1], providing deeper insights into its robustness and applicability under heterogeneous conditions.

[Continued below]

[Continued from above]

Response to W2

Empirical Contributions (See Section 1 Contribution Evaluation),

1. New Observations: Detailed Analysis on Overfitting in PFT As shown in Figure 2, we analyze commonly used fine-tuning strategies and demonstrate their tendency to overfit, even with carefully tuned hyperparameters, across three different types of distribution shifts. These findings highlight a critical gap in prior FL research, where such tendencies have not been systematically explored.

2. Comprehensive Benchmarking PFT Our experiments span seven diverse datasets (CIFAR-10-C, CIFAR-10, Digit5, DomainNet, Chexpert, CelebA, CIFAR-10), addressing four types of distribution shifts (input-level shift, feature shift, spurious correlation, and label shift, as detailed in Appendix Table 6). We evaluate five Personalized Fine-Tuning (PFT) methods using five robust metrics, providing a thorough and comprehensive analysis.

3. Establishing New Ablation Studies to Support Theoretical Results In Figure 4, we conduct detailed analysis on feature distortion relation to overfitting and in Table 2, we conduct large ranges of label flipping to support the theoretical analysis validating our theoretical insights in low heterogeneity regimes and further demonstrating LP-FT’s superior performance in scenarios with higher heterogeneity.

Reference

• [1] Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution. In ICLR. OpenReview.net, 2022.

Response to W2

We would like to clarify that the novelty and contribution of our work extend beyond simply adapting an existing method (LP-FT [1]) to a new scenario (PFT). Our key contributions lie in the methodological and theoretical novelty, as well as empirical findings, benchmarking, and analysis. These efforts provide deeper insights and fine-grained evaluations of diverse distribution shifts, demonstrating the effectiveness of LP-FT in the new PFT setting.

Methodological Novelties (See Section 1 Contribution Insight)

1. Simple yet Non-Trivial Method: Novel Focus on Final-Stage LP-FT in FL While our method may appear straightforward, its novelty lies in strategically shifting the insertion phase of LP-FT to the final local training stage in FL. Unlike most existing PFL methods that require modifications to the entire training process, including both General FL (GFL) training and local personalization, our approach focuses on the critical yet underexplored final local training stage. This makes our method a versatile solution, independent of the GFL training process. This shift not only introduces a fresh perspective on FL personalization but also demonstrates significant theoretical and empirical insights.

We also respectfully wish reviewers to agree that Simplicity does not diminish the novelty of our method; instead, it enhances its practical impact and potential for widespread applicability in real-world FL scenarios. A simple yet effective method, like ours, is a strength that enables broader adoption while providing substantial improvements over state-of-the-art techniques.

2. Pioneering LP-FT to FL As far as we know, this work is the first to systematically explore and leverage the potential of the simple yet effective LP-FT method within the FL framework. This represents a significant contribution, as it opens up a new line of research exploring the application and adaptation of LP-FT in addressing the unique challenges of FL, such as data heterogeneity and distribution shifts.

Theoretical Novelties (See Section 1 Contribution Theory),

1. Our theoretical setting and assumption differ significantly from that of [1]. While [1] assumes a centralized setting, we focus on personalized fine-tuning (PFT) in FL, a fundamentally different scenario. In our case, the pre-trained model is trained on data from all clients, reflecting a global client distribution. This contrasts with the centralized setting in [1], where the pre-trained model is trained without any notion of a global client distribution.

2. Our theoretical analysis diverges significantly from that of [1]. In [1], the assumption is that a single ground-truth data-generating function underlies all data points. This assumption makes it challenging to study LP-FT in different distribution shift settings in PFT for FL (e.g., concept shift). In contrast, we assume that each client has a distinct ground-truth function, enabling the exploration of personalized fine-tuning in FL under varying distribution shifts. Specifically, the way we define the data-generating function (Assumption 4.1), with unique linear heads before pre-training $V_i^* = \begin{bmatrix} {V_{com}^*}^T & \lambda e_i^T \end{bmatrix}^T$ for each client and the way we define data heterogeneity $\Bigl( x_i = e_i + \epsilon n , \text{ with } n \sim \mathcal{N}(0, I) \Bigr)$, allow us to study both concept shift and covariate shift using gradient descent to compare FT and LP-FT in terms of global performance. This approach provides new insights into the challenges of PFT in FL. In other words, the main challenge of our theoretical analysis lies in incorporating assumptions that align the theoretical setup with our specific framework. Our analysis connects the superior global performance of the LP-FT method to the concept of feature distortion. This novel perspective offers valuable insights into the effectiveness of LP-FT in FL through the lens of feature distortion.

3. In-depth Integrating Theoretical Analysis and Empirical Studies We have conducted additional empirical studies to substantiate the theoretical analysis on the heterogeneity level $\lambda$. These include a comprehensive investigation with diverse ablation studies across varying levels of label-flipping ratios (Section 5). This significantly extends the scope and granularity of distribution shifts explored in the original LP-FT paper [1], providing deeper insights into its robustness and applicability under heterogeneous conditions.

[Continued below]

Response to W1

1. Clarification: Our PFT setting is distinct from PFL methods. First, we would like to clarify our focus: it is essential to distinguish PFT from PFL, including methods like FedSelect mentioned by the reviewer. PFL focuses on comprehensive architectural or pipeline-level modifications to address data heterogeneity across the entire FL framework. In contrast, our focus PFT specifically targets the final localized fine-tuning stage, aiming to optimize the trade-off between local personalization and global generalization. These two approaches are complementary but distinct, addressing different aspects of Federated Learning. Our study centers on PFT as a modular, lightweight solution that complements broader PFL strategies without requiring significant system-level changes.

2. Clarification on the value of LP-FT research in FL We acknowledge the rapid advancements in Personalized Federated Learning (PFL) research. However, studying the effectiveness of LP-FT remains valuable because it offers a lightweight, modular solution that integrates seamlessly into existing FL frameworks without requiring architectural changes. Unlike complex PFL methods, LP-FT is a practical, plug-and-play approach adaptable to diverse FL applications. Investigating its effectiveness across various distribution shifts provides key insights and inspires future FL research, where LP-FT can serve as a foundational module for more advanced methods.

3. Clarification: Included advanced PFL baselines. We have included recent and advanced PFL baselines in Appendix Table 4 (updated to Appendix Table 5 in the revised version) to provide comprehensive comparisons. However, to maintain the focus on the primary contributions of this work, these baselines are not presented in the main text. Importantly, even these state-of-the-art PFL methods exhibit overfitting issues and relatively low global performance compared to LP-FT. This highlights LP-FT’s continued relevance and practical advantages, reinforcing its role as a reliable and effective strategy for addressing data heterogeneity in Federated Learning.

Thank you for your detailed response; it has resolved my confusion. I will maintain my positive rating.

Response to W3: Clarification on Computational Cost Concerns for LP-FT

We respectively disagree that LP-FT’s computation cost is overwhelming. We start by clarifying the PFT training pipeline. First, PFT has nothing to do with the GFL, thus LP-FT has the same computational cost compared to FedAvg for GFL. Second, our LP-FT is only integrated exclusively into the PFT stage.

Theoretically, compared to full fine-tuning (FT), LP-FT introduces an additional operation of linear probing (LP) with a computational complexity of $O(nmd)$ when applying gradient descent on a dataset of size $n$, using a linear head with input dimension $d$ and output dimension $m$. If the LP step is performed for $e$ epochs, the computational complexity becomes $O(e \cdot nmd)$. We have incorporated this theoretical analysis of the computational complexity of LP-FT compared to FT in the appendix of the revised paper.

Empirical Evaluation Empirically, we evaluate the FLOPs for one step (forward and backward for one batch) on Digit5 dataset. Specifically, we compare LP for 1 step with batch size of 128 and FT for 1 step with a batch size of 128. Since the computational cost of LP is only in the order of $O(10^{-3})$ of FT. LP is inherently more efficient in its operations. We measured the computational cost using FLOPs for one step (batch size of 128) on ResNet18 with the Digit5 dataset. Our analysis shows that FT requires $6.444 \times 10^{12}$ compared to $6.443 \times 10^{12}$ for LP.

Response to W4:Clarification on comprehensive experiments & Addressing misunderstanding on experimental settings

Clarification on the comprehensiveness of experiments:

We greatly appreciate the reviewer's feedback and would like to respectfully clarify that we believe our empirical studies are comprehensive, a perspective that has also been recognized by other reviewers. Our experiments cover seven diverse datasets (CIFAR-10-C, CIFAR-10, Digit5, DomainNet, Chexpert, CelebA, CIFAR10), four types of distribution shifts (input-level shift, feature shift, spurious correlation, label shift, as the label shift is detailed in Appendix Table 6), five Personalized Fine-Tuning (PFT) methods, and employ five evaluation metrics to provide a robust analysis. This comprehensiveness has been recognized by multiple reviewers:

• Reviewer 7VAu remarked that

  “Extensive experiments and theories were conducted.”

• Reviewer gFiS highlighted the

  “comprehensive experimental observation and theoretical analysis.”

• Reviewer sf8C noted the evaluation covers

  “multiple data heterogeneity scenarios.”

We believe these assessments strongly support the thoroughness of our empirical evaluations.

Regarding CIFAR-10-C:

CIFAR-10-C is a well-established benchmark in the research community for evaluating distribution shifts and federated learning settings. It comprises corruption patterns applied to CIFAR-10 to simulate real-world robustness challenges, making it a natural choice for our work. While it is unclear what is meant by “which CIFAR-10-C,” we have included detailed descriptions and references to CIFAR-10-C in Appendix Table 3 to assist readers who may be less familiar with this benchmark.

Clarification and Correction of Fundamental Factual Misunderstanding:

We respectfully disagree with the comments and clarify that we HAVE thoroughly addressed non-IID scenarios in FL. We emphasize that data heterogeneity and distribution shifts, inherent to non-IID settings, are central themes of our work. This is clearly highlighted throughout our paper:

• Abstract: Lines 16–19
• Section 1: Lines 75–79
• Section 1: Lines 97–99
• Section 1: Lines 103–106
• Section 2: Lines 110–125
• Section 3.1: Lines 162–176
• Section 3.2: Lines 181–192
• Section 3.3: Lines 214–235
• Section 3.4: Lines 315–317
• Section 4: Assumption 4.1 (Line 416)
• Section 4: Lines 477–492
• Section 4: Lines 500–521
• Section 6: Line 539

In summary. Our experiments encompass diverse Non-IID settings, including various distribution shift types and degrees of data heterogeneity. Throughout the original paper, we have clearly conveyed these Non-IID settings.

[Continued from above]

Response to W2 & Q2: Hilighting Contribution and Novelty

Empirical Contributions (See Section 1 Contribution Evaluation):,

1. New Observations: Detailed Analysis on Overfitting in PFT: As shown in Figure 2, we analyze commonly used fine-tuning strategies and demonstrate their tendency to overfit, even with carefully tuned hyperparameters, across three different types of distribution shifts. These findings highlight a critical gap in prior FL research, where such tendencies have not been systematically explored.

2. Comprehensive Benchmarking PFT: Our experiments span seven diverse datasets (CIFAR-10-C, CIFAR-10, Digit5, DomainNet, Chexpert, CelebA, CIFAR-10), addressing four types of distribution shifts (input-level shift, feature shift, spurious correlation, and label shift, as detailed in Appendix Table 6). We evaluate five Personalized Fine-Tuning (PFT) methods using five robust metrics, providing a thorough and comprehensive analysis.

3. Establishing New Ablation Studies to Support Theoretical Results: In Figure 4, we conduct detailed analysis on feature distortion relation to overfitting and in Table 2, we conduct large ranges of label flipping to support the theoretical analysis validating our theoretical insights in low heterogeneity regimes and further demonstrating LP-FT’s superior performance in scenarios with higher heterogeneity.

Reference

• [1] Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution. In ICLR. OpenReview.net, 2022.

Response to W2 & Q2: Highlighting Contribution and Novelty

We would like to clarify that the novelty and contribution of our work extend beyond simply adapting an existing method (LP-FT [1]) to a new scenario (PFT). Our key contributions lie in the methodological and theoretical novelty, as well as empirical findings, benchmarking, and analysis. These efforts provide deeper insights and fine-grained evaluations of diverse distribution shifts, demonstrating the effectiveness of LP-FT in the new PFT setting.

Methodological Novelties (See Section 1 Contribution Insight)

1. Simple yet Non-Trivial Method: Novel Focus on Final-Stage LP-FT in FL

While our method may appear straightforward, its novelty lies in strategically shifting the insertion phase of LP-FT to the final local training stage in FL. Unlike most existing PFL methods that require modifications to the entire training process, including both General FL (GFL) training and local personalization, our approach focuses on the critical yet underexplored final local training stage. This makes our method a versatile solution, independent of the GFL training process. This shift not only introduces a fresh perspective on FL personalization but also demonstrates significant theoretical and empirical insights.

We also respectfully wish reviewers to agree that Simplicity does not diminish the novelty of our method; instead, it enhances its practical impact and potential for widespread applicability in real-world FL scenarios. A simple yet effective method, like ours, is a strength that enables broader adoption while providing substantial improvements over state-of-the-art techniques.

2. Pioneering LP-FT to FL:

To the best of our knowledge, this work is the first to systematically explore and leverage the potential of the simple yet effective LP-FT method within the FL framework. This represents a significant contribution, as it opens up a new line of research exploring the application and adaptation of LP-FT in addressing the unique challenges of FL, such as data heterogeneity and distribution shifts.

Theoretical Novelties (See Section 1 Theory Contribution),

1. Our theoretical setting and assumption differ significantly from that of [1]. While [1] assumes a centralized setting, we focus on personalized fine-tuning (PFT) in FL, a fundamentally different scenario. In our case, the pre-trained model is trained on data from all clients, reflecting a global client distribution. This contrasts with the centralized setting in [1], where the pre-trained model is trained without any notion of a global client distribution.

2. Our theoretical analysis diverges significantly from that of [1]. In [1], the assumption is that a single ground-truth data-generating function underlies all data points. This assumption makes it challenging to study LP-FT in different distribution shift settings in PFT for FL (e.g., concept shift). In contrast, we assume that each client has a distinct ground-truth function, enabling the exploration of personalized fine-tuning in FL under varying distribution shifts. Specifically, the way we define the data-generating function (Assumption 4.1), with unique linear heads before pre-training $V_i^* = \begin{bmatrix} {V_{com}^*}^T & \lambda e_i^T \end{bmatrix}^T$ for each client and the way we define data heterogeneity $x_i = e_i + \epsilon n, \quad \text{with} \quad n \sim \mathcal{N}(0, I).$ allow us to study both concept shift and covariate shift using gradient descent to compare FT and LP-FT in terms of global performance. This approach provides new insights into the challenges of PFT in FL. In other words, the main challenge we address in our theoretical analysis lies in incorporating assumptions that align the theoretical framework with our specific setup. Our analysis connects the superior global performance of the LP-FT method to the concept of feature distortion. This novel perspective offers valuable insights into the effectiveness of LP-FT in FL through the lens of feature distortion.

3. In-depth Integrating Theoretical Analysis and Empirical Studies: We have conducted additional empirical studies to substantiate the theoretical analysis on the heterogeneity level $\lambda$. These include a comprehensive investigation with diverse ablation studies across varying levels of label-flipping ratios (Section 5). This significantly extends the scope and granularity of distribution shifts explored in the original LP-FT paper [1], providing deeper insights into its robustness and applicability under heterogeneous conditions.

[Continued below]

Response to W1.1 & Q1: Clarification on our focused PFT vs. reviewer-suggested PEFT

Thank you for your feedback. We would like to clarify that our paper primarily focuses on the personalized fine-tuning (PFT) rather than the parameter-efficient fine-tuning (PEFT) on FL.

PFT focuses on balancing the global and local risk during the final personalized fine-tuning stage of the FL (note that the fine-tuning here indicates different clients continue to train the global model on their individual local data). Such fine-tuning can be PEFT and non-PEFT. Given that our work is not centered on efficient fine-tuning foundation models, our primary focus is on effectively managing the trade-off between local and global risks.

Although our primary focus differs, we address the reviewer’s request by showing results for PEFT. We adopted two widely recognized PEFT methods: LoRA [1] an Adapter [2]. Both methods were fine-tuned with meticulous adjustments to their learning rates on the ViT, as detailed in Table below. It is worth noting that the PEFT methods are not specifically designed for PFT. As expected, their performance significantly underperforms our proposed LP-FL.

Table: Different PEFT compared on DomainNet with ViT.

Reference

• [1] Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models." arXiv preprint arXiv:2106.09685 (2021).
• [2] Houlsby, Neil, et al. "Parameter-efficient transfer learning for NLP." International conference on machine learning. PMLR, 2019.

Response to W1.2: Clarification on our focused PFT vs. other FL methods that consider distribution shift.

1. Clarification on our PFT’s focus and setting Although our method also considers the distribution shift problem, our focus is fundamentally different from the reviewer's pointed FL methods. Unlike most existing FL methods on distribution shift that address the issue through the whole FL training process, our approach focuses on the critical yet underexplored final local training stage after the global model is obtained. This makes our method a versatile solution, independent of the global model training process. This shift not only introduces a fresh perspective on FL personalization but also demonstrates significant theoretical and empirical insights.

Specifically, FedIIR is a GFL method that modifies the entire FL training pipeline across all stages, whereas our PFT setting focuses exclusively on the final personalized fine-tuning stage, operating in parallel to GFL methods. FedTHE addresses a completely different scenario, focusing on test-time adaptation to improve personalization. Our PFT setting, however, emphasizes balancing global and local performance during the final personalized fine-tuning stage.

**2. Discussion on related work ** Although the focus and settings differ, we have not omitted related works on FL for distribution shift, as evidenced by our extensive discussion on heterogeneous FL in the first paragraph of Section 2 (Related Work).

Response to W5: Results Explanation and Takeaway

We respectfully disagree that our experiments are not conclusive. We would like to start by clarifying our evaluation criteria, and then highlight the takeaways from our results.

Clarification of Evaluation Criteria:

LP-FT's superiority over baseline methods is primarily evaluated based on Global Accuracy, the Worst Accuracy, and Avg. Accuracy, as these align with our focus on feature distortion effects. We emphasize that LP-FT consistently outperforms the baselines across most scenarios under these criteria.

Highlighting Key Improvements on LP-FTs

To provide more concrete evidence of LP-FT's efficacy, we point out the following improvements: In CIFAR10-C (input shift), LP-FT improves upon the best baseline (Proximal FT) by 1.77% in Global Accuracy, 0.51% in Worst Accuracy, and 1.36% in Average over local and global accuracy. In DomainNet (feature shift), LP-FT achieves an improvement of 1.47% in Global Accuracy, 1.33% in Worst Accuracy, and 0.92% in Average over local and global accuracy over the strongest baseline. Across CelebA (spurious correlation), LP-FT demonstrates 2.21% gains in Global Accuracy, 4.39% in Worst Accuracy, and 2.04% in Average over local and global accuracy.

Justifications for Limited Cases Where LP-FT Did Not Outperform:

For the cases where LP-FT does not surpass the baselines, we have provided detailed explanations in Section 3.4 results analysis of the paper. LP-FT already surpasses 90% of the baselines. LP-FT's design prioritizes preventing feature distortion, which is deeply discussed in the theoretical analysis. This aligns with the overarching objectives of this study. While certain trade-offs (e.g., performance in niche scenarios) are observed, they are an intentional choice to achieve better overall robustness and adaptability across diverse datasets and clients.

Response to W6: Our Theoretical Analysis Exactly Addresses Data Heterogeneity

There appears to be a misunderstanding that the reviewer thought ``.theoretical analysis overlooks the data is heterogeneous”. We would like to clarify that our theoretical analysis IS CENTERED ON data heterogeneity. As stated in Line 466-469 of our original submission (Line 478–481 of our revised submission), we introduce covariate shift with the explicit assumption that data is heterogeneous.

Our theoretical analysis incorporates data heterogeneity based on the specific setting under consideration, whether it involves concept shift or a combination of concept and covariate shifts. In Section 4.1, we study the global performance of LP-FT under concept shift alone, assuming data is not heterogeneous (i.e., $\mathbb{E}_{x \sim \mathcal{D}_i}[x x^T] = I_d$ for all clients $i \in [C]$). In Section 4.2, however, we extend our analysis to encompass both concept and covariate shifts, incorporating data heterogeneity by allowing each client's data to be generated as $x_i = e_i + \epsilon n$, where $n \sim \mathcal{N}(0, I)$ and $e_i$ represents a client-specific shift. Therefore, our theoretical analysis considers or excludes data heterogeneity in accordance with each section’s focus.

Response to Q3

Thanks for the suggestions on exploring text data. In this work, we follow the common stream of heterogeneous FL studies and closely related work on vision tasks. It would be an interesting future extension work.

Our focus on vision tasks is consistent with the original LP-FT paper [1] and its extended analysis [2], both of which are centered on vision data. By following this established scope, we aim to build upon and extend the insights of prior research within the same modality.

Reference:

• [1] Kumar, A., Raghunathan, A., Jones, R., Ma, T., & Liang, P. (2022). Fine-tuning can distort pretrained features and underperform out-of-distribution. arXiv preprint arXiv:2202.10054.
• [2] Trivedi, P., Koutra, D., & Thiagarajan, J. J. (2023). A closer look at model adaptation using feature distortion and simplicity bias. arXiv preprint arXiv:2303.13500.

Response to W3

Thank you for the insightful suggestion. We have revised the section structure and naming to align more clearly with observations, methods, and experiments, improving readability and coherence. We believe this refinement enhances the clarity and flow of the paper.

Response to W4

1. Clarification: Our PFT setting is distinct from PFL methods. irst, we would like to clarify our focus: it is essential to distinguish PFT from PFL. PFL focuses on comprehensive architectural or pipeline-level modifications to address data heterogeneity across the entire FL framework. In contrast, our focus PFT specifically targets the final localized fine-tuning stage, aiming to optimize the trade-off between local personalization and global generalization. These two approaches are complementary but distinct, addressing different aspects of Federated Learning. Our study centers on PFT as a modular, lightweight solution that complements broader PFL strategies without requiring significant system-level changes.

2. Clarification on the value of LP-FT research in FL We acknowledge the rapid advancements in Personalized Federated Learning (PFL) research. However, studying the effectiveness of LP-FT remains valuable because it offers a lightweight, modular solution that integrates seamlessly into existing FL frameworks without requiring architectural changes. Unlike complex PFL methods, LP-FT is a practical, plug-and-play approach adaptable to diverse FL applications. Investigating its effectiveness across various distribution shifts provides key insights and inspires future FL research, where LP-FT can serve as a foundational module for more advanced methods.

3. Clarification: Included advanced PFL baselines. We have included recent and advanced PFL baselines in Appendix Table 4 (updated to Appendix Table 5 in the revised version) to provide comprehensive comparisons. However, to maintain the focus on the primary contributions of this work, these baselines are not presented in the main text. Importantly, even these state-of-the-art PFL methods exhibit overfitting issues and relatively low global performance compared to LP-FT. This highlights LP-FT’s continued relevance and practical advantages, reinforcing its role as a reliable and effective strategy for addressing data heterogeneity in Federated Learning.

[Continued from above]

Response to W2

Empirical Contributions (See Section 1 Contribution Evaluation),

1. New Observations: Detailed Analysis on Overfitting in PFT As shown in Figure 2, we analyze commonly used fine-tuning strategies and demonstrate their tendency to overfit, even with carefully tuned hyperparameters, across three different types of distribution shifts. These findings highlight a critical gap in prior FL research, where such tendencies have not been systematically explored.

2. Comprehensive Benchmarking PFT Our experiments span seven diverse datasets (CIFAR-10-C, CIFAR-10, Digit5, DomainNet, Chexpert, CelebA, CIFAR-10), addressing four types of distribution shifts (input-level shift, feature shift, spurious correlation, and label shift, as detailed in Appendix Table 6). We evaluate five Personalized Fine-Tuning (PFT) methods using five robust metrics, providing a thorough and comprehensive analysis.

3. Establishing New Ablation Studies to Support Theoretical Results In Figure 4, we conduct detailed analysis on feature distortion relation to overfitting and in Table 2, we conduct large ranges of label flipping to support the theoretical analysis validating our theoretical insights in low heterogeneity regimes and further demonstrating LP-FT’s superior performance in scenarios with higher heterogeneity.

Reference

• [1] Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution. In ICLR. OpenReview.net, 2022.

Response to W2

We would like to clarify that the novelty and contribution of our work extend beyond simply adapting an existing method (LP-FT [1]) to a new scenario (PFT). Our key contributions lie in the methodological and theoretical novelty, as well as empirical findings, benchmarking, and analysis. These efforts provide deeper insights and fine-grained evaluations of diverse distribution shifts, demonstrating the effectiveness of LP-FT in the new PFT setting.

Methodological Novelties (See Section 1 Contribution Insight) 1. Simple yet Non-Trivial Method: Novel Focus on Final-Stage LP-FT in FL While our method may appear straightforward, its novelty lies in strategically shifting the insertion phase of LP-FT to the final local training stage in FL. Unlike most existing PFL methods that require modifications to the entire training process, including both General FL (GFL) training and local personalization, our approach focuses on the critical yet underexplored final local training stage. This makes our method a versatile solution, independent of the GFL training process. This shift not only introduces a fresh perspective on FL personalization but also demonstrates significant theoretical and empirical insights.

We also respectfully wish reviewers to agree that Simplicity does not diminish the novelty of our method; instead, it enhances its practical impact and potential for widespread applicability in real-world FL scenarios. A simple yet effective method, like ours, is a strength that enables broader adoption while providing substantial improvements over state-of-the-art techniques.

2. Pioneering LP-FT to FL As far as we know, this work is the first to systematically explore and leverage the potential of the simple yet effective LP-FT method within the FL framework. This represents a significant contribution, as it opens up a new line of research exploring the application and adaptation of LP-FT in addressing the unique challenges of FL, such as data heterogeneity and distribution shifts.

Theoretical Novelties (See Section 1 Contribution Theory),

1. Our theoretical setting and assumption differ significantly from that of [1]. While [1] assumes a centralized setting, we focus on personalized fine-tuning (PFT) in FL, a fundamentally different scenario. In our case, the pre-trained model is trained on data from all clients, reflecting a global client distribution. This contrasts with the centralized setting in [1], where the pre-trained model is trained without any notion of a global client distribution.

2. Our theoretical analysis diverges significantly from that of [1]. In [1], the assumption is that a single ground-truth data-generating function underlies all data points. This assumption makes it challenging to study LP-FT in different distribution shift settings in PFT for FL (e.g., concept shift). In contrast, we assume that each client has a distinct ground-truth function, enabling the exploration of personalized fine-tuning in FL under varying distribution shifts. Specifically, the way we define the data-generating function (Assumption 4.1), with unique linear heads before pre-training $V_i^* = \begin{bmatrix} {V_{com}^*}^T & \lambda e_i^T \end{bmatrix}^T$ for each client and the way we define data heterogeneity $\Bigl( x_i = e_i + \epsilon n, \text{ with } n \sim \mathcal{N}(0, I) \Bigr)$, allow us to study both concept shift and covariate shift using gradient descent to compare FT and LP-FT in terms of global performance. This approach provides new insights into the challenges of PFT in FL. In other words, the main challenge of our theoretical analysis lies in incorporating assumptions that align the theoretical setup with our specific framework. Our analysis connects the superior global performance of the LP-FT method to the concept of feature distortion. This novel perspective offers valuable insights into the effectiveness of LP-FT in FL through the lens of feature distortion.

3. In-depth Integrating Theoretical Analysis and Empirical Studies We have conducted additional empirical studies to substantiate the theoretical analysis on the heterogeneity level $\lambda$. These include a comprehensive investigation with diverse ablation studies across varying levels of label-flipping ratios (Section 5). This significantly extends the scope and granularity of distribution shifts explored in the original LP-FT paper [1], providing deeper insights into its robustness and applicability under heterogeneous conditions.

[Continued below]

Response to W1

We would like to provide additional context and clarify our focus:

1. Importance of Feature Distribution Shifts Feature distribution shifts are critical in many real-world scenarios, such as FL applications in healthcare and autonomous systems, where input features (e.g., sensor readings or image characteristics) vary significantly across clients due to differences in equipment, environments, or acquisition conditions (added in Section 1 Line 77-78 in revised version). Addressing feature shifts is essential for robust personalization and generalization in such settings.

2. Limitations of PFL Methods for Label Distribution Skew While many PFL methods have been developed for label distribution skew, these approaches do not necessarily generalize to feature distribution shift scenarios. As highlighted in our revised related work [1, 2], some methods tailored for label skew fail to effectively mitigate the challenges posed by feature skew, such as maintaining global generalization or handling federated feature distortion.

3. Focus on Theoretical Grounding for Feature Shift Our theoretical analysis is grounded in feature shift assumptions to explain the phenomena studied and the effectiveness of our method. These assumptions include specific conditions on the data distribution of different clients (e.g., $\mathbb{E}_{x \sim \mathcal{D}_i}[x x^T] = I_d$ or $x_i = e_i + \epsilon n, \text{ where } n \sim \mathcal{N}(0, I)$) as well as assumptions on each client’s data-generating function (Assumption 4.1). These assumptions are considered in both concept shift and covariate shift scenarios. This framework enables us to analyze the global performance of LP-FT and FT through the lens of gradient descent under these distribution shifts. Overall, our theoretical analysis, recognized by the Reviewer 7VAu as a strength of our work, is particularly built upon feature shift assumptions. We chose to focus on feature shifts to ensure consistency with our theoretical framework and to maintain clarity for readers, thereby avoiding potential confusion.

4. Results on Label Shifts Nevertheless, we acknowledge the importance of label distribution shifts. To address your concern, we conducted additional experiments on label shifts and found that our method performs robustly in these settings as well, as seen in Table 6 in the Appendix.

Reference:

• [1] Guo, Yongxin, Xiaoying Tang, and Tao Lin. "FedRC: Tackling Diverse Distribution Shifts Challenge in Federated Learning by Robust Clustering." arXiv preprint arXiv:2301.12379 (2023).
• [2] Son, Ha Min, et al. "FedUV: Uniformity and Variance for Heterogeneous Federated Learning." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.