Shared Representations for Personalized Federated Learning, Explained

By the RunLocalModel editorial team · October 14, 2026 · ~30 minute read

This page reads the published paper: Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai, Exploiting Shared Representations for Personalized Federated Learning, ICML 2021, PMLR 139, pages 2089–2099. Collins and Mokhtari and Shakkottai were at the University of Texas at Austin. Hassani was at the University of Pennsylvania. The proceedings text is 11 pages and points to an appendix for proofs and experimental details. Those details are in the author version, arXiv:2102.07078. Where the two texts disagree, this page says which one it is using. The code the authors' repository describes as official is at LittleStory233/FedRep. The appendix prints a different GitHub URL, and that URL does not resolve.

The hole this paper is filling

Federated learning, in the form McMahan and coauthors gave it in 2017, trains one model. Each client takes a few steps on its own data. The server averages the resulting models. The data stays put. The average is a good idea when every client is drawing from the same distribution. It is a strained idea when they are not.

The paper's Figure 1 draws the strain on a linear problem with dimension d = 20, subspace dimension k = 2, and n = 100 clients. With few samples per client, a model trained only on local data has large error. With more samples, that local model improves. A single global model, trained by FedAvg, stays inaccurate even after each client has a full d samples. The clients disagree, and the average cannot represent the disagreement. FedRep's error is small at every sample size in that figure.

The authors' explanation is the one representation learning already uses for images and text. A lot of the statistical work is in a shared map from raw inputs to features. What changes from task to task, or from client to client, is the map from those features to labels. If that story is right, federated learning should share the first map and leave the second map at home.

A concrete picture: a hundred phones classify photos. Edges, textures, and object parts are useful on every phone. One phone's owner mostly labels birds. Another mostly labels food. A network forced to use one classifier for both owners learns a compromise. A shared eye, with a private last layer, can serve both. The CIFAR experiments later in the paper are a controlled version of this picture. They hand each client a few classes from a shared image set. They do not hand each client a different camera, a different hospital, or a different language.

One feature map, many classifiers

There are n clients. Client i draws labeled pairs from its own distribution D_i, and it holds only M_i of them. The M_i are small compared with the dimension of the input. The number of clients is large. That is the usual federated regime: nobody has enough data to train a large model alone, and pooling the data in one place is off the table.

The model for client i is a composition. A shared representation φ maps an input in R^d down to R^k, with k much smaller than d. A head h_i, private to that client, maps the k-dimensional features to a label. The predictor is

The training objective judges a representation by how well each client can do after it has fit its own head:

The inner minimum is the whole point. A representation is good when, for that fixed representation, each client's best small head has low loss. The clients cooperate on φ. They do not cooperate on the heads.

This is a different problem from FedAvg. FedAvg minimizes the average loss of one shared model. If the distributions D_i differ, the functions f_i have different minimizers, and the shared model is a compromise among them. The paper's question is how to use everyone's data anyway, so that each client ends with a model that is actually good for D_i.

Why a single model cannot name the subspace

Section 2.1 makes the failure precise in a linear model, before any algorithm appears. The representation is a d by k matrix B. Client i's head is a vector w_i in R^k. The prediction is a linear function of the input:

The ground truth has the same shape. There is one true subspace B* and a true head w*_i for each client, and the label is y = w*_i^T B*^T x. Every client's ideal predictor, the vector B* w*_i in R^d, lives in the column space of B*. The clients share a k-dimensional flat inside a d-dimensional world. They differ in where they stand on that flat.

Assume the inputs are mean zero with covariance equal to the identity. The expected squared error for client i then collapses to a distance between predictors:

FedRep's objective asks for a B and a separate w_i for each client. There is a solution with zero error: set each product B w_i equal to B* w*_i. If the heads w*_i point in enough different directions to span R^k, those n equations determine the column space of B*. You have seen every direction of the shared flat, because different clients stood in different places on it.

Standard federated learning asks for one pair (B, w) used by every client. The best that pair can do is match the average client:

The leftover error on client i is the distance between that client's true predictor and the average predictor. It grows as the heads spread out. And the geometry is worse than the error. One equation produces one vector in R^d. A single vector does not determine a k-dimensional subspace once k is greater than 1. Many subspaces contain the average client. FedAvg, in this model, cannot recover the representation no matter how many samples you give it. That is what Figure 1 is showing when FedAvg's curve stays high as the sample count grows.

One round of FedRep

The published Algorithm 1 alternates a long local solve with a short shared step. A fraction r of the clients take part in a round. In the experiments, r = 0.1.

1. The server sends the current representation φ to the selected clients.
2. Each selected client starts from the head it saved last time and takes τ gradient steps on that head. The representation is held fixed. Gradient descent, SGD, or SGD with momentum all fit the template. The linear analysis solves this step exactly.
3. The client then takes one gradient step on the representation, with the head held at the value it just found.
4. The client sends the updated representation to the server. The head stays on the device.
5. The server averages the representations it received. Clients who sat out keep their old heads.

The experiments stretch step 3 from one step into a few epochs. FedRep runs ten local epochs of SGD with momentum on the head, then one epoch on the representation for CIFAR-10 with 100 clients, and five epochs on the representation in the other settings. The author-version algorithm writes that second stage as its own loop, with τ_φ steps. The published algorithm box shows one representation step. The schedule in the tables is the longer one.

In the linear case the head update is a small least-squares problem. The batch loss is quadratic and strongly convex in w_i, so gradient descent reaches an ε-accurate head in order log(1/ε) steps, and a direct solve costs order m k² + k³ operations. The analysis assumes the client returns the exact minimizer, then takes one gradient step on B using the same m fresh samples, then the server averages those B updates over the participating clients.

Clients who miss a round are not stuck with a stale model forever. The next time they are selected, they refit the head against whatever representation the server is sending now. The head is cheap relative to the representation, which is why missing rounds is tolerable. The representation is what needed everyone's data.

Why the long head solve is the algorithm

In ordinary FedAvg, many local steps are dangerous under heterogeneity. Each client walks toward its own minimizer. The average of those walks is a worse estimate of the shared minimizer as the walks get longer. SCAFFOLD and FedProx exist to pull those walks back. They still aim at one model.

In FedRep the head is supposed to walk toward the local labels. A longer walk is a better inner solve. The representation step is then taken at a head that is close to the best head for the current features. Averaging those representation steps is an average of gradients of the outer objective, the one written with the inner minimum. A short walk on the head leaves the representation gradient at the wrong head, and averaging those gradients does not estimate the outer objective.

Figure 3 tests that distinction on the linear problem, with d = 10, k = 2, m = 5 samples per update, and r = 0.1. FedRep, which solves the head exactly, pulls the subspace error down faster as the number of clients grows from 10 to 100 to 1,000. Replacing the exact solve by 1 step of gradient descent, or by 10 steps, does not scale the same way. More clients help when each participating client has actually fit its head.

The author-version appendix adds two facts the proceedings figure does not. The synthetic clients reuse the same m samples for the whole run. The proof draws a fresh batch every round. And FedRep's trajectory looked the same from a random Gaussian start, while the version that only takes gradient steps on the head often failed to converge from a random start. Exactness of the head solve is doing stabilization work, alongside the statistical work.

What the theorem measures

Theorem 1 is about the linear model. Inputs are isotropic and sub-gaussian. Labels are exactly linear in the shared subspace, with no misspecification. Each selected client draws m fresh samples on each round it participates. The server and clients first build an initial subspace by a method of moments, then run the alternating updates.

The quantity that converges is the principal-angle distance between the column space of the current B and the column space of B*. It is the sine of the largest angle between the two k-dimensional subspaces. It is zero when the subspaces are the same. It is one when they are orthogonal.

Write E₀ = 1 − dist(B₀, B*)². Let σ̄_min and σ̄_max be the smallest and largest singular values of any participating block of the true head matrix, after a 1/√(rn) scaling. The condition number is κ = σ̄_max / σ̄_min. The author version states the theorem this way. The proceedings line that defines κ prints the maximum singular value over itself. That is a typesetting error. The ratio of maximum to minimum is the quantity the rest of the bound uses.

If the batch size is large enough, and the step size on B is at most 1/(4 σ̄_max²), then after T rounds

with probability at least 1 − T exp(−100 min(k² log(rn), d)). The distance is multiplied by a constant smaller than one at every round. After order log(1/ε) rounds it is below ε. Optimization calls this linear convergence. The paper also calls it exponentially fast. The two phrases describe the same contraction.

The contraction needs E₀ bounded away from zero, so the initial principal-angle distance has to sit a constant below 1. A start that is nearly orthogonal to the truth barely moves. The method of moments is how they get an initialization that good, up to log factors in the sample count. Each client sends the server the d by d matrix

The server averages those matrices and takes a rank-k singular value decomposition. The top k directions initialize B. The intuition is a spiked covariance. For isotropic inputs, the fourth-moment matrix of one client is large along that client's predictor B* w*_i. Average over clients whose heads span R^k, and the spike fills out the column space of B*.

How many samples each client needs

The batch size that makes the contraction hold is, in the author-version statement of Theorem 1,

for an absolute constant c. Because the distance then falls exponentially in the number of rounds, the total number of samples per client scales as m times log(1/ε):

The proceedings abstract and the contribution list state a different expression, O(κ² k³ (log n + κ² k d / n)). The powers of κ and k do not match Theorem 1. Use the theorem.

The shape that matters is easier to see if k, κ, and r are treated as constants, which is how the paper summarizes it. A client training alone on a d-dimensional linear problem needs order d samples. A client inside FedRep needs order d/n + log n. The d-dimensional object, the subspace, is paid for by all n clients together, so each client pays d/n. The log n is the cost of concentration over the participating set. Collaboration helps when d/n + log n is much smaller than d, which is the range 1 << n << exp(Θ(d)). For any realistic input dimension, the upper end of that range is enormous. Federation helps each participating client, on this bound, long before you run out of clients.

A new client who arrives after the subspace is learned does not pay for d at all. The proceedings text says its sample cost scales as Θ(k log(1/ε)) to drive the error below ε. It only has to solve a k-dimensional regression on features that already exist.

Figure 1 is the picture for the simplified scaling, with d = 20 and n = 100, so d/n + log n is a small number while d is 20. FedRep's error is already small at 0.25 d = 5 samples per client, which is just above k = 2. The figure is a simulation of this linear model. It is not a measurement of the leading constant in the theorem, and the full bound still multiplies by κ⁴ k² and by log(1/ε).

Two structural assumptions sit under the d/n gain.

• Client diversity. Every participating set of rn clients must have heads that span R^k. The minimum singular value of every such block is positive. If a direction of the subspace is used by only a handful of clients, a random 10 percent sample can miss it, the condition number κ blows up, and the sample bound follows it. Diversity of heads is the signal that reveals the subspace. If every client wants the same classifier, there is one direction to see, and the extra heads are idle.
• Client normalization. Each true head has length √k, and B* has orthonormal columns. This keeps the rows of the ground-truth matrix at comparable scale, the incoherence condition this kind of sensing needs. The paper says the exact length can relax to order √k.

Why this is a matrix-sensing problem with a hole in it

Stack the predictors B* w*_i as rows of one rank-k matrix. Each sample is a measurement of one row, because client i's data only involves client i's predictor. Measurements never mix two clients. Standard matrix sensing assumes measurements that touch the whole matrix and satisfy a restricted isometry property: every rank-k matrix is approximately preserved. Row-wise sparse measurements do not give that property at an efficient sample size. The operator does not concentrate to an identity on all rank-k matrices.

FedRep's alternating procedure is built for that hole. Solve the row, which is the head, exactly. Then take one descent step on the shared factor. The paper presents Theorem 1 as the first convergence result for an alternating minimization-descent method on a matrix-sensing problem, and the first sample-efficient linear rate for rank-one, row-wise sparse measurements. The closest earlier rate they cite is Zhong, Jain, and Dhillon (2015), for rank-1 Gaussian measurements that are not row-wise sparse. FedRep matches that rate up to a factor of order κ². A κ⁴ dependence has shown up before in sparse measurement problems, including matrix completion.

The same split, a different schedule

Several personalized methods were already aimed at heterogeneous clients. The paper's claim is that they still ask each client to solve a full-dimensional problem. FedRep is the one that reduces the local problem to k dimensions.

FedPer, from Arivazhagan and coauthors (2019), is the nearest neighbor. It also keeps a global body and a local head. On each local round it updates both, for the same number of steps. FedRep's specific move is to finish the head, or come much closer to finishing it, before the body moves. The experiments are the evidence that this schedule matters. They are also the place where the margin turns out to be small on some datasets.

LG-FedAvg, from Liang and coauthors (2020), learns local representations and one global head. That cut fits a world where the classifier means the same thing everywhere and the features do not. The CIFAR partitions in this paper are the other world: the images are ordinary CIFAR images, and the clients differ by which classes they hold. The classifier is exactly the personal piece. Sharing it, and keeping the features local, throws away the part that benefits from pooling.

The author version ties the split to a meta-learning observation. Raghu, Raghu, Bengio, and Vinyals (2019) argued that few-shot methods such as MAML succeed largely by reusing a representation, and that adapting the head does most of the per-task work. Tripuraneni, Jin, and Jordan (2020) and Du and coauthors (2020) give generalization bounds of the form "pay for the representation during training, then pay only for k on a new task." Theorem 1 is an optimization result that sits next to those statistical bounds: in the linear case, this alternating procedure actually finds the representation, at a sample cost with the d/n shape.

What the experiments actually show

The real-data section asks whether the linear story survives a 5-layer convolutional net on CIFAR-10 and CIFAR-100, a multilayer perceptron on handwritten letters, and a recurrent net on tweet sentiment. It does not ask whether the theorem's rate holds. There is no subspace distance to plot once the representation is a stack of convolutions.

CIFAR heterogeneity is a class partition, the same construction McMahan and coauthors used. The 50,000 training images are cut into shards of a single class, and each client receives S shards. With 100 clients and 2 classes, each client has 500 images. With 1,000 clients and 2 classes, each client has 50 images. With 100 clients and 20 classes of CIFAR-100, each client still has 500 images, now spread over 20 classes. The images themselves are the standard ones. This is label heterogeneity.

FEMNIST uses 150 clients, 10 letter classes, and a log-normal assignment of samples in the style of Li and coauthors (2019). The average client has 148 samples, and the minimum in the appendix table is 50. The table header writes this partition as (150, 3). Sent140 uses the natural split by Twitter author: 183 clients, binary positive or negative sentiment, 72 tweets per client on average, at least 50. Handwriting style and personal vocabulary are a milder kind of feature shift than a change of camera. The dominant designed heterogeneity in the image experiments is still which labels appear.

From the author-version appendix, the architecture split is concrete. CIFAR uses a 5-layer network with two convolutional layers and three fully connected layers. FEMNIST uses a multilayer perceptron with two hidden layers. Sent140 uses a frozen 300-dimensional GloVe embedding, an LSTM, and two fully connected decoder layers. For FedRep, the head is the weights and biases of the final fully connected layer in each of these models. Everything before that is the shared representation. The "low-dimensional head" of the theorem is, in these nets, one linear layer. For CIFAR-100 that layer still has 100 outputs. It is smaller than the body. It is not k = 2.

Training details that affect how you read the table: participation r = 0.1, random initialization, SGD with momentum 0.5, batch size 10 on the image tasks and 4 on Sent140, learning rate 0.01 unless a baseline was tuned otherwise. CIFAR runs for 100 communication rounds, FEMNIST for 200, Sent140 for 50. The reported number is the average of each client's local test accuracy, averaged again over the last 10 rounds. The appendix says the whole process was repeated five times and the table shows the averages. The table does not show the spread across those five runs.

One fairness detail is easy to miss. On CIFAR-10 with 100 clients, every method other than FedRep gets one local epoch. On the other settings, those methods get five. FedRep gets that budget on the representation, and an additional ten epochs on the head. The authors' defense is Figure 5. On CIFAR-10 with 100 clients and 2 classes, raising FedAvg's local epochs from 1 to 40 does not produce a steady gain. Raising FedRep's head epochs from 1 to 40 does, and the curve does not turn down through 40. Figure 5 uses batch size 50 and learning rate 0.1, a setting in which both methods score worse than they do in the table. The table itself does not rerun every baseline with a matched compute budget of ten extra epochs.

The numbers, read against the claim

The paper says FedRep is the best method in the table, or very close to the best, in every setting. That sentence survives if "personalized federated method" is the comparison class, and if gaps of a few tenths of a point count as a tie. It does not survive as "FedRep dominates local training and fine-tuning." Here are the four columns that decide it. Accuracy is in percent.

Four readings sit on top of those cells.

A single shared model collapses when the labels differ. FedAvg's accuracies run from 23.94 to 52.75. On FEMNIST, FedProx scores 18.89 and SCAFFOLD scores 17.65, below FedAvg. Those methods are built to make one global model converge despite client drift. Convergence to the wrong target is still the wrong target. Label heterogeneity is a modeling choice, and correcting the optimization does not remove it.

Most of the recovery is "let each client have its own head." Fine-tuning the head of a finished FedAvg model for ten epochs jumps CIFAR-10 with 100 clients and 2 classes from 42.65 to 87.65. FedRep's further gain on that split is 0.05 points. On CIFAR-100 with 5 classes per client, fine-tuned FedAvg leads FedRep by 0.19 points. On Sent140 the FedRep lead over fine-tuning is 0.49 points. The discussion section of the paper says this outright: fine-tuning global methods, especially FedAvg, also performs very well, and the authors planned to study why. A gap of two tenths of a point, on an average of five runs with no reported spread, is too small to rank the methods.

The shared representation pulls ahead when local data is scarce, or when the head of a global model is a bad starting point. With 50 images per client, FedRep leads local training by about 5 points (83.27 against 78.30) and leads fine-tuned FedAvg by about 1.2 points. On FEMNIST the lead over fine-tuned FedAvg is about 6 points, and the lead over local training is about 18 points. Those gaps are large next to the two-tenths gaps. They are the empirical case for training the body while the head is already personal, which is the algorithmic idea.

Local training wins when each client already has an easy problem. CIFAR-10, 100 clients, 2 classes, 500 images: local accuracy 89.79, FedRep 87.70. A two-class problem with 500 images does not need a feature extractor trained to also serve everyone else's classes. Sharing costs about two points. This is the setting where the paper's own Figure 1, on a linear model with very different heads, would have predicted a win for FedRep. The neural net does not follow that prediction once local data is plentiful and the task is tiny.

The full table is worth having, because the ranking among personalized methods is part of the claim. Bold marks the best number in each column. FedRep's row is shaded.

Against FedPer, the same cut with a symmetric local schedule, FedRep is ahead on every column. The margins are 0.57, 1.84, 1.54, 3.15, 0.42, 0.29, and 1.65 points. The 3-point gap on CIFAR-100 with 5 classes per client is the cleanest evidence that solving the head harder, then moving the body, beats updating both together. The 0.29-point gap on Sent140 is inside the region where five-run noise could reorder the methods. The paper does not give you that noise.

Against LG-FedAvg, the opposite cut, FedRep is ahead everywhere, and the gap is large where the label set per client is rich: 56.10 against 38.76 on CIFAR-100 with 20 classes. A shared classifier cannot absorb twenty different local label subsets. That comparison supports the paper's choice of which piece to share, on this particular kind of heterogeneity.

Ditto, which trains local models pulled toward a global model, leads the CIFAR-100 20-class column by 0.24 points. APFL, which mixes a local model and a global model, is within a point of FedRep on CIFAR-10 with 1,000 clients. Per-FedAvg is competitive on FEMNIST (71.51) and weak on the scarce-data CIFAR-10 split (67.36, behind even local training at 78.30). Meta-learning one initialization is a different bet from learning a representation that a new head can sit on.

A new client, and a new label set

The linear theory says a new client, given the subspace, only fits a k-dimensional head. Figure 4 tests the picture with n = 100, d = 20, k = 2. FedRep and FedAvg are trained on the original clients. A new client receives the learned representation, solves for the best head on m_new labeled samples, and is scored on its own test data. FedRep's error becomes small once m_new is at least k. Local training on those same samples stays inaccurate until m_new is on the order of d. FedAvg's representation does not transfer: the new client's error stays large. The single shared model never learned the subspace, which is the Section 2.1 calculation showing up in a plot.

Figure 6 is the same question for a network, with a harsher shift. Models train on FEMNIST letters. New clients bring handwritten digits. There are 100 new clients. Each has 500 samples from 5 digit classes for fine-tuning the head, and 100 test samples from those classes. The representation stays frozen. The appendix notes that this fine-tuning phase samples every new client each round, unlike the training phase at r = 0.1. FedRep's frozen representation produces higher digit accuracy than the representations from FedAvg, Per-FedAvg, LG-FedAvg, FedProx, L2GD, and APFL, across the fine-tuning epochs they plot.

That is the strongest empirical claim in the paper, and it is worth stating at the right strength. A body trained on letters, while each training client kept a personal head, was a better frozen feature map for digits than a body trained as part of one global letter model. The experiment uses 500 labeled digit images per new client, which is a lot of data for "only fit a head." It shows transfer of the representation. It does not show the few-shot regime of the linear theorem, where order k samples would suffice.

Where the argument stops

The useful idea is stable: if clients share the structure of their inputs and disagree about labels, average the feature map and leave the classifier at home. Spend local computation on the classifier. Take short steps on the feature map, because those steps are what you average. A new client inherits the feature map and fits a classifier.

The theorem supports that idea inside a linear model with a true shared subspace, quadratic loss, fresh sub-gaussian samples, heads that span the subspace on every participating subset, and an initialization that is already a constant angle away from orthogonal to the truth. Inside that model the subspace converges linearly, a lone client pays order d, a participating client pays order d/n + log n once k and κ are constant, and a later client pays order k log(1/ε). The abstract's polynomial in κ is not the theorem's polynomial. The proceedings definition of κ is mistyped. The method-of-moments matrix that initializes the proof is a function of raw data, and the privacy sentence attached to it is not a guarantee.

The experiments support a weaker, still useful, version. On class-partitioned CIFAR, on letter images, and on tweet sentiment, a shared body plus a private last layer is the best or nearly the best personalized method they tried. The margin over the same architecture with a symmetric update, FedPer, is real on some splits and tiny on others. The margin over "train FedAvg, then fine-tune the head" is about six points on FEMNIST and about a point or less on most of the CIFAR and sentiment splits. Local training wins when the local task is easy and the local data is ample. Nothing in the table is a feature-shift benchmark in the strong sense: different sensors, different hospitals, different languages. The theory's fresh-sample assumption is not how the synthetic plot was run, and it is not how a finite CIFAR shard works.

Communication is almost an afterthought in the paper, and it should stay one. Each round, selected clients download the representation and upload an updated representation. The head never moves. In these networks the head is the last layer, so the upload is most of the network. FedRep is a statistical schedule. It is not a compression scheme.

Secure federated transfer learning, the other essay in this series, is a different cut of "train without pooling the data." There, two organizations hold different columns, few rows coincide, the labels sit on one side, and the products that cross the boundary are encrypted. FedRep assumes many clients, the same kind of input, different label maps, and a server that averages representation parameters in the clear. The two protocols answer different obstacles. Using one where the other applies makes a simple problem expensive, or a privacy problem unprotected.

How LLMs Are Trained on this site is about a dataset that sits in one place: a model reads, then it is taught. FedRep starts from the other end. The examples already live on the clients, the disagreement is in the labels, and the thing worth averaging is the representation underneath those labels.