Multi-Level Additive Modeling for Structured Non-IID Federated Learning

TL;DR

This paper introduces Multi-level Additive Models (MAM) for federated learning, enabling flexible modeling of non-IID data across clients through a hierarchical structure that improves knowledge sharing and personalization.

Contribution

It proposes FeMAM, a novel federated learning framework that adaptively constructs multi-level models to better capture client heterogeneity and non-IID distributions.

Findings

FeMAM outperforms existing clustered and personalized FL methods.

The method effectively models complex non-IID structures.

FeMAM demonstrates superior accuracy across various non-IID scenarios.

Abstract

The primary challenge in Federated Learning (FL) is to model non-IID distributions across clients, whose fine-grained structure is important to improve knowledge sharing. For example, some knowledge is globally shared across all clients, some is only transferable within a subgroup of clients, and some are client-specific. To capture and exploit this structure, we train models organized in a multi-level structure, called ``Multi-level Additive Models (MAM)'', for better knowledge-sharing across heterogeneous clients and their personalization. In federated MAM (FeMAM), each client is assigned to at most one model per level and its personalized prediction sums up the outputs of models assigned to it across all levels. For the top level, FeMAM trains one global model shared by all clients as FedAvg. For every mid-level, it learns multiple models each assigned to a subgroup of clients, as clustered FL. Every bottom-level model is trained for one client only. In the training objective, each model aims to minimize the residual of the additive predictions by the other models assigned to each client. To approximate the arbitrary structure of non-IID across clients, FeMAM introduces more flexibility and adaptivity to FL by incrementally adding new models to the prediction of each client and reassigning another if necessary, automatically optimizing the knowledge-sharing structure. Extensive experiments show that FeMAM surpasses existing clustered FL and personalized FL methods in various non-IID settings. Our code is available at https://github.com/shutong043/FeMAM.