Building Large Models from Small Distributed Models: A Layer Matching Approach

Cross-device federated learning (FL) has emerged as a promising approach to leverage the collective power of numerous clients for training machine learning models using their local data and computational resources. Meanwhile, huge models have recently been developed to fully utilize such a large amount of data and achieve high performance on various tasks. However, training such large models requires substantial storage and computational resources, and resource limitations of the individual client devices in the FL systems often prevent the utilization of large modern deep learning models. In this paper, we address this challenge by proposing a novel federated layer matching (FLM) algorithm that enables the server to construct a deep server model by leveraging relatively shallow client models. The FLM algorithm dynamically matches and combines similar layers from the client models into the server model while incorporating dissimilar layers as new layers. By employing this strategy, clients can train smaller models that align with their device capacities, while the server can aggregate and derive a larger, more powerful server model by leveraging the distributed client data. Our empirical results demonstrate that the FLM algorithm enables the construction of a server model that is significantly larger than the individual client models, and the server model matches the performance of the same-sized model trained in the centralized setting. Further, the model constructed by FLM significantly outperforms the model obtained using the classical model averaging algorithm under the same amount of communication and client computational resources.