FedRA: A Random Allocation Strategy for Federated Tuning to Unleash the Power of Heterogeneous Clients

Shangchao Su,Bin Li,Xiangyang Xue
2024-03-12
Abstract:With the increasing availability of Foundation Models, federated tuning has garnered attention in the field of federated learning, utilizing data and computation resources from multiple clients to collaboratively fine-tune foundation models. However, in real-world federated scenarios, there often exist a multitude of heterogeneous clients with varying computation and communication resources, rendering them incapable of supporting the entire model fine-tuning process. In response to this challenge, we propose a novel federated tuning algorithm, FedRA. The implementation of FedRA is straightforward and can be seamlessly integrated into any transformer-based model without the need for further modification to the original model. Specifically, in each communication round, FedRA randomly generates an allocation matrix. For resource-constrained clients, it reorganizes a small number of layers from the original model based on the allocation matrix and fine-tunes using adapters. Subsequently, the server aggregates the updated adapter parameters from the clients according to the current allocation matrix into the corresponding layers of the original model. It is worth noting that FedRA also supports scenarios where none of the clients can support the entire global model, which is an impressive advantage. We conduct experiments on two large-scale image datasets, DomainNet and NICO++, under various non-iid settings. The results demonstrate that FedRA outperforms the compared methods significantly. The source code is available at \url{<a class="link-external link-https" href="https://github.com/leondada/FedRA" rel="external noopener nofollow">this https URL</a>}.
Machine Learning,Artificial Intelligence,Distributed, Parallel, and Cluster Computing
What problem does this paper attempt to address?
### Problems Addressed by the Paper The paper aims to address the issue of efficiently performing federated tuning of pre-trained foundation models in federated learning, especially in the presence of a large number of heterogeneous clients. Specifically, the paper focuses on how to effectively fine-tune models on resource-constrained clients to fully utilize their computational and data resources while avoiding the problem of insufficient resources to support the entire model fine-tuning process. ### Background and Challenges 1. **Limitations of Federated Learning**: - Traditional federated learning methods usually require training the entire global model parameters, which is not feasible on resource-constrained clients. - The parameter size of pre-trained foundation models (such as ViT, CLIP, GPT, etc.) is very large, and many resource-constrained devices cannot support complete model fine-tuning. 2. **Shortcomings of Existing Methods**: - **Width Pruning Methods**: These methods allocate smaller models to resource-constrained clients by pruning model channels, but this approach disrupts the layer structure of pre-trained models, making it difficult to apply directly to federated tuning. - **Depth Pruning Methods**: These methods allocate only the initial layers of the foundation model to resource-constrained clients. While preserving the integrity of model layers, this leads to a feature imbalance problem, where only a few resource-rich clients can train the higher layers of the model, thereby affecting overall performance. ### Proposed Method To address the above challenges, the paper proposes a new federated tuning algorithm—FedRA (Federated Random Allocation). The main features of FedRA are as follows: 1. **Random Allocation Strategy**: - In each communication round, the server randomly generates an allocation matrix, which is used to allocate sub-layers of the global model to different clients. - Each client receives sub-layers that include trainable adapter parameters and initially frozen model parameters to construct a new local model. 2. **Local Fine-Tuning**: - Clients perform local fine-tuning using the allocated sub-layers, adjusting the adapter parameters and classification head. - The fine-tuned adapter parameters are sent back to the server for aggregation. 3. **Aggregation Strategy**: - The server aggregates the adapter parameters collected from clients into the corresponding layers of the global model based on the current allocation matrix. - Random allocation ensures that each layer of the global model can learn information from all clients, thereby solving the feature imbalance problem. ### Experimental Results The paper conducts experiments on two large-scale image datasets (DomainNet and NICO++) to validate the performance of FedRA under different non-I.I.D. settings. The experimental results show that: - **Overall Performance Improvement**: FedRA significantly outperforms other comparison methods on the DomainNet and NICO++ datasets, especially in scenarios with clients of varying model sizes. - **Extreme Heterogeneous Scenarios**: Even when no client can support the complete global model, FedRA can still work effectively. - **Convergence Analysis**: Both experimental and theoretical analyses show that FedRA can gradually converge after multiple training rounds, and due to the random allocation strategy, its performance is superior to existing depth pruning methods. ### Main Contributions 1. **Introduction of Heterogeneous Client Federated Tuning Task (HeFT)**: HeFT requires federated tuning algorithms to adapt to the computational capacity differences of heterogeneous clients while maximizing the utilization of pre-trained knowledge. 2. **Proposing the FedRA Algorithm**: By using a random allocation strategy, FedRA addresses the feature imbalance problem and can work effectively even when all clients lack the computational capacity to train the complete global model. 3. **Extensive Experimental Validation**: The paper conducts experiments on multiple large-scale datasets and model architectures, demonstrating the superior performance of FedRA. In summary, the paper effectively addresses the challenges of federated tuning on heterogeneous clients by proposing the FedRA algorithm, providing new insights and methods for the application of federated learning in real-world scenarios.