TLDR: A new method called FedFD improves federated learning for systems with diverse models (Hetero-FL). Unlike older methods that transfer final predictions (logit distillation) and often become unstable, FedFD uses “feature distillation” to transfer knowledge from internal model layers. It employs hierarchical feature alignment and orthogonal projection to efficiently combine knowledge from different model architectures, leading to significantly higher accuracy and more stable training.
Federated Learning (FL) has emerged as a crucial approach for training artificial intelligence models collaboratively across many devices while keeping sensitive data private on each device. Traditionally, FL has focused on scenarios where all participating devices use the same type of model. However, in the real world, devices often have different computing capabilities, leading to a need for Model-Heterogeneous Federated Learning (Hetero-FL), where various models can work together.
A key challenge in Hetero-FL is effectively combining the knowledge from these diverse models. Knowledge distillation, a technique where a “teacher” model transfers its knowledge to a “student” model, is often used. Existing methods primarily rely on “logit distillation,” which involves transferring the final output predictions of the models. While this works well for homogeneous models, it struggles with heterogeneous models because their internal structures and how they represent data are different. This can lead to unstable training and less effective knowledge transfer, as the logit representation doesn’t fully capture the underlying differences in how models process information.
Researchers have identified that logit distillation often results in fuzzy classification boundaries and unstable training processes when applied to heterogeneous models. This is because it only focuses on the output layer and doesn’t properly align the representations in different internal “feature spaces” of the diverse models.
Introducing FedFD: A New Approach to Knowledge Distillation
To overcome these limitations, a new method called FedFD (Feature Distillation for model-heterogeneous Federated learning) has been proposed. FedFD shifts the focus from logit distillation to “feature distillation,” which involves transferring knowledge from the intermediate layers of the models, where the data’s “features” are represented. This approach is more closely related to the model’s internal structure and can better handle the variations between heterogeneous models.
FedFD introduces two main components: hierarchical feature alignment and parameter orthogonality. Instead of maintaining a separate projection layer for every single client, the server groups client models with similar architectures and maintains a projection layer for each group. This reduces complexity and ensures enough knowledge for distillation. To prevent conflicts between the knowledge from different model architectures, FedFD uses orthogonal projection techniques. This ensures that the features are mapped into separate, non-conflicting spaces, maximizing the transferred knowledge and maintaining the shape of the features during transformation.
The process involves clients training their local models and sending them to the server. The server then aggregates these models into a global model. During distillation, clients extract feature representations, which are then aggregated by the server based on model architecture. Orthogonal projection layers are trained to align these aggregated features with the global model’s features, using a technique that ensures stability and efficiency.
Also Read:
- Optimizing Federated Learning for Real-World Data and Device Variations
- FedAKD: A New Approach to Fair Federated Learning with Diverse Data
Key Benefits and Experimental Results
Extensive experiments across various datasets (CIFAR-10, CIFAR-100, Tiny-ImageNet) and different settings demonstrate that FedFD significantly improves model accuracy compared to state-of-the-art methods, achieving up to 16.09% better global model accuracy. It also shows superior communication efficiency, converging faster and more stably than logit-distillation based methods. The modular design of FedFD means it can be easily integrated with existing FL techniques, offering benefits in optimization, privacy, and flexibility.
This research highlights that feature distillation is indeed a more effective choice for model-heterogeneous federated learning, providing a stable and efficient way to aggregate knowledge from diverse models. For more technical details, you can refer to the full research paper here.
