PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploiting Shared Representations for Personalized Federated Learning

Liam Collins, Hamed Hassani, Aryan Mokhtari, Sanjay Shakkottai

arXiv:2102.07078v3cs.LGmath.OC

TL;DR

Federated learning with heterogeneous client distributions can make a single shared model perform poorly, motivating personalized solutions. The paper proposes jointly learning a global representation and local client heads, and proves efficient linear representation learning with near-optimal per-client sample complexity under a linear setting.

  • Problem

    Heterogeneous client distributions can make a single shared model perform poorly, motivating personalized federated solutions.

  • Method

    The framework learns a global representation from all clients while each client uses local information to learn a personalized head.

  • Results

    FedRep converges exponentially fast to the ground-truth representation in the linear setting and reaches an ϵ-accurate representation using O((d/n + log(n)) log(1/ϵ)) samples per client.

  • Takeaways & Limitations

    The alternating minimization-descent analysis is relevant beyond federated learning to linear multi-task representation learning.

  • Takeaways & Limitations

    Theorem 1 requires initialization with principal angle distance dist(B0, B∗) bounded away from 1, obtained using Method of Moments up to logarithmic factors.

Abstract

from arXiv · show

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.

1 Introduction

Federated learning must personalize models under client data heterogeneity while respecting distributed constraints. FedRep addresses this by learning a shared low-dimensional representation and client-specific heads, with theoretical and empirical evidence of improved efficiency and performance.

  • Client data distributions can differ substantially, making a single shared model perform poorly for many clients and fail to generalize across diverse or unseen clients.
  • FedRep treats heterogeneous federated learning as parallel tasks with common structure, learning a global representation from all clients and personalized low-dimensional heads from local data.
  • Theoretical analysis shows exponentially fast convergence to the ground-truth representation for linear regression tasks, requiring O((d/n + log(n)) log(1/ϵ)) samples per client for an ϵ-accurate representation.
  • Experiments on synthetic data, CIFAR10, CIFAR100, FEMNIST, and Sent140 report benefits from local updates, robustness to heterogeneity, generalization to new clients, and improvements over several baselines.
  • FedRep’s sample complexity per client scales as Θ(d/n + log(n)), compared with Θ(d) for local learning, so collaboration is expected to help when 1 ≪ n ≪ e^Θ(d).
  • FedRep reduces each client’s head-learning problem to dimension k, giving sample complexity Θ(k) instead of Θ(d) when no representation is learned.

2 Problem Formulation

The formulation addresses heterogeneous federated clients by learning a shared low-dimensional representation alongside personalized local heads. In a linear comparison, this formulation can recover the underlying representation and has lower per-client sample complexity than local learning.

  • 2 Problem Formulation: Client i’s supervised objective is an expected loss over its distribution D_i, estimated from M_i labeled samples used to train its model.The loss penalizes the discrepancy between predicted and true labels.
  • 2 Problem Formulation: Federated learning coordinates clients through a central server so they can use cumulative data when each client has relatively few labeled samples.Clients exchange messages while retaining local data and computation.
  • 2 Problem Formulation: Standard federated learning uses one shared model optimized for average client performance, which can perform poorly when client distributions differ.Different client error functions may have different minimizers, motivating personalized models.
  • 2 Problem Formulation: The proposed formulation treats heterogeneous federated learning as parallel tasks that may share a common representation despite having different labels.The goal is to exploit this shared structure to improve each client’s model.
  • 2 Problem Formulation: Each client model composes a global representation φ: R^d → R^k with a client-specific head h_i: R^k → Y, where k ≪ d.The low-dimensional representation reduces the number of parameters learned locally.
  • 2.1 Comparison with Standard Federated Learning: In the linear quadratic setting, the personalized formulation can achieve zero global error, whereas standard federated learning incurs error that grows with head heterogeneity.The comparison uses a shared projection matrix and client-specific head vectors.
  • 2.1 Comparison with Standard Federated Learning: The personalized formulation provides n matrix equations and can recover the column space of B* when the ground-truth heads span R^k.The standard formulation provides only one matrix equation and cannot recover that column space for k > 1.

3 FedRep Algorithm

FedRep alternates personalized head updates with representation updates across clients and a server. Clients perform multiple local updates, then the server averages their representation updates to produce the next global representation.

  • FedRep Algorithm: FedRep jointly learns a global representation while each client learns its unique local head.The method distributes computation across clients and separates shared from personalized parameters.
  • Client Update: On each communication round, a constant fraction r ∈ (0, 1] of clients is selected for local computation.Selected clients execute the client update before participating in aggregation.
  • Client Update: Each selected client performs τ_h gradient-based head updates to approach its optimal head for the current global representation.The update rule permits gradient descent, SGD, or SGD with momentum through the generic GRD operator.
  • Client Update: Larger τ_h makes the inner head minimization more accurate and improves the accuracy of subsequent representation updates.The paper typically chooses many local head epochs for this reason.
  • Client Update: After updating its head, each client performs τ_φ local representation updates starting from the global representation φ^{t−1}.These updates run for s = 1, …, τ_φ.
  • Server Update: The server averages clients’ locally updated representations to compute the next global representation.Clients send their locally updated representations after completing both local update stages.

4 Low-Dimensional Linear Representation

The linear setting models each client with a shared low-dimensional representation and a unique regressor, then analyzes FedRep’s alternating head and representation updates. The method linearly recovers the representation under stated assumptions while reducing per-client sample complexity through collaboration.

  • Model: Each client combines a shared projection B ∈ R^d×k with a unique regressor w_i ∈ R^k to map inputs to labels.All clients’ optimal solutions lie in the same k-dimensional subspace, with k assumed small.
  • Algorithm: FedRep alternates local head optimization with gradient-based global-representation updates, using client batches and server aggregation.The analysis assumes selected clients obtain the optimal local head before updating the representation.
  • Analysis: The objective has non-unique factor solutions, so recovery targets the ground-truth representation’s column space rather than the exact factors.Invertible transformations of the factors produce the same global minima; principal angle distance measures subspace recovery.
  • Theory: FedRep’s iterates converge linearly to the optimal representation in principal angle distance under client-diversity and related assumptions.The theorem analyzes convergence using singular-value quantities derived from participating clients’ ground-truth heads.
  • Sample complexity: Θ(d/n + log(n)) samples per client suffice for an ϵ-accurate representation, compared with Θ(d) without collaboration when k, κ, and r are treated as constants.New clients with the learned representation require only Θ(k) samples to solve a k-dimensional regression problem.
  • Broader connections: The analysis connects the problem to matrix sensing and extends the scope of the result to representation learning, meta-learning, and multi-task learning.The setting involves row-wise sparse measurements and a non-symmetric alternating procedure.

5 Experiments

Experiments evaluate FedRep’s local-head updates, representation quality, new-client adaptation, and performance across heterogeneous real datasets. FedRep generally benefits from more head computation and performs competitively or better than alternatives.

  • Experimental focus: FedRep experiments test local-head computation, learned-representation quality, and applicability across synthetic and real datasets.The real datasets include CIFAR10, CIFAR100, FEMNIST, and Sent140.
  • Synthetic data: FedRep converges faster as the number of clients increases, while exact local-head minimization outpaces one- and ten-step gradient updates.FedRep finds the optimal local head each round; GD-GD and 10GD-GD use one and ten updates, respectively.
  • Synthetic data: FedRep enables accurate new-client models when mnew ≥ k, whereas Local Only requires mnew = Ω(d), and FedSGD shows large error.The comparison uses a learned representation and locally optimized heads on new-client samples.
  • Real data: FedRep’s head performance is monotonically non-decreasing with more local epochs, unlike FedAvg, which is not necessarily improved by additional local updates.This result is reported for CIFAR10 with 100 clients and two classes per client.
  • Real data: Across varied heterogeneity, client counts, and samples per client, FedRep is top-performing or close to top-performing in all tested settings.The evaluation covers CIFAR10, CIFAR100, Sent140, and FEMNIST at participation rate r=0.1.
  • Real data: On new FEMNIST-digit clients adapted from FEMNIST-letter models, FedRep achieves significantly better performance than the compared baselines.The new clients use 500 samples to fine-tune the head.

6 Discussion

The discussion presents FedRep as a representation-learning framework that combines a global representation with personalized heads. It reports theoretical and empirical support for this design in federated and related multi-task settings.

  • 6 Discussion: The framework receives both theoretical and empirical justification for utility in federated settings.The empirical claims concern heterogeneous federated environments and the theoretical result concerns linear representations.
  • 6 Discussion: FedRep learns a global representation from all clients’ data and uses it to enhance each client’s model and generalize to new users.The framework also leverages client computation for multiple local updates of client heads.
  • 6 Discussion: The analysis shows that alternating minimization-descent efficiently learns linear representations, extending the result’s relevance beyond federated learning.The paper identifies meta-learning and multi-task learning as related settings in the abstracted contribution.

A.1 Synthetic Data: Further comparison with GD-GD

The supplementary synthetic experiment compares FedRep with GD-GD across client and sample configurations. FedRep’s relative advantage grows as the number of clients increases.

  • A.1 Synthetic Data: Further comparison with GD-GD: Figure 7 fixes m by row and n by column while plotting average function-value trajectories over 10 trials with 95% confidence intervals.The experiments use full participation, r=1.
  • A.1 Synthetic Data: Further comparison with GD-GD: FedRep’s relative improvement over GD-GD increases with n, highlighting its advantage when there are many clients.Principal-angle distances are omitted because their results are reported as very similar.
  • A.1 Synthetic Data: Further comparison with GD-GD: The synthetic setup generates ground-truth matrices from normalized Gaussian samples and uses noisy synthetic observations.Clients train on the same m samples throughout training, with noiseless test samples.

A.2 Real Data: Further experimental details

The supplementary real-data section documents datasets, model implementations, evaluation procedures, and hyperparameters. It also records implementation choices used to compare FedRep with other methods.

  • A.2 Real Data: Further experimental details: The real-data experiments use CIFAR10, CIFAR100, FEMNIST, and Sent140, with dataset-specific client partitions and sample counts.CIFAR datasets are shard-based, FEMNIST uses handwritten-letter subsets, and Sent140 assigns tweets by author.
  • A.2 Real Data: Further experimental details: All methods use tuned learning rates and a participation rate r=0.1, except fine-tuning phases in Figure 6, which sample all clients.The reported learning-rate search is {0.001, 0.01, 0.1}.
  • A.2 Real Data: Further experimental details: Evaluation reports average local test accuracy over the final ten communication rounds, except for fine-tuning results.Fine-tuning accuracy is measured after locally training each client’s head on the fully trained global model.
  • A.2 Real Data: Further experimental details: The implementation compares FedRep with personalized, global-model, and fine-tuned federated baselines using a shared PyTorch-based codebase.The codebase includes FedPer, SCAFFOLD, FedProx, APFL, Ditto, L2GD, and Per-FedAvg implementations.

B Proof of Main Result

The proof establishes notation and assumptions for analyzing FedRep as a shared low-dimensional representation with client-specific parameters. It models the updates through random linear operators and matrix-sensing quantities.

  • The analysis assumes sub-gaussian data and defines row-wise incoherence for low-rank matrices.Row-wise incoherence controls the norm of individual rows relative to the Frobenius norm.
  • The ground-truth representation and iterates are represented by orthonormal matrices.The proof uses hats to denote orthonormal matrices, including the ground-truth representation and current iterates.
  • Client subsets of size rn define restricted parameter matrices and their associated singular-value bounds.The quantities σ̄max,∗ and σ̄min,∗ summarize the extreme singular values over all rn-row restrictions of the ground-truth client matrix.
  • The random measurement operator maps a matrix to concatenated client-sample inner products and converts the loss into a matrix-sensing problem.The operator A is defined using standard basis vectors and sample vectors, while the labels form a concatenated target vector.
  • FedRep updates are analyzed through an operator recursion involving an adjoint, pseudoinverse, and QR factorization.The proof treats each iteration with a fresh batch of i.i.d. data and uses QR to maintain orthonormal representation iterates.

B.2 Auxilliary Lemmas

The auxiliary lemmas derive the linear FedRep update and establish concentration and norm bounds needed to control its stochastic error. Together, these bounds support the contraction argument for the representation.

  • The linear FedRep update for W is characterized explicitly through block matrices built from the current and ground-truth representations.The update is expressed using G, C, and D, with their block structure enabling row-wise decoupling.
  • The stochastic error term F is bounded by combining an inverse-Gram bound with a bound on (GD − C)w∗.The resulting estimate follows from the Cauchy–Schwarz inequality and a union bound over the relevant concentration events.
  • The update error is decomposed into independent row-level subvectors, allowing Frobenius-norm control through concentration inequalities.The proof stacks vectorized block expressions and identifies each resulting subvector with a row of the error matrix F.
  • The Gram matrix G is shown to have a positive minimum singular value with high probability.The argument uses sub-gaussian concentration, nets, and a union bound over client-related indices.
  • Concentration of the measurement operator around the identity controls the remaining update error.The proof bounds row norms and applies net arguments and Bernstein’s inequality to sums of independent, mean-zero sub-exponential variables.

B.3 Main Result

The main result uses the auxiliary bounds to show contraction of the representation distance under suitable step-size and initialization conditions. A method-of-moments procedure supplies the required initialization without more than logarithmic sample overhead.

  • The descent lemma splits the next representation error into a contracting term and a concentration-controlled term.The first term is strictly smaller than the current principal-angle distance, while the second is bounded through operator concentration.
  • The analysis assumes an initial distance bounded below one, equivalently an initialization quality E0 bounded below by a positive constant.This condition ensures the recursive contraction argument starts within its valid region.
  • The contraction argument requires a sufficiently small step size and conditions involving the singular-value parameters of the client matrix.The proof imposes η ≤ 1/(4σ̄max,∗^2) and maintains bounds ensuring the perturbation does not eliminate contraction.
  • Theorem 1 follows by recursively applying the descent lemma and taking a union bound over all iterations.The resulting probability guarantee is at least 1 − e^-100 min(k^2 log(rn), d).
  • Method-of-moments initialization achieves the required starting point without increasing overall sample complexity by more than logarithmic factors.The initialization result is stated under Gaussian inputs and uses client-provided moment matrices.

B.5 Proof Challenges

The proof faces two structural challenges: row-sparse measurements prevent standard RIP analysis, and FedRep uses non-symmetric updates. The paper addresses these through incoherence-based concentration and a mixed update analysis.

  • Row-wise sparse measurements: Row-sparse measurement operators distribute computation across clients but do not satisfy the RIP required by standard matrix-sensing analyses.The operators have nonzero entries only in client-specific rows, which prevents uniform RIP guarantees in the relevant regime.
  • Row-wise sparse measurements: Row-wise incoherence restricts the measured matrix class and reduces the required samples per user to Ω(k^3 log(rn) + k^2d/(rn)), ignoring dimension-independent constants.This targeted concentration avoids requiring RIP over all rank-k matrices.
  • Non-symmetric updates: FedRep combines alternating minimization for heads with gradient descent for representations, unlike analyses using symmetric updates for both factors.The proof establishes contraction in principal-angle distance rather than the standard Procrustes-distance result.
Loading 2102.07078v3…