Personalized Federated Learning via Low-Rank Matrix Optimization
Authors: Ali Dadras
Sebastian U Stich
Alp Yurtsever
Venue: Accepted by TMLR ()

Abstract: Personalized Federated Learning (pFL) has gained significant attention for building a suite of models tailored to different clients. In pFL, the challenge lies in balancing the reliance on local datasets, which may lack representativeness, against the diversity of other clients' models, whose quality and relevance are uncertain. Focusing on the clustered FL scenario, where devices are grouped based on similarities in their data distributions without prior knowledge of cluster memberships, we develop a mathematical model for pFL using low-rank matrix optimization. Building on this formulation, we propose a pFL approach leveraging the Burer-Monteiro factorization technique. We examine the convergence guarantees of the proposed method and present numerical experiments on training deep neural networks, demonstrating the empirical performance of the proposed method in scenarios where personalization is crucial.

Published in Transactions on Machine Learning Research (08/2025)
Personalized Federated Learning via Low-Rank Matrix Optimization
Ali Dadras
ali.dadms dadro@rumu.se
Umeâ University, Sweden
Scbastian U. Stich
stich&rcispa. de
CISPA Helmholtz Center, Germany
Alp Yurtsever
alp yurtseverrumu..se
Umeâ University, Sweden
Reviewed on OpenReview: https: // openreview. net/ forum? id= DFJa.QBNNrr
Abstract
Personalized Federated Learning (pFL) has gained significant altention for building a suite of models tailored different clients. In pFL, the challenge lies in balancing the reliance local datasets, which may lack representativeness, against the diversity of other clients' models, whose quality and relevance are uncertain. Focusing the clustered FL scenario, where devices are grouped based on similarities in their data distributions without prior knowledge of cluster memberships, we develop a mathematical model for pFL using low-rank matrix optimization. Building on this formulation, we propose a pFL approach leveraging the Burer- Monteiro factorization technique. We examine the convergence guarantees of the proposed method and present numerical experiments on training deep neural networks, demonstrating the empirical performance of the proposed method in scenarios where personalization is crucial.
Introduction
Federated Learning (FL) is an important paradigm in machine learning, holding a great promise for training machine learning models over a large network with restricted data sharing. is most suitable when clients require collaboration often due the absence can capture the diversity and variability of the underlying behavior but in an environment where sharing datasets with collaborators is prohibited-one often driven by concerns and regulations surrounding data sharing and storage. Consequently, FL research has been focused designing algorithms that solve optimization and learning problems on network without sharing essential data; but instead communicating decision variables (referred to as local models in machine learning) or other auxiliary variables like gradients or step-directions, alongside privacy improvement strategies such as encryption or noise injection. However, the restriction of data sharing presents challenges beyond the training process. Primarily, limitations on data sharing hinder effective control over the quality and relevance of the data provided by participating clients-a a major concern that led to the rise of personalized federated learning models (pFL).
large, representative dataset available at hand that
The primary goal of pFL is finding the right balance between the two conflicting forces: the reliability of local datasots which may lack representativencss, and the diversity of collaborators' models whose quality and relevance are uncertain. As a result, pFL lacks a clear definition and direction without explicit specifications regarding data distribution of the clients. Regrettably, it appears there is no universally accepted measurable quantitative goal to evaluate the success in personalization. Adding this concern, many existing pFL methods are tested in settings that are inherently unsuited for pFL, where either Federated Averaging (FedAvg) local training produces the best accuracies. It is evident. that when Fed yields optimal results,

---

Published in Transactions on Machine Learning Research (08/2025)
indicating uniform data distributions across clients, there is room for personalization. On the other hand, if local training performs well, suggesting that local datasets represent the problem sufficiently well, the necessity of FL is called into question.
Motivated by these observations, our first step is to formulate a mathematical problem that highlights the role and necessity of a pFT, approach. Suppose there are n clients collaborating on an FI. system, indexed by = I,...,,,, n, and assume that the data for each client comes from specific data distribution, denoted by Dz. The true objective function for each client is defined as:
(0.) == EE,-D, (,
(1)
where 4 R2 x I2 + R is a loss function. Here, & = (æi, y:) denotes a random data sample drawn from the probability distribution D;, where 2, represents the input features and denotes the corresponding response. When data distributions are known, solution to this problem can be found by minimizing ((,)) locally:
min (O.). 0,
However, the true distribution D, is unknown in practice. Instead, client é approximates its objective using an empirical dataset M, = f consisting of m, independent and identically distributed data samples drawn from D. This yields the empirical objective
1 Σ (.00..
f.,,))
mi E.EMM
Throughout, we operate under the assumption that the dataset M₁ is not large enough for clients to accurately approximate a solution to problem (2) locally on their own. Otherwise, FL. would not be required.
An effective solution to this problem is possible only if the distributions D; exhibit some correlation that we can exploit. At one extreme, when all distributions are the same, the standard FL template can be used, which can be formulated as
min 0 .......,.
fa(O.)
s.L. 0, -.. On,
(3)
1=1
equivalently as
min 1 Σ f((0..
(4)
71
A significant portion of existing pFL methods are designed by relaxing the equality constraint (also called consensus roistraint);; examples include Moreau envelope smoothing and quadratic penalty regularization (Dinh et al.) 2020 Li et al., 2020). However, these approaches that penalize model dissimilarity using a specific norm have limitations, as they rely on the (implicit) assumption that similarity in distributions Dr translates to the proximity of client models in given norm. The following simple examples demonstrate these limitations:
Example 1 (Label noise in classification). Consider a linear binory classification problem with two groups of clients that differ in their sign conventions. Specifically, these groups label the positive and negative classes in opposite ways due to a misalignment in how they interpret the binary outcomes. Clearly, the data distributions of these lwo groups are nearly identical, differing only by one bil. However, for this linear classifier, this difference results in the optimal models for the two groups having opposite signs, leading to solutions 0 groupp = - grupp which are distant in all norms.
Although presented as toy example here, mislabeling is common problem in classification tasks, especially in domains like healthcare where human experts are involved in data collection. In real-world FL systems, where data remains private, detecting or preventing such issues is challenging. Therefore, algorithms must be designed robust against these inconsistencies.
Examplc 2 (Clustered FL.. Suppose each client draws data from one of r distinct distributions, forming clusters of clients. We assume that cluster memberships are unknown, and the challenge is establish effective collaboration without knowing in advance which clients share similar data distributions.
2

---

Published in Transactions on Machine Learning Research (08/2025)
Low dimestiion
0 = a
e elac
(B)
03 ag
0+
0
05
bace
01 a]
1 a2
lac
(a) Local
(b) FodAvg
(c) Distance based
(d) Low-rank (ours)
Figure 1: Solutions to the minimization problem 1 E== |00 a:ll where o;, E R.2 The red points represent individual minimizers of the problem which are equivalent to 0: = a, FEDAVG solution is 0" n Σ;= a, which is shown in yellow. Considering ||0; 0,1l S constraint one can find (5) as shown in green. Solving this problem with the low-rank assumption (r 1) gives solutions 0lae blue points, lying in low-dimensional subspace.
A specific class of pFL methods for clustered FL problems focuses on identifying clusters using various similarity measures and then performing cluster-aware aggregation. However, in a real-world FL system where data remains private, estimating these clusters is challenging without exposing additional information. This difficulty persists unless we rely on the stringent assumption that similar distributions produce models that are close in certain norm, which could then be used to estimate clusters during training.
Example 3 (Collaborative filtering). Consider the classical problem of recommendation systems. Suppose there are n clients and p items. Let 0; E TR represent the relevance scores of client i for the items. The data consists of the actual scores rated by the clients, where euch client rales only a subset of the items, denoted by S; S PP.. We denote these scores by 07 for j E S.. The goal is predicting unknown scores that are not in S,, based on hidden patterns among clients
Models based on Euclidean distance regularizations known to fail in accurately predicting movie preferences. Instead, low-rank matrix factorization is among the most popular and successful approaches to collaborative filtering, as demonstrated by the Netflix competition 2008 progress prize winning team (Koren et al. 2009 The typical explanation for the empirical success of low-rank models in this problem is that movie preferences are well- pollmearametee by a few meaningful factors, such as genre, cast, language, and year. more nuanced argument generalizes this by noting that low-rank matrices naturally arisc in latent variable modols (LVMs). While this is standard for LVMs with linear parameterizations, (Udell & Townsend 2019) demonstrate that low-rank models are effective for broad class of (possibly high-dimensional) LVMs Ms parameterized by piecewise analytic function.
Inspired by these examples, we explore how to formulate pFL without relying on a specific distance metric. This leads us to investigate low-dimensional subspace formulations, where personalized models are related by their membership to a low-dimensional subspace rather than their proximity a distance metric. This approach allows us conceptualize pFL by focusing on the inherent structure of the model relationships rather than their spatial closeness, illustrated in Figure 1 Drawing parallels to collaborative filtering, we specifically focus on low-rank formulations.
We can now summarize our main contributions:
We introduce a new formulation for pFL based on low-rank matrix optimization in Section 3 Utilizing a nonconvex matrix factorization method applied to this formulation, we propose a new method called Personalized Federated Learning via. Matrix Factorization (pPI**).
We investigate the convergence guarantees of the proposed method in SScctin 4 For the smooth nonconvex minimization problem, we show that the proposed method converges to a first-order stationary point at a rate of O(1T)); with the stochastic gradients, the rate becomes O(1/TT).
The decision variable in matrix completion reveals the data, limiting privacy benefits of FL.. Nevertheless, problem highlights the challenge of distributed learning with personalized models.
3

---

Published in Transactions on Machine Learning Research (08/2025)
We present numerical experiments on training various types of neural networks in Section 5 We compare the performance of the proposed method against the baseline in scenarios where personalization is crucial, such as classification in clustered FL with label misalignment or non-homogeneous data distributions.
2 Related Work
Many pFL algorithms regularize client models to remain close; the main assumption is that personal models are close under a chosen metric (Dinh et al.; 2020). Learning a mixture of the global model and the local models is proposed in (Hanzely & Richtárik 2020), where these personalized models are encouraged to stay close to their average by incorporating quadratic penalty. Empirical evidence suggests that model similarities among different neural networks, particularly in their classifier layers, correlates with the similarity of the training data distributions (Tan et al. 2023). Thus, enforcing model closeness (e.g., in Euclidean distance) tends to favor clients with similar data distributions.
More recently, model decoupling methods have been proposed (Arivazhagan et al. 2019 Oh et al. 2021 Pillutla et al. 2022: Mishchenko et al. 2023), showing a better performance than distancc-regalarizatton---- pFL methods. The main idea is to decouple each local model into two blocks, a feature extractor block followed by classifier block. The feature extractor block is communicated and aggregated over clients and the classifier block is trained locally by each client. (Arivazhagan et al. 2019) considered a personalization of only selected layers of the neural network: all devicos share a common set of base layers with identical weights, while each device maintains its own personalization layers that adapt to its local data. The base layers are synchronized with the server while the personalization layers are kept private by each device. In (Oh et al.] 2021), the entire network is decomposed into the body (extractor), which related to universality, and the head (classifier), which is related to personalization. This reduces the update and aggregation parts from the entire model to the body of the model during federated training.
Perhaps the most relevant works to ours are (Collins et al. 2021 and (Thekumparampiil et al. 2021). The goal in (Thekumparampil et al., 2021) is to find a shared low-dimensions representation of the data features. While the primary focus of (Thekumparampil et al. 2021 | is on multi-task learning, it can be applied for personalized FL by treating each client's 's learning problem as a separate task, as considered in (Collins et al, 2021). In (Collins et al.] (021). the server learns the common cow-dimensional features of the data, and each client learns local features suited to its requirements. This method, Federated Representation Learning (FedRep), uses gradient-based updates to train global low-dimensional reprosentation and allows cach client to compute personalized low-dinensional classifier. The key difference between FedRep and our approach is in the low-rank assumption. FedRep assumes that each clients parameters are low rank. In contrast, we assume that the concatenation of all clients parameters lies in a low-dimensional space. In other words, their feature extractor extracts the features from a single shared global model while our method trains a set of models and each client uses a combination of these models as its personalized model. Put differently, FedRep assumes that O, are low rank, whereas we assume that O == [vec(O1). (o2))...... (e.( is low rank. Furthermore, FedRep focuses on the lincar represcntation setting with quadratic loss.
Beyond distance-based and decoupling approaches, a variety of other strategies have been explored for pFT.. (Zhang et al. 2024) introduces LR-BPFL, a Bayesian personalized federated learning method that learns a global deterministic model along with personalized low-rank Bayesian corrections. (Bao et al. 2023 proposes FEDCOLLAB, a clustered federated learning framework that mitigates negative transfer by partitioning clients into non-overlapping coalitions informed by both pairwise distribution distances and relative data quantities. (Prakash al.] 2023) processes hierarchical, tree-like data in federated learning by developing algorithm tailored to hyperbolic spaces. (Anelli et al. 2022) investigates federated pair-wise learning for factorization models in a recommendation scenario. (Huang et al.) 2022) proposes an FL framework for solving the Shint-of-Interest (POI) recommendation problem. (Ammad-Ud-Dis et al. 2019) introduces a federated implementalionn of collaborative filtering for recommendation systems. Liang el al. (2020) introduces LG-FEDAVG, which combines local representation learning with global model learning in an end-to-end manner. Each local device learns to extract higher-level representations from raw data. before a global model operates on the representations (rather than raw data) from all devices. (Tan et al., (023) proposes a decoupling algorithm that also personalizes feature extractors by adjusting aggregation weights based on
4

---

Published in Transactions on Machine Learning Research (08/2025)
classifier similarity. (Deng et al., 2020 introduces APFL algorithm, which learns personalized model for each user convex combination of local and global models, with the combination coefficients adaptively updated during training. (Hao et al. 2022) assumes factorized weights for neural network weights and, rather than learning a single global model, learns a dictionary of rank-1 weight factor matrices. Each client then assembles a personalized model from this dictionary to match its own data distribution. (Jeong Hwang 2022) considers factorization of the model parameters and allows clients to perform a selective aggregation scheme to utilizee only the knowledge from the relevant participants for each client. (Pal et al. 2024] introduces LRS, modeling user parameters as low-rank + sparse to capture shared structure and individual-specifie characteristics. They further develop AMHT-I.S, providing theoretical guarantees in the linear Gaussian setting, and extend it to a user-level differentially private version.
Finally, several works use low-rank matrix optimization in FL, but with different motivations and goals from ours. One line uses low-rank structure the model parameters: (Pinto et al., 2023 projects private datasets onto a low-dimensional subspace spanned by principal components estimated from public unlabeled data, then apply gradient-based private algorithms (e-g-; Noisy-SGD) to learn linear classifier; (Zhao
et al. 2016 Cai et al. 2014) assumees one of the neural network layers is low rank; (Liu et al. 2024a uses
homogencous pro-factorized low-rank layers across clients; (Tran ct al. 2025 factorizes local prompts into
two lower-rank components plus a residual; and (Nuu et al. 2023) introduces PriSM training, which assigns
resource-constrrinnel clients low-rank sub-models via importance-aware probabilistic sampling. second line uses low-rank structure in the gradients, such as (Kasiviswanathan 2021 Yu et al. 2021 Gooneratne et al. 2020). Yao et al.] 2021) proposes FedHM, which trains low-rank factorized neural networks of a specified size and reconstructs a full-rank global model on the server via model shape alignment method.
3 Algorithm
We propose a novel formulation for pFT, based on low-rank matrix optimization:
F(O) = 1 n (..).) s.t. (ank(O) Sr.
min @RRaXn
(5)
Here, Q E RRdx denoles the system-level decision variable obtained by concatenating clients' decision variables as o = [0, 02.0..... 0n]; and r is problem specific tuning parameter. Note that this formulation suits well for the examples we discussed in the introduction.
There exists rich literature on rank-constrained matrix optimization problems, approaches including hard thresholding algorithms (Jain et al. 2010 Goldfarb & Ma 2011 Kyrillidis & Cevher 2014), convex relaxation
methods (Candès & Recht, 2012
Recht et al. 2010). and nonconvex matrix factorization techniques (Burer & Bhojanapalli et al. 2016 Park et al. 2017). The first two class of algorithms require expensive spectral decomposition steps, hence they are not suitable for federated implementation unless the server possesses sufficient computational power to perform such decompositions. Consequently, we adopt the nonconvex matrix factorization technique, also known as the Burer-Montelt (BM) factorization.
Monteiro 2003 Sun & Luo 2016
BM factorization strategy replaces the system-level decision variable O R with factorized form of θ = UVV, This transformation leads to the following optimization problem:
min (U, V), where (U, F(UVT) UeRdar V
f((U.)).
(6)
We denote by VT = [v..... vnl E grxn, In this notation, personalized model parameters can be computed as 0, = Uv, E Rd. One can interpret U E gddrr as a shared feature representation in the FL problem, computed by the server, and v, E R' as the feature extractor specific to client i.
While various optimization techniques can address problem G we design our algorithm based on the simple block-coordinate gradient updates. We can compute the gradient of with respect U and v, as follows:
VU = Σ (Uv.)) vT and V. V) = n
(7)
5

---

Published in Transactions on Machine Learning Research (08/2025)
Algorithm Personalized Federated Learning via Matrix Factorization (pFL
set U E Rmxx, v E I2" Vi E [n] for round = 0, 0,1,,...... T do
Client-level local training for client E St do
set v 1.1 = v
for k .,.... K - do
1++t1 vi 11 U Vf((U'
1.k
Vi
end for +++
viKk
G (V)((UV))) v!"
Client communicates G the server.
end for
Server-level aggregation
U = U T Eies. G
Server communicates U to the clients.
end for
It is crucial that v is separable with respect to Vs, enabling clients to compute Vv, (U, V) in parallel without requiring access data or model parameters from other clients, given the features U. Consequently, for given step-size 17: 0, local training steps can be independently formulated and performed by each participating client
v+1 = Th 1 U V.((U'V)). n
(8)
On the other hand, db not separable with respect to the rows or columns of U, necessitating collaboration among clients for computing Vu V). Consequently, the gradient step in U requires communication and will be performed at the server, forming aggregation step:
n mh (V)(U'V))) V
U++1= = U
Algorithm depicts the pseudo-code of our algorithm. Here, K is the number of local passes each client performs, and the output of the algorithm is a set of personalized paramcters 0, Uv, that each client can compute locally using its feature extractors and the shared feature representation U.
4 Convergence Guarantees
Several works have studied the convergence for the problem under different. assumptions; we refer to (Chi et al. 2019) and references therein. (Bhojanapalli et al. 2016 Park et al. 2018) proved linear/sub-linear rates for smooth functions and smooth and strongly couvex functions, respectively. Due to the nonconvex nature of BM factorization, even in cases where f(.) is convex in O, it is not possible to prove a convergence theorem to the global minimum. For more specialized cases (c.g,, matrix sonsing problems under some technical assumptions called restricted isometry property), convergence to a global solution can be characterized with careful initialization procedures (Park et al.) 2018 JJann al. 2013 Zheng Lafferty 2016 Park al., 2016). Since our focus primarily on neural network applications, where objectives are already nonconvex in θ, derive convergence guarantees to stationary point.
We begin our presentation of the main convergence guarantee by first. listing our assumptions. We assume (Uv.)) are directionally smooth, and that we have access to unbiased stochastic gradients VvF(UVT) and V (UF( with bounded variance.

---

Published in Transactions on Machine Learning Research (08/2025)
Assumption 1 (Directional smoothness). We assume that F(UVT) is smooth with respect to U and V, Le.. there evist constonds Ly, Lv 2 O such that for all U1 U2 E Rdxr and VI V2 E IR
[uF(U.VI] VUF(U Lo ( U2l1 + +I V2llF) [vF(UIVI) - VVF(U.V)) S (W(11 U211 + |V - V2llr)
Assumption 2 (Stochastic gradients). We assume access to an unbiased stochastic gradient estimator with bounded variance, i.e., there erists < to0 such that for all U Rdxr and V R
E FIFFUUVV VOF(UVT] ||² 0
[VF(UVT) VF(UVT) =FFUUT)
and
E (VFFUUTT VVF(UVT) )1 s
Theorem 1. Consider problem 6 with smooth loss functions fi(-) in the sense that Assumption holds. Assume access to a stochastic grodient estimalor such that Assumption 2 holds. Furthermore, assume that every client porticipales in each round with probability p and performs K local steps per ileration. Then, the sequence U', V generated by pFL with step-sizes t7 = mas and The < z where == max(Ly, Lv satisfies the following bound:
VVF(U'V 5 2 ) - +
21
K
n (T(1
1 - 2n
Corollary 2. Choosing 1) 2 in Theorem yields rate of O(1/VT) in the stochastic setting. If full gradients are available (a 0), then ar resulss in convergence rate of O11/T).
5 Numerical Experiments
In this section, we evaluate the performance of pFL We compare the performance of pFL against several basclines, including LOCAL training, FEDAVG [McMahan et al. FEDREP (Collins et al. 2021), APFL (Deng et al. 2020), CFL (Sattler et al. 2021), FLUTE (Liu et al.
2017). FEDPER (Arivazhagan et al. 2019
2024b), FEDAS (Yang et al. )024). FEDALT (Pillutla et al.) 2022). and PFEDEDA (MeLaughlin & Su 2024
by implementing pFL in the FL-BBench benchmark (Tan & Wang 2024 Tan al. 2023). The source code
is available al http:://ggtthu..c///aaaaaaaaaa/ wwtp.//aaaaallllll/....... We conducted experiments in live different setups:
Setup (1) For the MNIST, CIFAR10, and CIFAR100 datasets, we split the data according to the Dirichlet distribution Dir(0.5) and Dir(1) across 100 clients. The labels' distribution is shown Figures 2c
and 2d Performance of the algorithuns is shown in Table
Sctup (2) For the CIFAR-100 dataset, we partitioned the 100 classes into 20 groups, each containing 5 distinct labels. Data was then distributed among 500 clients, with each client exclusively assigned data from a single group, resulting in highly heterogeneons data. Results are shown in Table 2
Setup (3) For the MNIST, we follow the experimental setup in (Sattler et al. 2019 and consider 1000 clients divided into 10 groups, and labels in each group are re-mapped (permuted) according to a random permutation map. In other words, clients in group one would have the same numbers (0...., 9} but labeled differently; group one may consider 0 with label 0, and group two may consider 0 with label 8. Figures 2a and 2b show the distribution of the labels before and after re-labeling, respectively. Results are shown in Table
Setup (4) We sampled a subset of clients, 30% of the total clients, from FEMNIST dataset without changing the underlying data distribution, then we removed clients with less than 10 data points. The remaining set has 1091 clients. We ran the experiments for and 5 numbers of local epochs. Results shown in Table

---

Published in Transactions on Machine Learning Research (08/2025)
(a) Uniform
(b) Permuted
(c) (0.5)
(d) Dir(1)
Figure 2: Distribution of the labels for MNIST dataset and 100 clients. The vertical and horizontal axes show clients and the size of each client's data, respectively.
MNIST
CIFARIO
CIFAR100
Dir(0.5)
Dir(1)
D.5(
Dir(1)
Dir(0.5)
Dir(1)
Local
92.12%
(=0.59)
89.15%
(#1.12)
59.141%
(4:374) 74)
48.53%
(#2.53)
16.09%
(=1.52)
10.66%
(#0.98)
FedAvg
96.92%
(=0.65)
97.01%
(#0.54)
65.21%
(42.11)
65.44%
(#1.68)
28.30% (#1.8))
28.36%
(=1.32)
FedPer
96.30%
(=0.26)
95.16%
(4:0.58)
66
86%
(4:3.19
58.25%
(4:2.22)
19.98% (#:1.6)
14.22% (#:1.01)
FedRep
95.04%
(=0.40)
93.33%
(4:0.44
65.16%
(4:3.44)
35.4%
(#22066
17.49% (#1.12)
12.14%
(#:1.05)
APFL
97.93%
(=0.51)
97.64%
(4:0.39)
65.99%
(4:2.6)
65.14%
(#:1.54)
27 27.07%
(#1.57)
27.07%
(#:1.36)
CFL
96.92%
(=0.72)
97.04%
(#0.5)
64.97%
(4:2.68)
65.98%
(#:1.70)
27 27.02%
1.48
24.84%%
(#:0.91)
FLUTE
76.72%
(=0.88)
73.25%
(#11.11)
48. 20%
(#:0.74
40.59%
(4:0.69)
12.58%
(#0,28)
7.28%
(#:0.13)
FedAlt
98.09%
(=0.05)
97.91%
(4:0066)
62.89%
(=:0.77)
58.09%
(#:0.43)
20.57%
(=0.46)
16,2098
(4:0.19)
FedAS
97.17%
(-0012)
97.11%
(+-0.39)
66.335
(+1.144
63.80%
(++0.54)
8.48%
(411.79
5.12% (+ (456)))
DFedFDA
97.23%
(+0.04)
97.05%
(+H0.08)
70.65%
(+1,,60)
66.67%
(+1.44)
26.15%
(++0.19)
19.10% (++0.17)
PFTAF
r:=
96.75% (=0.61)
96.53%
(+0055)
43.89%
(#3.49) 49)
61.03% (=:1.66
34.32% (111.6)
35.24%% (#:1.77)
F = 5
96.78% (=0.51)
96.35%
(10060)
60.73%
(12)86)
65.89% (11.88)
35.64%
(12.00)
35.
(#:1.44
= 10
96. 96.98% (=0.70)
96.84%
(10056)
65.10%
(4:2.30)
67. 67.68%
(+1.56)
35. 251.28% (#1.73)
36.86 36.84% (=1.50)
== 15
98. 24% (+00266
97.93% (+0.22)
63% (+:2.43)
65.88% (+1.62)
35.70% (+:1.77)
36% (+:1.46)
Table 1: Performance of the algorithms for Setup (1) second best is underlined.
The best accuracy is shown in boldface, and the
Setup (5) We examine the sensitivity of pF1MP to the choice of rank, ranging from 1 to 20, on a highly heterogeneous data split, (0r(0.1) and 100 clients. Figure 3 shows the average test accuracy and runtime versus rank for the MNIST and CIFAR-10 datasets.
Setup (6) We evaluated pFL with range of local update counts, Ky 10, and r {1, 10, 20}, a highly heterogeneous MNIST split generated by Dir(0.1). The experiment involved 200 clients, each communication round sampling 10% subset. F Figure 4 reports the resulting average test accuracy.
Model. We used three-layer neural network, consisting of three linear layers, on the MNIST and FEMNIST datasets and a four-layer convolutional neural network, consisting of two convolutional layers followed by two linear layers, on the CIFAR10 and CIFAR100 datasets. For FEDPER and FEDREP, we treated the last layer as the classifier, while in PFLMP we factorized the entire model.
Hyper-parameters. We consider partial participation with probability equal to 0.1. We set the batch size equal 256 for all algorithms. We triod parameter values of (6) is set to r E f1, 10, 15). All experiments have 75% train and 25% test data splits on each client's data. We chose the best step size for each algorithm from the set {10- 4,00,00 10 3, 10 11002,10') 10
Observations. In the heterogencous experiments, Sctup (1) pFL outperforms other pFL methods in most cases, though the performance across algorithis is comparable on the MNIST dataset. Remarkably, DFL achieves substantial improvements in average test accuracy when different client groups have similar intra-group distributions but differ significantly across groups (see Table 2 It important to note that the low test accuracy observed in the CIFAR-100 experiment, /Setup (2) is attributed to the simplicity of the neural network model rather than the algorithms themselves. Additionally, the results indicate that multiple
8

---

Published in Transactions on Machine Learning Research (08/2025)
MNIST (permuuee labels)
CIFARI00 (super groups)
FEMNIST
1000 clients
500 clients
1091 clients
local epoch
local epoch
1 loral epoch
local epochs
Local
25.36%
(10.013)
40% (100,95)
50.77% (1:0.053)
65. (0.012)
FedAvg
12.02%
(+00022)
36.40% (#: (1.1)))
65.40% (1:0.017)
77. 77.19%(0.013)
Ped
19.86%
(+0.141)
14.80% (+0.81)
66.05% (+000010)
67. 67.72%(00.00)
FedRop
21.30%
(10.148)
12.18% (10.80)
66.10% (1:0.013)
66. (6.20%((100010)
PFIAE
==1
14.70%
(+0.083)
42.94% (+1.22)
67.82%
(#0.134)
71.42%(0.044)
r: 5
23.75%
(+:0027)
44.70% (#1.91
69.99% (#F6.123
72.009%[
F 10
34.23%
(40.090)
45.57% (+1.97)
72.56%
(+:00223)
72.47%((0.000))
15
39.31% (20.042)
45.43% (+123)
73.59%
(0.092)
(6.400(1
Table 2: Performance of the algorithms for Setup (2)
Setup (3)
and Setup (4)
The best accuracy is
shown in boldface, and the second best is underlined.
-10
1.1
MNIST CIFAR1O
0.8
1
accuracy
runtime 5
0.9
0.6
0.8
0.4
0.7
0.2
MNIST CIFARIO
0.6
0.5
7
9
11
13
15
17
19
3
5
7
9
11
13
15
17
19
rank
rank
Figure 3: Sensitivity analysis of PFL with respect the factorization rank r. The left plot shows prediction accuracy versus rank, while the right plot shows runtime versus rank. Results are averaged over random trials. See Sctup (5) for details.
local updates in pFL contribute positively to performance. Overall, the findings highlight the practical significance of the proposed method.
6 Conclusions
We introduced new pFL formulation based low-rank matrix optimization and developed novel pFL algorithm utilizing Eurer-Monteiro factorization. We further established convergence guarantees for the proposed method: for minimizing smooth non-convex objective, the algorithm converges to stationary point at a rate of OL/TT) with full gradients; and O(1/VT) for the stochastic setting, Evaluations across four experimental setups highlight the practical significance of the proposed method, especially in scenarios where personalization is essential, and standard approaches are unable adequately capture the complexity of the underlying data distributions.
We conclude by listing some limitations and future directions. Our mumerical experiments demonstrate improved performance of DFL"P with multiple local steps; however, this enhancement is not reflected in our thcoretical convergence guarantees. Establishing stronger guarantees that reffect this behavior is a valuable direction for future research. Another notable limitation is that our formulation currently factorizes the entire model (decision variable), which can be computationally intensive in some cases, particularly in large-scale neural network applications. A more efficient approach might be to apply the BM factorization selectively, targeting only subset of the parameters, which could reduce overhead while maintaining its benefits. Exploring such partial factorizations promising direction for future research

---

Published in Transactions on Machine Learning Research (08/2025)
Ki
10
Kr
100
1
0.8
0.8
0.8
accuracy
accuray
accuracy
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
200
400
600
800
1,000
200
400
600
800
1,000
200
400
600
800
1,000
communication round
commuicatiin round
communication round
(a)
r=
(b) r
(c) r 20
Figure 4: Performance of pFL for different numbers steps and ranks. See Setup (6) for details.
Acknowledgments
This work was partially conducted while Ali Dadras at research visit CTSPA; sincerely appreciate their support and the stimulating research environment they provided. Alp Yurtsever and Ali Dadras were supported by the Wallenberg AI, Aulonomous Systems and Soflware Program (WASP) funded by the Knut and Alice Wallenberg Foundation. The computations were enabled by the supercomputing resource Berzelius provided by National Supercomputer Centre at Linköping University and the Knut and Alice Wallenberg foundation. Alp Yurtsever and Ali Dadras further acknowledge support from the Swedish Research Council, under registration number 2023-05476.
References
Muhammad Ammad- Ud-Din, Elena Ivannikova, Suleiman A. Khan, Were Oyomno, Qiang Fu, Kuan Eeik Tan, and Adrian Flanagan. Federated collaborative filtering for privacy-preserving personalized recommendation 2019.
system. arXiv preprint
Vito Walter Anelli, Yashar Deldjoo, Tommaso Di Noia, Antonio Ferrara, and Fedelucio Narducci. User- controlled federated matrix factorization for recommender systems. Journal of Intelligent Information Systems, 58)))2287,309, 2022.
Manoj Ghuhan Arvaazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arXiu preprint xrXim:1912. 000010088, 2019.
Wenxuan Bao, Haohan Wang, Jun Wu, and Jingrui He. Optimizing the collaboration structure in cross-silo federated learning. In International Conference Machine Learning, pp. 1738- PMLR, 2023.
Srinadh Bhojanapalli, Anastasios Kyrillidis, and Sujay Sanghavi. Dropping convexity for faster semi-definite optimization. In Conference on Learning Theory, pp. 530-582. PMT.R, 2016.
Samuel Burer and Renato D. C. Monteiro. A nonlinear programming algorithm for solving semidefinite programs via low-rank factorization. Mathematicul Programming, 95(2):329 357, 2003.
Chenghao Cai, Dengfeng Ke, Yanyan Xu, and Kaile Su. Fast. learning of deep neural net works via singular value decomposition. In PRICAI 2014; Trends in Artificial Intelligence, pp. 820- 826. Springer, 2014.
Emmanuel Candès and Benjamin Recht. Exact matrix completion via convex optimization. Communications of the ACM, 55(6):111-119, 2012.
Yuejie Chi, Y. M. Lu, and Yuxin Chen. Nonconvex optimization meets low-rank matrix factorization: an overview. JEFE Transactions Signal Processing, 67(20):5239 5269, 2019.
10

---

Published in Transactions on Machine Learning Research (08/2025)
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized lederated learning. In International Conference on Machine Learning, pp. 10899-0099 PMTRR 2021.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning.
arXiv preprint XrX::20003. 13461, 2020.
Canh T. Dinh, Nguyen H. Tran, and Josh Nguyen. Personalized federated learning with Moreau envelopes. In Advances in Neural Information Processing Systems, volume 33, pp. 21394-21405, 2020.
Donald Goldfarb and Shiqian Ma. Convergence of fixed-point contimuation algorithms for matrix rank minimization. Foundalions of Compulational Malhemalics, (1(2))183 210, 2011.
Mary Gooneratne, Khe Chai Sim, Petr Zadrazil, Andreas Kabel, Françoise Beaufays, and Giovanni Motta. Low-rank gradient approximation for memory-efficient on-device training of deep neural networks. In ICASSP 2020 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3017 3021. IEEE, 2020.
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. ar preprint WrXiv:2002005516 2020.
Weituo Hao, Nikhil Mehta, Kevin Liang, Pengyu Cheng, Mostafa El-Khany, and Lawrence Carin. WAFFLE: weight anonymized factorization for federated learning. IEEE Access, 10:49207-492188 2022.
Jiwei Huang, Zeyu Tong, and Zihan Feng. Geographical POI recommendation for internet of things: a federated learning approach using matrix factorization. International Journal of Communicntion Systems, pp. e5161, 2022.
Prateek Jain, Raghu Meka, and Inderjit Dhillon. Guaranteed rank minimization via singular value projection. In Advances in Neural Information Processing Systems, volume 23, 2010.
Prateek Jain, Praneeth Netrapalli, and Sujay Sanghavi. Low-rank matrix completion using alternating minimization. In ACM Symposium on Theory of Computing, pp. 665-674. ACM, 2013.
Wonyong Jeong and Sung Ju Hwang. Factorized-FL: personalized federated learning with parameter factorization and similarity matching. In Advances in Newral Information Processing Systems, volume 35, pp. 35684 35695, 2022.
Shiva Prasad Kasiviswanathan. SGD with low-dimensional gradients with applications to private and distributed learning. In Uncertainty in Artificial Intelligence, pp. 1905-1915. PMLR, 2021.
Yehuda Koren, Robert Bell, and Chris Volinsky, Matrix factorization techniques for recommender systems. IEEE Computer, 42(8):300 3208):30-37, 2009.
Anastasios Kyrillidis and Volkan Cevher. Matrix recipes for hard thresholding methods. Journal of Mathematical Imaging and Vision, 48:235 265, 2014.
Tian Li, Anit Kumar Sahu, Manzil Zahocr, Maziar Sanjabi, Amoct Talwalkar, and Virginia Smith. Fedorated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429-450 2020.
Paul Pu Liang, Terrance Liu, Ziyin Liu, Nicholas B. Allen, Randy P. Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: federated learning with local and global representations. arXiv preprint Xrx:222011.0 0.5522, 2020.
Jiahao Liu, Yipeng Zhou, Di Wu, Miao Hu, Mohsen Guizani, and Quan Z. Sheng. FedLMT: tackling system heterogeneity of federated learning via low-rank model training with theoretical guarantees. In International Conference on Machine Learning. PMLR, 2024a.
Renpu Liu, Cong Shen, and Jing Yang. Federated representation learning in the palleeearameterizedd regime. Xix preprint ar Xix:::200. 045649 2024b.
11

---

Published in Transactions on Machine Learning Research (08/2025)
Connor McLaughlin and Lili Su. Personalized federated learning via feature distribution adaptation. In Aduances in Newral Information Processing Systems, volume 37, pp.
2024.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication- efficient. learning of deep networks from decentralized data. In International Conference on Artificial Intelligence and Statislics, pp. 1273 1282. PMLR, 2017.
Konstantin Mishchenko, Rustem Islamov, Eduard Gorbunov, and Samuel Horváth. Partially personalized federated learning: breaking the curse of data heterogeneity. arXiv preprinl arXiv:2305. 1tt:2305518285, 18285, 2023.
Yue Niu, Saurav Prakash, Souvik Kundu, Sunwoo Lee, and Salman Avestimehr. Overcoming resource constraints in federated learning: large models can be trained with only weak clients. Transactions Machine Learning Research, 2023.
Jachoon Oh, Sangmook Kim, and Se- So-Young Yun. FedBABU: towards enhancod representation for foderated image classification. ar preprint urXiv:2106.060. 2021.
Soumyabrata Pal, Prateek Varshney, Gagan Madan, Prateck Jain, Abhradeep Thakurta, Gaurav Aggarwal, Pradeep Shenoy, and Gaurav Srivastava. Sample-efficient personalization: modeling user parameters low rank plus sparse components. In International Conference on Artificial Intelligence and Statistics, pp. 1702 1710. PMLR, 2024.
Dohyung Park, Anastasios Kyrillidis, Srinadh Bhojanapalli, Constantine Caramanis, and Sujay Sanghavi. Provable Burer Monteiro factorization for class of norn-constrained matrix problems. orXiv preprinl arXiv:1606.011 013161 2016.
Dohyung Park, Anastasios Kyrillidis, Constantine Caramanis, and Sujay Sanghavi. Non-square matrix sensing without spurious local minima via the Burer Monteiro approach. In Artificial Intelligence and Statistics, pp. 65-74. PMLR, 2017.
Dohyung Park, Anastasios Kyrillidis, Constantine Caramanis, and Sujay Sanghavi. Finding low-rank solutions via nonconvex matrix factorization, efficiently and provably. SIAM Journal on Imaging Sciences, 11(4): 2165-2200 2018.
Krishna Pillutla, Kshitiz Malik, Abdel-Rahman Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference Machine Learning, pp. 17716-17758. PMLR, 2022.
Francesco Pinto, Yaxi Hu, Fanny Yang, and Amartya Sanyal. PILLAR: how to make semi-private learning more effective. arXiv preprint arXi::2306.0300 2023.
Saurav Prakash, Jin Sima, Chao Pan, Eli Chien, and Olgica Milenkovic. Federated classification in hyperbolic spaces via secure aggregation of convex hulls. Xiv preprint arXio:2308.0655 2023.
Benjamin Recht, Maryam Fazel, and Pablo A. Parrilo. Guaranteed minimum-raude solutions of linear matrix equations via nuclear norm mininization. SIAM Review, 52(())471 501, 2010.
Felix Sattler, Klaus-Robert Müller, and Wojciech Samek. Clustered federated learning. In NewrIPS19 Workshop on Federaled Learning for Dala Privacy and Confidentialily, pp. 1-5, 2019.
Felix Sattler, Klaus-Robert Müller, and Wojeiech Samek. Clustered federated learning: model-agnostic distributed multitask optimization under privacy constraints. JEEE Transactions on Neural Networks and Learning Systems, 32(8):3710 3722, 2021.
Ruoyu Sun and Zhi-Quan Luo. Guaranteed matrix completion via non-convex factorization. IEEE Transactions on Information Theory, (2(1))66535 6579, 2016.
Jiahao Tan and Xinpeng Wang FL-bench: a foderated learning benchmark for solving image classification tasks. http::://ggthb...c/////ooa/aaaaa
2024.
12

---

Published in Transactions on Machine Learning Research (08/2025)
Jiahao Tan, Yipeng Zhou, Gang Liu, Jessic Hui Wang, and Shui Yu. pFedSim: similarity-aware model aggregation towards personalized federated learning. Xiv preprint arXin:2305.15700 2023.
Kiran K. Thekumparampil, Prateek Jain, Pranceth Netrapalli, and Sewoong Oh. Statistically and compu- tationally efficient. linear meta-ropressntation learning. In Aduances in Neural Information Processing Systems, volume 34, pp. 18487-18500, 2021.
Linh Tran, Wei Sun, Stacy Patterson, and Ana Milanova. learning for multimodal large language models. arXiv preprint 0rxv:25001.1900 2025.
personalized federated prompt
Madeleine Udell and Alex Townsend. Why are big data matrices approximately low rank? SIAM Journal on Mathemalics of Dala Science, 1(1):144 110, 2019.
Xiyuan Yang, Wenke Huang, and Mang Ye. FedAS: bridging inconsistency in personalized federated learning. In IEEE/CVP Conference on Compuler Vision and Pallern Recognition, pp- 11986-112 11995, 2024.
Dezhong Yao, Wanning Pan, Michael O'Neill, Yutong Dai, Yao Wan, Hai Jin, and Lichao Sun. FedHM: efficient federated learning for heterogeneous models via low-rank factorization. orXiv preprint arXiv:2111.14665 2021.
Da Yu, Huishuai Zhang, Wei Chen, and Tie- Yan Liu. Do not let privacy overbill utility: gradient embedding perturbation for private learning. arXiv preprint arXi::2102.122 12677, 2021.
Boning Zhang, Dongzhu Liu, Osvaldo Simeone, Guanchu Wang, Dimitrios Pezaros, and Guangxu Zhu. Personalizing low-rank Bayesian neural networks via federated learning, arXio preprinl 2rX.v22410.43390 2024.
Yong Zhao, Jinyu Li, and Yifan Gong. Low-rank plus diagonal adaptation for deep neural networks. In 2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5005-5009. IEEE, 2016.
Qinqing Zheng and John Lafferty. Convergence analysis for roctangular matrix completion using Burer- Monteiro factorization and gradient descent. arXiv preprint arXic:1605.0705. 2016.
13

---

Published in Transactions on Machine Learning Research (08/2025)
A Theoretical Results
A.1 Compact Notation
We assume that each client participates in the learning process independently with probability p. To model this, we define a partial participation matrix D, for each time step i as diagonal matrix, where each diagonal entry [Deja, represents the participation status of client i at time t. Specifically,
(o. 1,
with probability p. with probability 1 P.
[Deea ==
where each DDlaa an independent Bernoulli random variable with parameter p. This implies that for each client [Dr],s 1 if the client participates in the training process at time t, and [Del, = 0 otherwise. We write algorithm in the compact form follows:
V =
for k 0, ..... 1, do
Vk+1 V% V.
end for
U U² - V..FuuF(UU'
y*+ V K
We define expectations with respect to gradient noise as Moieel | and E E and expectation with respect to participation randomness as [t. For participation randomness, we assume that Es, [D]] = pl, where is the identity matrix and p the probability of client participation under independent sampling. We define the conditional expectation given all randomness before iteration / and local step l
E
E
noise
randomness before (t,k), Dt).
where the randomness includes all prior gradient noise and participation randomness up local step of iteration Additionally, we define the conditional expectation given all randomness in the algorithm before iteration and the final local step K as
E, [-]
randomness before (t, ()
Finally, we use E||] to denote the total expectation over all sources randomness in the algorithm, including gradient noise and client. participation.
A.2 Convergence Analysis
We start with proving some usefuu bounds. For any t,
E, (FF(U = = [ I1 + -oN
+2(VOF(U'V'") (UF(U'V) "UFF"""""-"V -UFLUMMI
V +o2.
(10)
14

---

Published in Transactions on Machine Learning Research (08/2025)
Similar to above, can write, for any t and l,
Ey, |D. [/]] B.t.
| D - }/(
+2 D.V.F(U'V ) D.V.F(UUV - }
Be DD]=0
D. -v. ) + +1D.1 B.,
FvF(U'VE VvF(U'VE {1]]]
- | II. o²
where in the third line, used the submultiplicative property of the Frobenius norm. Now take the expectation with respect to the participation probability
DD.Vv
312]1 Es. ||D.VFF(UVV T>1l +|Di1 0²]
Es, Et,
s Tv }|p +o²
(11)
where we used Lemma and
Lemma
(A) First, we use the smoothness of F(UV!T) with respect V and write
F(U'Vi K41 J) s F(U'Vk ) (VVF(U'V)) ) V V1 + 2 | Vill
Iv
F(U'V )) no (VVF(U'V) D. D.
1) n Ly
D D. )F
Taking conditional expectation, we get
Ess. By, F(U'' }]] s F(U'VE 7>P
VvF(U'V }1/ + %/
pplz
VVF(UVV
11
1le Iv
0²
= F(U'v T) 7oP (1- Lv
VVF(U'V ))
Lvo n
2
where we used
111 in the second line. We rearrange the inequality above and average over k to obtain
1 (1 1/
Ly
Σ
[VF(UUV 3/1
TeP
K
k=
[ Es, E., (( []])
2 Lvv?
< K
th
2
(F(U'V(T - -Es, Br, [(U'V++1)))) I]) n/ Lvoz
(12)
where, in the second line, we used v = V and Vi,k = V
(B) Now, we will use the smoothness again, but this time with respect to U:
F < F(U'VE 3) + (VoF(U'VE )) 1 - -ut--U') VoF(U'V( + nhe
Lu
10t1-Uip
2
s F(UVV The (VoF(U'V
- )
LU 2
Similar to the previous case, if we take expectation with respect to the randomness in U update at iteration f, we obtain the following bound by using 10):
E, [F(U+++V viTi 5) F(U'V ) (VOE(U'VE ", VoF(U'((T) + th Tu 2
VUr(U + Tyoo
15

---

Published in Transactions on Machine Learning Research (08/2025)
If we split the inner product term as
(VUF(U'VE ) - - - - -(VOR(UVVE I (VP(U()) Vor(U - + |
F VUF(U'VV VUF(U'V }) |
2-2 -
Male | || 27
I
where the last line follows from Young's inequality assumption, we have
with a = nabo. Moreover, by the smoothness
(VR(U'T -VuF(UVV )|p
Le |
Vill2 v nh Lu
Vv FvFUU
2L
21/u
The 2
Combining all these bounds, we get
E, [ } & =(U'vt 7h (1- tl
LU
/ 31
LU 2
+ The
2
Lu
k-1
[VF(UU'V
+ mh Luo"
+ t
Σ
(13)
2
2
1=0
Now we consider two cases and K in
1. For = 0, we have
E, [F s F(U'v()) nu (1 (1-1)-4 noLo) (UF(U'v + mh
Luo?
(14)
where we used V V
2. For k= K, we have
Es [F(U+IV(+) F(UUVCCtTT) The
Lu (1- 1u
FUF(U'Vt /²2
Lu
+ nu 2
2
Lu
K-1
Σ VvF(U'V T)
LUo2 + the 2
+ t/e
2
k=O
s F(U'V(+1)) nh Lu
VUF(U'veT]] 1|²
2 Lu
K
Luo² 2
VVF(U'Vt
+ +Th
2
Rearranging the terms we can write
12 VOFrU'ti Σ VvF(U'V T)
K
2
Lu
Nhe
+ +
Looz 2
$
(C) We once again use smoothness with respect to V:
FU 4+1 F J) + (VrF(U 1) V vi) + LV V Vill
=F(U++VVg] nu (VVF(UHVV) J), F.V.V J)) +h . 2 |D.
16

---

Published in Transactions on Machine Learning Research (08/2025)
We take the conditional expectation
F(U k+1 D)] < F(U
Bii ( []]]
Es.
Ei
Bs.
pLy
TvF(U'VV
)> }1 + L
+ nl
2
where we used 11
Focusing again on the inner product term, we obtain
E ( (VrF(U V " ), D. ">]]
Es.
E1, [(VUF(U** ) VVF(U'V "), V.VVF & 5)1
-Es
(VVE(U++VVE) )) VVF(U'VE )) D.V.E(U'VE 1) (VVF(U'VE T ". ".
Es.
(VYF(UL+IVE) VVF(U'VA VVFIU'VA 3) +p VVF(U'VV >1?
VV - VVFLU'VE | - VVFIU'V 3 +D V
BoPKLy
2
pLv |0+1
-UIE
TopKLy VVP(UVV + 112
21/1K
pLl pLy-11
V Jo
VVP(U'V | +D VVF(U'V ||;
21p,
2
where we used Young's inequality 17
in the fourth line with = KvL. Substituting back, we get
Es. Er,
F(
k+1
D)] s F( ) noP (1 (1-n (K + +1LL
VvF T|
2
pLv
+ + Lvo²
nh 2K
By summing over and appropriately rearranging the terms in the inequality, we obtain the following:
(K + L)Lv
VVF(U'V |/ Tu 2
2LLV
|VUF(UVV }|²
ToP 1- 77
2
S SF(U++Vv(T) -ss, [Ex [FUU+Iyet 3]] n/
KTyv02
2
where we used Vt.K_ Vi+1 and Vo. V We define
VUF |1]]
Au = E,
VVF(U 3113]
Av Ess.
Σ Eid
(15)
K k=
We summarize the resulting inequalities part (A) to (C) as
Lv [ (1-no- no Luv Av S (F(U'VCT) Es. EL,K F(U''(++T)]]] []]
Lv
ToP
1/
2
K
2
(1-1) DU < F(U'V(T) [.( + 2 the LUo2
VVF(U'V 5/| 11) S < - (E(F(UUUIYHIIT) + + Luoz
Ly
K-1
The A. + o
Σ
2
k=o
(K + 1v
2ph 22 2
Du F(U -Es, E,, F[U++I+(+T]]
2 KLvo²
noP
1 nl
KAv
17

---

Published in Transactions on Machine Learning Research (08/2025)
where we used the definitions
15
We rewrite four inequalities above using % K := P and defining
L := max (Tv, Lv> } as
mp²
pI. "2
Av SF(U'V'
- Es. [E,K [
Ioo2 Too2
3]] + +1²
2K
(1 SL) S F(U'V')) - EF(UU + t,2
2
17
L
r²p²
R-1
Σ VVF(U'VA 15
S - - EF + +12 م0
2
K²
2
k=
npp² (K + 1)L
(1 - 7p
KKv - "1 Do )((
PL
- Es.
ELK
+ 17 Lo?
K
2K
2
2K
Summing up the inequalities above, we gel
7(1 ((1-2n1) (Au + Av) ( -2 + mp 17 PL
Av
2K
<2 < - Es. Er,x
[(U + +12(1 + 1 )
x2 ((U'V()) Es Er,
(r(utiva I]] + 20)2
(16)
where we used the following inequality
2)
L
K-1
Σ
VVFUU'
3/
mppp"
(1 1 up 2K
(K + 1)L
K-1
|VFFUUV }1²
2K²
K
k=0
k=
1
L
K
K-1
Σ
VVF(U'V! k
5//²
?pp
( (1- (1-m mp
(K + 1)L
K
[vF(U'V:3]]
2K
K
2K
k=0
km
K-1 1)
mpp L
(1 - p) + 1 -1
Σ [vF(U'
K 2
1 < 0 (n - 2L
where in the second line, we used Jensen's inequality, see 18 of randomncss in the algorithm. Then, we average both sides of 16 over the iterations t, followed by dividing both sides by n (1 - 2nL), yielding the following expression:
Next, we take the expectation over all sources
VVF(U'V b11))
1 T
Σ t=
E VUE(U'V TT 1
1
Σ
+E
K
k=0
2 ((UUVOT) FTU
2nIT.0
+
nT (1- (1-2nL)
1 2n
2 ( - 17
2n
nI(1 (T(L- 2nL)
2n
provided that L = ma, Lv3, 1/- K k and 11 S 1
21. This completes the proof.
18

---

Published in Transactions on Machine Learning Research (08/2025)
A.3 Useful Inequalities
We list below lew elementary facts that we used in our analysis, included here only for completeness.
Lemma 3. For any malrices A E Rdixd2 xd2 and B E Rd2xdo, the Frobenius norm of the product AB salisfies
|BB| F S
where |1-1l2 is the spectral norm.
Lemma 4 (Young's inequality). Let θ Y E Rd xd2 and a > 0. Then, the following inequality holds:
(o, Y) $ |01|
2a LYP.
(17)
Lemma Let Or E xdd for e O,..., , K 1. Then, the following bound holds:
K Σ 1=0
S K K
1 2
(18)
Proof. This inequality follows directly from Jensen's inequality applied to the Frobenius norm.
Lemma 6. Let D be diagonal matria with diagonal entries that are 1 with probability p and 0 with probability P, and let A be an arbitrory matriz. Then the expectation of the squared Frobenius norm of the product DA is given by
ED DAA 1?
Proof. The Probensus norm squared of DA is defined as:
|DAI = E(DA)
Since D is diagonal, the product DA will zero out all rows of A where the corresponding diogonal entry in D is O. Let d, represent the i-th diagonal entry of D, where each d, is Bernoulli random variable with ED [d] p.
Thus, we can express
DAI 0s:
72
n A, 12
DA
j=1
where - Ai is the squared norm of the i-th row of A. Now, taking the expectation, we have:
ED DA 1
EDId
Σ
pl 12
i=1
Simplifying, we get:
ED [ 12 =P Σ
| 11 A
1=1
Therefore, the expectation of DAA is:
ED |DA 1 E.
This completes the proof.
19

---

Published in Transactions on Machine Learning Research (08/2025)
Lemma 7. Let D be diagonal matrix where each diagonal entry is independently with probability p and 0 with probability 1 p. Then the expected value of the spectral norm |D[|2 is given by
B(I)D = -1- 1.
Proof. Since D is diagonal, its spectral norm |D|[2 is the largest absolute value among its diagonal entries. Therefore, |DI1 if at least one diagonal entry is 1, and |Dll2 = 0 only if all diagonal entries are 0.
Define X as the event that all diagonal entries are 0. The probability of this event, Pr(X), is:
Pr(X) - (1-p)",
since each diagomal entry is 0 independently with probability 1 - p.
Thus, the probability that |D||2 = (i.e., the event X does not occur) is:
1 1- = 100-111-1 p)"
Therefore, the expected value of |D||| is:
D(I = 1 (1 (1- +0 (1-
-1- 11- 1.
This completes the proof.
20


Reviews:
[{"review_id": "48THdBEgTs", "reviewer": "Reviewer_BrW6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\nThe proposed framework for personalized FL based on low-rank factorization is new and interesting. The paper is generally well-written, the work is grounded in theory and experiments demonstrate the performance of the proposed algorithm in practice.\n\n**Weaknesses**\nMany relevant prior works are not discussed/compared with. For example, multiple prior works have addressed various notions of `noisy' labels in FL scenarios [1], [2]. With respect to baselines, why not consider FedHM for comparison? I understand that it is not specifically for personalization, but it would be good to see comparison results for the global model. In relation to that, the work of [3] is also quite relevant for comparison. \n\nAdditionally, have the authors tried an FL type update on $U$? I mean each client updates $U$ based on local information and then server averages the updates from the clients? \n\nFurthermore, I am not sure why the baselines considered in the experiments are so old. Why not compare with Mishchenko et al. (2023); Pillutla et al. (2022), which have been cited in the paper?\n\n[1] Bao, W., Wang, H., Wu, J., & He, J. (2023). Optimizing the collaboration structure in cross-silo federated learning. In ICML.\n[2] Prakash, S., Sima, J., Pan, C., Chien, E., & Milenkovic, O. (2023). Federated classification in hyperbolic spaces via secure aggregation of convex hulls. In ICLR.\n[3] Niu, Y., Prakash, S., Kundu, S., Lee, S., & Avestimehr, S. (2023). Overcoming resource constraints in federated learning: Large models can be trained with only weak clients. Transactions on Machine Learning Research.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper develops a personalized federated learning algorithm. Specifically, the paper focuses on the clustered FL setting where clients are grouped into clusters based on similarities of data distributions without having prior knowledge of cluster memberships. However, instead of directly clustering clients, the authors provide a new formulation based on low-rank matrix factorization, where the $n$ clients' models are assumed to lie in an $r$ dimensional space. The authors then proceed to develop an algorithm where factor models based federated learning is carried out across clients. Particularly, the local model at a client can be described as product between a shared low-rank matrix across clients, and a personal part corresponding to the specific client. The authors also provide a convergence analysis of their algorithm. Experiments with neural networks are also provided.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please make the changes required to address the Weaknesses mentioned above.", "broader_impact_concerns": "I do not have any broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "GNPZDOh16K", "reviewer": "Reviewer_eYV6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The experimental design is diverse and relevant, with careful attention to personalized label mappings.\n\nHowever, the work also has some areas where further clarity or extension would be helpful. First, while the factorization assumption is powerful, the approach may incur high computational or memory overhead when applied to deep models, especially if the entire model is factorized. The paper acknowledges this and suggests partial factorization as future work, so I think the paper is currently okay for TMLR.\n\nSecond, although theoretical guarantees are provided, they are somewhat detached from practical choices like the number of local steps K, which appear crucial for empirical success but are not yet reflected in the theory. Lastly, while the baselines are comprehensive, the paper could benefit from deeper ablation studies: such as sensitivity to the rank r or performance under increasing label corruption", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a well-motivated and technically grounded contribution to the field of Personalized Federated Learning (pFL) through novel low-rank matrix optimization framework. Specifically, the authors propose a new method based on Burer-Monteiro factorization, which assumes that the matrix formed by stacking all clients’ models lies in a low-dimensional subspace.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Here are a few comments:\n\nAbstract: The paper mentions a focus on “clustered FL,” where devices are grouped by data similarities. It would be helpful if you provide a short rationale for why low-rank matrix optimization is especially suitable for discovering or leveraging these clusters.\n\nYou mention using the Burer–Monteiro factorization technique and that convergence guarantees are examined.  It might be clearer to add one or two lines outlining how the method handles potential numerical instability. If available, highlight any theoretical upper bounds on time/space complexity. That can address practicality concerns right away.\n\nThe introduction references the difficulty of balancing local data reliability with other clients’ contributions. This is a strong motivator for your work. However, the text could benefit from a slightly more concrete example. Consider adding a real world example for your motivation\n\nBecause you are using a low-rank matrix factorization, the choice of rank can heavily influence both performance and computational complexity. I would like to see a sensitivity test, and comment on how r must be selected. I understand that the authors have tested several values of r, but it can be expanded with sensitivity plots and insights", "broader_impact_concerns": "-", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "m6ZR2PVhyf", "reviewer": "Reviewer_b2RG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well written but there might be a critical issue.\n\n1. Are the authors aware of the following paper and its follow up works on meta-learning?\n\nhttps://openreview.net/forum?id=48LtSAkxjiX - this paper studies the exact same problem with the same algorithm (modulo a few minor details) as well. The authors also provide stronger guarantees on convergence to the global optimal in stead of a stationary point.\n\nThe same problem formulation, in the context of meta-learning, has been studied already in details (please check the papers that cite the above one as well). \n\nThe authors need to differentiate clearly with this line of work. \n\n2. (Minor) x being the notation for parameters looks weird. Why not use \\theta? \n\n3. The notation f_i subsumes both the loss personalized for model i and the dataset i. It would be good to clarify this", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors have looked at the problem of Personalized Federated Learning (pFL) - the goal is to ensure improved model training taking into account relevance and diversity of samples via a low rank formulation of the underlying parameters", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please differentiate clearly with the existing theoretical work on meta-learning starting with the above one.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)