[Paper Review] Federated Mutual Learning

Tao Shen, Jie Zhang+7 more|arXiv (Cornell University)|Jun 27, 2020

Privacy-Preserving Technologies in Data29 references71 citations

TL;DR

Federated Mutual Learning (FML) enables simultaneous training of a generalized global model and personalized local models by using deep mutual learning, handling data, objective, and model heterogeneities in federated settings. It updates local meme and personalized models via mutual knowledge distillation and aggregates meme models to form the new global model.

ABSTRACT

Federated learning (FL) enables collaboratively training deep learning models on decentralized data. However, there are three types of heterogeneities in FL setting bringing about distinctive challenges to the canonical federated learning algorithm (FedAvg). First, due to the Non-IIDness of data, the global shared model may perform worse than local models that solely trained on their private data; Second, the objective of center server and clients may be different, where center server seeks for a generalized model whereas client pursue a personalized model, and clients may run different tasks; Third, clients may need to design their customized model for various scenes and tasks; In this work, we present a novel federated learning paradigm, named Federated Mutual Leaning (FML), dealing with the three heterogeneities. FML allows clients training a generalized model collaboratively and a personalized model independently, and designing their private customized models. Thus, the Non-IIDness of data is no longer a bug but a feature that clients can be personally served better. The experiments show that FML can achieve better performance than alternatives in typical FL setting, and clients can be benefited from FML with different models and tasks.

Motivation & Objective

• Motivate federated learning under three heterogeneities: data non-IIDness, server-client objective mismatch, and model heterogeneity.
• Propose a flexible FL framework (FML) that trains a generalized global model and personalized local models with customized architectures.
• Leverage deep mutual learning to enable two-way knowledge transfer between meme (global-anchored) and personalized models during local updates.
• Allow clients to design their own customized models while still benefiting from collaboration through a shared representation learned by the global meme model.

Proposed method

• Introduce a two-model local setup per client: a meme model forked from the current global model and a client-specific personalized model.
• Apply deep mutual learning (DML) as the local update mechanism, with two-way KL-divergence based distillation between meme and personalized models.
• Define local update losses: L_local = alpha * L_C_local + (1 - alpha) * KL(p_meme || p_local); L_meme = beta * L_C_meme + (1 - beta) * KL(p_local || p_meme).
• During each round, clients update their meme and personalized models via DML on private data, then send meme models to the server for global aggregation by simple averaging of meme models.
• At convergence, the global model represents a generalized model over joint data, while personalized models remain on clients, tailored to their local distributions.
• Discuss privacy and fairness considerations by removing sample-count weighting in aggregation and treating each client as equal.

Experimental results

Research questions

• RQ1 How can federated learning be robust to data non-IIDness while still serving personalized local models?
• RQ2 Can a two-model local update with mutual knowledge transfer improve both global generalization and local personalization compared to standard FedAvg variants?
• RQ3 How does model heterogeneity (different architectures across clients) impact performance under Federated Mutual Learning?
• RQ4 What are the optimal alpha and beta settings for balancing data-driven learning and knowledge transfer during training?
• RQ5 What are the privacy and fairness implications of abandoning sample-count weighting in aggregation?

Key findings

• FML generally outperforms FedAvg and FedProx in typical FL settings across IID and non-IID data across multiple model families.
• DML-based local updates provide a strong regularizer via KL divergence, yielding steadier training trajectories under non-IID data.
• Clients with different tasks or architectures can benefit from a shared representation learned by the global meme model, supporting model heterogeneity.
• Personalized models retain independence and continue learning on private data, while meme models aggregate globally to form the next generation of the global model.
• Removing the n_k/n weighting and using equal client treatment can enhance privacy and fairness considerations, with negligible impact on performance in many settings.