Federated Adapter on Foundation Models: An Out-Of-Distribution Approach
Yiyuan Yang, Guodong Long, Tianyi Zhou, Qinghua Lu, Shanshan Ye, Jing Jiang
TL;DR
The paper addresses OOD generalization in Federated Foundation Models (FedFM), where large parameter scales and adapter-based PEFT create unique heterogeneity and data-shift challenges. It theoretically shows that the conventional aggregated global FedFM model inherently possesses OOD generalization, while personalized adapters can be guided to improve OOD performance via a feature distance-based regularization that aligns personalized and global representations. The proposed FedOA framework employs adapter-based PEFT with a personalized adapter per client and a global adapter, optimized with a distance-based regularizer to yield invariant feature learning and provable convergence in general non-convex settings. Empirically, FedOA attains state-of-the-art OOD generalization on heterogeneous NLP tasks with leave-one-task-out evaluations, demonstrating robustness to cross-domain shifts and scalability across more clients and tasks. This work provides a principled, scalable approach to OOD generalization in FedFM and sets a foundation for future PEFT-guided federated learning in diverse downstream tasks.
Abstract
As foundation models gain prominence, Federated Foundation Models (FedFM) have emerged as a privacy-preserving approach to collaboratively fine-tune models in federated learning (FL) frameworks using distributed datasets across clients. A key challenge for FedFM, given the versatile nature of foundation models, is addressing out-of-distribution (OOD) generalization, where unseen tasks or clients may exhibit distribution shifts leading to suboptimal performance. Although numerous studies have explored OOD generalization in conventional FL, these methods are inadequate for FedFM due to the challenges posed by large parameter scales and increased data heterogeneity. To address these, we propose FedOA, which employs adapter-based parameter-efficient fine-tuning methods for efficacy and introduces personalized adapters with feature distance-based regularization to align distributions and guarantee OOD generalization for each client. Theoretically, we demonstrate that the conventional aggregated global model in FedFM inherently retains OOD generalization capabilities, and our proposed method enhances the personalized model's OOD generalization through regularization informed by the global model, with proven convergence under general non-convex settings. Empirically, the effectiveness of the proposed method is validated on benchmark datasets across various NLP tasks.