FedTAIL: A Federated Learning Approach with Trans-Architecture Intermediate Links

Federated Learning (FL) allows multiple clients to collaboratively train deep learning models while preserving the privacy of local data, which is achieved by aggregating the gradient information from clients' models instead of using raw data. The success of FL depends significantly on the coordinated scheduling among clients. However, heterogeneity among clients in FL typically poses challenges to generalization performance, such as unbalanced data distribution and different model architectures. Particularly, as we enter the era of giant AI models, it is vital to assess FL’s compatibility with both transformer models and traditional convolutional networks. Knowledge Distillation-based FL is considered a solution to address heterogeneity, but transferring knowledge within clients using only final predictions may impact efficiency. To enhance generalization, we propose a representation-based federated learning model with Trans-Architecture Intermediate Linking (FedTAIL). Our model enables clients to share representations from intermediate layers of their local models, aggregates them at the server as global representations, and uses them to regularize local training based on similarities for all potential links between the server and clients. Empirical results verify that our framework achieves the best performance across various architectures.