FedDr+: Stabilizing Dot-regression with Global Feature Distillation for Federated Learning

TL;DR

FedDr+ is a novel federated learning algorithm that stabilizes model training across heterogeneous clients by combining feature alignment with global feature distillation, improving both global and personalized models.

Contribution

The paper introduces FedDr+, a new method that enhances local model alignment and preserves unseen class information using dot-regression loss and feature distillation.

Findings

FedDr+ outperforms existing methods in diverse data distributions.

The algorithm improves both global and personalized model accuracy.

Empirical results validate the effectiveness of feature distillation in FL.

Abstract

Federated Learning (FL) has emerged as a pivotal framework for the development of effective global models (global FL) or personalized models (personalized FL) across clients with heterogeneous, non-iid data distribution. A key challenge in FL is client drift, where data heterogeneity impedes the aggregation of scattered knowledge. Recent studies have tackled the client drift issue by identifying significant divergence in the last classifier layer. To mitigate this divergence, strategies such as freezing the classifier weights and aligning the feature extractor accordingly have proven effective. Although the local alignment between classifier and feature extractor has been studied as a crucial factor in FL, we observe that it may lead the model to overemphasize the observed classes within each client. Thus, our objectives are twofold: (1) enhancing local alignment while (2) preserving the representation of unseen class samples. This approach aims to effectively integrate knowledge from individual clients, thereby improving performance for both global and personalized FL. To achieve this, we introduce a novel algorithm named FedDr+, which empowers local model alignment using dot-regression loss. FedDr+ freezes the classifier as a simplex ETF to align the features and improves aggregated global models by employing a feature distillation mechanism to retain information about unseen/missing classes. Consequently, we provide empirical evidence demonstrating that our algorithm surpasses existing methods that use a frozen classifier to boost alignment across the diverse distribution.