[Paper Review] Exploiting Shared Representations for Personalized Federated Learning

Liam Collins, Hamed Hassani+2 more|arXiv (Cornell University)|Feb 14, 2021

Domain Adaptation and Few-Shot Learning48 references69 citations

TL;DR

The paper introduces FedRep, a federated learning framework that learns a global low-dimensional representation shared across clients while allowing each client to have a personalized local head, with theoretical and empirical performance benefits under data heterogeneity.

ABSTRACT

Deep neural networks have shown the ability to extract universal feature representations from data such as images and text that have been useful for a variety of learning tasks. However, the fruits of representation learning have yet to be fully-realized in federated settings. Although data in federated settings is often non-i.i.d. across clients, the success of centralized deep learning suggests that data often shares a global feature representation, while the statistical heterogeneity across clients or tasks is concentrated in the labels. Based on this intuition, we propose a novel federated learning framework and algorithm for learning a shared data representation across clients and unique local heads for each client. Our algorithm harnesses the distributed computational power across clients to perform many local-updates with respect to the low-dimensional local parameters for every update of the representation. We prove that this method obtains linear convergence to the ground-truth representation with near-optimal sample complexity in a linear setting, demonstrating that it can efficiently reduce the problem dimension for each client. This result is of interest beyond federated learning to a broad class of problems in which we aim to learn a shared low-dimensional representation among data distributions, for example in meta-learning and multi-task learning. Further, extensive experimental results show the empirical improvement of our method over alternative personalized federated learning approaches in federated environments with heterogeneous data.

Motivation & Objective

• Motivate personalized federated learning when client data are heterogeneous and labels differ across clients.
• Propose a framework that learns a shared representation with client-specific heads to improve performance over standard FL.
• Provide theoretical guarantees showing linear convergence and near-optimal sample complexity in a linear setting.
• Demonstrate empirical improvements over baselines on synthetic and real datasets under data heterogeneity.

Proposed method

• Define a global low-dimensional representation phi and client-specific heads h_i so that each client model is q_i = h_i o phi.
• Develop the FedRep algorithm where a subset of clients perform multiple local updates for heads, followed by representation updates, with server-side aggregation of phi.
• Analyze a linear, two-layer network setting where f_i(w_i, B) = 0.5 E[(y_i - w_i^T B^T x_i)^2] and show convergence of alternating minimization-descent to the ground-truth representation.
• Show that FedRep requires O((d/n + log n) log(1/epsilon)) samples per client to reach epsilon accuracy in representation learning, under suitable assumptions.
• Demonstrate that federation reduces per-client sample complexity from Theta(d) to Theta(d/n + log n) in the linear setting and enables generalization to new clients with Theta(k) samples.

Experimental results

Research questions

• RQ1 Can a shared, low-dimensional representation learned across heterogeneous clients improve personalization when local labels vary by client?
• RQ2 What are the convergence and sample-complexity guarantees for learning a global representation with client-specific heads in a federated setting?
• RQ3 How does FedRep perform empirically versus standard federated learning baselines under data heterogeneity on real and synthetic data?
• RQ4 Does the learned representation generalize to new clients, reducing their downstream training requirements?

Key findings

• FedRep converges to the ground-truth representation at an exponentially fast rate in the linear two-layer setting.
• Per-client sample complexity scales as O((d/n + log n) log(1/epsilon)) under the proposed method.
• Compared to learning a single global model, FedRep achieves near-optimal sample complexity gains in heterogeneous settings.
• For new clients, the representation is already available, so only a low-dimensional head of dimension k needs to be learned, yielding Theta(k) samples.
• Empirical results on CIFAR10, CIFAR100, FEMNIST, and Sent140 show improvements over baselines in heterogeneous federated settings.
• The framework highlights benefits of more local updates and effective generalization to unseen clients.