FedDiT: Federated Learning by Distillation Token Enhanced Vision Transformer

Abstract

Federated learning (FL) is a promising approach for privacy-preserving machine learning, enabling collaborative model training across distributed devices without sharing raw data. However, FL faces significant challenges due to the nonindependent and identically distributed (non-IID) nature of data across devices, leading to difficulties in model convergence and generalization. In this paper, we propose FedDiT, a novel federated learning framework that combines knowledge distillation with vision transformers. FedDiT introduces the Distilled Vision Transformer (DTViT) model on the client side, incorporating a distillation token to enhance local learning and knowledge transfer. This approach significantly improves the robustness and performance of FL in non-IID environments. We validated FedDiT through extensive experiments on public datasets, and the results show that it outperforms existing FL methods in both accuracy and smoother convergence. Additionally, FedDiT achieves higher throughput compared to standard transformers and knowledge distillation methods, making it more efficient for practical deployment in federated learning scenarios.