FedIN: Federated Intermediate Layers Learning for Model Heterogeneity

TL;DR

FedIN introduces a federated learning approach that enables training with heterogeneous models by leveraging client feature knowledge and intermediate layer alignment, improving performance without requiring public datasets.

Contribution

The paper proposes FedIN, a novel federated learning method supporting heterogeneous models through feature-based knowledge exchange and a convex optimization to reduce gradient divergence.

Findings

FedIN outperforms existing algorithms in heterogeneous environments.

IN training effectively aligns intermediate layers with minimal overhead.

The convex optimization mitigates gradient conflicts, enhancing convergence.

Abstract

Federated learning (FL) facilitates edge devices to cooperatively train a global shared model while maintaining the training data locally and privately. However, a common assumption in FL requires the participating edge devices to have similar computation resources and train on an identical global model architecture. In this study, we propose an FL method called Federated Intermediate Layers Learning (FedIN), supporting heterogeneous models without relying on any public dataset. Instead, FedIN leverages the inherent knowledge embedded in client model features to facilitate knowledge exchange. The training models in FedIN are partitioned into three distinct components: an extractor, intermediate layers, and a classifier. We capture client features by extracting the outputs of the extractor and the inputs of the classifier. To harness the knowledge from client features, we propose IN training for aligning the intermediate layers based on features obtained from other clients. IN training only needs minimal memory and communication overhead by utilizing a single batch of client features. Additionally, we formulate and address a convex optimization problem to mitigate the challenge of gradient divergence caused by conflicts between IN training and local training. The experiment results demonstrate the superior performance of FedIN in heterogeneous model environments compared to state-of-the-art algorithms. Furthermore, our ablation study demonstrates the effectiveness of IN training and the proposed solution for alleviating gradient divergence.