HeteroTune: Efficient Federated Learning for Large Heterogeneous Models
This addresses the problem of efficient and scalable federated learning for large models in real-world applications with resource constraints, representing a novel method rather than an incremental improvement.
The paper tackles the challenge of deploying large pre-trained models in privacy-sensitive, distributed environments with heterogeneous client resources by proposing HeteroTune, a federated fine-tuning paradigm that reduces communication overhead by 99.5%, cuts peak memory usage by ~50%, and improves performance by 4.61% on LLaMA models.
While large pre-trained models have achieved impressive performance across AI tasks, their deployment in privacy-sensitive and distributed environments remains challenging. Federated learning (FL) offers a viable solution by enabling decentralized fine-tuning without data sharing, but real-world applications face significant obstacles due to heterogeneous client resources in compute and memory. To address this, we propose HeteroTune, a novel federated fine-tuning paradigm for large, heterogeneous models operating under limited communication and computation budgets. The core of our method lies in a novel architecture, DeMA (Dense Mixture of Adapters), which enables flexible and efficient aggregation of heterogeneous models by preserving their full representational capacity while facilitating seamless cross-model knowledge fusion. We further introduce CMGA (Cross-Model Gradient Alignment), a lightweight yet effective mechanism that enhances training stability by harmonizing gradient directions across heterogeneous client models during aggregation, mitigating update conflicts and promoting more consistent convergence in federated settings. We provide both theoretical analysis and empirical evidence showing that HeteroTune achieves state-of-the-art performance and efficiency across diverse tasks and model architectures. For example, on LLaMA models, it reduces communication overhead by 99.5%, cuts peak memory usage by ~50%, and improves performance by 4.61%.