Dual-Personalizing Adapter for Federated Foundation Models

W1: Cannot clearly see the significance of considering personalization and test-time adaptation at the same time.

Our proposed method is not doing test-time adaptation. In general, the test-time adaptation methods [1] usually involve fine-tuning steps to align the model parameters with new distributions in test time. In contrast, our proposed method is designed without any adaptation step, aiming instead for test-time generalization where clients must maintain performance on their main tasks while also demonstrating generalization capabilities on new tasks during the testing stage.

Moreover, our method addresses an innovative application scenario known as federated foundation models (FFM) [2], an emerging domain that significantly diverges from traditional federated learning methods. We are the first to explore the test-time generalization on personalized FFM. Unlike the works you mentioned, we refined a new setting (test-time personalization) in FFM, proposed a new training paradigm with self-reconstructed datasets to train personalized FFM in extreme scenes like cross-tasks, and also introduced a new dynamic weighting mechanism to determine the combination of different adapters for test-time personalization. We will add a discussion about the difference between our work and your mentioned works in related work.

[1] Liang, J., He, R., & Tan, T. (2024). A comprehensive survey on test-time adaptation under distribution shifts. International Journal of Computer Vision, 1-34.

[2] Ren, C., Yu, H., Peng, H., Tang, X., Li, A., Gao, Y., ... & Yang, Q. (2024). Advances and open challenges in federated learning with foundation models. arXiv preprint arXiv:2404.15381.

W2: Propose two methods and can not find one solution that fits both two metrics.

Our paper aims to improve the test-time generalization without sacrificing the performance on personalization metric in federated settings. In traditional machine learning methods, an algorithm should be evaluated on both validation dataset and test dataset, which in our context correspond to the personalization metric and test-time generalization metric, respectively.

Our proposed framework (ref to Section 4), model architecture (ref to Figure 1) and loss function (ref to Equation 1) are unified design. The mentioned two methods can be considered as two hyperparameters to control the updating strategy (sequentially or iteratively) of local and global adapters within the federated learning process. Specifically, FedDPA-F adopts a sequential approach, first optimizing the global adapter and then optimizing the local adapter. In contrast, FedDPA-T alternates between optimizing the global and local adapters iteratively during each communication round. To facilitate understanding, Figure 2 illustrates the two options of the hyperparameter in the learning process. In the distributed learning scenario, it is a common way to try different updating strategies in federated learning, such as [3].

[3] Tian Li, Shengyuan Hu, AhmadBeirami, and Virginia Smith, “Ditto: Fair and Robust Federated Learning Through Personalization”, ICML 2021

W3: Limited performance gain.

The objective of this paper is to improve the test-time generalization of federated foundation models without sacrificing personalization performance. In two main tables (Table 1 & 2), in the test-time generalization metric, our proposed method shows significant improvement over FedLoRA. For example, both FedDPA-F and FedDPA-T demonstrate average improvements of approximately 2.3% and 0.9% over FedLoRA. Notably, in the summarization task, our approaches achieve an 8.4% higher score than FedLoRA.

Similarly, in the personalization metric, our method achieved relatively higher performance than FedLoRA. For example, both FedDPA-F and FedDPA-T demonstrate average improvements of approximately 0.5% and 2.3% over FedLoRA. Notably, in the reading comprehension task, our methods outperform FedLoRA by 5.9% in terms of personalization.

Moreover, we will further explore one more analysis of the statistical significance. Specifically, it will involve assuming that our method can improve test-time generalization performance and subsequently employing statistical tests, such as the P-value approach, to verify the significance and validate the efficacy and confidence of our method.

Q1: Difference between FedIT and FedLoRA.

Compared to personalized federated learning methods, FedIT focuses on training a global model applicable across all clients, without incorporating personalization techniques. In each communication round of the convergence curve, FedIT evaluates the global model on test data. Conversely, FedLoRA, a personalized federated learning method, evaluates performance using fine-tuned models in each communication round. Therefore, FedIT shows a relatively lower accuracy in the early stage of the learning process. For a more detailed discussion, please refer to Appendix A.2.

W1: The paper would benefit for a much stronger motivation.

Unlike traditional FL (especially personalized FL) which primarily addresses heterogeneity among clients during training, our setting also considers heterogeneity within each client during testing, building on insights from previous works [1]. Since the learning process of foundation models will involve more heterogeneous datasets from different tasks compared with the traditional machine learning methods, our setting aims to consider more complex scenarios with various distribution shifts in federated foundation models. For example, in traditional personalized FL, each client's training and testing are restricted to the same distribution (e.g., client 1 trains and tests on task 1); in contrast, our setting involves clients training on one task and testing on multiple different tasks (e.g., client 1 trains on task 1 and tests on tasks 1, 2, 3...).

Moreover, imagine the on-device foundation model in the near future, the model should be lightweight and versatile enough to tackle many different tasks. To this motivation, this paper is the first step towards exploring a solution that can enable the lightweight on-device foundation models with desired performance on training data and also gain a better test-time generalization while the test data is from another domain or task.

[1] Tan, Y., Chen, C., Zhuang, W., Dong, X., Lyu, L., & Long, G. (2024). Is heterogeneity notorious? taming heterogeneity to handle test-time shift in federated learning. Advances in Neural Information Processing Systems, 36.

W2: Novelty of our methods.

As discussed in the recent literature on Federated Foundation Models (FFM) [2], there are many new challenges that need to be rethought. For example, training versatile foundation models requires incorporating more heterogeneous datasets within the federated learning framework. Moreover, the model communication-efficient and fine-tuning strategy also needs to be considered in the federated foundation models. Before the research community can go further to explore more exciting applications of federated foundation models, some base work must be explored. Our work is to rethink the problems of traditional FL (especially PFL) based on FFM and establish the benchmark of personalized FFM.

We are the first to explore the test-time generalization of personalized FFM. The overall design is new with refined setting (test-time personalization) in FFM and the technique framework is composed of multiple pieces including a new training paradigm with self-reconstructed datasets to train personalized FFM, and also a new dynamic weighting mechanism to determine the combination of different adapters for test-time personalization.

[2] Ren, C., Yu, H., Peng, H., Tang, X., Li, A., Gao, Y., ... & Yang, Q. (2024). Advances and open challenges in federated learning with foundation models. arXiv preprint arXiv:2404.15381.

W3: Different distributions within a device and how to detect shifts.

In federated foundation models, each client’s device will be deployed with a lightweight intelligent assistant to tackle a broad range of downstream tasks, thus it is very likely that the client needs to tackle unseen tasks with different distributions during testing.

In this paper, we make new assumptions for FFM: 1) for each device, the test distribution may differ from the train distribution, 2) different clients have different distributions as well, and 3) the test distribution of client A could be observed by another client in their training tasks. For example, in practice, a client accustomed to writing emails in English may require translation assistance when working on a new project in Chinese. Here, the client has ample historical data for training in English email writing, but none for Chinese translation, which only emerges during the testing/inference phase. Additionally, other clients may be specialized in translation tasks.

To detect the distribution shifts during the testing phase, we designed an Instance-wise Dynamic Weighting mechanism. Specifically, we measure the similarity between the representations of test and training instances to adjust the importance weight of local adapter and global adapter. More detailed discussion could be found in Section 4.3.

W4: Extreme evaluation setting.

The problem setting and application scenario in federated foundation models have been different from the traditional federated learning methods. Our research is focused on the challenges of FFM, wherein foundation models, pre-trained on extensive datasets, already exhibit a degree of generalization capability towards general non-IID problems —a topic thoroughly investigated within traditional FL. Consequently, in the realm of FFM, we aim to address more complex scenarios, such as cross-domain and cross-dataset challenges, where centralized foundation models typically underperform [3].

In our future work, we could add more experiments to include traditional federated learning settings by splitting a dataset into many pieces with a slight non-IID. We believe our proposed method can be easily adapted to this type of non-IID.

[3] Wang, Y., Ivison, H., Dasigi, P., Hessel, J., Khot, T., Chandu, K., ... & Hajishirzi, H. (2023). How far can camels go? exploring the state of instruction tuning on open resources. Advances in Neural Information Processing Systems, 36, 74764-74786.

W2 and W3: Design of trade-off mechanism and large-scale experiment.

Thanks for your insightful advice. We will further explore better mechanisms from the perspective of generalization theory and more datasets in real applications.

Q1: Choice of similarity function.

We select cosine similarity due to its better robustness and normalization with high-dimensional vectors than other metrics, and its superiority has been demonstrated in many NLP/CV works. Additionally, we conducted an ablation study comparing various similarity functions in Appendix B.2, and the results in Table 9 demonstrate cosine similarity outperforms other similarity functions.

W1 and Q2: Theoretical analysis of the design of alpha.

In section 3.2, we conducted a preliminary theoretical analysis of the discordance between personalization and test-time tasks, which motivated us to learn two adapters: a local adapter for personalization task (training distribution) and a global adapter for test-time tasks. Therefore, when a test instance is more similar to the client’s training instance, we intuitively infer that its distribution is more similar to the training data, and because the local adapter aimed for personalization contains more knowledge of training data distribution, we increase the weight of local adapter. We will further explore the theoretical relation between them.

Q3: Why FedDPA-F wins more time than FedDPA-T.

It is because that FedDPA-F maintains a more generalized local adapter than FedDPA-T. The difference between FedDPA-F and FedDPA-T lies in the local adapter: FedDPA-F utilizes the global adapter as initialization of the local adapter, while FedDPA-T randomly initializes the local adapter. Given that the global adapter aggregated on different distributions matin certain generalization capabilities, the local adapter of FedDPA-F has better generalization performance than that of FedDPA-T, which leads to better performance on most test-time tasks. We will add this discussion in our experiment analysis to clarify.

Q4: Other data sources.

Our method is designed with a high degree of flexibility for its adaption across various data sources. As illustrated in Appendix A.1, any public text data sources (e.g., Dolly), uniformed into the generative task with appropriate prompts like Table 6, can be used in this setting. In addition, other types of data can also be applied to our methods by replacing LLM with the corresponding foundation models. For example, for image datasets (e.g., ImageNet, MS COCO), ViT can be used as the foundation model in our methods for this setting. We will add a discussion about the adaptation to other data sources.

W1: Discussion of assumption.

Yes, our main experiments are based on the assumption that all possible distributions are included in all clients. We already discussed this in the first paragraph of section 4 and Appendix C. To enhance clarity, we will rewrite this part and emphasize keywords in bold to make it clearer.

W2 and Q1: Heavily rely on the LoRA and adapt to other data types.

Our method is designed with a high degree of flexibility, facilitating its adaption across various adapter-based PEFT methods and transformer-based foundation models of different data types. In this paper, we just use LLM and LoRA as examples to illustrate our method, and our method can easily be adapted to other frameworks by substituting the LoRA and LLM with alternative adapter-based PEFT methods and transformer-based foundation models. We will revise our paper by including a discussion about adaptation to other adapter-based PEFT methods (e.g., series adapter) and other foundation models (e.g., ViT for images and UNITER for multimodality data).

W3: Different with multi-task learning.

Multi-task learning is centralized, and tuning on the combination of multi-task data of a centralized foundation model can be taken as the multi-task learning baseline. As in the centralized foundation model, all tasks are standardized into a uniform format (refer to Appendix A.1 and Table 6), and the model already acquires task-agnostic token embedding by the extensive data pre-training; directly tuning on these multi-task data is the implementation of multi-task learning with foundation models [1]. We have included this baseline, referred to as "Centralized," in our experimental comparisons. In addition, our setting accounts for test-time distribution shifts, a factor not typically addressed in multi-task learning. Although our main experiments are under the assumption that all possible distributions are included in all clients, our setting can also allow for scenarios involving unseen distributions for all clients, and our methods have demonstrated robustness in these scenarios, as detailed in Appendix C.1. We will add a discussion in related work to clarify the difference with multi-task learning.

[1] Yu, J., Dai, Y., Liu, X., Huang, J., Shen, Y., Zhang, K., ... & Chen, Y. (2024). Unleashing the Power of Multi-Task Learning: A Comprehensive Survey Spanning Traditional, Deep, and Pretrained Foundation Model Eras. arXiv preprint arXiv:2404.18961.

Q2: Merge two federated datasets into one for 16 clients.

Yes, we can merge these two datasets into one for more clients. As explored in section 6.2 of the ablation study on the impact of client number, our methods could maintain performance when scaling up the clients (up to 40 clients) in Figure 4. Therefore, our methods are supposed to be effective when having more clients and more tasks.

Q3: How to fix the value of alpha.

In line 315, we aim to investigate the impact of the global adapter on the training of the local adapter in FedDPA-T. As illustrated on the left of Fig 2(b), during the local adapter training of FedDPA-T, we employ the frozen global adapter to expedite the learning of the local adapter. This is because the global adapter contains certain task-related knowledge to accelerate the local adapter’s learning. It is worth noting that during the training phase, there are no distribution shifts; such shifts occur solely in the inference/testing phase, and the dynamic weighting mechanism is only used in the inference phase to determine the value of alpha for addressing test-time distribution shifts. In contrast, in the training phase, alpha is only used to control the contribution of the frozen global adapter to the local adapter learning. Therefore, during the training of FedDPA-T, we can fix the value of alpha, as its function differs from that during inference. To enhance clarity, we will revise the symbol and use two distinct symbols to differentiate these uses.

Q4: Theoretical analysis and more datasets.

For theoretical analysis, we believe that further work could rethink this problem within the context of existing theoretical frameworks and establish a new theoretical framework of FFM. For example, we can start from the perspective of out-of-domain generalization theory to rethink the generalization analysis and establish a new generalization and convergence analysis framework of FFM.

One main challenge in applying to more datasets is the significant computational resources required. Given that foundation models often consist of billions of parameters, they necessitate considerable time and storage for tuning. Additionally, the computational demand varies significantly across data types; for example, a thousand text data in the Flan occupy only about 1MB, whereas the same number of images in ImageNet can require around 130MB. Therefore, we believe that future work should explore more efficient methods for tuning FFM on larger datasets.