FedUNet: A Lightweight Additive U-Net Module for Federated Learning with Heterogeneous Models

TL;DR

FedUNet introduces a lightweight, architecture-agnostic federated learning framework that uses a U-Net-inspired additive module to enable effective knowledge sharing among heterogeneous models with minimal communication overhead.

Contribution

The paper proposes FedUNet, a novel federated learning approach that allows heterogeneous models to collaborate using a shared U-Net-inspired additive module, reducing communication costs.

Findings

Achieves over 93% accuracy with VGG models.

Maintains high performance with only 0.89 MB communication overhead.

Enables effective knowledge transfer among diverse model architectures.

Abstract

Federated learning (FL) enables decentralized model training without sharing local data. However, most existing methods assume identical model architectures across clients, limiting their applicability in heterogeneous real-world environments. To address this, we propose FedUNet, a lightweight and architecture-agnostic FL framework that attaches a U-Net-inspired additive module to each client's backbone. By sharing only the compact bottleneck of the U-Net, FedUNet enables efficient knowledge transfer without structural alignment. The encoder-decoder design and skip connections in the U-Net help capture both low-level and high-level features, facilitating the extraction of clientinvariant representations. This enables cooperative learning between the backbone and the additive module with minimal communication cost. Experiment with VGG variants shows that FedUNet achieves 93.11% accuracy and 92.68% in compact form (i.e., a lightweight version of FedUNet) with only 0.89 MB low communication overhead.